[Mlir-commits] [mlir] [mlir][SPIR-V] Add OpExpectKHR op (PR #195842)

Arseniy Obolenskiy llvmlistbot at llvm.org
Tue May 5 06:42:20 PDT 2026


https://github.com/aobolensk updated https://github.com/llvm/llvm-project/pull/195842

>From 02624cc8ebb996e3c255fd82d52e7501337fc1c7 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Tue, 5 May 2026 14:18:52 +0200
Subject: [PATCH 1/4] [mlir][SPIR-V] Add OpExpectKHR op

Complete SPV_KHR_expect_assume extension support
---
 .../mlir/Dialect/SPIRV/IR/SPIRVBase.td        |  3 +-
 .../mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td     | 52 +++++++++++++++++++
 mlir/test/Dialect/SPIRV/IR/misc-ops.mlir      | 52 +++++++++++++++++++
 3 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index f870e6d920c8c..b33a84d093fb9 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4682,6 +4682,7 @@ def SPIRV_OC_OpSetMeshOutputsEXT              : I32EnumAttrCase<"OpSetMeshOutput
 def SPIRV_OC_OpSubgroupBlockReadINTEL         : I32EnumAttrCase<"OpSubgroupBlockReadINTEL", 5575>;
 def SPIRV_OC_OpSubgroupBlockWriteINTEL        : I32EnumAttrCase<"OpSubgroupBlockWriteINTEL", 5576>;
 def SPIRV_OC_OpAssumeTrueKHR                  : I32EnumAttrCase<"OpAssumeTrueKHR", 5630>;
+def SPIRV_OC_OpExpectKHR                      : I32EnumAttrCase<"OpExpectKHR", 5631>;
 def SPIRV_OC_OpAtomicFAddEXT                  : I32EnumAttrCase<"OpAtomicFAddEXT", 6035>;
 def SPIRV_OC_OpConvertFToBF16INTEL            : I32EnumAttrCase<"OpConvertFToBF16INTEL", 6116>;
 def SPIRV_OC_OpConvertBF16ToFINTEL            : I32EnumAttrCase<"OpConvertBF16ToFINTEL", 6117>;
@@ -4806,7 +4807,7 @@ def SPIRV_OpcodeAttr :
       SPIRV_OC_OpSpecConstantCompositeReplicateEXT,
       SPIRV_OC_OpEmitMeshTasksEXT, SPIRV_OC_OpSetMeshOutputsEXT,
       SPIRV_OC_OpSubgroupBlockReadINTEL, SPIRV_OC_OpSubgroupBlockWriteINTEL,
-      SPIRV_OC_OpAssumeTrueKHR, SPIRV_OC_OpAtomicFAddEXT,
+      SPIRV_OC_OpAssumeTrueKHR, SPIRV_OC_OpExpectKHR, SPIRV_OC_OpAtomicFAddEXT,
       SPIRV_OC_OpConvertFToBF16INTEL, SPIRV_OC_OpConvertBF16ToFINTEL,
       SPIRV_OC_OpControlBarrierArriveINTEL, SPIRV_OC_OpControlBarrierWaitINTEL,
       SPIRV_OC_OpGroupIMulKHR, SPIRV_OC_OpGroupFMulKHR,
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td
index d828e118757e2..8efd9ee691113 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td
@@ -56,6 +56,58 @@ def SPIRV_KHRAssumeTrueOp : SPIRV_KhrVendorOp<"AssumeTrue", []> {
 
 // -----
 
+def SPIRV_KHRExpectOp : SPIRV_KhrVendorOp<"Expect",
+    [Pure, AllTypesMatch<["value", "expectedValue", "result"]>]> {
+  let summary = "Provides a hint to the optimizer that the value is likely "
+                "to equal the expected value.";
+
+  let description = [{
+    Result Type must be a scalar or vector of integer type or Boolean type.
+
+    Value and Expected Value must be of Result Type. Result is the same as
+    Value.
+
+    <!-- End of AutoGen section -->
+
+    ```
+    integer-scalar-vector-type ::= integer-type |
+                                   `vector<` integer-literal `x` integer-type `>`
+    bool-scalar-vector-type ::= `i1` |
+                                `vector<` integer-literal `x` `i1` `>`
+    expect-op ::= ssa-id `=` `spirv.KHR.Expect` ssa-use `,` ssa-use
+                  `:` (integer-scalar-vector-type | bool-scalar-vector-type)
+    ```
+
+    #### Example:
+
+    ```mlir
+    %2 = spirv.KHR.Expect %0, %1 : i32
+    %5 = spirv.KHR.Expect %3, %4 : vector<4xi1>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPIRV_V_1_0>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[SPV_KHR_expect_assume]>,
+    Capability<[SPIRV_C_ExpectAssumeKHR]>
+  ];
+
+  let arguments = (ins
+    SPIRV_ScalarOrVectorOf<AnyTypeOf<[SPIRV_Integer, SPIRV_Bool]>>:$value,
+    SPIRV_ScalarOrVectorOf<AnyTypeOf<[SPIRV_Integer, SPIRV_Bool]>>:$expectedValue
+  );
+
+  let results = (outs
+    SPIRV_ScalarOrVectorOf<AnyTypeOf<[SPIRV_Integer, SPIRV_Bool]>>:$result
+  );
+
+  let hasVerifier = 0;
+  let assemblyFormat = "operands attr-dict `:` type($result)";
+}
+
+// -----
+
 def SPIRV_UndefOp : SPIRV_Op<"Undef", [Pure]> {
   let summary = "Make an intermediate object whose value is undefined.";
 
diff --git a/mlir/test/Dialect/SPIRV/IR/misc-ops.mlir b/mlir/test/Dialect/SPIRV/IR/misc-ops.mlir
index 182b661035a61..cfd19a683bebf 100644
--- a/mlir/test/Dialect/SPIRV/IR/misc-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/misc-ops.mlir
@@ -44,3 +44,55 @@ func.func @assume_true(%arg : f32) -> () {
   spirv.KHR.AssumeTrue %arg
   spirv.Return
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.KHR.Expect
+//===----------------------------------------------------------------------===//
+
+func.func @expect_scalar_int(%val : i32, %expected : i32) -> i32 {
+  // CHECK: %{{.*}} = spirv.KHR.Expect %{{.*}}, %{{.*}} : i32
+  %0 = spirv.KHR.Expect %val, %expected : i32
+  spirv.ReturnValue %0 : i32
+}
+
+// -----
+
+func.func @expect_scalar_bool(%val : i1, %expected : i1) -> i1 {
+  // CHECK: %{{.*}} = spirv.KHR.Expect %{{.*}}, %{{.*}} : i1
+  %0 = spirv.KHR.Expect %val, %expected : i1
+  spirv.ReturnValue %0 : i1
+}
+
+// -----
+
+func.func @expect_vector_int(%val : vector<4xi32>, %expected : vector<4xi32>) -> vector<4xi32> {
+  // CHECK: %{{.*}} = spirv.KHR.Expect %{{.*}}, %{{.*}} : vector<4xi32>
+  %0 = spirv.KHR.Expect %val, %expected : vector<4xi32>
+  spirv.ReturnValue %0 : vector<4xi32>
+}
+
+// -----
+
+func.func @expect_vector_bool(%val : vector<4xi1>, %expected : vector<4xi1>) -> vector<4xi1> {
+  // CHECK: %{{.*}} = spirv.KHR.Expect %{{.*}}, %{{.*}} : vector<4xi1>
+  %0 = spirv.KHR.Expect %val, %expected : vector<4xi1>
+  spirv.ReturnValue %0 : vector<4xi1>
+}
+
+// -----
+
+func.func @expect_type_mismatch(%val : i32, %expected : i64) -> i32 {
+  // expected-error @+1 {{op failed to verify that all of {value, expectedValue, result} have same type}}
+  %0 = "spirv.KHR.Expect"(%val, %expected) : (i32, i64) -> i32
+  spirv.ReturnValue %0 : i32
+}
+
+// -----
+
+func.func @expect_float_invalid(%val : f32, %expected : f32) -> f32 {
+  // expected-error @+1 {{op operand #0 must be}}
+  %0 = "spirv.KHR.Expect"(%val, %expected) : (f32, f32) -> f32
+  spirv.ReturnValue %0 : f32
+}

>From 28e2f27a7a2f413c25e1dbb543ec73ecdfed86bb Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Tue, 5 May 2026 14:56:55 +0200
Subject: [PATCH 2/4] Add target tests

---
 mlir/test/Target/SPIRV/expect-assume-ops.mlir | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 mlir/test/Target/SPIRV/expect-assume-ops.mlir

diff --git a/mlir/test/Target/SPIRV/expect-assume-ops.mlir b/mlir/test/Target/SPIRV/expect-assume-ops.mlir
new file mode 100644
index 0000000000000..664646efbd9ed
--- /dev/null
+++ b/mlir/test/Target/SPIRV/expect-assume-ops.mlir
@@ -0,0 +1,41 @@
+// RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip \
+// RUN:  --split-input-file %s | FileCheck %s
+
+spirv.module Logical GLSL450 requires
+  #spirv.vce<v1.0, [Shader, Linkage, ExpectAssumeKHR], [SPV_KHR_expect_assume]> {
+
+  // CHECK-LABEL: @assume_true
+  spirv.func @assume_true(%arg : i1) "None" {
+    // CHECK: spirv.KHR.AssumeTrue %{{.*}}
+    spirv.KHR.AssumeTrue %arg
+    spirv.Return
+  }
+
+  // CHECK-LABEL: @expect_scalar_int
+  spirv.func @expect_scalar_int(%val : i32, %expected : i32) -> i32 "None" {
+    // CHECK: {{%.+}} = spirv.KHR.Expect %{{.*}}, %{{.*}} : i32
+    %0 = spirv.KHR.Expect %val, %expected : i32
+    spirv.ReturnValue %0 : i32
+  }
+
+  // CHECK-LABEL: @expect_scalar_bool
+  spirv.func @expect_scalar_bool(%val : i1, %expected : i1) -> i1 "None" {
+    // CHECK: {{%.+}} = spirv.KHR.Expect %{{.*}}, %{{.*}} : i1
+    %0 = spirv.KHR.Expect %val, %expected : i1
+    spirv.ReturnValue %0 : i1
+  }
+
+  // CHECK-LABEL: @expect_vector_int
+  spirv.func @expect_vector_int(%val : vector<4xi32>, %expected : vector<4xi32>) -> vector<4xi32> "None" {
+    // CHECK: {{%.+}} = spirv.KHR.Expect %{{.*}}, %{{.*}} : vector<4xi32>
+    %0 = spirv.KHR.Expect %val, %expected : vector<4xi32>
+    spirv.ReturnValue %0 : vector<4xi32>
+  }
+
+  // CHECK-LABEL: @expect_vector_bool
+  spirv.func @expect_vector_bool(%val : vector<4xi1>, %expected : vector<4xi1>) -> vector<4xi1> "None" {
+    // CHECK: {{%.+}} = spirv.KHR.Expect %{{.*}}, %{{.*}} : vector<4xi1>
+    %0 = spirv.KHR.Expect %val, %expected : vector<4xi1>
+    spirv.ReturnValue %0 : vector<4xi1>
+  }
+}

>From 4070458ae29684ee5f6b42810c7102531c9dd9cf Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Tue, 5 May 2026 15:18:00 +0200
Subject: [PATCH 3/4] Address review comments

---
 mlir/test/Target/SPIRV/expect-assume-ops.mlir | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Target/SPIRV/expect-assume-ops.mlir b/mlir/test/Target/SPIRV/expect-assume-ops.mlir
index 664646efbd9ed..e742835462227 100644
--- a/mlir/test/Target/SPIRV/expect-assume-ops.mlir
+++ b/mlir/test/Target/SPIRV/expect-assume-ops.mlir
@@ -1,5 +1,9 @@
-// RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip \
-// RUN:  --split-input-file %s | FileCheck %s
+// RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip %s | FileCheck %s
+
+// RUN: %if spirv-tools %{ rm -rf %t %}
+// RUN: %if spirv-tools %{ mkdir %t %}
+// RUN: %if spirv-tools %{ mlir-translate --no-implicit-module --serialize-spirv --split-input-file --spirv-save-validation-files-with-prefix=%t/module %s %}
+// RUN: %if spirv-tools %{ spirv-val %t %}
 
 spirv.module Logical GLSL450 requires
   #spirv.vce<v1.0, [Shader, Linkage, ExpectAssumeKHR], [SPV_KHR_expect_assume]> {

>From a1d27a8fd7c5eee7d2a56ac481b19efda567d44d Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Tue, 5 May 2026 15:42:05 +0200
Subject: [PATCH 4/4] Address review comments

---
 mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td
index 8efd9ee691113..786dc23c2bdfe 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVMiscOps.td
@@ -26,10 +26,6 @@ def SPIRV_KHRAssumeTrueOp : SPIRV_KhrVendorOp<"AssumeTrue", []> {
 
     <!-- End of AutoGen section -->
 
-    ```
-    assumetruekhr-op ::= `spirv.KHR.AssumeTrue` ssa-use
-    ```
-
     #### Example:
 
     ```mlir
@@ -69,15 +65,6 @@ def SPIRV_KHRExpectOp : SPIRV_KhrVendorOp<"Expect",
 
     <!-- End of AutoGen section -->
 
-    ```
-    integer-scalar-vector-type ::= integer-type |
-                                   `vector<` integer-literal `x` integer-type `>`
-    bool-scalar-vector-type ::= `i1` |
-                                `vector<` integer-literal `x` `i1` `>`
-    expect-op ::= ssa-id `=` `spirv.KHR.Expect` ssa-use `,` ssa-use
-                  `:` (integer-scalar-vector-type | bool-scalar-vector-type)
-    ```
-
     #### Example:
 
     ```mlir



More information about the Mlir-commits mailing list