[mlir][DataLayout] Add IsolatedFromAbove to DataLayoutOpInterface

[fabianmcg]

This patch adds the IsolatedFromAbove trait as a dependent trait to the DataLayoutOpInterface op interface.

The motivation behind this change comes from the implementation of the ptr dialect, specifically the ptr.type_offset op. This op produces an int-like value that equates to the size of a memory element. This is useful for ptr arithmetic and indexing arrays. For example:

%f32_off = ptr.type_offset f32 : index
%addr = ptr.ptradd %ptr, %f32_off : !ptr, index
%x = ptr.load %addr : !ptr -> f32 // Read ptr[1]

Without the IsolatedFromAvobe trait in the DL interface, the ptr.type_offset cannot be ConstantLike. Why?
Take the example:

op {DL1} {
  %f32_off0 = ptr.type_offset f32 : index
  op {DL2} {
    %f32_off1 = ptr.type_offset f32 : index
  }
}

If ptr.type_offset were to be ConstantLike then canonicalize would hoist and unique the value. However, that could be wrong as DL2 could have an entry to specify the size that's different from the size in DL1.

The best solution to the above problem is to make DataLayoutOpInterface require the IsolatedFromAbove trait, as it preserves the constness of values in the DL with respect to the canonicalizer.

[github-actions]

⚠️ We detected that you are using a GitHub private e-mail address to contribute to the repo.
Please turn off Keep my email addresses private setting in your account.
See LLVM Discourse for more information.

[llvmbot]

@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir

Author: Fabian Mora (fabianmcg)

Changes

This patch adds the IsolatedFromAbove trait as a dependent trait to the DataLayoutOpInterface op interface.

The motivation behind this change comes from the implementation of the ptr dialect, specifically the ptr.type_offset op. This op produces an int-like value that equates to the size of a memory element. This is useful for ptr arithmetic and indexing arrays. For example:

%f32_off = ptr.type_offset f32 : index
%addr = ptr.ptradd %ptr, %f32_off : !ptr, index
%x = ptr.load %addr : !ptr -> f32 // Read ptr[1]

Without the IsolatedFromAvobe trait in the DL interface, the ptr.type_offset cannot be ConstantLike. Why?
Take the example:

op {DL1} {
  %f32_off0 = ptr.type_offset f32 : index
  op {DL2} {
    %f32_off1 = ptr.type_offset f32 : index
  }
}

If ptr.type_offset were to be ConstantLike then canonicalize would hoist and unique the value. However, that would be wrong as DL2 can have an entry to specify the size that's different from the size in DL1.

The best solution to the above problem is to make DataLayoutOpInterface require the IsolatedFromAbove trait, as it preserves the constness of values in the DL with respect to the canonicalizer.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/GPU/IR/GPUOps.td (+1-1)
• (modified) mlir/include/mlir/Interfaces/DataLayoutInterfaces.td (+1-1)
• (modified) mlir/test/lib/Dialect/Test/TestOps.td (+1-1)

diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 2b1ce573effd0..3241aff4b683c 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -1388,7 +1388,7 @@ def GPU_BarrierOp : GPU_Op<"barrier"> {
 }
 
 def GPU_GPUModuleOp : GPU_Op<"module", [
-      DataLayoutOpInterface, HasDefaultDLTIDataLayout, IsolatedFromAbove,
+      IsolatedFromAbove, DataLayoutOpInterface, HasDefaultDLTIDataLayout,
       NoRegionArguments, SymbolTable, Symbol] # GraphRegionNoTerminator.traits> {
   let summary = "A top level compilation unit containing code to be run on a GPU.";
   let description = [{
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 818b441b9a770..fc701b20cb007 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -565,7 +565,7 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
       }]
     >
   ];
-
+  let dependentTraits = [IsolatedFromAbove];
   let verify = [{ return ::mlir::detail::verifyDataLayoutOp($_op); }];
 }
 
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 94c722038f1cc..f653e4465cfef 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2555,7 +2555,7 @@ def MakeTupleOp: TEST_Op<"make_tuple"> {
 //===----------------------------------------------------------------------===//
 
 def OpWithDataLayoutOp : TEST_Op<"op_with_data_layout",
-                                 [HasDefaultDLTIDataLayout, DataLayoutOpInterface]> {
+                                 [IsolatedFromAbove, HasDefaultDLTIDataLayout, DataLayoutOpInterface]> {
   let summary =
       "An op that uses DataLayout implementation from the Target dialect";
   let regions = (region VariadicRegion<AnyRegion>:$regions);

[joker-eph]

Having an SSA value crossing a boundary with a different data layout seems iffy indeed, but that brings the question of things like gpu.launch right?

[fabianmcg]

but that brings the question of things like gpu.launch right?

Indeed. IMO, I think it also exposes the need to have a clear documented divide between target dependent an independent ops. This divide would ideally also mean tighter verification constraints and the idea of target aware and independent passes. Where for example ptr.type_offset expansion to a constant would fail if there is no target.

Under such system I would classify gpu.launch as target independent, and gpu.module as the target dependent version. Which would align with the status quo, where only gpu.module can have a DL or gpu target attrs. But this needs more discussion.

[rolfmorel]

This makes sense to me. However, currently the DLTI attributes can occur on any op, also those not implementing DataLayoutOpInterface (if I am not mistaken) and will be picked up/respected. Hence, if ptr.type_offset is ConstLike, I think there might still be situations where the canonicalizer might do something you don't want.

@fabianmcg, if this is the case it might be worth opening an issue for it.