[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