[flang][cuf] Add to cuf.alloc/cuf.allocate mem alloc effect

[SusanTan]

Add MemAlloc effect to the result so that cuf.alloc/cuf.allocate can be recognized by FIR alias analysis.

[llvmbot]

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

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

Changes

Add MemAlloc effect to the result so that cuf.alloc/cuf.allocate can be recognized by FIR alias analysis.

---

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

3 Files Affected:

• (modified) flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td (+8-6)
• (added) flang/test/Analysis/AliasAnalysis/cuf-alloc-source-kind.mlir (+22)
• (added) flang/test/Analysis/AliasAnalysis/cuf-allocate-source-kind.mlir (+24)

diff --git a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
index e38738230ffbc..43b383057aa6d 100644
--- a/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
+++ b/flang/include/flang/Optimizer/Dialect/CUF/CUFOps.td
@@ -26,8 +26,7 @@ include "mlir/IR/BuiltinAttributes.td"
 class cuf_Op<string mnemonic, list<Trait> traits>
     : Op<CUFDialect, mnemonic, traits>;
 
-def cuf_AllocOp : cuf_Op<"alloc", [AttrSizedOperandSegments,
-    MemoryEffects<[MemAlloc]>]> {
+def cuf_AllocOp : cuf_Op<"alloc", [AttrSizedOperandSegments]> {
   let summary = "Allocate an object on device";
 
   let description = [{
@@ -47,7 +46,9 @@ def cuf_AllocOp : cuf_Op<"alloc", [AttrSizedOperandSegments,
     cuf_DataAttributeAttr:$data_attr
   );
 
-  let results = (outs fir_ReferenceType:$ptr);
+  // Value-scoped Allocate on the returned reference
+  let results = (outs Res<fir_ReferenceType, "",
+                          [MemAlloc<DefaultResource>]>:$ptr);
 
   let assemblyFormat = [{
     $in_type (`(` $typeparams^ `:` type($typeparams) `)`)?
@@ -84,8 +85,7 @@ def cuf_FreeOp : cuf_Op<"free", [MemoryEffects<[MemFree]>]> {
   let hasVerifier = 1;
 }
 
-def cuf_AllocateOp : cuf_Op<"allocate", [AttrSizedOperandSegments,
-    MemoryEffects<[MemAlloc<DefaultResource>]>]> {
+def cuf_AllocateOp : cuf_Op<"allocate", [AttrSizedOperandSegments]> {
   let summary = "Perform the device allocation of data of an allocatable";
 
   let description = [{
@@ -94,7 +94,9 @@ def cuf_AllocateOp : cuf_Op<"allocate", [AttrSizedOperandSegments,
     is initialized before with the standard flang runtime calls.
   }];
 
-  let arguments = (ins Arg<fir_ReferenceType, "", [MemRead, MemWrite]>:$box,
+  // Value-scoped Allocate on the descriptor being allocated
+  let arguments = (ins Arg<fir_ReferenceType, "",
+                           [MemAlloc<DefaultResource>, MemRead, MemWrite]>:$box,
       Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$errmsg,
       Optional<fir_ReferenceType>:$stream,
       Arg<Optional<AnyRefOrBoxType>, "", [MemWrite]>:$pinned,
diff --git a/flang/test/Analysis/AliasAnalysis/cuf-alloc-source-kind.mlir b/flang/test/Analysis/AliasAnalysis/cuf-alloc-source-kind.mlir
new file mode 100644
index 0000000000000..f062dcb3a3360
--- /dev/null
+++ b/flang/test/Analysis/AliasAnalysis/cuf-alloc-source-kind.mlir
@@ -0,0 +1,22 @@
+// REQUIRES: asserts
+// RUN: fir-opt %s -pass-pipeline='builtin.module(func.func(test-fir-alias-analysis))' -debug-only=fir-alias-analysis --mlir-disable-threading 2>&1 | FileCheck %s
+
+// Verify that a CUF allocation is recognized as SourceKind::Allocate by
+// fir::AliasAnalysis::getSource.
+
+module {
+  func.func @_QQmain() attributes {fir.bindc_name = "TEST"} {
+    // Allocate two independent device arrays and tag the results; with
+    // value-scoped MemAlloc handling in AA, these should be classified as
+    // Allocate and not alias.
+    %a = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a1", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa1", test.ptr = "cuf_alloc_a"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+    %b = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a2", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa2", test.ptr = "cuf_alloc_b"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+    return
+  }
+}
+
+// CHECK-LABEL: Testing : "_QQmain"
+// Distinct allocations should not alias.
+// CHECK: cuf_alloc_a#0 <-> cuf_alloc_b#0: NoAlias
+
+
diff --git a/flang/test/Analysis/AliasAnalysis/cuf-allocate-source-kind.mlir b/flang/test/Analysis/AliasAnalysis/cuf-allocate-source-kind.mlir
new file mode 100644
index 0000000000000..201d841b5cde0
--- /dev/null
+++ b/flang/test/Analysis/AliasAnalysis/cuf-allocate-source-kind.mlir
@@ -0,0 +1,24 @@
+// REQUIRES: asserts
+// RUN: fir-opt %s -pass-pipeline='builtin.module(func.func(test-fir-alias-analysis))' -debug-only=fir-alias-analysis --mlir-disable-threading 2>&1 | FileCheck %s
+
+// Verify that CUF allocation via cuf.allocate is recognized as
+// SourceKind::Allocate by fir::AliasAnalysis::getSource on the box value.
+
+module {
+  func.func @_QQmain() attributes {fir.bindc_name = "TEST"} {
+    // Create two independent device boxes and tag their refs.
+    %a = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a1", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa1", test.ptr = "cuf_allocate_a"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+    %b = cuf.alloc !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a2", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa2", test.ptr = "cuf_allocate_b"} -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+    // Allocate device data for each descriptor; AA should classify the box
+    // values (tagged above) as Allocate sources and not alias.
+    %sa = cuf.allocate %a : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+    %sb = cuf.allocate %b : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {data_attr = #cuf.cuda<device>} -> i32
+    return
+  }
+}
+
+// CHECK-LABEL: Testing : "_QQmain"
+// Distinct allocations via cuf.allocate should not alias.
+// CHECK: cuf_allocate_a#0 <-> cuf_allocate_b#0: NoAlias
+
+

[vzakhari]

I am not sure, but I think MemAlloc<DefaultResource> does not do anything on the operand, and we have to keep the MemAlloc on the operation itself.

Unfortunately, "indirect" allocations are not that easy to express.

[jeanPerier]

An option here could be to revamp this operation a bit so that it returns a fir.box and that a fir.store is generated on top of it to store this result in the original fir.ref<box>.
This decomposition would allow having MemAlloc on the result fir.box and could allow mem2reg to directly link fir.load of the allocatable to this allocatable when possible (not sure how useful that is in the context of CUF).
The downside would be that the "naive" codegen would create an extra desc -> temp desc (new cuf.allocate codegen needs to generate a load to produce the new fir.box result) and temp desc -> copy (the extra fir.store from lowering). This should however rather be easy to remove if LLVM is not already doing it for us.

[SusanTan]

good point! let me try that

[SusanTan]

We decided to revert the changes for cuf.allocate, because it's only used for fortran allocatables which always allocates data for the descriptor, and the descriptor can be distinguished by cuf.alloc and that should be enough for AA to figure out the allocation site.

[clementval]

Is this change done by clang-format? If not, can you revert it?

[vzakhari]

nit: I think this test is redundant, because cuf-alloc-source-kind.mlir above already tests the same thing.