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

Rajat Bajpai llvmlistbot at llvm.org
Mon Feb 23 23:50:53 PST 2026


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

>From fa87158b043d306eb846f0fbd28920242d590ed5 Mon Sep 17 00:00:00 2001
From: rbajpai <rbajpai at nvidia.com>
Date: Sat, 21 Feb 2026 12:08:01 +0530
Subject: [PATCH] [MLIR][ODS][NVVM] Add EnumAttrIsOneOf/IsNoneOf constraints
 and use them in NVVM fence ops

Add `EnumAttrIsOneOf` and `EnumAttrIsNoneOf` to `EnumAttr.td`, a generic
`AttrConstraint`s that restrict an `EnumAttr` to an allowlist or denylist
of cases respectively. These are the enum analogs of the existing
`IntIsOneOf` constraint for integer attributes.

`EnumAttrIsOneOf` restricts the attribute to a closed set of permitted
cases. `EnumAttrIsNoneOf` excludes specific cases while allowing all
others, which is more ergonomic when the disallowed set is smaller than
the allowed set and more future-proof when new enum cases are added.

This replaces hand-written verifier logic with ODS-generated
verification, eliminating the `FenceSyncRestrictOp` verifier entirely
and simplifying both `FenceProxyOp` and `FenceProxySyncRestrictOp`
verifiers.

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]>;
```
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td   | 17 +++++++---
 mlir/include/mlir/IR/EnumAttr.td              | 34 +++++++++++++++++++
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp    | 15 --------
 .../Target/LLVMIR/nvvm/fence-invalid.mlir     | 10 +++---
 4 files changed, 51 insertions(+), 25 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 9725cdb6a4c4d..6ff3ef45110a5 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"> {
@@ -1360,8 +1363,12 @@ def ProxyKindAttr : EnumAttr<NVVM_Dialect, ProxyKind, "proxy_kind"> {
   let assemblyFormat = "`<` $value `>`";
 }
 
+defvar ProxyKindNotTensormapOrGeneric =
+  ConfinedAttr<ProxyKindAttr,
+    [EnumAttrIsNoneOf<ProxyKindAttr, [ProxyTensorMap, ProxyGeneric]>]>;
+
 def NVVM_FenceProxyOp : NVVM_Op<"fence.proxy">,
-  Arguments<(ins ProxyKindAttr:$kind,
+  Arguments<(ins ProxyKindNotTensormapOrGeneric:$kind,
                  OptionalAttr<SharedSpaceAttr>:$space)> {
   let description = [{
     Fence operation with proxy to establish an ordering between memory accesses
@@ -1439,7 +1446,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..4c5438b52229e 100644
--- a/mlir/include/mlir/IR/EnumAttr.td
+++ b/mlir/include/mlir/IR/EnumAttr.td
@@ -750,4 +750,38 @@ 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), ", ") # "}">;
+
+/// Attribute constraint restricting the enum attribute to exclude a subset of
+/// cases. `enumAttr` is the `EnumAttr` attribute and `disallowedCases` is a
+/// list of `EnumCase`s that are not permitted.
+///
+/// Example:
+/// ```
+/// ConfinedAttr<MyEnumAttr, [EnumAttrIsNoneOf<MyEnumAttr, [CaseA, CaseB]>]>
+/// ```
+class EnumAttrIsNoneOf<EnumAttr enumAttr, list<EnumCase> disallowedCases>
+  : AttrConstraint<
+      And<!foreach(case, disallowedCases,
+        Neg<CPred<"::llvm::cast<" # enumAttr.cppNamespace # "::"
+              # enumAttr.cppClassName # ">($_self).getValue() == "
+              # enumAttr.enum.cppType # "::" # case.symbol>>)>,
+      "whose value is none of {"
+        # !interleave(!foreach(case, disallowedCases, 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..1a9169dc8464b 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -2731,18 +2731,7 @@ 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";
-  if (getKind() == NVVM::ProxyKind::GENERIC)
-    return emitOpError() << "generic proxy not a supported proxy kind";
   if (getKind() == NVVM::ProxyKind::async_shared && !getSpace().has_value()) {
     return emitOpError() << "async_shared fence requires space attribute";
   }
@@ -2775,10 +2764,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..d2e14637f094f 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
 }
@@ -17,7 +17,7 @@ llvm.func @fence_sync_restrict() {
 // -----
 
 llvm.func @fence_proxy() {
-  // expected-error @below {{tensormap proxy is not a supported proxy kind}}
+  // expected-error @below {{attribute 'kind' failed to satisfy constraint: Proxy kind whose value is none of {tensormap, generic}}}
   nvvm.fence.proxy {kind = #nvvm.proxy_kind<tensormap>}
   llvm.return
 }
@@ -25,7 +25,7 @@ llvm.func @fence_proxy() {
 // -----
 
 llvm.func @fence_proxy() {
-  // expected-error @below {{generic proxy not a supported proxy kind}}
+  // expected-error @below {{attribute 'kind' failed to satisfy constraint: Proxy kind whose value is none of {tensormap, generic}}}
   nvvm.fence.proxy {kind = #nvvm.proxy_kind<generic>}
   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
 }



More information about the Mlir-commits mailing list