[Mlir-commits] [mlir] [MLIR][ODS][NVVM] Add EnumAttrIsOneOf constraint and use it in NVVM fence ops (PR #182662)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Rajat Bajpai (rajatbajpai)

<details>
<summary>Changes</summary>

Add `EnumAttrIsOneOf` to `EnumAttr.td`, a generic `AttrConstraint` that restricts an `EnumAttr` to a subset of allowed cases. This is the enum analog of the existing `IntIsOneOf` constraint for integer attributes.

Use it in the NVVM dialect to declaratively constrain the `MemOrderKind` attribute to `{acquire, release}` for `FenceSyncRestrictOp` and `FenceProxySyncRestrictOp`, replacing hand-written verifier logic with ODS-generated verification.

This eliminates the `FenceSyncRestrictOp` verifier method entirely and simplifies `FenceProxySyncRestrictOp` verifier by removing its `MemOrderKind` check.

It can also be parameterized through base class hierarchies, allowing derived ops to specify their own allowed subsets:

```
class MyBase<string mnemonic, list<EnumCase> allowed>
    : NVVM_Op<mnemonic>,
      Arguments<(ins
        ConfinedAttr<MyEnumAttr,
          [EnumAttrIsOneOf<MyEnumAttr, allowed>]>:$kind)>;

def Op1 : MyBase<"op1", [CaseA, CaseB]>;
def Op2 : MyBase<"op2", [CaseB, CaseC, CaseD]>;
```

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


4 Files Affected:

- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+7-4) 
- (modified) mlir/include/mlir/IR/EnumAttr.td (+17) 
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (-11) 
- (modified) mlir/test/Target/LLVMIR/nvvm/fence-invalid.mlir (+3-3) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 9725cdb6a4c4d..ddbf7c22d8c62 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1307,8 +1307,13 @@ def NVVM_FenceScClusterOp : NVVM_Op<"fence.sc.cluster"> {
   let assemblyFormat = "attr-dict";
 }
 
+defvar MemOrderAcquireOrRelease =
+  ConfinedAttr<MemOrderKindAttr,
+    [EnumAttrIsOneOf<MemOrderKindAttr,
+      [MemOrderKindAcquire, MemOrderKindRelease]>]>;
+
 def NVVM_FenceSyncRestrictOp : NVVM_Op<"fence.sync_restrict">,
-                               Arguments<(ins MemOrderKindAttr:$order)> {
+                               Arguments<(ins MemOrderAcquireOrRelease:$order)> {
   let summary = "Uni-directional thread fence operation";
   let description = [{
     The `nvvm.fence.sync_restrict` Op restricts the class of memory
@@ -1322,8 +1327,6 @@ def NVVM_FenceSyncRestrictOp : NVVM_Op<"fence.sync_restrict">,
   let llvmBuilder = [{
     createIntrinsicCall(builder, getFenceSyncRestrictID($order));
   }];
-
-  let hasVerifier = 1;
 }
 
 def NVVM_FenceMbarrierInitOp : NVVM_Op<"fence.mbarrier.init"> {
@@ -1439,7 +1442,7 @@ def NVVM_FenceProxyReleaseOp : NVVM_Op<"fence.proxy.release">,
 }
 
 def NVVM_FenceProxySyncRestrictOp : NVVM_Op<"fence.proxy.sync_restrict">,
-      Arguments<(ins MemOrderKindAttr:$order,
+      Arguments<(ins MemOrderAcquireOrRelease:$order,
                      DefaultValuedAttr<ProxyKindAttr, "ProxyKind::GENERIC">:$fromProxy,
                      DefaultValuedAttr<ProxyKindAttr, "ProxyKind::async">:$toProxy)> {
   let summary = "Uni-directional proxy fence operation with sync_restrict";
diff --git a/mlir/include/mlir/IR/EnumAttr.td b/mlir/include/mlir/IR/EnumAttr.td
index ec57626ebde65..97241f8a12e90 100644
--- a/mlir/include/mlir/IR/EnumAttr.td
+++ b/mlir/include/mlir/IR/EnumAttr.td
@@ -750,4 +750,21 @@ class ConstantEnumCase<Attr attribute, string case>
     "attribute must be one of 'EnumAttr' or 'EnumInfo'";
 }
 
+/// Attribute constraint restricting the enum attribute to a subset of allowed
+/// cases. `enumAttr` is the `EnumAttr` attribute and `allowedCases` is a list
+/// of `EnumCase`s that are permitted.
+///
+/// Example:
+/// ```
+/// ConfinedAttr<MyEnumAttr, [EnumAttrIsOneOf<MyEnumAttr, [CaseA, CaseB]>]>
+/// ```
+class EnumAttrIsOneOf<EnumAttr enumAttr, list<EnumCase> allowedCases>
+  : AttrConstraint<
+      Or<!foreach(case, allowedCases,
+        CPred<"::llvm::cast<" # enumAttr.cppNamespace # "::"
+              # enumAttr.cppClassName # ">($_self).getValue() == "
+              # enumAttr.enum.cppType # "::" # case.symbol>)>,
+      "whose value is one of {"
+        # !interleave(!foreach(case, allowedCases, case.str), ", ") # "}">;
+
 #endif // ENUMATTR_TD
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index b51eb3408df00..dc0abae881465 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -2731,13 +2731,6 @@ bool NVVM::WgmmaMmaAsyncOp::getAsmValues(
   return true; // Has manual mapping
 }
 
-LogicalResult NVVM::FenceSyncRestrictOp::verify() {
-  if (getOrder() != NVVM::MemOrderKind::ACQUIRE &&
-      getOrder() != NVVM::MemOrderKind::RELEASE)
-    return emitOpError("only acquire and release semantics are supported");
-  return success();
-}
-
 LogicalResult NVVM::FenceProxyOp::verify() {
   if (getKind() == NVVM::ProxyKind::TENSORMAP)
     return emitOpError() << "tensormap proxy is not a supported proxy kind";
@@ -2775,10 +2768,6 @@ LogicalResult NVVM::FenceProxyReleaseOp::verify() {
 }
 
 LogicalResult NVVM::FenceProxySyncRestrictOp::verify() {
-  if (getOrder() != NVVM::MemOrderKind::ACQUIRE &&
-      getOrder() != NVVM::MemOrderKind::RELEASE)
-    return emitOpError("only acquire and release semantics are supported");
-
   if (getFromProxy() != NVVM::ProxyKind::GENERIC)
     return emitOpError("only generic is support for from_proxy attribute");
 
diff --git a/mlir/test/Target/LLVMIR/nvvm/fence-invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/fence-invalid.mlir
index 22578b5581da4..606b3fcffb272 100644
--- a/mlir/test/Target/LLVMIR/nvvm/fence-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/fence-invalid.mlir
@@ -1,7 +1,7 @@
 // RUN: mlir-translate --mlir-to-llvmir -verify-diagnostics -split-input-file %s
 
 llvm.func @fence_sync_restrict() {
-  // expected-error @below {{only acquire and release semantics are supported}}
+  // expected-error @below {{attribute 'order' failed to satisfy constraint: NVVM Memory Ordering kind whose value is one of {acquire, release}}}
   nvvm.fence.sync_restrict {order = #nvvm.mem_order<weak>}
   llvm.return
 }
@@ -9,7 +9,7 @@ llvm.func @fence_sync_restrict() {
 // -----
 
 llvm.func @fence_sync_restrict() {
-  // expected-error @below {{only acquire and release semantics are supported}}
+  // expected-error @below {{attribute 'order' failed to satisfy constraint: NVVM Memory Ordering kind whose value is one of {acquire, release}}}
   nvvm.fence.sync_restrict {order = #nvvm.mem_order<mmio>}
   llvm.return
 }
@@ -65,7 +65,7 @@ llvm.func @fence_proxy_release() {
 // -----
 
 llvm.func @fence_proxy_sync_restrict() {
-  // expected-error @below {{only acquire and release semantics are supported}}
+  // expected-error @below {{attribute 'order' failed to satisfy constraint: NVVM Memory Ordering kind whose value is one of {acquire, release}}}
   nvvm.fence.proxy.sync_restrict {order = #nvvm.mem_order<mmio>}
   llvm.return
 }

``````````

</details>


https://github.com/llvm/llvm-project/pull/182662