[MLIR] Add stage to side effect

    [MLIR] Add stage and effectOnFullRegion to side effect

    This patch add stage and effectOnFullRegion to side effect for optimization pass
    to obtain more accurate information.
    Stage uses numbering to track the side effects's stage of occurrence.
    EffectOnFullRegion indicates if effect act on every single value of resource.

    RFC disscussion: https://discourse.llvm.org/t/rfc-add-effect-index-in-memroy-effect/72235
    Differential Revision: https://reviews.llvm.org/D156087

Reviewed By: mehdi_amini, Mogball

Differential Revision: https://reviews.llvm.org/D156087

Author: cxy-1993

mlir/docs/Rationale/SideEffectsAndSpeculation.md

@@ -47,7 +47,10 @@ Operations with implicit behaviors can be broadly categorized as follows:
    `longjmp`, operations that throw exceptions.
 
 Finally, a given operation may have a combination of the above implicit
-behaviors.
+behaviors. The combination of implicit behaviors during the execution of the
+operation may be ordered. We use 'stage' to label the order of implicit
+behaviors during the execution of 'op'. Implicit behaviors with a lower stage
+number happen earlier than those with a higher stage number.
 
 ## Modeling
 
@@ -76,6 +79,10 @@ When adding a new op, ask:
 
 1. Does it read from or write to the heap or stack? It should probably implement
    `MemoryEffectsOpInterface`.
+1. Does these side effects ordered? It should probably set the stage of
+   side effects to make analysis more accurate.
+1. Does These side effects act on every single value of resource? It probably
+   should set the FullEffect on effect.
 1. Does it have side effects that must be preserved, like a volatile store or a
    syscall? It should probably implement `MemoryEffectsOpInterface` and model
    the effect as a read from or write to an abstract `Resource`. Please start an
@@ -91,3 +98,83 @@ When adding a new op, ask:
 1. Is your operation free of side effects and can be freely hoisted, introduced
    and eliminated? It should probably be marked `Pure`. (TODO: revisit this name
    since it has overloaded meanings in C++.)
+
+## Examples
+
+This section describes a few very simple examples that help understand how to
+add side effect correctly.
+
+### SIMD compute operation
+
+If we have a SIMD backend dialect with a "simd.abs" operation, which reads all
+values from the source memref, calculates their absolute values, and writes them
+to the target memref.
+
+```mlir
+  func.func @abs(%source : memref<10xf32>, %target : memref<10xf32>) {
+    simd.abs(%source, %target) : memref<10xf32> to memref<10xf32>
+    return
+  }
+```
+
+The abs operation reads each individual value from the source resource and then
+writes these values to each corresponding value in the target resource.
+Therefore, we need to specify a read side effect for the source and a write side
+effect for the target. The read side effect occurs before the write side effect,
+so we need to mark the read stage as earlier than the write stage. Additionally,
+we need to indicate that these side effects apply to each individual value in
+the resource.
+
+A typical approach is as follows:
+``` mlir
+  def AbsOp : SIMD_Op<"abs", [...] {
+    ...
+
+    let arguments = (ins Arg<AnyRankedOrUnrankedMemRef, "the source memref",
+                             [MemReadAt<0, FullEffect>]>:$source,
+                         Arg<AnyRankedOrUnrankedMemRef, "the target memref",
+                             [MemWriteAt<1, FullEffect>]>:$target);
+
+    ...
+  }
+```
+
+In the above example, we attach the side effect [MemReadAt<0, FullEffect>] to
+the source, indicating that the abs operation reads each individual value from
+the source during stage 0. Likewise, we attach the side effect
+[MemWriteAt<1, FullEffect>] to the target, indicating that the abs operation
+writes to each individual value within the target during stage 1 (after reading
+from the source).
+
+### Load like operation
+
+Memref.load is a typical load like operation:
+```mlir
+  func.func @foo(%input : memref<10xf32>, %index : index) -> f32 {
+    %result = memref.load  %input[index] : memref<10xf32>
+    return %result : f32
+  }
+```
+
+The load like operation reads a single value from the input memref and returns
+it. Therefore, we need to specify a partial read side effect for the input
+memref, indicating that not every single value is used.
+
+A typical approach is as follows:
+``` mlir
+  def LoadOp : MemRef_Op<"load", [...] {
+    ...
+
+    let arguments = (ins Arg<AnyMemRef, "the reference to load from",
+                             [MemReadAt<0, PartialEffect>]>:$memref,
+                         Variadic<Index>:$indices,
+                         DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
+
+    ...
+  }
+```
+
+In the above example, we attach the side effect [MemReadAt<0, PartialEffect>] to
+the source, indicating that the load operation reads parts of values from the
+memref during stage 0. Since side effects typically occur at stage 0 and are
+partial by default, we can abbreviate it as "[MemRead]".

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

@@ -383,7 +383,8 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
   }];
 
   let arguments = (ins Arg<AnyRankedOrUnrankedMemRef,
-                       "the reference to load from", [MemRead]>:$memref,
+                           "the reference to load from",
+                           [MemReadAt<0, FullEffect>]>:$memref,
                        UnitAttr:$restrict, UnitAttr:$writable);
   let results = (outs AnyTensor:$result);
 

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

@@ -1228,7 +1228,7 @@ def GPU_AllocOp : GPU_Op<"alloc", [
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    Variadic<Index>:$dynamicSizes, Variadic<Index>:$symbolOperands,
                    UnitAttr:$hostShared);
-  let results = (outs Res<AnyMemRef, "", [MemAlloc]>:$memref,
+  let results = (outs Res<AnyMemRef, "", [MemAllocAt<0, FullEffect>]>:$memref,
                  Optional<GPU_AsyncToken>:$asyncToken);
 
   let extraClassDeclaration = [{
@@ -1268,7 +1268,7 @@ def GPU_DeallocOp : GPU_Op<"dealloc", [GPU_AsyncOpInterface]> {
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   Arg<AnyMemRef, "", [MemFree]>:$memref);
+                   Arg<AnyMemRef, "", [MemFreeAt<0, FullEffect>]>:$memref);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
@@ -1299,8 +1299,8 @@ def GPU_MemcpyOp : GPU_Op<"memcpy", [GPU_AsyncOpInterface]> {
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   Arg<AnyMemRef, "", [MemWrite]>:$dst,
-                   Arg<AnyMemRef, "", [MemRead]>:$src);
+                   Arg<AnyMemRef, "", [MemWriteAt<0, FullEffect>]>:$dst,
+                   Arg<AnyMemRef, "", [MemReadAt<0, FullEffect>]>:$src);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
@@ -1335,7 +1335,7 @@ def GPU_MemsetOp : GPU_Op<"memset",
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
-                   Arg<AnyMemRef, "", [MemWrite]>:$dst,
+                   Arg<AnyMemRef, "", [MemWriteAt<0, FullEffect>]>:$dst,
                    Arg<AnyType, "">:$value);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
@@ -1390,7 +1390,8 @@ def GPU_SubgroupMmaLoadMatrixOp : GPU_Op<"subgroup_mma_load_matrix",
     ```
   }];
 
-  let arguments = (ins Arg<GPU_MMAMemRef, "", [MemRead]>:$srcMemref,
+  let arguments = (ins Arg<GPU_MMAMemRef, "",
+                          [MemReadAt<0, FullEffect>]>:$srcMemref,
                   Variadic<Index>:$indices,
                   IndexAttr:$leadDimension,
                   OptionalAttr<UnitAttr>:$transpose);
@@ -1431,7 +1432,7 @@ def GPU_SubgroupMmaStoreMatrixOp : GPU_Op<"subgroup_mma_store_matrix",
   }];
 
   let arguments = (ins Arg<MMAMatrixOf<[SI8, UI8, I32, F16, F32]>>:$src,
-                  Arg<GPU_MMAMemRef, "",[MemWrite]>:$dstMemref,
+                  Arg<GPU_MMAMemRef, "",[MemWriteAt<0, FullEffect>]>:$dstMemref,
                   Variadic<Index>:$indices,
                   IndexAttr:$leadDimension,
                   OptionalAttr<UnitAttr>:$transpose);

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

@@ -72,7 +72,8 @@ class AllocLikeOp<string mnemonic,
                        Variadic<Index>:$symbolOperands,
                        ConfinedAttr<OptionalAttr<I64Attr>,
                                 [IntMinValue<0>]>:$alignment);
-  let results = (outs Res<AnyMemRef, "", [MemAlloc<resource>]>:$memref);
+  let results = (outs Res<AnyMemRef, "",
+                          [MemAlloc<resource, 0, FullEffect>]>:$memref);
 
   let builders = [
     OpBuilder<(ins "MemRefType":$memrefType,
@@ -276,12 +277,15 @@ def MemRef_ReallocOp : MemRef_Op<"realloc"> {
   // memref and allocating the outcoming memref, even though this may not
   // physically happen on each execution.
 
-  let arguments = (ins Arg<MemRefRankOf<[AnyType], [1]>, "", [MemFree]>:$source,
+  let arguments = (ins Arg<MemRefRankOf<[AnyType], [1]>, "",
+                                        [MemFreeAt<0, FullEffect>]>:$source,
                    Optional<Index>:$dynamicResultSize,
                    ConfinedAttr<OptionalAttr<I64Attr>,
                                 [IntMinValue<0>]>:$alignment);
 
-  let results = (outs Res<MemRefRankOf<[AnyType], [1]>, "", [MemAlloc<DefaultResource>]>);
+  let results = (outs Res<MemRefRankOf<[AnyType], [1]>, "",
+                                       [MemAlloc<DefaultResource, 1,
+                                                 FullEffect>]>);
 
   let builders = [
     OpBuilder<(ins "MemRefType":$resultType,
@@ -532,9 +536,9 @@ def CopyOp : MemRef_Op<"copy", [CopyOpInterface, SameOperandsElementType,
   }];
 
   let arguments = (ins Arg<AnyRankedOrUnrankedMemRef, "the memref to copy from",
-                           [MemRead]>:$source,
+                           [MemReadAt<0, FullEffect>]>:$source,
                        Arg<AnyRankedOrUnrankedMemRef, "the memref to copy to",
-                           [MemWrite]>:$target);
+                           [MemWriteAt<0, FullEffect>]>:$target);
 
   let assemblyFormat = [{
     $source `,` $target attr-dict `:` type($source) `to` type($target)
@@ -564,7 +568,8 @@ def MemRef_DeallocOp : MemRef_Op<"dealloc", [MemRefsNormalizable]> {
     ```
   }];
 
-  let arguments = (ins Arg<AnyRankedOrUnrankedMemRef, "", [MemFree]>:$memref);
+  let arguments = (ins Arg<AnyRankedOrUnrankedMemRef, "",
+                           [MemFreeAt<0, FullEffect>]>:$memref);
 
   let hasFolder = 1;
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
@@ -2107,7 +2112,8 @@ def TensorStoreOp : MemRef_Op<"tensor_store",
   }];
 
   let arguments = (ins AnyTensor:$tensor, Arg<AnyRankedOrUnrankedMemRef,
-                       "the reference to store to", [MemWrite]>:$memref);
+                       "the reference to store to",
+                       [MemWriteAt<0, FullEffect>]>:$memref);
 
   let assemblyFormat = "$tensor `,` $memref attr-dict `:` type($memref)";
 }

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

@@ -239,7 +239,7 @@ def NVGPU_LdMatrixOp : NVGPU_Op<"ldmatrix", [
     ```
   }];
 
-  let arguments = (ins Arg<AnyMemRef, "", [MemRead]>:$srcMemref,
+  let arguments = (ins Arg<AnyMemRef, "", [MemReadAt<0, FullEffect>]>:$srcMemref,
                            Variadic<Index>:$indices, BoolAttr:$transpose,
                            I32Attr:$numTiles);
   let results = (outs AnyVector:$res);
@@ -423,9 +423,9 @@ def NVGPU_DeviceAsyncCopyOp : NVGPU_Op<"device_async_copy", [
     ```
   }];
   let results = (outs NVGPU_DeviceAsyncToken:$asyncToken);
-  let arguments = (ins Arg<AnyMemRef, "", [MemWrite]>:$dst,
+  let arguments = (ins Arg<AnyMemRef, "", [MemWriteAt<0, FullEffect>]>:$dst,
                        Variadic<Index>:$dstIndices,
-                       Arg<AnyMemRef, "", [MemRead]>:$src,
+                       Arg<AnyMemRef, "", [MemReadAt<0, FullEffect>]>:$src,
                        Variadic<Index>:$srcIndices,
                        IndexAttr:$dstElements,
                        Optional<Index>:$srcElements,
@@ -630,7 +630,7 @@ def NVGPU_TmaAsyncLoadOp : NVGPU_Op<"tma.async.load", []> {
 
     The Op uses `$barrier` mbarrier based completion mechanism. 
   }];  
-  let arguments = (ins Arg<AnyMemRef, "", [MemWrite]>:$dst,
+  let arguments = (ins Arg<AnyMemRef, "", [MemWriteAt<0, FullEffect>]>:$dst,
                        NVGPU_MBarrierGroup:$barriers,
                        NVGPU_TensorMapDescriptor:$tensorMapDescriptor,
                        Variadic<Index>:$coordinates, 

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

@@ -149,10 +149,19 @@ class EffectOpInterfaceBase<string name, string baseEffect>
   string baseEffectName = baseEffect;
 }
 
+
+class EffectRange <bits<1> val> {
+  bits<1> Value = val;
+}
+
+def FullEffect : EffectRange<1>;
+def PartialEffect : EffectRange<0>;
+
 // This class is the general base side effect class. This is used by derived
 // effect interfaces to define their effects.
 class SideEffect<EffectOpInterfaceBase interface, string effectName,
-                 Resource resourceReference> : OpVariableDecorator {
+                 Resource resourceReference, int effectStage, EffectRange range>
+    : OpVariableDecorator {
   /// The name of the base effects class.
   string baseEffectName = interface.baseEffectName;
 
@@ -167,6 +176,13 @@ class SideEffect<EffectOpInterfaceBase interface, string effectName,
 
   /// The resource that the effect is being applied to.
   string resource = resourceReference.name;
+
+  /// The stage of side effects, we use it to describe the sequence in which
+  /// effects occur.
+  int stage = effectStage;
+
+  // Does this side effect act on every single value of resource.
+  bit effectOnFullRegion = range.Value;
 }
 
 // This class is the base used for specifying effects applied to an operation.