[Mlir-commits] [flang] [mlir] [flang][acc] Introduce interface for rematerializable ops (PR #174467)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-openacc

Author: Razvan Lupusoru (razvanlupusoru)

<details>
<summary>Changes</summary>

During the outlining process when offloading acc regions, the body of the compute kernel is separated from its original location and live-in values are handled in various ways including becoming function arguments.

However, some operations are purely synthetic and only make sense when included with another operation (usually such operations exist to simplify IR design). Bounds and shapes are examples where during outlining they should be recreated inside to capture the full information.

Therefore, introduce a new operation interface named OutlineRematerializationOpInterface meant to be attached to such operations. It is currently expected that all such operations are memory effect free to ensure there are no considerations needed when moving or cloning them into outlined regions.

The interface is attached to the following operations:
- acc.bounds (directly in TableGen)
- fir.shape (via external model)
- fir.shape_shift (via external model)
- fir.shift (via external model)
- fir.field_index (via external model)

The pass that will use this interface and associated testing will follow in another pull request.

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


4 Files Affected:

- (modified) flang/include/flang/Optimizer/OpenACC/Support/FIROpenACCOpsInterfaces.h (+13) 
- (modified) flang/lib/Optimizer/OpenACC/Support/RegisterOpenACCExtensions.cpp (+12) 
- (modified) mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td (+2-1) 
- (modified) mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td (+15) 


``````````diff
diff --git a/flang/include/flang/Optimizer/OpenACC/Support/FIROpenACCOpsInterfaces.h b/flang/include/flang/Optimizer/OpenACC/Support/FIROpenACCOpsInterfaces.h
index d7f8f87ccb8bf..28e110bb913ca 100644
--- a/flang/include/flang/Optimizer/OpenACC/Support/FIROpenACCOpsInterfaces.h
+++ b/flang/include/flang/Optimizer/OpenACC/Support/FIROpenACCOpsInterfaces.h
@@ -18,7 +18,11 @@
 namespace fir {
 class AddrOfOp;
 class DeclareOp;
+class FieldIndexOp;
 class GlobalOp;
+class ShapeOp;
+class ShapeShiftOp;
+class ShiftOp;
 } // namespace fir
 
 namespace hlfir {
@@ -77,6 +81,15 @@ struct IndirectGlobalAccessModel
                             mlir::SymbolTable *symbolTable) const;
 };
 
+/// External model for OutlineRematerializationOpInterface.
+/// This interface marks operations that are candidates for rematerialization
+/// during outlining. These operations produce synthetic types or values
+/// that cannot be passed as arguments to outlined regions.
+template <typename Op>
+struct OutlineRematerializationModel
+    : public mlir::acc::OutlineRematerializationOpInterface::ExternalModel<
+          OutlineRematerializationModel<Op>, Op> {};
+
 } // namespace fir::acc
 
 #endif // FLANG_OPTIMIZER_OPENACC_FIROPENACC_OPS_INTERFACES_H_
diff --git a/flang/lib/Optimizer/OpenACC/Support/RegisterOpenACCExtensions.cpp b/flang/lib/Optimizer/OpenACC/Support/RegisterOpenACCExtensions.cpp
index acd1d01ef1e87..60be72d8a7103 100644
--- a/flang/lib/Optimizer/OpenACC/Support/RegisterOpenACCExtensions.cpp
+++ b/flang/lib/Optimizer/OpenACC/Support/RegisterOpenACCExtensions.cpp
@@ -61,6 +61,18 @@ void registerOpenACCExtensions(mlir::DialectRegistry &registry) {
         *ctx);
     fir::TypeDescOp::attachInterface<
         IndirectGlobalAccessModel<fir::TypeDescOp>>(*ctx);
+
+    // Attach OutlineRematerializationOpInterface to FIR operations that
+    // produce synthetic types (shapes, field indices) which cannot be passed
+    // as arguments to outlined regions and must be rematerialized inside.
+    fir::ShapeOp::attachInterface<
+        OutlineRematerializationModel<fir::ShapeOp>>(*ctx);
+    fir::ShapeShiftOp::attachInterface<
+        OutlineRematerializationModel<fir::ShapeShiftOp>>(*ctx);
+    fir::ShiftOp::attachInterface<
+        OutlineRematerializationModel<fir::ShiftOp>>(*ctx);
+    fir::FieldIndexOp::attachInterface<
+        OutlineRematerializationModel<fir::FieldIndexOp>>(*ctx);
   });
 
   // Register HLFIR operation interfaces
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 2375104eadb85..73ca362c6dc3d 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -471,7 +471,8 @@ def OpenACC_VarNameAttr : OpenACC_Attr<"VarName", "var_name"> {
 // Used for data specification in data clauses (2.7.1).
 // Either (or both) extent and upperbound must be specified.
 def OpenACC_DataBoundsOp : OpenACC_Op<"bounds",
-    [AttrSizedOperandSegments, NoMemoryEffect]> {
+    [AttrSizedOperandSegments, NoMemoryEffect,
+     OutlineRematerializationOpInterface]> {
   let summary = "Represents normalized bounds information for acc data clause.";
 
   let description = [{
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
index d958006d58bad..bef228fd5769a 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOpsInterfaces.td
@@ -100,4 +100,19 @@ def IndirectGlobalAccessOpInterface : OpInterface<"IndirectGlobalAccessOpInterfa
   ];
 }
 
+def OutlineRematerializationOpInterface : OpInterface<"OutlineRematerializationOpInterface"> {
+  let cppNamespace = "::mlir::acc";
+
+  let description = [{
+    An interface for operations that are candidates for rematerialization
+    during outlining. These operations produce synthetic types or values
+    that cannot be passed as arguments to outlined regions and must be
+    rematerialized inside the region instead.
+
+    Operations implementing this interface are expected to be memory effect
+    free. Their results are typically used for type construction or providing
+    metadata (such as shape information for arrays).
+  }];
+}
+
 #endif // OPENACC_OPS_INTERFACES

``````````

</details>


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