[flang-commits] [flang] [flang][cuda] Avoid crash when source is a block argument (PR #217459)

[via flang-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

Author: Valentin Clement (バレンタイン クレメン) (clementval)

<details>
<summary>Changes</summary>

When the value is a block argument, getDefiningOp() returns null. Avoid crash when this is the case by checking for the defining op first. 

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


2 Files Affected:

- (modified) flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp (+6-2) 
- (modified) flang/test/Fir/CUDA/cuda-data-transfer.fir (+16) 


``````````diff
diff --git a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
index 8e16473eaf982..e9e5989bee5b4 100644
--- a/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/CUDA/CUFOpConversion.cpp
@@ -337,14 +337,18 @@ struct CUFDataTransferOpConversion
     }
 
     auto materializeBoxIfNeeded = [&](mlir::Value val) -> mlir::Value {
-      if (mlir::isa<fir::EmboxOp, fir::ReboxOp>(val.getDefiningOp())) {
+      // val can be a block argument and therefore has no defining operation.
+      mlir::Operation *defOp = val.getDefiningOp();
+      if (!defOp)
+        return val;
+      if (mlir::isa<fir::EmboxOp, fir::ReboxOp>(defOp)) {
         // Materialize the box to memory to be able to call the runtime.
         mlir::Value box = builder.createTemporary(loc, val.getType());
         fir::StoreOp::create(builder, loc, val, box);
         return box;
       }
       if (mlir::isa<fir::BaseBoxType>(val.getType()))
-        if (auto loadOp = mlir::dyn_cast<fir::LoadOp>(val.getDefiningOp()))
+        if (auto loadOp = mlir::dyn_cast<fir::LoadOp>(defOp))
           return loadOp.getMemref();
       return val;
     };
diff --git a/flang/test/Fir/CUDA/cuda-data-transfer.fir b/flang/test/Fir/CUDA/cuda-data-transfer.fir
index 235bcae1ccf98..37f5448ba8b89 100644
--- a/flang/test/Fir/CUDA/cuda-data-transfer.fir
+++ b/flang/test/Fir/CUDA/cuda-data-transfer.fir
@@ -725,5 +725,21 @@ func.func @_QQmain() attributes {fir.bindc_name = "T"} {
 // CHECK-LABEL: func.func @_QQmain() 
 // CHECK: fir.call @_FortranACUFDataTransferDescDesc
 
+func.func @_QPdesc_block_arg(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>> {cuf.data_attr = #cuf.cuda<device>}, %arg1: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>> {cuf.data_attr = #cuf.cuda<device>}, %arg2: !fir.heap<!fir.array<?xf64>>, %arg3: index, %arg4: i1) {
+  cf.cond_br %arg4, ^bb1, ^bb2
+^bb1:
+  cf.br ^bb3(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>)
+^bb2:
+  cf.br ^bb3(%arg1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>)
+^bb3(%0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>):
+  %1 = fir.shape %arg3 : (index) -> !fir.shape<1>
+  %2 = fir.embox %arg2(%1) : (!fir.heap<!fir.array<?xf64>>, !fir.shape<1>) -> !fir.box<!fir.array<?xf64>>
+  cuf.data_transfer %0 to %2 {transfer_kind = #cuf.cuda_transfer<device_device>} : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf64>>>>, !fir.box<!fir.array<?xf64>>
+  return
+}
+
+// CHECK-LABEL: func.func @_QPdesc_block_arg
+// CHECK: fir.call @_FortranACUFDataTransferDescDesc
+
 } // end of module
 

``````````

</details>


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