Skip to content

Commit 4c0e8e8

Browse files
committed
Use rebox, assumed-rank handling, expand tests
Signed-off-by: Kajetan Puchalski <[email protected]>
1 parent 837d45b commit 4c0e8e8

File tree

2 files changed

+75
-6
lines changed

2 files changed

+75
-6
lines changed

flang/lib/Optimizer/HLFIR/Transforms/InlineHLFIRCopyIn.cpp

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ InlineCopyInConversion::matchAndRewrite(hlfir::CopyInOp copyIn,
4848
fir::FirOpBuilder builder(rewriter, copyIn.getOperation());
4949
mlir::Location loc = copyIn.getLoc();
5050
hlfir::Entity inputVariable{copyIn.getVar()};
51+
mlir::Type resultAddrType = copyIn.getCopiedIn().getType();
5152
if (!fir::isa_trivial(inputVariable.getFortranElementType()))
5253
return rewriter.notifyMatchFailure(copyIn,
5354
"CopyInOp's data type is not trivial");
@@ -66,6 +67,10 @@ InlineCopyInConversion::matchAndRewrite(hlfir::CopyInOp copyIn,
6667
return rewriter.notifyMatchFailure(copyIn,
6768
"CopyInOp has no direct CopyOut");
6869

70+
if (mlir::cast<fir::BaseBoxType>(resultAddrType).isAssumedRank())
71+
return rewriter.notifyMatchFailure(copyIn,
72+
"The result array is assumed-rank");
73+
6974
// Only inline the copy_in when copy_out does not need to be done, i.e. in
7075
// case of intent(in).
7176
if (copyOut.getVar())
@@ -85,12 +90,9 @@ InlineCopyInConversion::matchAndRewrite(hlfir::CopyInOp copyIn,
8590
.genThen([&]() {
8691
mlir::Value result = inputVariable;
8792
if (fir::isPointerType(inputVariable.getType())) {
88-
auto boxAddr = builder.create<fir::BoxAddrOp>(loc, inputVariable);
89-
fir::ReferenceType refTy = fir::ReferenceType::get(sequenceType);
90-
mlir::Value refVal = builder.createConvert(loc, refTy, boxAddr);
91-
mlir::Value shape = hlfir::genShape(loc, builder, inputVariable);
92-
result = builder.create<fir::EmboxOp>(loc, resultBoxType, refVal,
93-
shape);
93+
result = builder.create<fir::ReboxOp>(
94+
loc, resultBoxType, inputVariable, mlir::Value{},
95+
mlir::Value{});
9496
}
9597
builder.create<fir::ResultOp>(
9698
loc, mlir::ValueRange{result, builder.createBool(loc, false)});
@@ -142,10 +144,13 @@ InlineCopyInConversion::matchAndRewrite(hlfir::CopyInOp copyIn,
142144
mlir::OpResult resultBox = results[0];
143145
mlir::OpResult needsCleanup = results[1];
144146

147+
// Prepare the corresponding copyOut to free the temporary if it is required
145148
auto alloca = builder.create<fir::AllocaOp>(loc, resultBox.getType());
146149
auto store = builder.create<fir::StoreOp>(loc, resultBox, alloca);
150+
rewriter.startOpModification(copyOut);
147151
copyOut->setOperand(0, store.getMemref());
148152
copyOut->setOperand(1, needsCleanup);
153+
rewriter.finalizeOpModification(copyOut);
149154

150155
rewriter.replaceOp(copyIn, {resultBox, builder.genNot(loc, isContiguous)});
151156
return mlir::success();

flang/test/HLFIR/inline-hlfir-copy-in.fir

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ func.func private @_test_inline_copy_in(%arg0: !fir.box<!fir.array<?x?x?xf64>> {
3434
return
3535
}
3636

37+
// CHECK: #loop_vectorize = #llvm.loop_vectorize<disable = true>
38+
// CHECK: #loop_annotation = #llvm.loop_annotation<vectorize = #loop_vectorize>
3739
// CHECK-LABEL: func.func private @_test_inline_copy_in(
3840
// CHECK-SAME: %[[VAL_0:.*]]: !fir.box<!fir.array<?x?x?xf64>> {fir.bindc_name = "x"},
3941
// CHECK-SAME: %[[VAL_1:.*]]: !fir.ref<i32> {fir.bindc_name = "i"},
@@ -140,3 +142,65 @@ func.func private @_test_no_inline_copy_in(%arg0: !fir.box<!fir.array<?x?x?xf64>
140142
// CHECK: hlfir.end_associate %[[VAL_22:.*]]#1, %[[VAL_22:.*]]#2 : !fir.ref<i32>, i1
141143
// CHECK: return
142144
// CHECK: }
145+
146+
// Test not inlining optional dummy arguments (no direct copy-out)
147+
func.func @_QPoptional_copy_in_out(%arg0: !fir.box<!fir.array<?xf32>> {fir.bindc_name = "x", fir.optional}) {
148+
%false = arith.constant false
149+
%0 = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
150+
%1 = fir.dummy_scope : !fir.dscope
151+
%2:2 = hlfir.declare %arg0 dummy_scope %1 {fortran_attrs = #fir.var_attrs<optional>, uniq_name = "_QFoptional_copy_in_outEx"} : (!fir.box<!fir.array<?xf32>>, !fir.dscope) -> (!fir.box<!fir.array<?xf32>>, !fir.box<!fir.array<?xf32>>)
152+
%3 = fir.is_present %2#0 : (!fir.box<!fir.array<?xf32>>) -> i1
153+
%4:2 = fir.if %3 -> (!fir.ref<!fir.array<?xf32>>, i1) {
154+
%5:2 = hlfir.copy_in %2#0 to %0 : (!fir.box<!fir.array<?xf32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.box<!fir.array<?xf32>>, i1)
155+
%6 = fir.box_addr %5#0 : (!fir.box<!fir.array<?xf32>>) -> !fir.ref<!fir.array<?xf32>>
156+
fir.result %6, %5#1 : !fir.ref<!fir.array<?xf32>>, i1
157+
} else {
158+
%5 = fir.absent !fir.ref<!fir.array<?xf32>>
159+
fir.result %5, %false : !fir.ref<!fir.array<?xf32>>, i1
160+
}
161+
fir.call @_QPtakes_optional_explicit(%4#0) fastmath<contract> : (!fir.ref<!fir.array<?xf32>>) -> ()
162+
hlfir.copy_out %0, %4#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, i1) -> ()
163+
return
164+
}
165+
166+
// CHECK-LABEL: func.func @_QPoptional_copy_in_out(
167+
// CHECK-SAME: %[[ARG_0:.*]]: !fir.box<!fir.array<?xf32>> {fir.bindc_name = "x", fir.optional}) {
168+
// CHECK: %false = arith.constant false
169+
// CHECK: %[[VAL_0:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>>
170+
// CHECK: %[[VAL_1:.*]] = fir.dummy_scope : !fir.dscope
171+
// CHECK: %[[VAL_2:.*]]:2 = hlfir.declare %[[ARG_0:.*]] dummy_scope %[[VAL_1:.*]] {fortran_attrs = #fir.var_attrs<optional>, uniq_name = "_QFoptional_copy_in_outEx"} : (!fir.box<!fir.array<?xf32>>, !fir.dscope) -> (!fir.box<!fir.array<?xf32>>, !fir.box<!fir.array<?xf32>>)
172+
// CHECK: %[[VAL_3:.*]] = fir.is_present %[[VAL_2:.*]]#0 : (!fir.box<!fir.array<?xf32>>) -> i1
173+
// CHECK: %[[VAL_4:.*]]:2 = fir.if %[[VAL_3:.*]] -> (!fir.ref<!fir.array<?xf32>>, i1) {
174+
// CHECK: %[[VAL_5:.*]]:2 = hlfir.copy_in %[[VAL_2:.*]]#0 to %[[VAL_0:.*]] : (!fir.box<!fir.array<?xf32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.box<!fir.array<?xf32>>, i1)
175+
// CHECK: %[[VAL_6:.*]] = fir.box_addr %[[VAL_5:.*]]#0 : (!fir.box<!fir.array<?xf32>>) -> !fir.ref<!fir.array<?xf32>>
176+
// CHECK: fir.result %[[VAL_6:.*]], %[[VAL_5:.*]]#1 : !fir.ref<!fir.array<?xf32>>, i1
177+
// CHECK: } else {
178+
// CHECK: %[[VAL_5:.*]] = fir.absent !fir.ref<!fir.array<?xf32>>
179+
// CHECK: fir.result %[[VAL_5:.*]], %false : !fir.ref<!fir.array<?xf32>>, i1
180+
// CHECK: }
181+
// CHECK: fir.call @_QPtakes_optional_explicit(%[[VAL_4:.*]]#0) fastmath<contract> : (!fir.ref<!fir.array<?xf32>>) -> ()
182+
// CHECK: hlfir.copy_out %[[VAL_0:.*]], %[[VAL_4:.*]]#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, i1) -> ()
183+
// CHECK: return
184+
// CHECK: }
185+
186+
// Test not inlining of assumed-rank arrays
187+
func.func @_QPtest_copy_in_out_2(%arg0: !fir.box<!fir.array<*:f32>> {fir.bindc_name = "x"}) {
188+
%0 = fir.alloca !fir.box<!fir.heap<!fir.array<*:f32>>>
189+
%1 = fir.dummy_scope : !fir.dscope
190+
%2:2 = hlfir.declare %arg0 dummy_scope %1 {uniq_name = "_QFtest_copy_in_out_2Ex"} : (!fir.box<!fir.array<*:f32>>, !fir.dscope) -> (!fir.box<!fir.array<*:f32>>, !fir.box<!fir.array<*:f32>>)
191+
%3:2 = hlfir.copy_in %2#0 to %0 : (!fir.box<!fir.array<*:f32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<*:f32>>>>) -> (!fir.box<!fir.array<*:f32>>, i1)
192+
fir.call @_QPtakes_contiguous_intentin(%3#0) fastmath<contract> : (!fir.box<!fir.array<*:f32>>) -> ()
193+
hlfir.copy_out %0, %3#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<*:f32>>>>, i1) -> ()
194+
return
195+
}
196+
197+
// CHECK-LABEL: func.func @_QPtest_copy_in_out_2(
198+
// CHECK-SAME: %[[ARG_0]]: !fir.box<!fir.array<*:f32>> {fir.bindc_name = "x"}) {
199+
// CHECK: %[[VAL_0]] = fir.alloca !fir.box<!fir.heap<!fir.array<*:f32>>>
200+
// CHECK: %[[VAL_1]] = fir.dummy_scope : !fir.dscope
201+
// CHECK: %[[VAL_2]]:2 = hlfir.declare %[[ARG_0]] dummy_scope %[[VAL_1]] {uniq_name = "_QFtest_copy_in_out_2Ex"} : (!fir.box<!fir.array<*:f32>>, !fir.dscope) -> (!fir.box<!fir.array<*:f32>>, !fir.box<!fir.array<*:f32>>)
202+
// CHECK: %[[VAL_3]]:2 = hlfir.copy_in %[[VAL_2]]#0 to %[[VAL_0]] : (!fir.box<!fir.array<*:f32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<*:f32>>>>) -> (!fir.box<!fir.array<*:f32>>, i1)
203+
// CHECK: fir.call @_QPtakes_contiguous_intentin(%[[VAL_3]]#0) fastmath<contract> : (!fir.box<!fir.array<*:f32>>) -> ()
204+
// CHECK: hlfir.copy_out %[[VAL_0]], %[[VAL_3]]#1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<*:f32>>>>, i1) -> ()
205+
// CHECK: return
206+
// CHECK: }

0 commit comments

Comments
 (0)