Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[mlir] NamedAttribute utility generator #75118

Closed
wants to merge 1 commit into from
Closed

Conversation

sjw36
Copy link
Contributor

@sjw36 sjw36 commented Dec 11, 2023

All attributes in MLIR are named, inherent attributes have unscoped names and arbitrary attributes should be scoped with a dialect. Current usage is ad-hoc and much of the codebase is sprinkled with constant strings used to lookup and set attributes, leading to potential bugs when names are not updated in all usages.

This PR adds a tablegen'd utility wrapper for a NamedAttribute that manages scoped/unscoped name lookup for accessing the attribute on an Operation based on inherentness.

@llvmbot
Copy link
Collaborator

llvmbot commented Dec 11, 2023

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-ods
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir-gpu

Author: SJW (sjw36)

Changes

All attributes in MLIR are named, inherent attributes have unscoped names and arbitrary attributes should be scoped with a dialect. Current usage is ad-hoc and much of the codebase is sprinkled with constant strings used to lookup and set attributes, leading to potential bugs when names are not updated in all usages.

This PR adds a tablegen'd utility wrapper for a NamedAttribute that manages scoped/unscoped name lookup for accessing the attribute on an Operation based on inherentness.


Full diff: https://github.com/llvm/llvm-project/pull/75118.diff

5 Files Affected:

  • (modified) mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td (+16-10)
  • (modified) mlir/include/mlir/IR/AttrTypeBase.td (+72)
  • (modified) mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp (+3-6)
  • (modified) mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp (+2-2)
  • (modified) mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp (+4-6)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 48b830ae34f29..a40599e91e4b5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -28,16 +28,6 @@ def ROCDL_Dialect : Dialect {
   let hasOperationAttrVerify = 1;
 
   let extraClassDeclaration = [{
-    /// Get the name of the attribute used to annotate external kernel
-    /// functions.
-    static StringRef getKernelFuncAttrName() { return "rocdl.kernel"; }
-    static constexpr ::llvm::StringLiteral getFlatWorkGroupSizeAttrName() {
-      return ::llvm::StringLiteral("rocdl.flat_work_group_size");
-    }
-    static constexpr ::llvm::StringLiteral getReqdWorkGroupSizeAttrName() {
-      return ::llvm::StringLiteral("rocdl.reqd_work_group_size");
-    }
-
     /// The address space value that represents global memory.
     static constexpr unsigned kGlobalMemoryAddressSpace = 1;
     /// The address space value that represents shared memory.
@@ -58,6 +48,22 @@ class ROCDL_Attr<string attrName, string attrMnemonic, list<Trait> traits = []>
   let mnemonic = attrMnemonic;
 }
 
+//===----------------------------------------------------------------------===//
+// ROCDL named attr definitions
+//===----------------------------------------------------------------------===//
+
+class ROCDL_NamedAttr<string name, string userName, string baseAttrType = "::mlir::Attribute"> :
+  NamedAttrDef<ROCDL_Dialect, name, userName, baseAttrType>;
+
+def ROCDL_KernelAttr : ROCDL_NamedAttr<"Kernel", "kernel", "::mlir::UnitAttr">;
+def ROCDL_ReqdWorkGroupSizeAttr :
+    ROCDL_NamedAttr<"ReqdWorkGroupSize", "reqd_work_group_size", "::mlir::DenseI32ArrayAttr">;
+def ROCDL_FlatWorkGroupSizeAttr :
+    ROCDL_NamedAttr<"FlatWorkGroupSize", "flat_work_group_size", "::mlir::StringAttr">;
+def ROCDL_MaxFlatWorkGroupSizeAttr :
+    ROCDL_NamedAttr<"MaxFlatWorkGroupSize", "max_flat_work_group_size", "::mlir::IntegerAttr">;
+def ROCDL_WavesPerEuAttr :
+    ROCDL_NamedAttr<"WavesPerEu", "waves_per_eu", "::mlir::IntegerAttr">;
 
 //===----------------------------------------------------------------------===//
 // ROCDL op definitions
diff --git a/mlir/include/mlir/IR/AttrTypeBase.td b/mlir/include/mlir/IR/AttrTypeBase.td
index 91c9283de8bd4..eeabbf7e06471 100644
--- a/mlir/include/mlir/IR/AttrTypeBase.td
+++ b/mlir/include/mlir/IR/AttrTypeBase.td
@@ -283,6 +283,78 @@ class AttrDef<Dialect dialect, string name, list<Trait> traits = [],
                                  "::" # cppClassName # ">($_self)">;
 }
 
+// Define a StringAttr wrapper for the NamedAttribute `name`
+// - `name` is dialect-scoped when not-inherent.
+// - Utilities to is/has/get/set/lookup/create typed Attr on an Operation
+//   including typed `value` attribute
+class NamedAttrDef<Dialect dialect, string name, string userName,
+    string valueAttrType = "::mlir::Attribute">
+    : AttrDef<dialect, name, [], "::mlir::StringAttr"> {
+  let mnemonic = userName;
+
+  string scopedName = dialect.name # "." # mnemonic;
+  code typedefValueAttr = "typedef " # valueAttrType # " ValueAttrType;\n";
+  code getNameFunc = "static constexpr llvm::StringLiteral getScopedName() { return \""
+      # scopedName # "\"; }\n";
+
+  code namedAttrFuncs = !strconcat(typedefValueAttr, getNameFunc, [{
+    // Get name based on inherentness
+    static llvm::StringLiteral getName(Operation *op = nullptr) {
+      if (op && op->getPropertiesStorageSize()) {
+       auto mnemonic = getMnemonic();
+       if (op->getInherentAttr(mnemonic))
+         return mnemonic;
+      }
+      return getScopedName();
+    }
+    // Is or Has
+    static bool is(::mlir::NamedAttribute attr) {
+      return attr.getName() == getScopedName();
+    }
+    static bool isInherent(::mlir::NamedAttribute attr) {
+      return attr.getName() == getMnemonic();
+    }
+    static bool has(::mlir::Operation *op) {
+      return op->hasAttr(getName(op));
+    }
+    // Get Name
+    static ::mlir::StringAttr get(::mlir::MLIRContext *ctx, ::mlir::Operation *op = nullptr) {
+      return ::mlir::StringAttr::get(ctx, getName(op));
+    }
+    // Get Value
+    static ValueAttrType getValue(::mlir::Operation *op) {
+      return op->getAttrOfType<ValueAttrType>(getName(op));
+    }
+    // Scoped lookup for inheritance
+    static ValueAttrType lookupValue(::mlir::Operation *op) {
+      if (auto attr = getValue(op))
+        return attr;
+      std::optional<RegisteredOperationName> opInfo = op->getRegisteredInfo();
+      if (!opInfo || !opInfo->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
+        if (auto *par = op->getParentOp())
+          return lookupValue(par);
+      }
+      return ValueAttrType();
+    }
+    // Set Value on Op
+    static void setValue(::mlir::Operation *op, ValueAttrType val) {
+      assert(op);
+      op->setAttr(getName(op), val);
+    }
+    // Remove Value from Op
+    static void removeValue(::mlir::Operation *op) {
+      assert(op);
+      op->removeAttr(getName(op));
+    }
+    // Create (scoped) NamedAttribute
+    static ::mlir::NamedAttribute create(::mlir::Builder &b, ValueAttrType val) {
+      return b.getNamedAttr(getScopedName(), val);
+    }
+  }]);
+
+  let extraClassDeclaration = namedAttrFuncs;
+}
+
 // Define a new type, named `name`, belonging to `dialect` that inherits from
 // the given C++ base class.
 class TypeDef<Dialect dialect, string name, list<Trait> traits = [],
diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index 599bb13190f12..81342e0679a7b 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -291,8 +291,7 @@ struct LowerGpuOpsToROCDLOpsPass
     m.walk([ctx](LLVM::LLVMFuncOp op) {
       if (auto blockSizes = dyn_cast_or_null<DenseI32ArrayAttr>(
               op->removeAttr(gpu::GPUFuncOp::getKnownBlockSizeAttrName()))) {
-        op->setAttr(ROCDL::ROCDLDialect::getReqdWorkGroupSizeAttrName(),
-                    blockSizes);
+        ROCDL::ReqdWorkGroupSizeAttr::setValue(op, blockSizes);
         // Also set up the rocdl.flat_work_group_size attribute to prevent
         // conflicting metadata.
         uint32_t flatSize = 1;
@@ -301,8 +300,7 @@ struct LowerGpuOpsToROCDLOpsPass
         }
         StringAttr flatSizeAttr =
             StringAttr::get(ctx, Twine(flatSize) + "," + Twine(flatSize));
-        op->setAttr(ROCDL::ROCDLDialect::getFlatWorkGroupSizeAttrName(),
-                    flatSizeAttr);
+        ROCDL::FlatWorkGroupSizeAttr::setValue(op, flatSizeAttr);
       }
     });
   }
@@ -355,8 +353,7 @@ void mlir::populateGpuToROCDLConversionPatterns(
       converter,
       /*allocaAddrSpace=*/ROCDL::ROCDLDialect::kPrivateMemoryAddressSpace,
       /*workgroupAddrSpace=*/ROCDL::ROCDLDialect::kSharedMemoryAddressSpace,
-      StringAttr::get(&converter.getContext(),
-                      ROCDL::ROCDLDialect::getKernelFuncAttrName()));
+      ROCDL::KernelAttr::get(&converter.getContext()));
   if (Runtime::HIP == runtime) {
     patterns.add<GPUPrintfOpToHIPLowering>(converter);
   } else if (Runtime::OpenCL == runtime) {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
index 26e46b31ddc01..078d026ac5222 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
@@ -253,9 +253,9 @@ void ROCDLDialect::initialize() {
 LogicalResult ROCDLDialect::verifyOperationAttribute(Operation *op,
                                                      NamedAttribute attr) {
   // Kernel function attribute should be attached to functions.
-  if (attr.getName() == ROCDLDialect::getKernelFuncAttrName()) {
+  if (ROCDL::KernelAttr::is(attr)) {
     if (!isa<LLVM::LLVMFuncOp>(op)) {
-      return op->emitError() << "'" << ROCDLDialect::getKernelFuncAttrName()
+      return op->emitError() << "'" << ROCDL::KernelAttr::getName()
                              << "' attribute attached to unexpected op";
     }
   }
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
index 5ab70280f6c81..0942bd2b5f3b3 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
@@ -83,7 +83,7 @@ class ROCDLDialectLLVMIRTranslationInterface
   LogicalResult
   amendOperation(Operation *op, NamedAttribute attribute,
                  LLVM::ModuleTranslation &moduleTranslation) const final {
-    if (attribute.getName() == ROCDL::ROCDLDialect::getKernelFuncAttrName()) {
+    if (ROCDL::KernelAttr::is(attribute)) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();
@@ -105,7 +105,7 @@ class ROCDLDialectLLVMIRTranslationInterface
     // Override flat-work-group-size
     // TODO: update clients to rocdl.flat_work_group_size instead,
     // then remove this half of the branch
-    if ("rocdl.max_flat_work_group_size" == attribute.getName()) {
+    if (ROCDL::MaxFlatWorkGroupSizeAttr::is(attribute)) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();
@@ -120,8 +120,7 @@ class ROCDLDialectLLVMIRTranslationInterface
       attrValueStream << "1," << value.getInt();
       llvmFunc->addFnAttr("amdgpu-flat-work-group-size", llvmAttrValue);
     }
-    if (ROCDL::ROCDLDialect::getFlatWorkGroupSizeAttrName() ==
-        attribute.getName()) {
+    if (ROCDL::FlatWorkGroupSizeAttr::is(attribute)) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();
@@ -137,8 +136,7 @@ class ROCDLDialectLLVMIRTranslationInterface
     }
 
     // Set reqd_work_group_size metadata
-    if (ROCDL::ROCDLDialect::getReqdWorkGroupSizeAttrName() ==
-        attribute.getName()) {
+    if (ROCDL::ReqdWorkGroupSizeAttr::is(attribute)) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();

@llvmbot
Copy link
Collaborator

llvmbot commented Dec 11, 2023

@llvm/pr-subscribers-mlir-llvm

Author: SJW (sjw36)

Changes

All attributes in MLIR are named, inherent attributes have unscoped names and arbitrary attributes should be scoped with a dialect. Current usage is ad-hoc and much of the codebase is sprinkled with constant strings used to lookup and set attributes, leading to potential bugs when names are not updated in all usages.

This PR adds a tablegen'd utility wrapper for a NamedAttribute that manages scoped/unscoped name lookup for accessing the attribute on an Operation based on inherentness.


Full diff: https://github.com/llvm/llvm-project/pull/75118.diff

5 Files Affected:

  • (modified) mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td (+16-10)
  • (modified) mlir/include/mlir/IR/AttrTypeBase.td (+72)
  • (modified) mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp (+3-6)
  • (modified) mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp (+2-2)
  • (modified) mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp (+4-6)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 48b830ae34f29..a40599e91e4b5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -28,16 +28,6 @@ def ROCDL_Dialect : Dialect {
   let hasOperationAttrVerify = 1;
 
   let extraClassDeclaration = [{
-    /// Get the name of the attribute used to annotate external kernel
-    /// functions.
-    static StringRef getKernelFuncAttrName() { return "rocdl.kernel"; }
-    static constexpr ::llvm::StringLiteral getFlatWorkGroupSizeAttrName() {
-      return ::llvm::StringLiteral("rocdl.flat_work_group_size");
-    }
-    static constexpr ::llvm::StringLiteral getReqdWorkGroupSizeAttrName() {
-      return ::llvm::StringLiteral("rocdl.reqd_work_group_size");
-    }
-
     /// The address space value that represents global memory.
     static constexpr unsigned kGlobalMemoryAddressSpace = 1;
     /// The address space value that represents shared memory.
@@ -58,6 +48,22 @@ class ROCDL_Attr<string attrName, string attrMnemonic, list<Trait> traits = []>
   let mnemonic = attrMnemonic;
 }
 
+//===----------------------------------------------------------------------===//
+// ROCDL named attr definitions
+//===----------------------------------------------------------------------===//
+
+class ROCDL_NamedAttr<string name, string userName, string baseAttrType = "::mlir::Attribute"> :
+  NamedAttrDef<ROCDL_Dialect, name, userName, baseAttrType>;
+
+def ROCDL_KernelAttr : ROCDL_NamedAttr<"Kernel", "kernel", "::mlir::UnitAttr">;
+def ROCDL_ReqdWorkGroupSizeAttr :
+    ROCDL_NamedAttr<"ReqdWorkGroupSize", "reqd_work_group_size", "::mlir::DenseI32ArrayAttr">;
+def ROCDL_FlatWorkGroupSizeAttr :
+    ROCDL_NamedAttr<"FlatWorkGroupSize", "flat_work_group_size", "::mlir::StringAttr">;
+def ROCDL_MaxFlatWorkGroupSizeAttr :
+    ROCDL_NamedAttr<"MaxFlatWorkGroupSize", "max_flat_work_group_size", "::mlir::IntegerAttr">;
+def ROCDL_WavesPerEuAttr :
+    ROCDL_NamedAttr<"WavesPerEu", "waves_per_eu", "::mlir::IntegerAttr">;
 
 //===----------------------------------------------------------------------===//
 // ROCDL op definitions
diff --git a/mlir/include/mlir/IR/AttrTypeBase.td b/mlir/include/mlir/IR/AttrTypeBase.td
index 91c9283de8bd4..eeabbf7e06471 100644
--- a/mlir/include/mlir/IR/AttrTypeBase.td
+++ b/mlir/include/mlir/IR/AttrTypeBase.td
@@ -283,6 +283,78 @@ class AttrDef<Dialect dialect, string name, list<Trait> traits = [],
                                  "::" # cppClassName # ">($_self)">;
 }
 
+// Define a StringAttr wrapper for the NamedAttribute `name`
+// - `name` is dialect-scoped when not-inherent.
+// - Utilities to is/has/get/set/lookup/create typed Attr on an Operation
+//   including typed `value` attribute
+class NamedAttrDef<Dialect dialect, string name, string userName,
+    string valueAttrType = "::mlir::Attribute">
+    : AttrDef<dialect, name, [], "::mlir::StringAttr"> {
+  let mnemonic = userName;
+
+  string scopedName = dialect.name # "." # mnemonic;
+  code typedefValueAttr = "typedef " # valueAttrType # " ValueAttrType;\n";
+  code getNameFunc = "static constexpr llvm::StringLiteral getScopedName() { return \""
+      # scopedName # "\"; }\n";
+
+  code namedAttrFuncs = !strconcat(typedefValueAttr, getNameFunc, [{
+    // Get name based on inherentness
+    static llvm::StringLiteral getName(Operation *op = nullptr) {
+      if (op && op->getPropertiesStorageSize()) {
+       auto mnemonic = getMnemonic();
+       if (op->getInherentAttr(mnemonic))
+         return mnemonic;
+      }
+      return getScopedName();
+    }
+    // Is or Has
+    static bool is(::mlir::NamedAttribute attr) {
+      return attr.getName() == getScopedName();
+    }
+    static bool isInherent(::mlir::NamedAttribute attr) {
+      return attr.getName() == getMnemonic();
+    }
+    static bool has(::mlir::Operation *op) {
+      return op->hasAttr(getName(op));
+    }
+    // Get Name
+    static ::mlir::StringAttr get(::mlir::MLIRContext *ctx, ::mlir::Operation *op = nullptr) {
+      return ::mlir::StringAttr::get(ctx, getName(op));
+    }
+    // Get Value
+    static ValueAttrType getValue(::mlir::Operation *op) {
+      return op->getAttrOfType<ValueAttrType>(getName(op));
+    }
+    // Scoped lookup for inheritance
+    static ValueAttrType lookupValue(::mlir::Operation *op) {
+      if (auto attr = getValue(op))
+        return attr;
+      std::optional<RegisteredOperationName> opInfo = op->getRegisteredInfo();
+      if (!opInfo || !opInfo->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
+        if (auto *par = op->getParentOp())
+          return lookupValue(par);
+      }
+      return ValueAttrType();
+    }
+    // Set Value on Op
+    static void setValue(::mlir::Operation *op, ValueAttrType val) {
+      assert(op);
+      op->setAttr(getName(op), val);
+    }
+    // Remove Value from Op
+    static void removeValue(::mlir::Operation *op) {
+      assert(op);
+      op->removeAttr(getName(op));
+    }
+    // Create (scoped) NamedAttribute
+    static ::mlir::NamedAttribute create(::mlir::Builder &b, ValueAttrType val) {
+      return b.getNamedAttr(getScopedName(), val);
+    }
+  }]);
+
+  let extraClassDeclaration = namedAttrFuncs;
+}
+
 // Define a new type, named `name`, belonging to `dialect` that inherits from
 // the given C++ base class.
 class TypeDef<Dialect dialect, string name, list<Trait> traits = [],
diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index 599bb13190f12..81342e0679a7b 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -291,8 +291,7 @@ struct LowerGpuOpsToROCDLOpsPass
     m.walk([ctx](LLVM::LLVMFuncOp op) {
       if (auto blockSizes = dyn_cast_or_null<DenseI32ArrayAttr>(
               op->removeAttr(gpu::GPUFuncOp::getKnownBlockSizeAttrName()))) {
-        op->setAttr(ROCDL::ROCDLDialect::getReqdWorkGroupSizeAttrName(),
-                    blockSizes);
+        ROCDL::ReqdWorkGroupSizeAttr::setValue(op, blockSizes);
         // Also set up the rocdl.flat_work_group_size attribute to prevent
         // conflicting metadata.
         uint32_t flatSize = 1;
@@ -301,8 +300,7 @@ struct LowerGpuOpsToROCDLOpsPass
         }
         StringAttr flatSizeAttr =
             StringAttr::get(ctx, Twine(flatSize) + "," + Twine(flatSize));
-        op->setAttr(ROCDL::ROCDLDialect::getFlatWorkGroupSizeAttrName(),
-                    flatSizeAttr);
+        ROCDL::FlatWorkGroupSizeAttr::setValue(op, flatSizeAttr);
       }
     });
   }
@@ -355,8 +353,7 @@ void mlir::populateGpuToROCDLConversionPatterns(
       converter,
       /*allocaAddrSpace=*/ROCDL::ROCDLDialect::kPrivateMemoryAddressSpace,
       /*workgroupAddrSpace=*/ROCDL::ROCDLDialect::kSharedMemoryAddressSpace,
-      StringAttr::get(&converter.getContext(),
-                      ROCDL::ROCDLDialect::getKernelFuncAttrName()));
+      ROCDL::KernelAttr::get(&converter.getContext()));
   if (Runtime::HIP == runtime) {
     patterns.add<GPUPrintfOpToHIPLowering>(converter);
   } else if (Runtime::OpenCL == runtime) {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
index 26e46b31ddc01..078d026ac5222 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
@@ -253,9 +253,9 @@ void ROCDLDialect::initialize() {
 LogicalResult ROCDLDialect::verifyOperationAttribute(Operation *op,
                                                      NamedAttribute attr) {
   // Kernel function attribute should be attached to functions.
-  if (attr.getName() == ROCDLDialect::getKernelFuncAttrName()) {
+  if (ROCDL::KernelAttr::is(attr)) {
     if (!isa<LLVM::LLVMFuncOp>(op)) {
-      return op->emitError() << "'" << ROCDLDialect::getKernelFuncAttrName()
+      return op->emitError() << "'" << ROCDL::KernelAttr::getName()
                              << "' attribute attached to unexpected op";
     }
   }
diff --git a/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
index 5ab70280f6c81..0942bd2b5f3b3 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
@@ -83,7 +83,7 @@ class ROCDLDialectLLVMIRTranslationInterface
   LogicalResult
   amendOperation(Operation *op, NamedAttribute attribute,
                  LLVM::ModuleTranslation &moduleTranslation) const final {
-    if (attribute.getName() == ROCDL::ROCDLDialect::getKernelFuncAttrName()) {
+    if (ROCDL::KernelAttr::is(attribute)) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();
@@ -105,7 +105,7 @@ class ROCDLDialectLLVMIRTranslationInterface
     // Override flat-work-group-size
     // TODO: update clients to rocdl.flat_work_group_size instead,
     // then remove this half of the branch
-    if ("rocdl.max_flat_work_group_size" == attribute.getName()) {
+    if (ROCDL::MaxFlatWorkGroupSizeAttr::is(attribute)) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();
@@ -120,8 +120,7 @@ class ROCDLDialectLLVMIRTranslationInterface
       attrValueStream << "1," << value.getInt();
       llvmFunc->addFnAttr("amdgpu-flat-work-group-size", llvmAttrValue);
     }
-    if (ROCDL::ROCDLDialect::getFlatWorkGroupSizeAttrName() ==
-        attribute.getName()) {
+    if (ROCDL::FlatWorkGroupSizeAttr::is(attribute)) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();
@@ -137,8 +136,7 @@ class ROCDLDialectLLVMIRTranslationInterface
     }
 
     // Set reqd_work_group_size metadata
-    if (ROCDL::ROCDLDialect::getReqdWorkGroupSizeAttrName() ==
-        attribute.getName()) {
+    if (ROCDL::ReqdWorkGroupSizeAttr::is(attribute)) {
       auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
       if (!func)
         return failure();

Copy link
Contributor

@Dinistro Dinistro left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems a helpful extension, especially for attributes that are added to LLVM dialect operations. I would suggest to add a dedicated test for this within the test dialect, to ensure that the correctness of this is checked explicitly.

@sjw36
Copy link
Contributor Author

sjw36 commented Dec 12, 2023

This seems a helpful extension, especially for attributes that are added to LLVM dialect operations. I would suggest to add a dedicated test for this within the test dialect, to ensure that the correctness of this is checked explicitly.

Thanks @Dinistro , will do.

@sjw36
Copy link
Contributor Author

sjw36 commented Dec 12, 2023

@joker-eph does this approach work going forward with properties? I see comments on getInherentAttr about it's transient nature.

Copy link
Contributor

@krzysz00 krzysz00 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor code-grammar nitpicks, but this is a solid utility.

static bool isInherent(::mlir::NamedAttribute attr) {
return attr.getName() == getMnemonic();
}
static bool has(::mlir::Operation *op) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The grammar of this method name feels weird. Would it make more sense as something like .isPresent or .isSetOn or some such? Asking if an attribute has an operation feels backwards

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, to follow the rest of the interface it could be hasValue.

return getScopedName();
}
// Is or Has
static bool is(::mlir::NamedAttribute attr) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wouldn't it make some sense to generate is[AttrName] out into the wider namespace? But that'd require more tablegen hackery.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we have is instead of == ? Does this mimic any existing convention on MLIR objects?

return getScopedName();
}
// Is or Has
static bool is(::mlir::NamedAttribute attr) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we have is instead of == ? Does this mimic any existing convention on MLIR objects?

@@ -283,6 +283,78 @@ class AttrDef<Dialect dialect, string name, list<Trait> traits = [],
"::" # cppClassName # ">($_self)">;
}

// Define a StringAttr wrapper for the NamedAttribute `name`
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The "NamedAttribute" part isn't clear to me.

Also this is more than a "wrapper", it seems to me this is actually registering its own full fledge attribute?
I'm not sure I follow the definition, but I suspect this won't be uniqued with StringAttr.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, this is no different than a StringAttr, so the usual interface works, but the new get only takes MLIRContext*. And once it is created, there is no way to get back to the wrapper class but it is mostly static lookup methods anyway.


code namedAttrFuncs = !strconcat(typedefValueAttr, getNameFunc, [{
// Get name based on inherentness
static llvm::StringLiteral getName(Operation *op = nullptr) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand the intended semantics of all this right now.

op->setAttr(getName(op), val);
}
// Remove Value from Op
static void removeValue(::mlir::Operation *op) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The terminology is confusing: why are you using "Value" for referring to "Attributes"?

@joker-eph
Copy link
Collaborator

joker-eph commented Dec 12, 2023

This is an interesting problem to solve, and there are a few aspects to it:

  1. In terms of performance, we'd want to use StringAttr as much as possible to avoid string hashing and string comparisons everywhere. But I don't think you're achieving this right now?
  2. In terms of correctness, there is the question of inherent vs discardable attributes. However this is an individual property of each operation, and the way you try to "infer" this does not seem right to me: if I have a dialect NamedAttribute that I intend to be mydialect.foo, I shouldn't lookup the inherent attribute of an operation to see if foo is a valid key before preprinting the mydialect. prefix: they are two different things!

@sjw36
Copy link
Contributor Author

sjw36 commented Dec 12, 2023

Thanks for the feedback @joker-eph .

In short, it is a wrapper that is does 2 things:

  1. Named lookup on objects depending on inherentness and Value Type
  2. Get/Set the Typed Value for consistency
  1. In terms of performance, we'd want to use StringAttr as much as possible to avoid string hashing and string comparisons everywhere. But I don't think you're achieving this right now?

No, that is true but most usages are StringLiteral except when copying from a NamedAttribute. It is complicated since it is context dependent. Do you have any suggestion?

  1. In terms of correctness, there is the question of inherent vs discardable attributes. However this is an individual property of each operation, and the way you try to "infer" this does not seem right to me: if I have a dialect NamedAttribute that I intend to be mydialect.foo, I shouldn't lookup the inherent attribute of an operation to see if foo is a valid key before preprinting the mydialect. prefix: they are two different things!

Perhaps that is a bridge too far. We have scenarios where an inherent attribute on an op is copied to the enclosing func or module and should then be dialect-qualified. I could live with a DialectNamedAttrDef and an NamedAttrDef (unqualified).

@joker-eph
Copy link
Collaborator

No, that is true but most usages are StringLiteral except when copying from a NamedAttribute. It is complicated since it is context dependent. Do you have any suggestion?

What I have been doing in the past is defining these as StringAttr on the dialect instance instead.

Perhaps that is a bridge too far. We have scenarios where an inherent attribute on an op is copied to the enclosing func or module and should then be dialect-qualified. I could live with a DialectNamedAttrDef and an NamedAttrDef (unqualified).

The whole concept of using an "AttrDef" here is still unclear to me right now to begin with.
I'm also not sure I'm making a direct connection between your usage of inherent attribute and a propagation as a discardable one.

In general with properties we should be relying more and more on Interfaces to access informations about an operation instead of relying on "name lookup": in this model the name of the inherent attribute shouldn't matter, and shouldn't even be exposed (it is an implementation detail of the storage).

@sjw36
Copy link
Contributor Author

sjw36 commented Dec 13, 2023

No, that is true but most usages are StringLiteral except when copying from a NamedAttribute. It is complicated since it is context dependent. Do you have any suggestion?

What I have been doing in the past is defining these as StringAttr on the dialect instance instead.

Can you point me to an example. I see getConsumedAttrName() in the Transform dialect, but it is just using a StringLiteral under the hood and the usage seems cumbersome having to get loaded dialect first. Is there a better example that is storing these on the dialect?

Perhaps that is a bridge too far. We have scenarios where an inherent attribute on an op is copied to the enclosing func or module and should then be dialect-qualified. I could live with a DialectNamedAttrDef and an NamedAttrDef (unqualified).

The whole concept of using an "AttrDef" here is still unclear to me right now to begin with. I'm also not sure I'm making a direct connection between your usage of inherent attribute and a propagation as a discardable one.

If we set aside the inherent/property (see below) case for now and just consider a DialectNamedAttr then:

  • the ODS class can be used throughout the code instead of "arbitrary_name" strings
  • and the ODS class can be used to Get/Set the Typed Value on arbitrary Operations

In this way it formalizes a specific NamedAttribute in a Dialect with a specific Typed Value. See ROCDL example usage.

In general with properties we should be relying more and more on Interfaces to access informations about an operation instead of relying on "name lookup": in this model the name of the inherent attribute shouldn't matter, and shouldn't even be exposed (it is an implementation detail of the storage).

This makes sense and certainly more optimal. I had a few use cases for promoting inherent attributes to parent ops but we are migrating them to properties in any case.

@joker-eph
Copy link
Collaborator

Can you point me to an example. I see getConsumedAttrName() in the Transform dialect, but it is just using a StringLiteral under the hood and the usage seems cumbersome having to get loaded dialect first. Is there a better example that is storing these on the dialect?

I don't have anything at hand: I've written this myself multiple times though. Yes you need to have a dialect pointer to get these then by nature, which is an extra step but I can't really see how to avoid though.

the ODS class can be used throughout the code instead of "arbitrary_name" strings

How is this class any different than an "arbitrary name" though? You built a string wrapper: isn't it just like a handle to a string?

and the ODS class can be used to Get/Set the Typed Value on arbitrary Operations

I don't quite get this part? Is this just making the setDiscardableAttr() API typed?
Is the whole thing here trying to make it so that the string attribute name key can only be set with a specific attribute type?

I would need to build the branch to looks a bit more, but from my read you're actually defining more than a "convenient type-safe API" and instead it is a new storage for each of these "AttrDef", even though they will be initialized with a single entry ever?

Copy link

github-actions bot commented Dec 14, 2023

:white_check_mark: With the latest revision this PR passed the C/C++ code formatter.

@sjw36
Copy link
Contributor Author

sjw36 commented Dec 14, 2023

I don't quite get this part? Is this just making the setDiscardableAttr() API typed? Is the whole thing here trying to make it so that the string attribute name key can only be set with a specific attribute type?

Yes, this is a wrapper around a specific Key for the NamedAttribute and static typed accessors to the Value on an Operation.
Unfortunately I don't see a way to enforce that a particular NameAttribute must have a specified Value Type.

I cleaned it up a bit to remove the inherent check and so it is now only dialect-qualified keys for discardable attributes.
Also added unit test.

@sjw36 sjw36 force-pushed the named-attr-ods branch 2 times, most recently from e49df21 to ecacfe2 Compare December 14, 2023 20:38
All attributes in MLIR are named, inherent attributes have unscoped names
and discardable attributes should be scoped with a dialect. Current usage is
ad-hoc and much of the codebase is sprinkled with constant strings used to
lookup and set attributes, leading to potential bugs when names are not
updated in all usages.

This PR adds a tablegen'd utility wrapper for a NamedAttribute that manages
scoped/unscoped name lookup for consistent typed access the attribute on an
Operation.
@sjw36
Copy link
Contributor Author

sjw36 commented Dec 14, 2023

Removed an unused file.

@joker-eph
Copy link
Collaborator

I checked locally and this matches my expectation: it defines a full-fledge new attribute.

I guess I still don't quite get why we would do that, when it seems like some very standard C++ can achieve the same thing.
(I'll try to draft something later)

@sjw36
Copy link
Contributor Author

sjw36 commented Dec 15, 2023

I checked locally and this matches my expectation: it defines a full-fledge new attribute.

I guess I still don't quite get why we would do that, when it seems like some very standard C++ can achieve the same thing. (I'll try to draft something later)

Sure, it is simple code, could easily be a template type (however, embedding the name is challenging).
I like being able to define a specific Key and typed Value attribute that is specific to my dialect in the tablegen file with a single line.

def ROCDL_ReqdWorkGroupSizeAttr :
    ROCDL_NamedAttr<"ReqdWorkGroupSize", "reqd_work_group_size", "::mlir::DenseI32ArrayAttr">;

I did consider a new tablegen type instead of overloading AttrDef, but that seemed overly burdensome, and since the AttrDef could service the key aspects reasonably.

Certainly welcome a better solution. Thanks.

@krzysz00
Copy link
Contributor

I think there's an argument for defining a new record type here. To give it a provisional name, something like DialectScopedAttr, so that you can generate accessors and the like in mlir-tblgen.

My sense here is that there're two related usecases, which may or may not need the same abstraction:

  1. Dialect-specific attributes that need to live on other people's operations. For example, that there reqd_work_group_size, which lives on llvm.func. Or gpu.container_module, which lives on builtin.module. It'll be good to have those documented and to have getter/setter functions that check invariants (since we can't inject code into other people's verify() methods). This sort of thing could also allow you to specify properties of the operation you put it on - for example, it could be known at tablegen time that gpu.container_module is supposed to go on module-like operations only.
  2. Dialect-wide attribute grab bags that aren't repeated as properties on every operation. For example, stuff like llvm.align, llvm.invariant_load, and so on and so forth. In principle, we could go clean all that up and give every LLVM dialect operation a long list of properties matching exactly the kind of metadata or attributes you can put on it in LLVM IR ... or you can just have a cleaner interface to these attributes by routing them through the mechanisms for case (1) above

Making this its own tablegen record type would have the advantage of allowing you to generate documentation automatically as well, for instance, but it is a lot more work.

@joker-eph
Copy link
Collaborator

since we can't inject code into other people's verify() methods

Can you clarify what you have in mind exactly? (I don't know if you're familiar with https://mlir.llvm.org/docs/DefiningDialects/#discardable-attribute-verification ; I would assume so since you give the example of gpu.container_modulewhich is exactly doing that).

Dialect-wide attribute grab bags that aren't repeated as properties on every operation. For example, stuff like llvm.align, llvm.invariant_load, and so on and so forth. In principle, we could go clean all that up and give every LLVM dialect operation a long list of properties matching exactly the kind of metadata or attributes you can put on it in LLVM IR ...

There is an semantic difference between "discardable" and "inherent" attributes: it's not clear to me if you're mixing these up together like if an attribute can be either of these as only a "cosmetic" thing.

Making this its own tablegen record type would have the advantage of allowing you to generate documentation automatically as well, for instance, but it is a lot more work.

"TableGen record" is not something well defined without connecting this to a backend.
I'm mostly pointing out that using AttrDef just does not seem right to me here: I haven't perceived a requirement that would justify doing this (other than it's there and you can tweak it to get something that kind of look like providing some of the feature you want).

@krzysz00
Copy link
Contributor

I wasn't aware of that dialect attribute verification hook - I think that's relatively new.

What I was trying to say about the LLVM dialect is that there're a lot of llvm.* discardable attributes these days, and it'd be good to either have them defined in a less ad-hoc way or to make them inherent attributes of relevant operations.

I'm also not sure AttrDef is the right mechanism here, I was pointing out motivations for describing "well-known" discardable attributes in tablegen.

@joker-eph
Copy link
Collaborator

@joker-eph
Copy link
Collaborator

I'll try to draft something later

Couldn't get to it before the break, but here is a draft: #77024

@sjw36
Copy link
Contributor Author

sjw36 commented Jan 9, 2024

Closing in favor of #77024.

@sjw36 sjw36 closed this Jan 9, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants