[MLIR] Add stage to side effect

[cxy-1993]

Summary

[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

[cxy-1993]

This is the RFC corresponding to this patch：https://discourse.llvm.org/t/rfc-add-effect-index-in-memroy-effect/72235

[cxy-1993]

Update the patch according to the discussion in https://discourse.llvm.org/t/rfc-add-effect-index-in-memroy-effect/72235/4

[cxy-1993]

Update patch fixes coding style

[mehdi_amini]

There is likely some documentation to update.

Also I'm not sure that "index" is expressive enough to carry the intent here. The closest I can relate here is pipeline stages in VLIW for example, so maybe the word "stage" should appear in the API?

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

284

We're losing the "DefaultResource" here, how does the "*At" compose with various resources?

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

173

Why a string and not an int?

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

218

Where is the notion of "index" defined?

[cxy-1993]

Updated the code according to the review comments

[cxy-1993]

In D156087#4544491, @mehdi_amini wrote:

There is likely some documentation to update.

Also I'm not sure that "index" is expressive enough to carry the intent here. The closest I can relate here is pipeline stages in VLIW for example, so maybe the word "stage" should appear in the API?

Thank you for your insightful comments. Indeed, the index is not accurate in this scenario. I have made the necessary changes to the patch based on your feedback. Please review it again.

[cxy-1993]

Rename SideEffectOrderAttr -> SideEffectStageAttr

[cxy-1993]

Could you please help me review this piece of code again @mehdi_amini . Also, I am unable to trigger the tests for this commit again. I found online that it's currently only allowed for maintainers to trigger the tests. Could you please help me trigger the tests? Thank you very much for your assistance.

[cxy-1993]

Update test

[mehdi_amini]

I still think documentation is missing, I don't see a definition of "Stage" yet?

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

225

It goes with my comment below, but attributes are suboptimal here. If stages is important why not add an Optional<int> stage as member for this class?

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

3205

Using attributes does not seem optimal to me: it's far from free and we shouldn't put these creation here.

[cxy-1993]

In D156087#4549079, @mehdi_amini wrote:

I still think documentation is missing, I don't see a definition of "Stage" yet?

When I first implemented it, I added a member variable in this class, but during the discussion with Alex Zinenko, he thought using properties would be more appropriate. What do you think about this issue? Looking forward to your reply.

I

Relate disscussion: https://discourse.llvm.org/t/rfc-add-effect-index-in-memroy-effect/72235/4

[cxy-1993]

Updated the code according to the latest review comments

[cxy-1993]

I agree with your suggestion of making 'stage' as a member variable in the side effect class. I have already updated the code and added the necessary documentation based on your advice. I kindly request you to review it once again. Thank you so much for your valuable comments! @mehdi_amini

[mehdi_amini]

Seems better with the int :)

We should still provide more doc, and an example in particular.

Also update the title and the description to reflect the current version, and add a link to the RFC in the description please.

Finally, something that makes all this still not totally complete for safely do the kind of things you described in the RFC is the notion of May/Must associated with the effect, as well as the "partial/full region" aspect of the effect.

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

172

Nit: this isn't a sentence?

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

84

Seems like the stage isn't optional here: that is, it is implicitly 0 which I can imagine being different from "undefined".

It may be OK, but it can also be surprising mixing implicit/explicit effects.

[cxy-1993]

Add effectOnFullRegion to side effect and update docs according to commentsx

[cxy-1993]

Update style

[mehdi_amini]

mlir/docs/Rationale/SideEffectsAndSpeculation.md

174

Thanks for providing a good doc for the "effectOnFullRegion", what about stage? Seems like stages aren't really needed in any of these examples?

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

1226 ↗

(On Diff #545955)

I'm not sure how to interpret the stages here: do we guarantee that all of the read happen before the write?
What about a loop that copy one element at a time?

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

38

Do we have enums in TableGen? (I don't remember).

I rather read "Full" and "Partial" than true/false if possible.

[cxy-1993]

Thanks for your advice.
The "partial/full region" aspect of the effect is also an important feature in side effect, i have already added them to the member variable of EffectInstance and they can be specified via the 'td' parameter.
Regarding May/Must associated with the effect, I believe you are referring to the relationship between memrefs. If that's the case, we will leave this issue to the pass to invoke alias analysis. If it's not, please let me know.

[mehdi_amini]

In D156087#4549709, @cxy-1993 wrote:

Regarding May/Must associated with the effect, I believe you are referring to the relationship between memrefs. If that's the case, we will leave this issue to the pass to invoke alias analysis. If it's not, please let me know.

No I wasn't referring to alias analysis (where "May alias"/"Must alias" is important) but to the read/write effects: how do we model an op that may write the full resource, or may not?
We can say that conservatively we say "write full" when it may do so, but then you can't use this information for kill the liveness of the entire resource.

[cxy-1993]

In D156087#4549716, @mehdi_amini wrote:

In D156087#4549709, @cxy-1993 wrote:

Regarding May/Must associated with the effect, I believe you are referring to the relationship between memrefs. If that's the case, we will leave this issue to the pass to invoke alias analysis. If it's not, please let me know.

No I wasn't referring to alias analysis (where "May alias"/"Must alias" is important) but to the read/write effects: how do we model an op that may write the full resource, or may not?
We can say that conservatively we say "write full" when it may do so, but then you can't use this information for kill the liveness of the entire resource.

I believe that when we say “write full” to a particular resource, this write is certain to occur. Therefore, any writes to this resource before that point can be killed by this operation. If the write is only possible, and we shouldn't explicitly mark it as writing all resources, do you think this approach is feasible?

[cxy-1993]

Sorry, I don't know how to reply inline comment...
The copy like example read source on stage 0 and write target on stage 1. By using stages, we indicate the sequence in which these two side effects occur. Is it enough for "stage" or we need to provide more exmaple?.

[cxy-1993]

Do not set side effects stage for operations where we cannot determine the order of side effects on all values.

[cxy-1993]

Thank you for your guidance. I have update the patch according to your comment. @mehdi_amini

mlir/docs/Rationale/SideEffectsAndSpeculation.md

174

I rewrote the first example and used a virtual SIMD operation to illustrate the effect of stages.

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

1226 ↗

(On Diff #545955)

NICE CATCH! Stage modeling indeed not be suitable for operations that can lower to loops, as its primary application lies in SIMD-like operations where hardware guarantees the order of reads and writes. I have updated existing dialect's side effects and the descriptions in the documentation.

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

284

DONE

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

172

I refactored this comment for better clarity.

173

I have changed it to int.

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

218

Changed to "stage" now.

225

DONE

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

38

I can not found any enum in TableGen except EnumAttr...

84

I have replace optional<int> member variable with default value null with int member variable with default value 0.

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

3205

DONE

[mehdi_amini]

I believe that when we say “write full” to a particular resource, this write is certain to occur. Therefore, any writes to this resource before that point can be killed by this operation. If the write is only possible, and we shouldn't explicitly mark it as writing all resources, do you think this approach is feasible?

It might be ok, I am not sure it is conservative in every situation? Assuming it is, it should likely be the default, right now your MemAlloc default to full.

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

38

I think it should look like this in TableGen:

class EffectRange <bits<1> val> {
  bits<1> Value = val;
}

def FullEffect : EffectStage<1>;
def PartialEffect : EffectStage<0>;

class SideEffect<EffectOpInterfaceBase interface, string effectName,
                 Resource resourceReference, int effectStage, EffectRange range> {
   ...

  // Does this side effect act on every single value of resource.
  EffectRange effectOnFullRegion = range;
}

But it may be possible to write also here:

// Does this side effect act on every single value of resource.
bit effectOnFullRegion = range.Value;

[cxy-1993]

In D156087#4551469, @mehdi_amini wrote:

I believe that when we say “write full” to a particular resource, this write is certain to occur. Therefore, any writes to this resource before that point can be killed by this operation. If the write is only possible, and we shouldn't explicitly mark it as writing all resources, do you think this approach is feasible?

It might be ok, I am not sure it is conservative in every situation? Assuming it is, it should likely be the default, right now your MemAlloc default to full.

I agree that treating side effects as default on partially memrefs is a cautious approach. Currently, almost all visible ops have side effects on the entire memref, so can we consider this common but less safe behavior as the default and require the definers of specific ops to understand the scope of side effects and make modifications accordingly?Otherwise, I would have to modify almost all existing side effects to become FullEffect.

[mehdi_amini]

I see where you're coming from, but I'm much more on the side of avoiding footgun in our APIs: especially this kind, where it won't be caught until there is a subtle miscompute which requires long hours of debug.

[cxy-1993]

Update patch according to review comments

[cxy-1993]

In D156087#4552053, @mehdi_amini wrote:

I see where you're coming from, but I'm much more on the side of avoiding footgun in our APIs: especially this kind, where it won't be caught until there is a subtle miscompute which requires long hours of debug.

I have already changed the default range of side effects to "partial" and made corresponding modifications to all feasible ops. Thank you for your valuable comments, please review it again.

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

38

DONE

[cxy-1993]

I have already modified the default range of side effects to "partial" and made necessary adjustments to various ops based on your feedback. I am reaching out to kindly request another code review for the changes I have made. @mehdi_amini

[Mogball]

+1. This patch makes sense to me. Having this level of granularity is useful.

Do any of the upstream transformations need to be taught about this?

[Mogball]

Please wait for @mehdi_amini and others to chime in first

[cxy-1993]

In D156087#4555106, @Mogball wrote:

+1. This patch makes sense to me. Having this level of granularity is useful.

Do any of the upstream transformations need to be taught about this?

Thank you for the recognition. I believe that currently, the existing passes in MLIR don't have an immediate use for this feature. The motivation behind this feature is to improve the modeling of operation read and write order, as well as their ranges, for certain SIMD backends. With the introduction of this feature, I am confident that more analysis and optimizations will gradually be contributed to the MLIR community.

[cxy-1993]

Update docs

[cxy-1993]

Ping @rriddle Could you please help me review this piece of code? Thank you! :)

[cxy-1993]

ping @mehdi_amini @rriddle Could you please take some time to help me with the review and merging of this patch? Thank you. :)

[cxy-1993]

@rriddle @mehdi_amini Could you kindly spare a moment to review the updated code at your convenience? Thank you for your time and consideration.

[mehdi_amini]

That seems fine to me, but I wanted to hear if @rriddle does not have any opposition here!

[cxy-1993]

@rriddle Ping :)

[cxy-1993]

@rriddle Please help me review this piece of code. Thank you for your time. :)

[cxy-1993]

@rriddle ping :) Could you please spare some time to review this patch, thank you!

[Mogball]

@cxy-1993 please just go ahead and land this. River is indisposed wrt MLIR reviews at the moment

[ftynse]

DIdn't the forum RFC conclude on using the parameters attribute for this instead of adding new fields?

[cxy-1993]

In D156087#4647473, @ftynse wrote:

DIdn't the forum RFC conclude on using the parameters attribute for this instead of adding new fields?

Sorry for not inform you and updating the RFC in a timely manner. During our discussions with amini, we concluded that it would be more appropriate to use a new field to represent the stages at which side effects occur, rather than directly using the attribute. As a result, I have replaced the usage of the parameter with two new fields: "stage" and "effectOnFullRegion." You can refer to the comments above for the discussion details.

Personally, I believe that adding a new field instead of using parameters would make this more straightforward and easier to comprehend.

[ftynse]

Then I'll block this until I had time to go through the discussion that should have happened on the RFC.

[cxy-1993]

rebase main

[cxy-1993]

This patch has already been rebased onto the main branch and passed precheckin. I don't have write access to the LLVM repo. Could you please help me to push this code into the master, Thanks. @Mogball

[cxy-1993]

rebase main

[mehdi_amini]

ftynse acknowledged my rationale with a "like" on Discourse, so I suspect this is good to go?

[cxy-1993]

In D156087#4652100, @mehdi_amini wrote:

ftynse acknowledged my rationale with a "like" on Discourse, so I suspect this is good to go?

Thank you for your recognition.I have already retested this code locally and i don't have write access to the LLVM repository. Could you please help me push this patch to the repository？ @mehdi_amini

mlir/docs/Rationale/SideEffectsAndSpeculation.md

1. Operations with undefined behavior. These operations are not defined on
    loading from a pointer that has been freed.
 1. Operations that don't terminate. E.g. an `scf.while` where the condition is
    always true.
 1. Operations with non-local control flow. These operations may pop their
    current frame of execution and return directly to an older frame. E.g.
    `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
 
 Modeling these behaviors has to walk a fine line -- we need to empower more
 complicated passes to reason about the nuances of such behaviors while
 simultaneously not overburdening simple passes that only need a coarse grained
 "can this op be freely moved" query.
 

 
 We don't have proper modeling yet to fully capture non-local control flow
 semantics.
 
 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? If not, it
+   should set the effectOnFullRegion field with false.
 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
    RFC if your operation has a novel side effect that cannot be adequately
    captured by `MemoryEffectsOpInterface`.
 1. Is it well defined in all inputs or does it assume certain runtime
    restrictions on its inputs, e.g. the pointer operand must point to valid
    memory? It should probably implement `ConditionallySpeculatable`.
 1. Can it infinitely loop on certain inputs? It should probably implement
    `ConditionallySpeculatable`.
 1. Does it have non-local control flow (e.g. `longjmp`)? We don't have proper
    modeling for these yet, patches welcome!
 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 and have a "simd.abs" operation, which read
+all value in source memref, calculate its absolute value and write to 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 every single value from the source resource and
+then writes these values to every single 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 mark these side effects as acting on every single 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>]>:$source,
+                         Arg<AnyRankedOrUnrankedMemRef, "the target memref",
+                             [MemWriteAt<1>]>:$target);
+
+    ...
+  }
+```
+
+In the above example, we added the side effect [MemReadAt<0>] to the source,
+indicating abs operation reads every single value from source in stage 0.
+[MemReadAt<0>] is a shorthand notation for [MemReadAt<0, true>]. We added the
+side effect [MemWriteAt<0>] to the target, indicating abs operation writes on
+every single value inside the target on stage 1(after read from 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 read one value from input memref and return it.
+Therefore, we needs to specify a read side effect for input memref, and mark
+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, false>]>:$memref,
+                         Variadic<Index>:$indices,
+                         DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
+
+    ...
+  }
+```
+
+In the above example, we added the side effect [MemReadAt<0, false>] to the
+source, indicating load operation read parts of value from memref at stage 0.

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

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,
                    Optional<Index>:$dynamicResultSize,
                    ConfinedAttr<OptionalAttr<I64Attr>,
                                 [IntMinValue<0>]>:$alignment);
 
   let results = (outs Res<MemRefRankOf<[AnyType], [1]>, "", [MemAlloc<DefaultResource>]>);
 
   let builders = [
     OpBuilder<(ins "MemRefType":$resultType,
                   "Value":$source,
                   CArg<"Value", "Value()">:$dynamicResultSize), [{
       return build($_builder, $_state, resultType, source, dynamicResultSize,
                    IntegerAttr());
     }]>];
let description = [{
     chains (e.g. through [`affine.apply`](Affine.md/#affineapply-affineapplyop)
     operations) to precisely analyze references at compile-time using polyhedral
     techniques. This is possible because of the
     [restrictions on dimensions and symbols](Affine.md/#restrictions-on-dimensions-and-symbols)
     in these contexts.
   }];
 
   let arguments = (ins Arg<AnyMemRef, "the reference to load from",
-                           [MemRead]>:$memref,
+                           [MemReadAt<0, false>]>:$memref,
                        Variadic<Index>:$indices,
                        DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
   let results = (outs AnyType:$result);
 
   let extraClassDeclaration = [{
     Value getMemRef() { return getOperand(0); }
     void setMemRef(Value value) { setOperand(0, value); }
     MemRefType getMemRefType() {
let description = [{
     operations) to precisely analyze references at compile-time using polyhedral
     techniques. This is possible because of the
     [restrictions on dimensions and symbols](Affine.md/#restrictions-on-dimensions-and-symbols)
     in these contexts.
   }];
 
   let arguments = (ins AnyType:$value,
                        Arg<AnyMemRef, "the reference to store to",
-                           [MemWrite]>:$memref,
+                           [MemWriteAt<0, false>]>:$memref,
                        Variadic<Index>:$indices,
                        DefaultValuedOptionalAttr<BoolAttr, "false">:$nontemporal);
 
   let builders = [
     OpBuilder<(ins "Value":$valueToStore, "Value":$memref), [{
       $_state.addOperands(valueToStore);
       $_state.addOperands(memref);
     }]>];

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

class EffectOpInterfaceBase<string name, string baseEffect>
 
   // The base effect name of this interface.
   string baseEffectName = baseEffect;
 }
 
 // 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, bit fullRegion>
+    : OpVariableDecorator {
   /// The name of the base effects class.
   string baseEffectName = interface.baseEffectName;
 
   /// The parent interface that the effect belongs to.
   string interfaceTrait = interface.trait;
 
   /// The cpp namespace of the interface trait.
   string cppNamespace = interface.cppNamespace;
 
   /// The derived effect that is being applied.
   string effect = 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 = fullRegion;
 }
 
 // This class is the base used for specifying effects applied to an operation.
 class SideEffectsTraitBase<EffectOpInterfaceBase parentInterface,
                            list<SideEffect> staticEffects>
     : OpInterfaceTrait<""> {
   /// The name of the interface trait to use.
   let trait = parentInterface.trait;

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

 /// effect being applied, a resource that corresponds to where the effect is
 /// applied, and an optional symbol reference or value(either operand, result,
 /// or region entry argument) that the effect is applied to, and an optional
 /// parameters attribute further specifying the details of the effect.
 template <typename EffectT>
 class EffectInstance {
 public:
   EffectInstance(EffectT *effect, Resource *resource = DefaultResource::get())
-      : effect(effect), resource(resource) {}
+      : effect(effect), resource(resource), stage(0), effectOnFullRegion(true) {
+  }
+  EffectInstance(EffectT *effect, int stage, bool effectOnFullRegion,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), stage(stage),
+        effectOnFullRegion(effectOnFullRegion) {}
   EffectInstance(EffectT *effect, Value value,
                  Resource *resource = DefaultResource::get())
-      : effect(effect), resource(resource), value(value) {}
+      : effect(effect), resource(resource), value(value), stage(0),
+        effectOnFullRegion(true) {}
+  EffectInstance(EffectT *effect, Value value, int stage,
+                 bool effectOnFullRegion,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), value(value), stage(stage),
+        effectOnFullRegion(effectOnFullRegion) {}
   EffectInstance(EffectT *effect, SymbolRefAttr symbol,
                  Resource *resource = DefaultResource::get())
-      : effect(effect), resource(resource), value(symbol) {}
+      : effect(effect), resource(resource), value(symbol), stage(0),
+        effectOnFullRegion(true) {}
+  EffectInstance(EffectT *effect, SymbolRefAttr symbol, int stage,
+                 bool effectOnFullRegion,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), value(symbol), stage(stage),
+        effectOnFullRegion(effectOnFullRegion) {}
   EffectInstance(EffectT *effect, Attribute parameters,
                  Resource *resource = DefaultResource::get())
-      : effect(effect), resource(resource), parameters(parameters) {}
+      : effect(effect), resource(resource), parameters(parameters), stage(0),
+        effectOnFullRegion(true) {}
+  EffectInstance(EffectT *effect, Attribute parameters, int stage,
+                 bool effectOnFullRegion,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), parameters(parameters),
+        stage(stage), effectOnFullRegion(effectOnFullRegion) {}
   EffectInstance(EffectT *effect, Value value, Attribute parameters,
                  Resource *resource = DefaultResource::get())
       : effect(effect), resource(resource), value(value),
-        parameters(parameters) {}
+        parameters(parameters), stage(0), effectOnFullRegion(true) {}
+  EffectInstance(EffectT *effect, Value value, Attribute parameters, int stage,
+                 bool effectOnFullRegion,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), value(value),
+        parameters(parameters), stage(stage),
+        effectOnFullRegion(effectOnFullRegion) {}
   EffectInstance(EffectT *effect, SymbolRefAttr symbol, Attribute parameters,
                  Resource *resource = DefaultResource::get())
       : effect(effect), resource(resource), value(symbol),
-        parameters(parameters) {}
+        parameters(parameters), stage(0), effectOnFullRegion(true) {}
+  EffectInstance(EffectT *effect, SymbolRefAttr symbol, Attribute parameters,
+                 int stage, bool effectOnFullRegion,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), value(symbol),
+        parameters(parameters), stage(stage),
+        effectOnFullRegion(effectOnFullRegion) {}
 
   /// Return the effect being applied.
   EffectT *getEffect() const { return effect; }
 
   /// Return the value the effect is applied on, or nullptr if there isn't a
   /// known value being affected.
-  Value getValue() const { return value ? llvm::dyn_cast_if_present<Value>(value) : Value(); }
+  Value getValue() const {
+    return value ? llvm::dyn_cast_if_present<Value>(value) : Value();
+  }
 
   /// Return the symbol reference the effect is applied on, or nullptr if there
   /// isn't a known smbol being affected.
   SymbolRefAttr getSymbolRef() const {
-    return value ? llvm::dyn_cast_if_present<SymbolRefAttr>(value) : SymbolRefAttr();
+    return value ? llvm::dyn_cast_if_present<SymbolRefAttr>(value)
+                 : SymbolRefAttr();
   }
 
   /// Return the resource that the effect applies to.
   Resource *getResource() const { return resource; }
 
   /// Return the parameters of the effect, if any.
   Attribute getParameters() const { return parameters; }
 
+  /// Return the effect happen stage.
+  int getStage() const { return stage; }
+
+  /// Return if this side effect act on every single value of resource.
+  bool getEffectOnFullRegion() const { return effectOnFullRegion; }
+
 private:
   /// The specific effect being applied.
   EffectT *effect;
 
   /// The resource that the given value resides in.
   Resource *resource;
 
   /// The Symbol or Value that the effect applies to. This is optionally null.
   PointerUnion<SymbolRefAttr, Value> value;
 
   /// Additional parameters of the effect instance. An attribute is used for
   /// type-safe structured storage and context-based uniquing. Concrete effects
   /// can use this at their convenience. This is optionally null.
   Attribute parameters;
+
+  // The stage side effect happen. Side effect with a lower stage
+  // number happen earlier than those with a higher stage number
+  int stage;
+
+  // Does this side effect act on every single value of resource.
+  bool effectOnFullRegion;
 };
 } // namespace SideEffects
 
 namespace Speculation {
 /// This enum is returned from the `getSpeculatability` method in the
 /// `ConditionallySpeculatable` op interface.
 enum class Speculatability {
   /// The Operation in question cannot be speculatively executed.  This could be

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

def MemoryEffectsOpInterface
   let description = [{
     An interface used to query information about the memory effects applied by
     an operation.
   }];
   let cppNamespace = "::mlir";
 }
 
 // The base class for defining specific memory effects.
-class MemoryEffect<string effectName, Resource resource>
-  : SideEffect<MemoryEffectsOpInterface, effectName, resource>;
+class MemoryEffect<string effectName, Resource resource, int stage,
+                   bit fullRegion>
+  : SideEffect<MemoryEffectsOpInterface, effectName, resource, stage,
+               fullRegion>;
 
 // This class represents the trait for memory effects that may be placed on
 // operations.
 class MemoryEffects<list<MemoryEffect> effects = []>
   : SideEffectsTraitBase<MemoryEffectsOpInterface, effects>;
 
 //===----------------------------------------------------------------------===//
 // Effects
 
 // The following effect indicates that the operation allocates from some
 // resource. An 'allocate' effect implies only allocation of the resource, and
 // not any visible mutation or dereference.
-class MemAlloc<Resource resource>
-  : MemoryEffect<"::mlir::MemoryEffects::Allocate", resource>;
-def MemAlloc : MemAlloc<DefaultResource>;
+class MemAlloc<Resource resource, int stage = 0, bit fullRegion = 1>
+  : MemoryEffect<"::mlir::MemoryEffects::Allocate", resource, stage,
+                 fullRegion>;
+def MemAlloc : MemAlloc<DefaultResource, 0, 1>;
+class MemAllocAt<int stage, bit fullRegion = 1>
+  : MemAlloc<DefaultResource, stage, fullRegion>;
 
 // The following effect indicates that the operation frees some resource that
 // has been allocated. A 'free' effect implies only de-allocation of the
 // resource, and not any visible allocation, mutation or dereference.
-class MemFree<Resource resource>
-  : MemoryEffect<"::mlir::MemoryEffects::Free", resource>;
-def MemFree : MemFree<DefaultResource>;
+class MemFree<Resource resource, int stage = 0, bit fullRegion = 1>
+  : MemoryEffect<"::mlir::MemoryEffects::Free", resource, stage, fullRegion>;
+def MemFree : MemFree<DefaultResource, 0, 1>;
+class MemFreeAt<int stage, bit fullRegion = 1>
+  : MemFree<DefaultResource, stage, fullRegion>;
 
 // The following effect indicates that the operation reads from some
 // resource. A 'read' effect implies only dereferencing of the resource, and
 // not any visible mutation.
-class MemRead<Resource resource>
-  : MemoryEffect<"::mlir::MemoryEffects::Read", resource>;
-def MemRead : MemRead<DefaultResource>;
+class MemRead<Resource resource, int stage = 0, bit fullRegion = 1>
+  : MemoryEffect<"::mlir::MemoryEffects::Read", resource, stage, fullRegion>;
+def MemRead : MemRead<DefaultResource, 0, 1>;
+class MemReadAt<int stage, bit fullRegion = 1>
+  : MemRead<DefaultResource, stage, fullRegion>;
 
 // The following effect indicates that the operation writes to some
 // resource. A 'write' effect implies only mutating a resource, and not any
 // visible dereference or read.
-class MemWrite<Resource resource>
-  : MemoryEffect<"::mlir::MemoryEffects::Write", resource>;
-def MemWrite : MemWrite<DefaultResource>;
+class MemWrite<Resource resource, int stage = 0, bit fullRegion = 1>
+  : MemoryEffect<"::mlir::MemoryEffects::Write", resource, stage, fullRegion>;
+def MemWrite : MemWrite<DefaultResource, 0, 1>;
+class MemWriteAt<int stage, bit fullRegion = 1>
+  : MemWrite<DefaultResource, stage, fullRegion>;
 
 //===----------------------------------------------------------------------===//
 // Effect Traits
 //===----------------------------------------------------------------------===//
 
 // Op has no effect on memory but may have undefined behavior.
 def NoMemoryEffect : MemoryEffects<[]>;
 

mlir/include/mlir/TableGen/SideEffects.h

public:
   StringRef getBaseEffectName() const;
 
   // Return the name of the Interface that the effect belongs to.
   std::string getInterfaceTrait() const;
 
   // Return the name of the resource class.
   StringRef getResource() const;
 
+  // Return the stage of the effect happen.
+  int64_t getStage() const;
+
+  // Return if this side effect act on every single value of resource.
+  bool getEffectOnfullRegion() const;
+
   static bool classof(const Operator::VariableDecorator *var);
 };
 
 // This class represents an instance of a side effect interface applied to an
 // operation. This is a wrapper around an OpInterfaceTrait that also includes
 // the effects that are applied.
 class SideEffectTrait : public InterfaceTrait {
 public:

mlir/lib/TableGen/SideEffects.cpp

std::string SideEffect::getInterfaceTrait() const {
   return cppNamespace.empty() ? trait.str()
                               : (cppNamespace + "::" + trait).str();
 }
 
 StringRef SideEffect::getResource() const {
   return def->getValueAsString("resource");
 }
 
+int64_t SideEffect::getStage() const { return def->getValueAsInt("stage"); }
+
+bool SideEffect::getEffectOnfullRegion() const {
+  return def->getValueAsBit("effectOnFullRegion");
+}
+
 bool SideEffect::classof(const Operator::VariableDecorator *var) {
   return var->getDef().isSubClassOf("SideEffect");
 }
 
 //===----------------------------------------------------------------------===//
 // SideEffectsTrait
 //===----------------------------------------------------------------------===//
 

mlir/test/lib/Dialect/Test/TestInterfaces.td

 
 def TestEffectOpInterface
     : EffectOpInterfaceBase<"TestEffectOpInterface",
                             "::mlir::TestEffects::Effect"> {
   let cppNamespace = "::mlir";
 }
 
 class TestEffect<string effectName>
-    : SideEffect<TestEffectOpInterface, effectName, DefaultResource>;
+    : SideEffect<TestEffectOpInterface, effectName, DefaultResource, 0, 1>;
 
 class TestEffects<list<TestEffect> effects = []>
    : SideEffectsTraitBase<TestEffectOpInterface, effects>;
 
 def TestConcreteEffect : TestEffect<"TestEffects::Concrete">;
 
 #endif // MLIR_TEST_DIALECT_TEST_INTERFACES

mlir/test/lib/Dialect/Test/TestOps.td

 def TestEffectsOpA : TEST_Op<"op_with_effects_a"> {
   let arguments = (ins
     Arg<Variadic<AnyMemRef>, "", [MemRead]>,
     Arg<FlatSymbolRefAttr, "", [MemRead]>:$first,
     Arg<SymbolRefAttr, "", [MemWrite]>:$second,
     Arg<OptionalAttr<SymbolRefAttr>, "", [MemRead]>:$optional_symbol
   );
 
-  let results = (outs Res<AnyMemRef, "", [MemAlloc<TestResource>]>);
+  let results = (outs Res<AnyMemRef, "", [MemAlloc<TestResource, 0>]>);
 }
 
 def TestEffectsOpB : TEST_Op<"op_with_effects_b",
-    [MemoryEffects<[MemWrite<TestResource>]>]>;
+    [MemoryEffects<[MemWrite<TestResource, 0>]>]>;
 
 def TestEffectsRead : TEST_Op<"op_with_memread",
     [MemoryEffects<[MemRead]>]> {
   let results = (outs AnyInteger);
 }
 
 def TestEffectsWrite : TEST_Op<"op_with_memwrite",
     [MemoryEffects<[MemWrite]>]>;

mlir/test/mlir-tblgen/op-side-effects.td

 // RUN: mlir-tblgen -gen-op-defs -I %S/../../include %s | FileCheck %s
 
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def TEST_Dialect : Dialect {
   let name = "test";
 }
 class TEST_Op<string mnemonic, list<Trait> traits = []> :
     Op<TEST_Dialect, mnemonic, traits>;
 
 def CustomResource : Resource<"CustomResource">;
 
 def SideEffectOpA : TEST_Op<"side_effect_op_a"> {
   let arguments = (ins
     Arg<Variadic<AnyMemRef>, "", [MemRead]>,
+    Arg<AnyMemRef, "", [MemWriteAt<1, false>]>,
     Arg<SymbolRefAttr, "", [MemRead]>:$symbol,
     Arg<FlatSymbolRefAttr, "", [MemWrite]>:$flat_symbol,
     Arg<OptionalAttr<SymbolRefAttr>, "", [MemRead]>:$optional_symbol
   );
-  let results = (outs Res<AnyMemRef, "", [MemAlloc<CustomResource>]>);
+  let results = (outs Res<AnyMemRef, "", [MemAlloc<CustomResource, 0>]>);
 }
 
 def SideEffectOpB : TEST_Op<"side_effect_op_b",
-    [MemoryEffects<[MemWrite<CustomResource>]>]>;
+    [MemoryEffects<[MemWrite<CustomResource, 0>]>]>;
 
 // CHECK: void SideEffectOpA::getEffects
 // CHECK:   for (::mlir::Value value : getODSOperands(0))
-// CHECK:     effects.emplace_back(::mlir::MemoryEffects::Read::get(), value, ::mlir::SideEffects::DefaultResource::get());
-// CHECK:   effects.emplace_back(::mlir::MemoryEffects::Read::get(), getSymbolAttr(), ::mlir::SideEffects::DefaultResource::get());
-// CHECK:   effects.emplace_back(::mlir::MemoryEffects::Write::get(), getFlatSymbolAttr(), ::mlir::SideEffects::DefaultResource::get());
+// CHECK:     effects.emplace_back(::mlir::MemoryEffects::Read::get(), value, 0, true, ::mlir::SideEffects::DefaultResource::get());
+// CHECK:   for (::mlir::Value value : getODSOperands(1))
+// CHECK:     effects.emplace_back(::mlir::MemoryEffects::Write::get(), value, 1, false, ::mlir::SideEffects::DefaultResource::get());
+// CHECK:   effects.emplace_back(::mlir::MemoryEffects::Read::get(), getSymbolAttr(), 0, true, ::mlir::SideEffects::DefaultResource::get());
+// CHECK:   effects.emplace_back(::mlir::MemoryEffects::Write::get(), getFlatSymbolAttr(), 0, true, ::mlir::SideEffects::DefaultResource::get());
 // CHECK:   if (auto symbolRef = getOptionalSymbolAttr())
-// CHECK:     effects.emplace_back(::mlir::MemoryEffects::Read::get(), symbolRef, ::mlir::SideEffects::DefaultResource::get());
+// CHECK:     effects.emplace_back(::mlir::MemoryEffects::Read::get(), symbolRef, 0, true, ::mlir::SideEffects::DefaultResource::get());
 // CHECK:   for (::mlir::Value value : getODSResults(0))
-// CHECK:     effects.emplace_back(::mlir::MemoryEffects::Allocate::get(), value, CustomResource::get());
+// CHECK:     effects.emplace_back(::mlir::MemoryEffects::Allocate::get(), value, 0, true, CustomResource::get());
 
 // CHECK: void SideEffectOpB::getEffects
-// CHECK:   effects.emplace_back(::mlir::MemoryEffects::Write::get(), CustomResource::get());
+// CHECK:   effects.emplace_back(::mlir::MemoryEffects::Write::get(), 0, true, CustomResource::get());

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

void OpEmitter::genSideEffectInterfaceMethods() {
   }
   /// Results.
   for (unsigned i = 0, e = op.getNumResults(); i != e; ++i)
     resolveDecorators(op.getResultDecorators(i), i, EffectKind::Result);
 
   // The code used to add an effect instance.
   // {0}: The effect class.
   // {1}: Optional value or symbol reference.
-  // {1}: The resource class.
+  // {2}: The side effect stage.
+  // {3}: Does this side effect act on every single value of resource.
+  // {4}: The resource class.
   const char *addEffectCode =
-      "  effects.emplace_back({0}::get(), {1}{2}::get());\n";
+      "  effects.emplace_back({0}::get(), {1}{2}, {3}, {4}::get());\n";
 
   for (auto &it : interfaceEffects) {
     // Generate the 'getEffects' method.
     std::string type = llvm::formatv("::llvm::SmallVectorImpl<::mlir::"
                                      "SideEffects::EffectInstance<{0}>> &",
                                      it.first())
                            .str();
     auto *getEffects = opClass.addMethod("void", "getEffects",
                                          MethodParameter(type, "effects"));
     ERROR_IF_PRUNED(getEffects, "getEffects", op);
     auto &body = getEffects->body();
 
     // Add effect instances for each of the locations marked on the operation.
     for (auto &location : it.second) {
       StringRef effect = location.effect.getName();
       StringRef resource = location.effect.getResource();
+      int stage = (int)location.effect.getStage();
+      bool effectOnFullRegion = (int)location.effect.getEffectOnfullRegion();
       if (location.kind == EffectKind::Static) {
         // A static instance has no attached value.
-        body << llvm::formatv(addEffectCode, effect, "", resource).str();
+        body << llvm::formatv(addEffectCode, effect, "", stage,
+                              effectOnFullRegion, resource)
+                    .str();
       } else if (location.kind == EffectKind::Symbol) {
         // A symbol reference requires adding the proper attribute.
         const auto *attr = op.getArg(location.index).get<NamedAttribute *>();
         std::string argName = op.getGetterName(attr->name);
         if (attr->attr.isOptional()) {
           body << "  if (auto symbolRef = " << argName << "Attr())\n  "
-               << llvm::formatv(addEffectCode, effect, "symbolRef, ", resource)
+               << llvm::formatv(addEffectCode, effect, "symbolRef, ", stage,
+                                effectOnFullRegion, resource)
                       .str();
         } else {
           body << llvm::formatv(addEffectCode, effect, argName + "Attr(), ",
-                                resource)
+                                stage, effectOnFullRegion, resource)
                       .str();
         }
       } else {
         // Otherwise this is an operand/result, so we need to attach the Value.
         body << "  for (::mlir::Value value : getODS"
              << (location.kind == EffectKind::Operand ? "Operands" : "Results")
              << "(" << location.index << "))\n  "
-             << llvm::formatv(addEffectCode, effect, "value, ", resource).str();
+             << llvm::formatv(addEffectCode, effect, "value, ", stage,
+                              effectOnFullRegion, resource)
+                    .str();
       }
     }
   }
 }
 
 void OpEmitter::genTypeInterfaceMethods() {
   if (!op.allResultTypesKnown())
     return;