[Mlir-commits] [mlir] 0c63122 - [MLIR] Add stage to side effect
Mehdi Amini
llvmlistbot at llvm.org
Fri Sep 29 17:52:48 PDT 2023
Author: cxy
Date: 2023-09-29T17:47:13-07:00
New Revision: 0c63122713c2d719789aef4bdfaf4e0b29c3b79e
URL: https://github.com/llvm/llvm-project/commit/0c63122713c2d719789aef4bdfaf4e0b29c3b79e
DIFF: https://github.com/llvm/llvm-project/commit/0c63122713c2d719789aef4bdfaf4e0b29c3b79e.diff
LOG: [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
Added:
Modified:
mlir/docs/Rationale/SideEffectsAndSpeculation.md
mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td
mlir/include/mlir/Interfaces/SideEffectInterfaces.h
mlir/include/mlir/Interfaces/SideEffectInterfaces.td
mlir/include/mlir/TableGen/SideEffects.h
mlir/lib/TableGen/SideEffects.cpp
mlir/test/lib/Dialect/Test/TestInterfaces.td
mlir/test/lib/Dialect/Test/TestOps.td
mlir/test/mlir-tblgen/op-side-effects.td
mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
Removed:
################################################################################
diff --git a/mlir/docs/Rationale/SideEffectsAndSpeculation.md b/mlir/docs/Rationale/SideEffectsAndSpeculation.md
index 95320feaebd2d83..0cd5ea59822a0fa 100644
--- a/mlir/docs/Rationale/SideEffectsAndSpeculation.md
+++ b/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]".
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 9761ab12134ad28..db93a51775ffcd7 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/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);
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 34cc129053c3259..6375d35f4311295 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/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);
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index ea6e363a6c3257f..6b0ccbe37e89e9c 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/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)";
}
diff --git a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
index 31b137160545772..ec250b1b65899d9 100644
--- a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
+++ b/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,
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td b/mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td
index 8b8904b3c5a3fc8..45a9ffa94363ef3 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td
+++ b/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.
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index 74fb96662934a9e..ec4e36263bbe6d1 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -139,36 +139,74 @@ template <typename EffectT>
class EffectInstance {
public:
EffectInstance(EffectT *effect, Resource *resource = DefaultResource::get())
- : effect(effect), resource(resource) {}
+ : effect(effect), resource(resource), stage(0),
+ effectOnFullRegion(false) {}
+ 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(false) {}
+ 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(false) {}
+ 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(false) {}
+ 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(false) {}
+ 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), stage(0), effectOnFullRegion(false) {}
EffectInstance(EffectT *effect, SymbolRefAttr symbol, Attribute parameters,
+ int stage, bool effectOnFullRegion,
Resource *resource = DefaultResource::get())
: effect(effect), resource(resource), value(symbol),
- parameters(parameters) {}
+ 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.
@@ -177,6 +215,12 @@ class EffectInstance {
/// 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;
@@ -191,6 +235,13 @@ class EffectInstance {
/// 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
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
index eb810c2cfada124..b2ab4fee9d29c03 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
@@ -34,8 +34,9 @@ def MemoryEffectsOpInterface
}
// 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,
+ EffectRange range>
+ : SideEffect<MemoryEffectsOpInterface, effectName, resource, stage, range>;
// This class represents the trait for memory effects that may be placed on
// operations.
@@ -48,30 +49,42 @@ class MemoryEffects<list<MemoryEffect> 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,
+ EffectRange range = PartialEffect>
+ : MemoryEffect<"::mlir::MemoryEffects::Allocate", resource, stage, range>;
+def MemAlloc : MemAlloc<DefaultResource, 0, PartialEffect>;
+class MemAllocAt<int stage, EffectRange range = PartialEffect>
+ : MemAlloc<DefaultResource, stage, range>;
// 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,
+ EffectRange range = PartialEffect>
+ : MemoryEffect<"::mlir::MemoryEffects::Free", resource, stage, range>;
+def MemFree : MemFree<DefaultResource, 0, PartialEffect>;
+class MemFreeAt<int stage, EffectRange range = PartialEffect>
+ : MemFree<DefaultResource, stage, range>;
// 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,
+ EffectRange range = PartialEffect>
+ : MemoryEffect<"::mlir::MemoryEffects::Read", resource, stage, range>;
+def MemRead : MemRead<DefaultResource, 0, PartialEffect>;
+class MemReadAt<int stage, EffectRange range = PartialEffect>
+ : MemRead<DefaultResource, stage, range>;
// 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,
+ EffectRange range = PartialEffect>
+ : MemoryEffect<"::mlir::MemoryEffects::Write", resource, stage, range>;
+def MemWrite : MemWrite<DefaultResource, 0, PartialEffect>;
+class MemWriteAt<int stage, EffectRange range = PartialEffect>
+ : MemWrite<DefaultResource, stage, range>;
//===----------------------------------------------------------------------===//
// Effect Traits
diff --git a/mlir/include/mlir/TableGen/SideEffects.h b/mlir/include/mlir/TableGen/SideEffects.h
index 4dcc6c90fbeeb47..5a9a34d4e427ccf 100644
--- a/mlir/include/mlir/TableGen/SideEffects.h
+++ b/mlir/include/mlir/TableGen/SideEffects.h
@@ -35,6 +35,12 @@ class SideEffect : public Operator::VariableDecorator {
// 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);
};
diff --git a/mlir/lib/TableGen/SideEffects.cpp b/mlir/lib/TableGen/SideEffects.cpp
index a635f198c595ae1..55ad59d3d0d01a8 100644
--- a/mlir/lib/TableGen/SideEffects.cpp
+++ b/mlir/lib/TableGen/SideEffects.cpp
@@ -36,6 +36,12 @@ 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");
}
diff --git a/mlir/test/lib/Dialect/Test/TestInterfaces.td b/mlir/test/lib/Dialect/Test/TestInterfaces.td
index 79a6c86493d418c..dea26b8dda62a0b 100644
--- a/mlir/test/lib/Dialect/Test/TestInterfaces.td
+++ b/mlir/test/lib/Dialect/Test/TestInterfaces.td
@@ -127,7 +127,8 @@ def TestEffectOpInterface
}
class TestEffect<string effectName>
- : SideEffect<TestEffectOpInterface, effectName, DefaultResource>;
+ : SideEffect<TestEffectOpInterface, effectName, DefaultResource, 0,
+ PartialEffect>;
class TestEffects<list<TestEffect> effects = []>
: SideEffectsTraitBase<TestEffectOpInterface, effects>;
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 6887f151eef7695..edb63924b3553f2 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2519,11 +2519,11 @@ def TestEffectsOpA : TEST_Op<"op_with_effects_a"> {
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]>]> {
diff --git a/mlir/test/mlir-tblgen/op-side-effects.td b/mlir/test/mlir-tblgen/op-side-effects.td
index b8e5f6a0f72d911..09612db905899fd 100644
--- a/mlir/test/mlir-tblgen/op-side-effects.td
+++ b/mlir/test/mlir-tblgen/op-side-effects.td
@@ -13,25 +13,28 @@ def CustomResource : Resource<"CustomResource">;
def SideEffectOpA : TEST_Op<"side_effect_op_a"> {
let arguments = (ins
Arg<Variadic<AnyMemRef>, "", [MemRead]>,
+ Arg<AnyMemRef, "", [MemWriteAt<1, FullEffect>]>,
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, false, ::mlir::SideEffects::DefaultResource::get());
+// CHECK: for (::mlir::Value value : getODSOperands(1))
+// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), value, 1, true, ::mlir::SideEffects::DefaultResource::get());
+// CHECK: effects.emplace_back(::mlir::MemoryEffects::Read::get(), getSymbolAttr(), 0, false, ::mlir::SideEffects::DefaultResource::get());
+// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), getFlatSymbolAttr(), 0, false, ::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, false, ::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, false, 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, false, CustomResource::get());
diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 985d82d0f0eb9d7..5f6a4e3bc52a840 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
@@ -3211,9 +3211,11 @@ void OpEmitter::genSideEffectInterfaceMethods() {
// 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.
@@ -3230,20 +3232,25 @@ void OpEmitter::genSideEffectInterfaceMethods() {
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 {
@@ -3251,7 +3258,9 @@ void OpEmitter::genSideEffectInterfaceMethods() {
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();
}
}
}
More information about the Mlir-commits
mailing list