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

[Susan Tan ス-ザン タン via flang-commits]

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

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: ?>>'`


>From 24c036b767c5617707bc90e0517faff075bf239e Mon Sep 17 00:00:00 2001
From: Susan Tan <zujunt at nvidia.com>
Date: Thu, 5 Mar 2026 13:28:44 -0800
Subject: [PATCH 1/2] add cuf block arg support

---
 .../lib/Optimizer/Transforms/FIRToMemRef.cpp  | 27 +++++++++++++------
 1 file changed, 19 insertions(+), 8 deletions(-)

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) {

>From 12be91903f74f848c8cb20ce4312d40a7897f0c7 Mon Sep 17 00:00:00 2001
From: Susan Tan <zujunt at nvidia.com>
Date: Thu, 5 Mar 2026 13:57:11 -0800
Subject: [PATCH 2/2] add unit test

---
 .../FIRToMemRef/array-coor-block-arg.mlir     | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

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
+