[mlir] Introduce support for parametric side-effects

The side effect infrastructure is based on the Effect and Resource class
templates, instances of instantiations of which are constructed as
thread-local singletons. With this scheme, it is impossible to further
parameterize either of those, or the EffectInstance class that contains
pointers to an Effect and Resource instances. Such a parameterization is
necessary to express more detailed side effects, e.g. those of a loop or
a function call with affine operations inside where it is possible to
precisely specify the slices of accessed buffers.

Include an additional Attribute to EffectInstance class for further
parameterization. This allows to leverage the dialect-specific
registration and uniquing capabilities of the attribute infrastructure
without requiring Effect or Resource instantiations to be attached to a
dialect themselves.

Split out the generic part of the side effect Tablegen classes into a
separate file to avoid generating built-in MemoryEffect interfaces when
processing any .td file that includes SideEffectInterfaceBase.td.

Reviewed By: rriddle

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

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

@@ -0,0 +1,166 @@
+//===-- SideEffectInterfaceBase.td - Side Effect Base ------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains base class definitions for side effect interfaces, i.e.
+// the customizable interfaces that provide information about which effects are
+// applied by an operation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_SIDEEFFECTS_BASE
+#define MLIR_INTERFACES_SIDEEFFECTS_BASE
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// Resource Bindings
+//===----------------------------------------------------------------------===//
+
+// A generic resource that can be attached to a general base side effect.
+class Resource<string resourceName> {
+  /// The resource that the associated effect is being applied to.
+  string name = resourceName;
+}
+
+// An intrinsic resource that lives in the ::mlir::SideEffects namespace.
+class IntrinsicResource<string resourceName> :
+  Resource<!strconcat("::mlir::SideEffects::", resourceName)> {
+}
+
+// A link to the DefaultResource class.
+def DefaultResource : IntrinsicResource<"DefaultResource">;
+// A link to the AutomaticAllocationScopeResource class.
+def AutomaticAllocationScopeResource :
+  IntrinsicResource<"AutomaticAllocationScopeResource">;
+
+//===----------------------------------------------------------------------===//
+// EffectOpInterface
+//===----------------------------------------------------------------------===//
+
+// A base interface used to query information about the side effects applied to
+// an operation. This template class takes the name of the derived interface
+// class, as well as the name of the base effect class.
+class EffectOpInterfaceBase<string name, string baseEffect>
+    : OpInterface<name> {
+  let methods = [
+    InterfaceMethod<[{
+        Collects all of the operation's effects into `effects`.
+      }],
+      "void", "getEffects",
+         (ins "SmallVectorImpl<::mlir::SideEffects::EffectInstance<"
+              # baseEffect # ">> &":$effects)
+    >,
+  ];
+
+  let extraClassDeclaration = [{
+    /// Collect all of the effect instances that correspond to the given
+    /// `Effect` and place them in 'effects'.
+    template <typename Effect> void getEffects(
+      SmallVectorImpl<::mlir::SideEffects::EffectInstance<
+                                              }] # baseEffect # [{>> &effects) {
+      getEffects(effects);
+      llvm::erase_if(effects, [&](auto &it) {
+        return !llvm::isa<Effect>(it.getEffect());
+      });
+    }
+
+    /// Returns true if this operation exhibits the given effect.
+    template <typename Effect> bool hasEffect() {
+      SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
+      getEffects(effects);
+      return llvm::any_of(effects, [](const auto &it) {
+        return llvm::isa<Effect>(it.getEffect());
+      });
+    }
+
+    /// Returns true if this operation only has the given effect.
+    template <typename Effect> bool onlyHasEffect() {
+      SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
+      getEffects(effects);
+      return !effects.empty() && llvm::all_of(effects, [](const auto &it) {
+        return isa<Effect>(it.getEffect());
+      });
+    }
+
+    /// Returns true if this operation has no effects.
+    bool hasNoEffect() {
+      SmallVector<::mlir::SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
+      getEffects(effects);
+      return effects.empty();
+    }
+
+    /// Returns true if the given operation has no effects for this interface.
+    static bool hasNoEffect(Operation *op) {
+      if (auto interface = dyn_cast<}] # name # [{>(op))
+        return interface.hasNoEffect();
+      return op->hasTrait<::mlir::OpTrait::HasRecursiveSideEffects>();
+    }
+
+    /// Collect all of the effect instances that operate on the provided value
+    /// and place them in 'effects'.
+    void getEffectsOnValue(::mlir::Value value,
+              llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<
+              }] # baseEffect # [{>> & effects) {
+      getEffects(effects);
+      llvm::erase_if(effects, [&](auto &it) { return it.getValue() != value; });
+    }
+
+    /// Collect all of the effect instances that operate on the provided
+    /// resource and place them in 'effects'.
+    void getEffectsOnValue(::mlir::SideEffects::Resource *resource,
+              llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<
+              }] # baseEffect # [{>> & effects) {
+      getEffects(effects);
+      llvm::erase_if(effects, [&](auto &it) {
+        return it.getResource() != resource;
+      });
+    }
+  }];
+
+  // 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 {
+  /// 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;
+}
+
+// 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;
+
+  /// The cpp namespace of the interface trait.
+  string cppNamespace = parentInterface.cppNamespace;
+
+  /// The name of the base effects class.
+  string baseEffectName = parentInterface.baseEffectName;
+
+  /// The derived effects being applied.
+  list<SideEffect> effects = staticEffects;
+}
+
+#endif // MLIR_INTERFACES_SIDEEFFECTS_BASE

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

@@ -131,8 +131,9 @@ struct AutomaticAllocationScopeResource
 
 /// This class represents a specific instance of an effect. It contains the
 /// effect being applied, a resource that corresponds to where the effect is
-/// applied, and an optional value(either operand, result, or region entry
-/// argument) that the effect is applied to.
+/// applied, an optional 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())
@@ -140,6 +141,13 @@ public:
   EffectInstance(EffectT *effect, Value value,
                  Resource *resource = DefaultResource::get())
       : effect(effect), resource(resource), value(value) {}
+  EffectInstance(EffectT *effect, Attribute parameters,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), parameters(parameters) {}
+  EffectInstance(EffectT *effect, Value value, Attribute parameters,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), value(value),
+        parameters(parameters) {}
 
   /// Return the effect being applied.
   EffectT *getEffect() const { return effect; }
@@ -151,6 +159,9 @@ public:
   /// 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; }
+
 private:
   /// The specific effect being applied.
   EffectT *effect;
@@ -160,6 +171,11 @@ private:
 
   /// The value that the effect applies to. This is optionally null.
   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;
 };
 } // namespace SideEffects
 

mlir/include/mlir/Interfaces/SideEffectInterfaces.td

@@ -1,4 +1,4 @@
-//===-- SideEffectInterfaces.td - Side Effect Interfaces ------------*- tablegen -*-===//
+//===-- SideEffectInterfaces.td - Side Effect Interfaces ---*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -14,153 +14,7 @@
 #ifndef MLIR_INTERFACES_SIDEEFFECTS
 #define MLIR_INTERFACES_SIDEEFFECTS
 
-include "mlir/IR/OpBase.td"
-
-//===----------------------------------------------------------------------===//
-// Resource Bindings
-//===----------------------------------------------------------------------===//
-
-// A generic resource that can be attached to a general base side effect.
-class Resource<string resourceName> {
-  /// The resource that the associated effect is being applied to.
-  string name = resourceName;
-}
-
-// An intrinsic resource that lives in the ::mlir::SideEffects namespace.
-class IntrinsicResource<string resourceName> :
-  Resource<!strconcat("::mlir::SideEffects::", resourceName)> {
-}
-
-// A link to the DefaultResource class.
-def DefaultResource : IntrinsicResource<"DefaultResource">;
-// A link to the AutomaticAllocationScopeResource class.
-def AutomaticAllocationScopeResource :
-  IntrinsicResource<"AutomaticAllocationScopeResource">;
-
-//===----------------------------------------------------------------------===//
-// EffectOpInterface
-//===----------------------------------------------------------------------===//
-
-// A base interface used to query information about the side effects applied to
-// an operation. This template class takes the name of the derived interface
-// class, as well as the name of the base effect class.
-class EffectOpInterfaceBase<string name, string baseEffect>
-    : OpInterface<name> {
-  let methods = [
-    InterfaceMethod<[{
-        Collects all of the operation's effects into `effects`.
-      }],
-      "void", "getEffects",
-         (ins "SmallVectorImpl<::mlir::SideEffects::EffectInstance<"
-              # baseEffect # ">> &":$effects)
-    >,
-  ];
-
-  let extraClassDeclaration = [{
-    /// Collect all of the effect instances that correspond to the given
-    /// `Effect` and place them in 'effects'.
-    template <typename Effect> void getEffects(
-      SmallVectorImpl<::mlir::SideEffects::EffectInstance<
-                                              }] # baseEffect # [{>> &effects) {
-      getEffects(effects);
-      llvm::erase_if(effects, [&](auto &it) {
-        return !llvm::isa<Effect>(it.getEffect());
-      });
-    }
-
-    /// Returns true if this operation exhibits the given effect.
-    template <typename Effect> bool hasEffect() {
-      SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
-      getEffects(effects);
-      return llvm::any_of(effects, [](const auto &it) {
-        return llvm::isa<Effect>(it.getEffect());
-      });
-    }
-
-    /// Returns true if this operation only has the given effect.
-    template <typename Effect> bool onlyHasEffect() {
-      SmallVector<SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
-      getEffects(effects);
-      return !effects.empty() && llvm::all_of(effects, [](const auto &it) {
-        return isa<Effect>(it.getEffect());
-      });
-    }
-
-    /// Returns true if this operation has no effects.
-    bool hasNoEffect() {
-      SmallVector<::mlir::SideEffects::EffectInstance<}] # baseEffect # [{>, 4> effects;
-      getEffects(effects);
-      return effects.empty();
-    }
-
-    /// Returns true if the given operation has no effects for this interface.
-    static bool hasNoEffect(Operation *op) {
-      if (auto interface = dyn_cast<}] # name # [{>(op))
-        return interface.hasNoEffect();
-      return op->hasTrait<::mlir::OpTrait::HasRecursiveSideEffects>();
-    }
-
-    /// Collect all of the effect instances that operate on the provided value
-    /// and place them in 'effects'.
-    void getEffectsOnValue(::mlir::Value value,
-              llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<
-              }] # baseEffect # [{>> & effects) {
-      getEffects(effects);
-      llvm::erase_if(effects, [&](auto &it) { return it.getValue() != value; });
-    }
-
-    /// Collect all of the effect instances that operate on the provided
-    /// resource and place them in 'effects'.
-    void getEffectsOnValue(::mlir::SideEffects::Resource *resource,
-              llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<
-              }] # baseEffect # [{>> & effects) {
-      getEffects(effects);
-      llvm::erase_if(effects, [&](auto &it) {
-        return it.getResource() != resource;
-      });
-    }
-  }];
-
-  // 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 {
-  /// 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;
-}
-
-// 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;
-
-  /// The cpp namespace of the interface trait.
-  string cppNamespace = parentInterface.cppNamespace;
-
-  /// The name of the base effects class.
-  string baseEffectName = parentInterface.baseEffectName;
-
-  /// The derived effects being applied.
-  list<SideEffect> effects = staticEffects;
-}
+include "mlir/Interfaces/SideEffectInterfaceBase.td"
 
 //===----------------------------------------------------------------------===//
 // MemoryEffects

mlir/test/IR/test-side-effects.mlir

@@ -18,3 +18,10 @@
 %3 = "test.side_effect_op"() {effects = [
   {effect="allocate", on_result, test_resource}
 ]} : () -> i32
+
+// No _memory_ effects, but a parametric test effect.
+// expected-remark@+2 {{operation has no memory effects}}
+// expected-remark@+1 {{found a parametric effect with affine_map<(d0, d1) -> (d1, d0)>}}
+%4 = "test.side_effect_op"() {
+  effect_parameter = affine_map<(i, j) -> (j, i)>
+} : () -> i32

mlir/test/lib/Dialect/Test/CMakeLists.txt

@@ -7,6 +7,8 @@ set(LLVM_OPTIONAL_SOURCES
 set(LLVM_TARGET_DEFINITIONS TestInterfaces.td)
 mlir_tablegen(TestTypeInterfaces.h.inc -gen-type-interface-decls)
 mlir_tablegen(TestTypeInterfaces.cpp.inc -gen-type-interface-defs)
+mlir_tablegen(TestOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(TestOpInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRTestInterfaceIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestTypeDefs.td)
@@ -29,6 +31,7 @@ add_public_tablegen_target(MLIRTestOpsIncGen)
 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestDialect
   TestDialect.cpp
+  TestInterfaces.cpp
   TestPatterns.cpp
   TestTraits.cpp
   TestTypes.cpp

mlir/test/lib/Dialect/Test/TestDialect.cpp

@@ -758,6 +758,15 @@ void SideEffectOp::getEffects(
   }
 }
 
+void SideEffectOp::getEffects(
+    SmallVectorImpl<TestEffects::EffectInstance> &effects) {
+  auto effectsAttr = getAttrOfType<AffineMapAttr>("effect_parameter");
+  if (!effectsAttr)
+    return;
+
+  effects.emplace_back(TestEffects::Concrete::get(), effectsAttr);
+}
+
 //===----------------------------------------------------------------------===//
 // StringAttrPrettyNameOp
 //===----------------------------------------------------------------------===//
@@ -911,6 +920,7 @@ void RegionIfOp::getSuccessorRegions(
 }
 
 #include "TestOpEnums.cpp.inc"
+#include "TestOpInterfaces.cpp.inc"
 #include "TestOpStructs.cpp.inc"
 #include "TestTypeInterfaces.cpp.inc"
 

mlir/test/lib/Dialect/Test/TestDialect.h

@@ -14,6 +14,7 @@
 #ifndef MLIR_TESTDIALECT_H
 #define MLIR_TESTDIALECT_H
 
+#include "TestInterfaces.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Dialect.h"
@@ -30,7 +31,7 @@
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "TestOpEnums.h.inc"
-
+#include "TestOpInterfaces.h.inc"
 #include "TestOpStructs.h.inc"
 #include "TestOpsDialect.h.inc"
 

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

@@ -0,0 +1,8 @@
+#include "TestInterfaces.h"
+
+using namespace mlir;
+
+bool mlir::TestEffects::Effect::classof(
+    const mlir::SideEffects::Effect *effect) {
+  return isa<mlir::TestEffects::Concrete>(effect);
+}

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

@@ -0,0 +1,37 @@
+//===- TestInterfaces.h - MLIR interfaces for testing -----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares interfaces for the 'test' dialect that can be used for
+// testing the interface infrastructure.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TEST_LIB_DIALECT_TEST_TESTINTERFACES_H
+#define MLIR_TEST_LIB_DIALECT_TEST_TESTINTERFACES_H
+
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+namespace mlir {
+namespace TestEffects {
+struct Effect : public SideEffects::Effect {
+  using SideEffects::Effect::Effect;
+
+  template <typename Derived>
+  using Base = SideEffects::Effect::Base<Derived, Effect>;
+
+  static bool classof(const SideEffects::Effect *effect);
+};
+
+using EffectInstance = SideEffects::EffectInstance<Effect>;
+
+struct Concrete : public Effect::Base<Concrete> {};
+
+} // namespace TestEffects
+} // namespace mlir
+
+#endif // MLIR_TEST_LIB_DIALECT_TEST_TESTINTERFACES_H

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

@@ -6,10 +6,11 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef TEST_INTERFACES
-#define TEST_INTERFACES
+#ifndef MLIR_TEST_DIALECT_TEST_INTERFACES
+#define MLIR_TEST_DIALECT_TEST_INTERFACES
 
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/SideEffectInterfaceBase.td"
 
 // A type interface used to test the ODS generation of type interfaces.
 def TestTypeInterface : TypeInterface<"TestTypeInterface"> {
@@ -52,4 +53,18 @@ def TestTypeInterface : TypeInterface<"TestTypeInterface"> {
   }];
 }
 
-#endif // TEST_INTERFACES
+def TestEffectOpInterface
+    : EffectOpInterfaceBase<"TestEffectOpInterface",
+                            "::mlir::TestEffects::Effect"> {
+  let cppNamespace = "::mlir";
+}
+
+class TestEffect<string effectName>
+    : SideEffect<TestEffectOpInterface, effectName, DefaultResource>;
+
+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

@@ -18,6 +18,7 @@ include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "TestInterfaces.td"
 
 def Test_Dialect : Dialect {
   let name = "test";
@@ -1721,7 +1722,8 @@ def FormatTypesMatchAttrOp : TEST_Op<"format_types_match_attr", [
 //===----------------------------------------------------------------------===//
 
 def SideEffectOp : TEST_Op<"side_effect_op",
-    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     DeclareOpInterfaceMethods<TestEffectOpInterface>]> {
   let results = (outs AnyType:$result);
 }
 

mlir/test/lib/IR/TestSideEffects.cpp

@@ -47,6 +47,20 @@ struct SideEffectsPass
         diag << " on resource '" << instance.getResource()->getName() << "'";
       }
     });
+
+    SmallVector<TestEffects::EffectInstance, 1> testEffects;
+    module.walk([&](TestEffectOpInterface op) {
+      testEffects.clear();
+      op.getEffects(testEffects);
+
+      if (testEffects.empty())
+        return;
+
+      for (const TestEffects::EffectInstance &instance : testEffects) {
+        op.emitRemark() << "found a parametric effect with "
+                        << instance.getParameters();
+      }
+    });
   }
 };
 } // end anonymous namespace