diff --git a/mlir/docs/Rationale/SideEffectsAndSpeculation.md b/mlir/docs/Rationale/SideEffectsAndSpeculation.md --- a/mlir/docs/Rationale/SideEffectsAndSpeculation.md +++ b/mlir/docs/Rationale/SideEffectsAndSpeculation.md @@ -76,6 +76,11 @@ 1. Does it read from or write to the heap or stack? It should probably implement `MemoryEffectsOpInterface`. +1. Does these side effects ordered? It may be necessary to add the stage of + side effects to make the analysis more accurate. We can use the `MemxxAt` + method to mark the stage of side effects in the TD file or add the + SideEffectStageAttr in the constructor of EffectInstance to specify the stage + of side effects. 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 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/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td @@ -186,7 +186,7 @@ OptionalAttr:$elem_type, UnitAttr:$inalloca); let results = (outs Res]>:$res); + [MemAlloc]>:$res); string llvmInstName = "Alloca"; string llvmBuilder = [{ auto addrSpace = $_resultType->getPointerAddressSpace(); 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 @@ -72,7 +72,7 @@ Variadic:$symbolOperands, ConfinedAttr, [IntMinValue<0>]>:$alignment); - let results = (outs Res]>:$memref); + let results = (outs Res]>:$memref); let builders = [ OpBuilder<(ins "MemRefType":$memrefType, @@ -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, "", [MemAlloc]>); 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) 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, int effectStage> + : 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 stage. + int stage = effectStage; } // 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 @@ -163,12 +163,15 @@ /// 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 getValue() const { + return value ? llvm::dyn_cast_if_present(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(value) : SymbolRefAttr(); + return value ? llvm::dyn_cast_if_present(value) + : SymbolRefAttr(); } /// Return the resource that the effect applies to. @@ -177,6 +180,20 @@ /// Return the parameters of the effect, if any. Attribute getParameters() const { return parameters; } + /// Return the effect happen stage. + unsigned getStage() const { + auto dicAttr = llvm::dyn_cast_or_null(parameters); + if (!dicAttr) + return 0; + + auto sideEffectOrder = + dicAttr.template getAs("SideEffectStageAttr"); + if (sideEffectOrder) { + return sideEffectOrder.getValue().getZExtValue(); + } + return 0; + } + private: /// The specific effect being applied. EffectT *effect; 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. @@ -48,30 +48,34 @@ // 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 - : MemoryEffect<"::mlir::MemoryEffects::Allocate", resource>; -def MemAlloc : MemAlloc; +class MemAlloc + : MemoryEffect<"::mlir::MemoryEffects::Allocate", resource, stage>; +def MemAlloc : MemAlloc; +class MemAllocAt : MemAlloc; // 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>; -def MemFree : MemFree; +class MemFree + : MemoryEffect<"::mlir::MemoryEffects::Free", resource, stage>; +def MemFree : MemFree; +class MemFreeAt : MemFree; // 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>; -def MemRead : MemRead; +class MemRead + : MemoryEffect<"::mlir::MemoryEffects::Read", resource, stage>; +def MemRead : MemRead; +class MemReadAt : MemRead; // 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>; -def MemWrite : MemWrite; +class MemWrite + : MemoryEffect<"::mlir::MemoryEffects::Write", resource, stage>; +def MemWrite : MemWrite; +class MemWriteAt : MemWrite; //===----------------------------------------------------------------------===// // 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 stage of the effect happen. + int64_t getStage() const; + static bool classof(const Operator::VariableDecorator *var); }; 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,8 @@ return def->getValueAsString("resource"); } +int64_t SideEffect::getStage() const { return def->getValueAsInt("stage"); } + 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/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/mlir/test/lib/Dialect/Test/TestOps.td @@ -2407,11 +2407,11 @@ Arg, "", [MemRead]>:$optional_symbol ); - let results = (outs Res]>); + let results = (outs Res]>); } def TestEffectsOpB : TEST_Op<"op_with_effects_b", - [MemoryEffects<[MemWrite]>]>; + [MemoryEffects<[MemWrite]>]>; 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 --- a/mlir/test/mlir-tblgen/op-side-effects.td +++ b/mlir/test/mlir-tblgen/op-side-effects.td @@ -13,25 +13,43 @@ def SideEffectOpA : TEST_Op<"side_effect_op_a"> { let arguments = (ins Arg, "", [MemRead]>, + Arg]>, Arg:$symbol, Arg:$flat_symbol, Arg, "", [MemRead]>:$optional_symbol ); - let results = (outs Res]>); + let results = (outs Res]>); } def SideEffectOpB : TEST_Op<"side_effect_op_b", - [MemoryEffects<[MemWrite]>]>; + [MemoryEffects<[MemWrite]>]>; // 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()); +// CEHCK: mlir::StringAttr attrName = mlir::StringAttr::get(getContext(), "SideEffectOrderAttr"); +// CHECK: mlir::IntegerAttr stageAttr = mlir::IntegerAttr::get(mlir::IndexType::get(getContext()), 0); +// CHECK: SmallVector attrs; +// CHECK: attrs.push_back(mlir::NamedAttribute(attrName, stageAttr)); +// CHECK: mlir::DictionaryAttr paramAttr = mlir::DictionaryAttr::get(getContext(), attrs); +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Read::get(), value, paramAttr, ::mlir::SideEffects::DefaultResource::get()); +// CHECK: for (::mlir::Value value : getODSOperands(1)) +// CEHCK: mlir::StringAttr attrName = mlir::StringAttr::get(getContext(), "SideEffectOrderAttr"); +// CHECK: mlir::IntegerAttr stageAttr = mlir::IntegerAttr::get(mlir::IndexType::get(getContext()), 1); +// CHECK: SmallVector attrs; +// CHECK: attrs.push_back(mlir::NamedAttribute(attrName, stageAttr)); +// CHECK: mlir::DictionaryAttr paramAttr = mlir::DictionaryAttr::get(getContext(), attrs); +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), value, paramAttr, ::mlir::SideEffects::DefaultResource::get()); +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Read::get(), getSymbolAttr(), paramAttr, ::mlir::SideEffects::DefaultResource::get()); +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), getFlatSymbolAttr(), paramAttr, ::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, paramAttr, ::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, paramAttr, CustomResource::get()); // CHECK: void SideEffectOpB::getEffects -// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), CustomResource::get()); +// CEHCK: mlir::StringAttr attrName = mlir::StringAttr::get(getContext(), "SideEffectOrderAttr"); +// CHECK: mlir::IntegerAttr stageAttr = mlir::IntegerAttr::get(mlir::IndexType::get(getContext()), 0); +// CHECK: SmallVector attrs; +// CHECK: attrs.push_back(mlir::NamedAttribute(attrName, stageAttr)); +// CHECK: mlir::DictionaryAttr paramAttr = mlir::DictionaryAttr::get(getContext(), attrs); +// CHECK: effects.emplace_back(::mlir::MemoryEffects::Write::get(), paramAttr, CustomResource::get()); 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 @@ -3192,11 +3192,22 @@ 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. + // {0}: The effect happen stage. + // {1}: The effect class. + // {2}: Optional value or symbol reference. + // {3}: The resource class. const char *addEffectCode = - " effects.emplace_back({0}::get(), {1}{2}::get());\n"; + " {\n" + " mlir::StringAttr attrName = mlir::StringAttr::get(getContext(), " + "\"SideEffectOrderAttr\");\n" + " mlir::IntegerAttr stageAttr = " + "mlir::IntegerAttr::get(mlir::IndexType::get(getContext()), {0});\n" + " SmallVector attrs;\n" + " attrs.push_back(mlir::NamedAttribute(attrName, stageAttr));\n" + " mlir::DictionaryAttr paramAttr = " + "mlir::DictionaryAttr::get(getContext(), attrs);\n" + " effects.emplace_back({1}::get(), {2}paramAttr, {3}::get());\n" + " }\n"; for (auto &it : interfaceEffects) { // Generate the 'getEffects' method. @@ -3213,20 +3224,22 @@ for (auto &location : it.second) { StringRef effect = location.effect.getName(); StringRef resource = location.effect.getResource(); + int stage = location.effect.getStage(); if (location.kind == EffectKind::Static) { // A static instance has no attached value. - body << llvm::formatv(addEffectCode, effect, "", resource).str(); + body << llvm::formatv(addEffectCode, stage, effect, "", resource).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, stage, effect, "symbolRef, ", + resource) .str(); } else { - body << llvm::formatv(addEffectCode, effect, argName + "Attr(), ", - resource) + body << llvm::formatv(addEffectCode, stage, effect, + argName + "Attr(), ", resource) .str(); } } else { @@ -3234,7 +3247,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, stage, effect, "value, ", resource) + .str(); } } }