diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td @@ -1222,8 +1222,8 @@ }]; let arguments = (ins Variadic:$asyncDependencies, - Arg:$dst, - Arg:$src); + Arg]>:$dst, + Arg]>:$src); let results = (outs Optional:$asyncToken); let assemblyFormat = [{ diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td @@ -212,8 +212,8 @@ DeclareOpInterfaceMethods, DeclareOpInterfaceMethods], /*requiresAccessGroup=*/1, /*requiresAliasAnalysis=*/1> { - dag args = (ins Arg:$dst, - Arg:$src, + dag args = (ins Arg]>:$dst, + Arg]>:$src, AnySignlessInteger:$len, I1Attr:$isVolatile); // Append the alias attributes defined by LLVM_IntrOpBase. let arguments = !con(args, aliasAttrs); @@ -253,8 +253,8 @@ DeclareOpInterfaceMethods, DeclareOpInterfaceMethods], /*requiresAccessGroup=*/1, /*requiresAliasAnalysis=*/1> { - dag args = (ins Arg:$dst, - Arg:$src, + dag args = (ins Arg]>:$dst, + Arg]>:$src, APIntAttr:$len, I1Attr:$isVolatile); // Append the alias attributes defined by LLVM_IntrOpBase. let arguments = !con(args, aliasAttrs); diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td @@ -276,12 +276,12 @@ // memref and allocating the outcoming memref, even though this may not // physically happen on each execution. - let arguments = (ins Arg, "", [MemFree]>:$source, + let arguments = (ins Arg, "", [MemFreeAt<"0">]>:$source, Optional:$dynamicResultSize, ConfinedAttr, [IntMinValue<0>]>:$alignment); - let results = (outs Res, "", [MemAlloc]>); + let results = (outs Res, "", [MemAllocAt<"1">]>); let builders = [ OpBuilder<(ins "MemRefType":$resultType, @@ -532,9 +532,9 @@ }]; let arguments = (ins Arg:$source, + [MemReadAt<"0">]>:$source, Arg:$target); + [MemWriteAt<"1">]>:$target); let assemblyFormat = [{ $source `,` $target attr-dict `:` type($source) `to` type($target) @@ -794,11 +794,11 @@ SmallVectorImpl> & effects) { effects.emplace_back(MemoryEffects::Read::get(), getSrcMemRef(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 0); effects.emplace_back(MemoryEffects::Write::get(), getDstMemRef(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 1); effects.emplace_back(MemoryEffects::Read::get(), getTagMemRef(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 0); } }]; let hasCustomAssemblyFormat = 1; diff --git a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td --- a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td +++ b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td @@ -381,9 +381,9 @@ ``` }]; let results = (outs NVGPU_DeviceAsyncToken:$asyncToken); - let arguments = (ins Arg:$dst, + let arguments = (ins Arg]>:$dst, Variadic:$dstIndices, - Arg:$src, + Arg]>:$src, Variadic:$srcIndices, IndexAttr:$dstElements, Optional:$srcElements, diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td b/mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td --- a/mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td +++ b/mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td @@ -152,7 +152,8 @@ // This class is the general base side effect class. This is used by derived // effect interfaces to define their effects. class SideEffect : OpVariableDecorator { + Resource resourceReference, string indexString> + : OpVariableDecorator { /// The name of the base effects class. string baseEffectName = interface.baseEffectName; @@ -167,6 +168,9 @@ /// The resource that the effect is being applied to. string resource = resourceReference.name; + + /// The effect happen index. + string index = indexString; } // 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 --- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h +++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h @@ -138,25 +138,32 @@ template class EffectInstance { public: - EffectInstance(EffectT *effect, Resource *resource = DefaultResource::get()) - : effect(effect), resource(resource) {} + EffectInstance(EffectT *effect, Resource *resource = DefaultResource::get(), + unsigned index = 0) + : effect(effect), resource(resource), index(index) {} EffectInstance(EffectT *effect, Value value, - Resource *resource = DefaultResource::get()) - : effect(effect), resource(resource), value(value) {} + Resource *resource = DefaultResource::get(), + unsigned index = 0) + : effect(effect), resource(resource), value(value), index(index) {} EffectInstance(EffectT *effect, SymbolRefAttr symbol, - Resource *resource = DefaultResource::get()) - : effect(effect), resource(resource), value(symbol) {} + Resource *resource = DefaultResource::get(), + unsigned index = 0) + : effect(effect), resource(resource), value(symbol), index(index) {} EffectInstance(EffectT *effect, Attribute parameters, - Resource *resource = DefaultResource::get()) - : effect(effect), resource(resource), parameters(parameters) {} + Resource *resource = DefaultResource::get(), + unsigned index = 0) + : effect(effect), resource(resource), parameters(parameters), + index(index) {} EffectInstance(EffectT *effect, Value value, Attribute parameters, - Resource *resource = DefaultResource::get()) + Resource *resource = DefaultResource::get(), + unsigned index = 0) : effect(effect), resource(resource), value(value), - parameters(parameters) {} + parameters(parameters), index(index) {} EffectInstance(EffectT *effect, SymbolRefAttr symbol, Attribute parameters, - Resource *resource = DefaultResource::get()) + Resource *resource = DefaultResource::get(), + unsigned index = 0) : effect(effect), resource(resource), value(symbol), - parameters(parameters) {} + parameters(parameters), index(index) {} /// Return the effect being applied. EffectT *getEffect() const { return effect; } @@ -177,6 +184,9 @@ /// Return the parameters of the effect, if any. Attribute getParameters() const { return parameters; } + /// Return the index when effect happen. + unsigned getIndex() const { return index; } + private: /// The specific effect being applied. EffectT *effect; @@ -191,6 +201,9 @@ /// type-safe structured storage and context-based uniquing. Concrete effects /// can use this at their convenience. This is optionally null. Attribute parameters; + + /// The index when memroy side effect happen. + unsigned index; }; } // namespace SideEffects diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td --- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td +++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td @@ -34,8 +34,8 @@ } // The base class for defining specific memory effects. -class MemoryEffect - : SideEffect; +class MemoryEffect + : SideEffect; // This class represents the trait for memory effects that may be placed on // operations. @@ -49,30 +49,42 @@ // resource. An 'allocate' effect implies only allocation of the resource, and // not any visible mutation or dereference. class MemAlloc - : MemoryEffect<"::mlir::MemoryEffects::Allocate", resource>; + : MemoryEffect<"::mlir::MemoryEffects::Allocate", resource, "0">; def MemAlloc : MemAlloc; +class MemAllocAt + : MemoryEffect<"::mlir::MemoryEffects::Allocate", DefaultResource, index>; + // 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 - : MemoryEffect<"::mlir::MemoryEffects::Free", resource>; + : MemoryEffect<"::mlir::MemoryEffects::Free", resource, "0">; def MemFree : MemFree; +class MemFreeAt + : MemoryEffect<"::mlir::MemoryEffects::Free", DefaultResource, index>; + // 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 - : MemoryEffect<"::mlir::MemoryEffects::Read", resource>; + : MemoryEffect<"::mlir::MemoryEffects::Read", resource, "0">; def MemRead : MemRead; +class MemReadAt + : MemoryEffect<"::mlir::MemoryEffects::Read", DefaultResource, index>; + // 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 - : MemoryEffect<"::mlir::MemoryEffects::Write", resource>; + : MemoryEffect<"::mlir::MemoryEffects::Write", resource, "0">; def MemWrite : MemWrite; +class MemWriteAt + : MemoryEffect<"::mlir::MemoryEffects::Write", DefaultResource, index>; + //===----------------------------------------------------------------------===// // Effect Traits //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/TableGen/SideEffects.h b/mlir/include/mlir/TableGen/SideEffects.h --- a/mlir/include/mlir/TableGen/SideEffects.h +++ b/mlir/include/mlir/TableGen/SideEffects.h @@ -35,6 +35,9 @@ // Return the name of the resource class. StringRef getResource() const; + // Return the index of the effect happen. + StringRef getIndex() const; + static bool classof(const Operator::VariableDecorator *var); }; diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -1787,11 +1787,11 @@ SmallVectorImpl> &effects) { effects.emplace_back(MemoryEffects::Read::get(), getSrcMemRef(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 0); effects.emplace_back(MemoryEffects::Write::get(), getDstMemRef(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 1); effects.emplace_back(MemoryEffects::Read::get(), getTagMemRef(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 0); } //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp --- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp @@ -493,11 +493,11 @@ SmallVectorImpl> &effects) { effects.emplace_back(MemoryEffects::Read::get(), getInput(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 0); effects.emplace_back(MemoryEffects::Write::get(), getOutput(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 2); effects.emplace_back(MemoryEffects::Allocate::get(), getOutput(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 1); } OpFoldResult CloneOp::fold(FoldAdaptor adaptor) { diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -977,15 +977,15 @@ if (!llvm::isa(operand->get().getType())) continue; effects.emplace_back(MemoryEffects::Read::get(), operand->get(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 0); } for (auto *operand : outputOperands) { if (!llvm::isa(operand->get().getType())) continue; effects.emplace_back(MemoryEffects::Read::get(), operand->get(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 0); effects.emplace_back(MemoryEffects::Write::get(), operand->get(), - SideEffects::DefaultResource::get()); + SideEffects::DefaultResource::get(), 1); } } diff --git a/mlir/lib/TableGen/SideEffects.cpp b/mlir/lib/TableGen/SideEffects.cpp --- a/mlir/lib/TableGen/SideEffects.cpp +++ b/mlir/lib/TableGen/SideEffects.cpp @@ -36,6 +36,10 @@ return def->getValueAsString("resource"); } +StringRef SideEffect::getIndex() const { + return def->getValueAsString("index"); +} + 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 --- a/mlir/test/lib/Dialect/Test/TestInterfaces.td +++ b/mlir/test/lib/Dialect/Test/TestInterfaces.td @@ -127,7 +127,7 @@ } class TestEffect - : SideEffect; + : SideEffect; class TestEffects effects = []> : SideEffectsTraitBase; diff --git a/mlir/test/mlir-tblgen/op-side-effects.td b/mlir/test/mlir-tblgen/op-side-effects.td --- a/mlir/test/mlir-tblgen/op-side-effects.td +++ b/mlir/test/mlir-tblgen/op-side-effects.td @@ -13,6 +13,7 @@ def SideEffectOpA : TEST_Op<"side_effect_op_a"> { let arguments = (ins Arg, "", [MemRead]>, + Arg]>, Arg:$symbol, Arg:$flat_symbol, Arg, "", [MemRead]>:$optional_symbol @@ -25,13 +26,15 @@ // 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, ::mlir::SideEffects::DefaultResource::get(), 0); +// CHECK: for (::mlir::Value value : getODSOperands(1)) +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), value, ::mlir::SideEffects::DefaultResource::get(), 1); +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Read::get(), getSymbolAttr(), ::mlir::SideEffects::DefaultResource::get(), 0); +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), getFlatSymbolAttr(), ::mlir::SideEffects::DefaultResource::get(), 0); // 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, ::mlir::SideEffects::DefaultResource::get(), 0); // 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, CustomResource::get(), 0); // CHECK: void SideEffectOpB::getEffects -// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), CustomResource::get()); +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), CustomResource::get(), 0); diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp --- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp +++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp @@ -2964,7 +2964,7 @@ // {1}: Optional value or symbol reference. // {1}: The resource class. const char *addEffectCode = - " effects.emplace_back({0}::get(), {1}{2}::get());\n"; + " effects.emplace_back({0}::get(), {1}{2}::get(), {3});\n"; for (auto &it : interfaceEffects) { // Generate the 'getEffects' method. @@ -2981,20 +2981,22 @@ for (auto &location : it.second) { StringRef effect = location.effect.getName(); StringRef resource = location.effect.getResource(); + StringRef index = location.effect.getIndex(); if (location.kind == EffectKind::Static) { // A static instance has no attached value. - body << llvm::formatv(addEffectCode, effect, "", resource).str(); + body << llvm::formatv(addEffectCode, effect, "", resource, index).str(); } else if (location.kind == EffectKind::Symbol) { // A symbol reference requires adding the proper attribute. const auto *attr = op.getArg(location.index).get(); 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, ", resource, + index) .str(); } else { body << llvm::formatv(addEffectCode, effect, argName + "Attr(), ", - resource) + resource, index) .str(); } } else { @@ -3002,7 +3004,8 @@ 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, ", resource, index) + .str(); } } }