[mlir] Add support for referencing a SymbolRefAttr in a SideEffectInstance

[rriddle]

Summary

This allows for operations that exclusively affect symbol operations to better describe their side effects.

mlir/include/mlir/Interfaces/SideEffectInterfaceBase.td

let extraClassDeclaration = [{
     /// 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 symbol
+    /// reference and place them in 'effects'.
+    void getEffectsOnSymbol(::mlir::SymbolRefAttr value,
+              llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<
+              }] # baseEffect # [{>> & effects) {
+      getEffects(effects);
+      llvm::erase_if(effects, [&](auto &it) {
+        return it.getSymbolRef() != value;
+      });
+    }
+
     /// Collect all of the effect instances that operate on the provided
     /// resource and place them in 'effects'.
-    void getEffectsOnValue(::mlir::SideEffects::Resource *resource,
+    void getEffectsOnResource(::mlir::SideEffects::Resource *resource,
               llvm::SmallVectorImpl<::mlir::SideEffects::EffectInstance<
               }] # baseEffect # [{>> & effects) {
       getEffects(effects);
       llvm::erase_if(effects, [&](auto &it) {
         return it.getResource() != resource;
       });
     }
   }];

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

 /// parent AutomaticAllocationScope trait.
 struct AutomaticAllocationScopeResource
     : public Resource::Base<AutomaticAllocationScopeResource> {
   StringRef getName() final { return "AutomaticAllocationScope"; }
 };
 
 /// 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, 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.
+/// 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) {}
   EffectInstance(EffectT *effect, Value value,
                  Resource *resource = DefaultResource::get())
       : effect(effect), resource(resource), value(value) {}
+  EffectInstance(EffectT *effect, SymbolRefAttr symbol,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), value(symbol) {}
   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) {}
+  EffectInstance(EffectT *effect, SymbolRefAttr symbol, Attribute parameters,
+                 Resource *resource = DefaultResource::get())
+      : effect(effect), resource(resource), value(symbol),
+        parameters(parameters) {}
 
   /// 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; }
+  Value getValue() const { return value ? value.dyn_cast<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 ? value.dyn_cast<SymbolRefAttr>() : 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; }
 
 private:
   /// The specific effect being applied.
   EffectT *effect;
 
   /// The resource that the given value resides in.
   Resource *resource;
 
-  /// The value that the effect applies to. This is optionally null.
-  Value value;
+  /// 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;
 };
 } // namespace SideEffects
 

mlir/include/mlir/TableGen/Attribute.h

public:
   // Returns true if this attribute is a derived attribute (i.e., a subclass
   // of `DerivedAttr`).
   bool isDerivedAttr() const;
 
   // Returns true if this attribute is a type attribute (i.e., a subclass
   // of `TypeAttrBase`).
   bool isTypeAttr() const;
 
+  // Returns true if this attribute is a symbol reference attribute (i.e., a
+  // subclass of `SymbolRefAttr` or `FlatSymbolRefAttr`).
+  bool isSymbolRefAttr() const;
+
   // Returns true if this attribute is an enum attribute (i.e., a subclass of
   // `EnumAttrInfo`)
   bool isEnumAttr() const;
 
   // Returns this attribute's TableGen def name. If this is an `OptionalAttr`
   // or `DefaultValuedAttr` without explicit name, returns the base attribute's
   // name.
   StringRef getAttrDefName() const;

mlir/lib/TableGen/Attribute.cpp

 }
 
 Attribute::Attribute(const DefInit *init) : Attribute(init->getDef()) {}
 
 bool Attribute::isDerivedAttr() const { return isSubClassOf("DerivedAttr"); }
 
 bool Attribute::isTypeAttr() const { return isSubClassOf("TypeAttrBase"); }
 
+bool Attribute::isSymbolRefAttr() const {
+  StringRef defName = def->getName();
+  if (defName == "SymbolRefAttr" || defName == "FlatSymbolRefAttr")
+    return true;
+  return isSubClassOf("SymbolRefAttr") || isSubClassOf("FlatSymbolRefAttr");
+}
+
 bool Attribute::isEnumAttr() const { return isSubClassOf("EnumAttrInfo"); }
 
 StringRef Attribute::getStorageType() const {
   const auto *init = def->getValueInit("storageType");
   auto type = getValueAsString(init);
   if (type.empty())
     return "Attribute";
   return type;

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

%2 = "test.side_effect_op"() {effects = [
   {effect="write", test_resource}
 ]} : () -> i32
 
 // expected-remark@+1 {{found an instance of 'allocate' on a value, on resource '<Test>'}}
 %3 = "test.side_effect_op"() {effects = [
   {effect="allocate", on_result, test_resource}
 ]} : () -> i32
 
+// expected-remark@+1 {{found an instance of 'read' on a symbol '@foo_ref', on resource '<Test>'}}
+"test.side_effect_op"() {effects = [
+  {effect="read", on_reference = @foo_ref, 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/TestDialect.cpp

for (Attribute element : effectsAttr) {
     MemoryEffects::Effect *effect =
         StringSwitch<MemoryEffects::Effect *>(
             effectElement.get("effect").cast<StringAttr>().getValue())
             .Case("allocate", MemoryEffects::Allocate::get())
             .Case("free", MemoryEffects::Free::get())
             .Case("read", MemoryEffects::Read::get())
             .Case("write", MemoryEffects::Write::get());
 
-    // Check for a result to affect.
-    Value value;
-    if (effectElement.get("on_result"))
-      value = getResult();
-
     // Check for a non-default resource to use.
     SideEffects::Resource *resource = SideEffects::DefaultResource::get();
     if (effectElement.get("test_resource"))
       resource = TestResource::get();
 
-    effects.emplace_back(effect, value, resource);
+    // Check for a result to affect.
+    if (effectElement.get("on_result"))
+      effects.emplace_back(effect, getResult(), resource);
+    else if (Attribute ref = effectElement.get("on_reference"))
+      effects.emplace_back(effect, ref.cast<SymbolRefAttr>(), resource);
+    else
+      effects.emplace_back(effect, resource);
   }
 }
 
 void SideEffectOp::getEffects(
     SmallVectorImpl<TestEffects::EffectInstance> &effects) {
   auto effectsAttr = getAttrOfType<AffineMapAttr>("effect_parameter");
   if (!effectsAttr)
     return;

mlir/test/lib/IR/TestSideEffects.cpp

module.walk([&](MemoryEffectOpInterface op) {
           diag << "'free'";
         else if (isa<MemoryEffects::Read>(instance.getEffect()))
           diag << "'read'";
         else if (isa<MemoryEffects::Write>(instance.getEffect()))
           diag << "'write'";
 
         if (instance.getValue())
           diag << " on a value,";
+        else if (SymbolRefAttr symbolRef = instance.getSymbolRef())
+          diag << " on a symbol '" << symbolRef << "',";
 
         diag << " on resource '" << instance.getResource()->getName() << "'";
       }
     });
 
     SmallVector<TestEffects::EffectInstance, 1> testEffects;
     module.walk([&](TestEffectOpInterface op) {
       testEffects.clear();

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<OpTrait> 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]>);
+  let arguments = (ins
+    Arg<Variadic<AnyMemRef>, "", [MemRead]>,
+    Arg<SymbolRefAttr, "", [MemRead]>:$symbol,
+    Arg<FlatSymbolRefAttr, "", [MemWrite]>:$flat_symbol,
+    Arg<OptionalAttr<SymbolRefAttr>, "", [MemRead]>:$optional_symbol
+  );
   let results = (outs Res<AnyMemRef, "", [MemAlloc<CustomResource>]>);
 }
 
 def SideEffectOpB : TEST_Op<"side_effect_op_b",
     [MemoryEffects<[MemWrite<CustomResource>]>]>;
 
 // CHECK: void SideEffectOpA::getEffects
 // CHECK:   for (::mlir::Value value : getODSOperands(0))
 // CHECK:     effects.emplace_back(MemoryEffects::Read::get(), value, ::mlir::SideEffects::DefaultResource::get());
+// CHECK:   effects.emplace_back(MemoryEffects::Read::get(), symbol(), ::mlir::SideEffects::DefaultResource::get());
+// CHECK:   effects.emplace_back(MemoryEffects::Write::get(), flat_symbol(), ::mlir::SideEffects::DefaultResource::get());
+// CHECK:   if (auto symbolRef = optional_symbolAttr())
+// CHECK:     effects.emplace_back(MemoryEffects::Read::get(), symbolRef, ::mlir::SideEffects::DefaultResource::get());
 // CHECK:   for (::mlir::Value value : getODSResults(0))
 // CHECK:     effects.emplace_back(MemoryEffects::Allocate::get(), value, CustomResource::get());
 
 // CHECK: void SideEffectOpB::getEffects
 // CHECK:   effects.emplace_back(MemoryEffects::Write::get(), CustomResource::get());

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

void OpEmitter::genOpInterfaceMethods() {
   for (const auto &trait : op.getTraits()) {
     if (const auto *opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait))
       if (opTrait->shouldDeclareMethods())
         genOpInterfaceMethod(opTrait);
   }
 }
 
 void OpEmitter::genSideEffectInterfaceMethods() {
-  enum EffectKind { Operand, Result, Static };
+  enum EffectKind { Operand, Result, Symbol, Static };
   struct EffectLocation {
     /// The effect applied.
     SideEffect effect;
 
-    /// The index if the kind is either operand or result.
+    /// The index if the kind is not static.
     unsigned index : 30;
 
     /// The kind of the location.
     unsigned kind : 2;
   };
 
   StringMap<SmallVector<EffectLocation, 1>> interfaceEffects;
   auto resolveDecorators = [&](Operator::var_decorator_range decorators,
for (const auto &trait : op.getTraits()) {
     const auto *opTrait = dyn_cast<tblgen::SideEffectTrait>(&trait);
     if (!opTrait)
       continue;
     auto &effects = interfaceEffects[opTrait->getBaseEffectName()];
     for (auto decorator : opTrait->getEffects())
       effects.push_back(EffectLocation{cast<SideEffect>(decorator),
                                        /*index=*/0, EffectKind::Static});
   }
-  /// Operands.
+  /// Attributes and Operands.
   for (unsigned i = 0, operandIt = 0, e = op.getNumArgs(); i != e; ++i) {
-    if (op.getArg(i).is<NamedTypeConstraint *>()) {
+    Argument arg = op.getArg(i);
+    if (arg.is<NamedTypeConstraint *>()) {
       resolveDecorators(op.getArgDecorators(i), operandIt, EffectKind::Operand);
       ++operandIt;
+      continue;
     }
+    const NamedAttribute *attr = arg.get<NamedAttribute *>();
+    if (attr->attr.getBaseAttr().isSymbolRefAttr())
+      resolveDecorators(op.getArgDecorators(i), i, EffectKind::Symbol);
   }
   /// 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.
+  const char *addEffectCode =
+      "  effects.emplace_back({0}::get(), {1}{2}::get());\n";
+
   for (auto &it : interfaceEffects) {
     // Generate the 'getEffects' method.
     std::string type = llvm::formatv("::mlir::SmallVectorImpl<::mlir::"
                                      "SideEffects::EffectInstance<{0}>> &",
                                      it.first())
                            .str();
     auto *getEffects =
         opClass.addMethodAndPrune("void", "getEffects", type, "effects");
     auto &body = getEffects->body();
 
     // Add effect instances for each of the locations marked on the operation.
     for (auto &location : it.second) {
-      if (location.kind != EffectKind::Static) {
+      StringRef effect = location.effect.getName();
+      StringRef resource = location.effect.getResource();
+      if (location.kind == EffectKind::Static) {
+        // A static instance has no attached value.
+        body << llvm::formatv(addEffectCode, 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<NamedAttribute *>();
+        if (attr->attr.isOptional()) {
+          body << "  if (auto symbolRef = " << attr->name << "Attr())\n  "
+               << llvm::formatv(addEffectCode, effect, "symbolRef, ", resource)
+                      .str();
+        } else {
+          body << llvm::formatv(addEffectCode, effect, attr->name + "(), ",
+                                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  ";
+             << "(" << location.index << "))\n  "
+             << llvm::formatv(addEffectCode, effect, "value, ", resource).str();
       }
-
-      body << "  effects.emplace_back(" << location.effect.getName()
-           << "::get()";
-
-      // If the effect isn't static, it has a specific value attached to it.
-      if (location.kind != EffectKind::Static)
-        body << ", value";
-      body << ", " << location.effect.getResource() << "::get());\n";
     }
   }
 }
 
 void OpEmitter::genTypeInterfaceMethods() {
   if (!op.allResultTypesKnown())
     return;