[flang-commits] [flang] [flang] materialize fir.box when it is from a block argument (PR #184898)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Susan Tan (ス-ザン　タン) (SusanTan)

<details>
<summary>Changes</summary>

We have to materialize `fir.box` before adding a `fir.convert` to a memref type. Otherwise we get:
`'fir.convert' op invalid type conversion'!fir.box<!fir.array<?xi32>>' / 'memref<?xi32, strided<[?], offset: ?>>'`


---
Full diff: https://github.com/llvm/llvm-project/pull/184898.diff


2 Files Affected:

- (modified) flang/lib/Optimizer/Transforms/FIRToMemRef.cpp (+19-8) 
- (modified) flang/test/Transforms/FIRToMemRef/array-coor-block-arg.mlir (+26) 


``````````diff
diff --git a/flang/lib/Optimizer/Transforms/FIRToMemRef.cpp b/flang/lib/Optimizer/Transforms/FIRToMemRef.cpp
index e8c30e62430c0..90b4b01e30f62 100644
--- a/flang/lib/Optimizer/Transforms/FIRToMemRef.cpp
+++ b/flang/lib/Optimizer/Transforms/FIRToMemRef.cpp
@@ -172,6 +172,9 @@ class FIRToMemRef : public fir::impl::FIRToMemRefBase<FIRToMemRef> {
   bool memrefIsDeviceData(Operation *memref) const;
 
   mlir::Attribute findCudaDataAttr(Value val) const;
+
+  Value materializeBoxAddressIfNeeded(Value basePtr, PatternRewriter &rewriter,
+                                      Location loc) const;
 };
 
 void FIRToMemRef::populateShapeAndShift(SmallVectorImpl<Value> &shapeVec,
@@ -251,6 +254,18 @@ mlir::Attribute FIRToMemRef::findCudaDataAttr(Value val) const {
   return nullptr;
 }
 
+Value FIRToMemRef::materializeBoxAddressIfNeeded(Value basePtr,
+                                                 PatternRewriter &rewriter,
+                                                 Location loc) const {
+  if (!isa<fir::BoxType>(basePtr.getType()))
+    return basePtr;
+
+  auto boxAddrOp = fir::BoxAddrOp::create(rewriter, loc, basePtr);
+  if (auto cudaAttr = findCudaDataAttr(basePtr))
+    boxAddrOp->setAttr(cuf::getDataAttrName(), cudaAttr);
+  return boxAddrOp.getResult();
+}
+
 void FIRToMemRef::populateShift(SmallVectorImpl<Value> &vec,
                                 fir::ShiftOp shift) const {
   vec.append(shift.getOrigins().begin(), shift.getOrigins().end());
@@ -468,9 +483,10 @@ FIRToMemRef::convertArrayCoorOp(Operation *memOp, fir::ArrayCoorOp arrayCoorOp,
   FailureOr<Value> converted;
   if (auto blockArg = dyn_cast<BlockArgument>(firMemref)) {
     rewriter.setInsertionPoint(arrayCoorOp);
-    Type memrefTy = typeConverter.convertMemrefType(blockArg.getType());
+    Value basePtr = materializeBoxAddressIfNeeded(blockArg, rewriter, loc);
+    Type memrefTy = typeConverter.convertMemrefType(basePtr.getType());
     converted =
-        fir::ConvertOp::create(rewriter, loc, memrefTy, blockArg).getResult();
+        fir::ConvertOp::create(rewriter, loc, memrefTy, basePtr).getResult();
     rewriter.setInsertionPointAfter(arrayCoorOp);
   } else if ((memref = firMemref.getDefiningOp()) &&
              enableFIRConvertOptimizations && isMarshalLike(memref) &&
@@ -639,12 +655,7 @@ FIRToMemRef::getFIRConvert(Operation *memOp, Operation *op,
 
   if (isa<fir::BoxType>(basePtr.getType())) {
     Operation *baseOp = basePtr.getDefiningOp();
-    auto boxAddrOp = fir::BoxAddrOp::create(rewriter, loc, basePtr);
-
-    if (auto cudaAttr = findCudaDataAttr(basePtr))
-      boxAddrOp->setAttr(cuf::getDataAttrName(), cudaAttr);
-
-    basePtr = boxAddrOp;
+    basePtr = materializeBoxAddressIfNeeded(basePtr, rewriter, loc);
     memrefTy = typeConverter.convertMemrefType(basePtr.getType());
 
     if (baseOp) {
diff --git a/flang/test/Transforms/FIRToMemRef/array-coor-block-arg.mlir b/flang/test/Transforms/FIRToMemRef/array-coor-block-arg.mlir
index 7e0acbae78af4..de0a1832de4c4 100644
--- a/flang/test/Transforms/FIRToMemRef/array-coor-block-arg.mlir
+++ b/flang/test/Transforms/FIRToMemRef/array-coor-block-arg.mlir
@@ -26,3 +26,29 @@ func.func @block_arg_memref(%arg0: !fir.ref<!fir.array<32xi32>>) {
 // CHECK:       memref.store {{%.+}}, [[BASE]][[[IDX]]] : memref<32xi32>
 // CHECK-NOT:   fir.array_coor
 
+// Verify fir.array_coor lowering when the base is a boxed SCF region block
+// argument. The pass must materialize fir.box_addr
+// first and must not emit an illegal direct box->memref fir.convert.
+func.func @block_arg_boxed_array(%arg0: !fir.box<!fir.array<?xi32>>) {
+  %c1 = arith.constant 1 : index
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %loop_res = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg1 = %arg0) -> (!fir.box<!fir.array<?xi32>>) {
+    %idxp1 = arith.addi %i, %c1 : index
+    %val = arith.index_cast %idxp1 : index to i32
+    %idx64 = arith.extsi %val : i32 to i64
+    %elt = fir.array_coor %arg1 %idx64 : (!fir.box<!fir.array<?xi32>>, i64) -> !fir.ref<i32>
+    fir.store %val to %elt : !fir.ref<i32>
+    scf.yield %arg1 : !fir.box<!fir.array<?xi32>>
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @block_arg_boxed_array
+// CHECK:         scf.for {{.*}} iter_args([[BOXARG:%.+]] = %arg0) -> (!fir.box<!fir.array<?xi32>>) {
+// CHECK:           [[BOXADDR:%.+]] = fir.box_addr [[BOXARG]]
+// CHECK:           [[BASE:%.+]] = fir.convert [[BOXADDR]] : (!fir.ref<!fir.array<?xi32>>) -> memref<?xi32{{.*}}>
+// CHECK-NOT:       fir.convert [[BOXARG]] : (!fir.box<!fir.array<?xi32>>) -> memref
+// CHECK:           memref.store {{%.+}}, {{%.+}}[{{%.+}}] : memref<?xi32, strided<[?], offset: ?>>
+// CHECK-NOT:       fir.array_coor
+

``````````

</details>


https://github.com/llvm/llvm-project/pull/184898