[Dialect] Cleanup and improve granularity of side effects (triton-lang#6476)

This PR cleans up side effect specification to use the `Arg<Type, desc,
[effects...]>` trick in ODS, which is a more concise way to specify
memory effects. At the same time, many ops are specifying coarser side
effects than necessary, and this PR changes them to target a specific
SSA operand or result where application.

Most notably, the many ops that access TensorMemory or GlobalMemory were
made to have finer grained side effects.

Author: Mogball

include/triton/Dialect/Triton/IR/TritonOps.td

@@ -301,7 +301,6 @@ def TT_LoadOp : TT_Op<"load", [
 def TT_StoreOp : TT_Op<"store", [
   SameLoadStoreOperandsShape,
   SameLoadStoreOperandsEncoding,
-  MemoryEffects<[MemWrite<GlobalMemory>]>,
   TypesMatchWith<"value type matches ptr type", "ptr", "value",
                  "getPointeeType($_self)">,
   TypesMatchWith<"mask type matches ptr type", "ptr", "mask",
@@ -312,7 +311,7 @@ def TT_StoreOp : TT_Op<"store", [
 
     let arguments = (
       ins
-      AnyTypeOf<[TT_PtrLike, TT_TensorPtr]>:$ptr,
+      Arg<AnyTypeOf<[TT_PtrLike, TT_TensorPtr]>, "", [MemWrite<GlobalMemory>]>:$ptr,
       TT_Type:$value,
       Optional<TT_BoolLike>:$mask,
       DefaultValuedAttr<DenseI32ArrayAttr, "::llvm::ArrayRef<int32_t>{}">:$boundaryCheck,
@@ -352,8 +351,6 @@ def TT_StoreOp : TT_Op<"store", [
 def TT_AtomicRMWOp : TT_Op<"atomic_rmw", [
   SameOperandsAndResultShape,
   SameOperandsAndResultEncoding,
-  MemoryEffects<[MemRead<GlobalMemory>]>,
-  MemoryEffects<[MemWrite<GlobalMemory>]>,
   TypesMatchWith<"ptr type matches value type", "val", "ptr",
                  "getPointerTypeSameShape($_self)">,
   TypesMatchWith<"mask type matches value type",
@@ -368,9 +365,14 @@ def TT_AtomicRMWOp : TT_Op<"atomic_rmw", [
         return old value at $ptr
     }];
 
-    let arguments = (ins TT_AtomicRMWAttr:$atomic_rmw_op, TT_PtrLike:$ptr,
-                         TT_Type:$val, Optional<TT_BoolLike>:$mask,
-                         TT_MemSemanticAttr:$sem, TT_MemSyncScopeAttr:$scope);
+    let arguments = (ins
+      TT_AtomicRMWAttr:$atomic_rmw_op,
+      Arg<TT_PtrLike, "", [MemWrite<GlobalMemory>, MemRead<GlobalMemory>]>:$ptr,
+      TT_Type:$val,
+      Optional<TT_BoolLike>:$mask,
+      TT_MemSemanticAttr:$sem,
+      TT_MemSyncScopeAttr:$scope
+    );
 
     let results = (outs TT_Type:$result);
 
@@ -382,9 +384,7 @@ def TT_AtomicRMWOp : TT_Op<"atomic_rmw", [
     }];
 }
 
-def TT_AtomicCASOp : TT_Op<"atomic_cas", [MemoryEffects<[MemRead<GlobalMemory>]>,
-                                          MemoryEffects<[MemWrite<GlobalMemory>]>,
-                                          SameOperandsAndResultShape,
+def TT_AtomicCASOp : TT_Op<"atomic_cas", [SameOperandsAndResultShape,
                                           SameOperandsAndResultEncoding]> {
     let summary = "atomic cas";
 
@@ -398,8 +398,12 @@ def TT_AtomicCASOp : TT_Op<"atomic_cas", [MemoryEffects<[MemRead<GlobalMemory>]>
         return $old
     }];
 
-    let arguments = (ins TT_PtrLike:$ptr, TT_Type:$cmp, TT_Type:$val,
-                     TT_MemSemanticAttr:$sem, TT_MemSyncScopeAttr:$scope);
+    let arguments = (ins
+      Arg<TT_PtrLike, "", [MemWrite<GlobalMemory>, MemRead<GlobalMemory>]>:$ptr,
+      TT_Type:$cmp,
+      TT_Type:$val,
+      TT_MemSemanticAttr:$sem,
+      TT_MemSyncScopeAttr:$scope);
 
     let results = (outs TT_Type:$result);
 
@@ -489,8 +493,7 @@ def TT_BroadcastOp : TT_Op<"broadcast", [Pure,
     let hasVerifier = 1;
 }
 
-// cat is not `pure` because it may reorder elements
-def TT_CatOp : TT_Op<"cat", [NoMemoryEffect,
+def TT_CatOp : TT_Op<"cat", [Pure,
                              SameTypeOperands,
                              SameOperandsAndResultElementType]> {
     let summary = "concatenate 2 tensors";
@@ -502,8 +505,7 @@ def TT_CatOp : TT_Op<"cat", [NoMemoryEffect,
     let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `->` type($result)";
 }
 
-def TT_JoinOp : TT_Op<"join", [
-    NoMemoryEffect, SameTypeOperands]> {
+def TT_JoinOp : TT_Op<"join", [Pure, SameTypeOperands]> {
     let summary = "join two tensors along a new, minor dimension";
     let description = [{
         For example, if the two input tensors are 4x8xf32, returns a tensor of
@@ -523,7 +525,7 @@ def TT_JoinOp : TT_Op<"join", [
 }
 
 def TT_SplitOp : TT_Op<"split", [
-  NoMemoryEffect,
+  Pure,
   InferTypeOpWithLayoutEquivalence,
   TypesMatchWith<"outLHS and outRHS types match",
                   "outLHS", "outRHS", "$_self">,
@@ -1256,18 +1258,15 @@ def ReturnOp : TT_Op<"return", [Pure, HasParent<"FuncOp">, /*MemRefsNormalizable
 }
 
 
-def TT_DescriptorLoadOp : TT_Op<"descriptor_load", [MemoryEffects<[MemRead<GlobalMemory>]>]> {
+def TT_DescriptorLoadOp : TT_Op<"descriptor_load"> {
   let summary = "Load from descriptor";
   let description = [{
     This operation will be lowered to Nvidia TMA load operation on targets supporting it.
     `desc` is a tensor descriptor object.
     The destination tensor type and shape must match the descriptor otherwise the result is undefined.
-
-    This is an escape hatch and is only there for testing/experimenting.
-    This op will be removed in the future.
   }];
   let arguments = (ins
-    TT_TensorDescType:$desc,
+    Arg<TT_TensorDescType, "", [MemRead<GlobalMemory>]>:$desc,
     Variadic<I32>:$indices,
     DefaultValuedAttr<TT_CacheModifierAttr, "::mlir::triton::CacheModifier::NONE">:$cache,
     DefaultValuedAttr<TT_EvictionPolicyAttr, "::mlir::triton::EvictionPolicy::NORMAL">:$evict
@@ -1287,20 +1286,15 @@ def TT_DescriptorLoadOp : TT_Op<"descriptor_load", [MemoryEffects<[MemRead<Globa
   let hasVerifier = 1;
 }
 
-def TT_DescriptorStoreOp : TT_Op<"descriptor_store", [
-    MemoryEffects<[MemRead<GlobalMemory>, MemWrite<GlobalMemory>]>,
-]> {
+def TT_DescriptorStoreOp : TT_Op<"descriptor_store"> {
   let summary = "store value based on descriptor";
   let description = [{
     This operation will be lowered to Nvidia TMA store operation on targets supporting it.
     `desc` is a tensor descriptor object.
     The shape and types of `src` must match the descriptor otherwise the result is undefined.
-
-    This is an escape hatch and is only there for testing/experimenting.
-    This op will be removed in the future.
   }];
   let arguments = (ins
-    TT_TensorDescType:$desc,
+    Arg<TT_TensorDescType, "", [MemRead<GlobalMemory>, MemWrite<GlobalMemory>]>:$desc,
     TT_Tensor:$src,
     Variadic<I32>:$indices
   );
@@ -1313,22 +1307,19 @@ def TT_DescriptorStoreOp : TT_Op<"descriptor_store", [
   let hasVerifier = 1;
 }
 
-def TT_DescriptorGatherOp : TT_Op<"descriptor_gather", [MemoryEffects<[MemRead<GlobalMemory>]>]> {
+def TT_DescriptorGatherOp : TT_Op<"descriptor_gather"> {
   let summary = "gather multiple rows from a descriptor into a single tensor";
   let description = [{
     The `tt.descriptor_gather` op will be lowered to NVIDIA TMA
-    load operations on targets that support it.
+    gather operations on targets that support it.
 
     `desc_ptr` is a pointer to the TMA descriptor allocated in global memory.
     The descriptor block must have 1 row and the indices must be a 1D tensor.
     Accordingly, the result is a 2D tensor multiple rows.
-
-    This is an escape hatch and is only there for testing/experimenting. This
-    op will be removed in the future.
   }];
 
   let arguments = (ins
-    TT_TensorDescType:$desc,
+    Arg<TT_TensorDescType, "", [MemRead<GlobalMemory>]>:$desc,
     RankedTensorOf<[I32]>:$x_offsets,
     I32:$y_offset
   );
@@ -1349,11 +1340,19 @@ def TT_DescriptorGatherOp : TT_Op<"descriptor_gather", [MemoryEffects<[MemRead<G
   }];
 }
 
-def TT_DescriptorScatterOp : TT_Op<"descriptor_scatter", [
-    MemoryEffects<[MemRead<GlobalMemory>, MemWrite<GlobalMemory>]>,
-]> {
+def TT_DescriptorScatterOp : TT_Op<"descriptor_scatter"> {
+  let summary = "scatter multiple rows from a tensor into a descriptor";
+  let description = [{
+    The `tt.descriptor_scatter` op will be lowered to NVIDIA TMA
+    scatter operations on targets that support it.
+
+    `desc_ptr` is a pointer to the TMA descriptor allocated in global memory.
+    The descriptor block must have 1 row and the indices must be a 1D tensor.
+    Accordingly, the result is a 2D tensor multiple rows.
+  }];
+
   let arguments = (ins
-    TT_TensorDescType:$desc,
+    Arg<TT_TensorDescType, "", [MemRead<GlobalMemory>, MemWrite<GlobalMemory>]>:$desc,
     RankedTensorOf<[I32]>:$x_offsets,
     I32:$y_offset,
     TT_Tensor:$src

include/triton/Dialect/TritonGPU/IR/TritonGPUOps.td

@@ -77,7 +77,6 @@ def TTG_AsyncCommitGroupOp : TTG_Op<"async_commit_group"> {
 
 def TTG_AsyncCopyGlobalToLocalOp : TTG_Op<"async_copy_global_to_local", [
   AttrSizedOperandSegments,
-  DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
   TypesMatchWith<"infer mask type from src type",
                  "src", "mask", "getI1SameShape($_self)",
                  "($_op.getOperands().size() <= 3) || std::equal_to<>()">,
@@ -95,9 +94,9 @@ def TTG_AsyncCopyGlobalToLocalOp : TTG_Op<"async_copy_global_to_local", [
     operands are the same as tt.load.
   }];
 
-  let arguments = (
-    ins TT_PtrTensor:$src,
-    TTG_MemDescType:$result,
+  let arguments = (ins
+    Arg<TT_PtrTensor, "", [MemRead<GlobalMemory>]>:$src,
+    Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$result,
     Optional<I1Tensor>:$mask,
     Optional<TT_Type>:$other,
     DefaultValuedAttr<TT_CacheModifierAttr, "triton::CacheModifier::NONE">:$cache,
@@ -146,11 +145,11 @@ def TTG_LocalAllocOp : TTG_Op<"local_alloc", [DeclareOpInterfaceMethods<MemoryEf
 
     Explicitly deallocating a buffer is optional; see local_dealloc.
   }];
-  let arguments = (
-    ins
+  let arguments = (ins
     Optional<TT_Tensor>:$src,
     OptionalAttr<I32Attr>:$alignment
   );
+  let results = (outs TTG_MemDescType:$result);
 
   let builders = [
     OpBuilder<(ins "Type":$result),
@@ -172,13 +171,12 @@ def TTG_LocalAllocOp : TTG_Op<"local_alloc", [DeclareOpInterfaceMethods<MemoryEf
     ($src^)? attr-dict `:` functional-type(operands, results)
   }];
 
-  let results = (outs TTG_MemDescType:$result);
   let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 // Deallocate shared memory
-def TTG_LocalDeallocOp : TTG_Op<"local_dealloc", [MemoryEffects<[MemFree<SharedMemory>]>]> {
+def TTG_LocalDeallocOp : TTG_Op<"local_dealloc"> {
   let summary = "dealloc buffer";
 
   let description = [{
@@ -195,7 +193,7 @@ def TTG_LocalDeallocOp : TTG_Op<"local_dealloc", [MemoryEffects<[MemFree<SharedM
     operand.
   }];
 
-  let arguments = (ins TTG_MemDescType:$src);
+  let arguments = (ins Arg<TTG_MemDescType, "", [MemFree<SharedMemory>]>:$src);
 
   // Use qualified() otherwise "!ttg.memdesc<X>" is printed as "<X>".
   let assemblyFormat = [{$src attr-dict `:` qualified(type($src))}];
@@ -251,13 +249,17 @@ def TTG_MemDescTransOp : TTG_Op<"memdesc_trans", [Pure,
   let hasFolder = 1;
 }
 
-def TTG_LocalLoadOp : TTG_Op<"local_load", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+def TTG_LocalLoadOp : TTG_Op<"local_load"> {
   let summary = "Load a buffer from local memory into a distributed tensor";
 
   let description = [{
     Load a tensor from the local memory descriptor into a distributed tensor.
   }];
-  let arguments = (ins TTG_MemDescType:$src, Optional<TTG_AsyncToken> :$token);
+  let arguments = (ins
+    Arg<TTG_MemDescType, "", [MemRead<SharedMemory>]>:$src,
+    Optional<TTG_AsyncToken>:$token
+  );
+  let results = (outs TT_Tensor:$result);
 
   let builders = [
       OpBuilder<(ins "Type":$retType, "Value":$src),
@@ -268,16 +270,18 @@ def TTG_LocalLoadOp : TTG_Op<"local_load", [DeclareOpInterfaceMethods<MemoryEffe
   // Use qualified() otherwise "!ttg.memdesc<X>" is printed as "<X>".
   let assemblyFormat = [{$src (`token` $token^)? attr-dict `:` qualified(type($src)) `->` type($result)}];
 
-  let results = (outs TT_Tensor:$result);
 }
 
-def TTG_LocalStoreOp : TTG_Op<"local_store", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+def TTG_LocalStoreOp : TTG_Op<"local_store"> {
   let summary = "Store a distributed tensor into a buffer in local memory";
 
   let description = [{
     Store a distributed tensor into a buffer in local memory.
   }];
-  let arguments = (ins TT_Tensor:$src, TTG_MemDescType:$dst);
+  let arguments = (ins
+    TT_Tensor:$src,
+    Arg<TTG_MemDescType, "", [MemWrite<SharedMemory>]>:$dst
+  );
 
   let hasVerifier = 1;
   // Use qualified() otherwise "!ttg.memdesc<X>" is printed as "<X>".
@@ -313,7 +317,7 @@ def TTG_Fp4ToFpOp : TTG_Op<"fp4_to_fp", [Pure]> {
 }
 
 // Allocate global memory
-def TTG_GlobalScratchAllocOp : TTG_Op<"global_scratch_alloc", [MemoryEffects<[MemAlloc<GlobalMemory>]>]> {
+def TTG_GlobalScratchAllocOp : TTG_Op<"global_scratch_alloc"> {
   let summary = "allocate a global memory buffer";
   let description = [{
     This operation allocates a buffer in global memory that is private to the current program.
@@ -323,7 +327,7 @@ def TTG_GlobalScratchAllocOp : TTG_Op<"global_scratch_alloc", [MemoryEffects<[Me
     I32Attr:$nbytes,
     I32Attr:$alignment
   );
-  let results = (outs TT_Ptr:$result);
+  let results = (outs Arg<TT_Ptr, "", [MemAlloc<GlobalMemory>]>:$result);
 
   let builders = [
     OpBuilder<(ins "Type":$result, "int32_t":$nbytes, "int32_t":$alignment),