[Mlir-commits] [mlir] [mlir][spirv] Allow composite types in SelectOp (PR #196105)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Igor Wodiany (IgWod)

<details>
<summary>Changes</summary>

This is allowed from the version 1.4 of the spec.

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


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td (+2-2) 
- (modified) mlir/test/Dialect/SPIRV/IR/logical-ops.mlir (+18) 
- (modified) mlir/test/Target/SPIRV/logical-ops.mlir (+25) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index 8badb84a879fa..379022ba23bdb 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4352,8 +4352,8 @@ def SPIRV_IntVec4 : SPIRV_Vec4<SPIRV_Integer>;
 def SPIRV_IOrUIVec4 : SPIRV_Vec4<SPIRV_SignlessOrUnsignedInt>;
 def SPIRV_Int32Vec4 : SPIRV_Vec4<AnyI32>;
 
-// TODO: From 1.4, this should also include Composite type.
-def SPIRV_SelectType : AnyTypeOf<[SPIRV_Scalar, SPIRV_Vector, SPIRV_AnyPtr]>;
+def SPIRV_SelectType : AnyTypeOf<[SPIRV_Scalar, SPIRV_Vector, SPIRV_AnyPtr,
+                                  SPIRV_AnyMatrix, SPIRV_AnyArray, SPIRV_AnyStruct]>;
 
 //===----------------------------------------------------------------------===//
 // SPIR-V OpTrait definitions
diff --git a/mlir/test/Dialect/SPIRV/IR/logical-ops.mlir b/mlir/test/Dialect/SPIRV/IR/logical-ops.mlir
index 1018751cf65e0..019ce8121947a 100644
--- a/mlir/test/Dialect/SPIRV/IR/logical-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/logical-ops.mlir
@@ -251,6 +251,24 @@ func.func @select_op_vec_condn_vec(%arg0: vector<3xi1>) -> () {
   return
 }
 
+func.func @select_op_array(%arg0: i1, %arg1: !spirv.array<4 x i32>, %arg2: !spirv.array<4 x i32>) -> () {
+  // CHECK: spirv.Select {{%.*}}, {{%.*}}, {{%.*}} : i1, !spirv.array<4 x i32>
+  %0 = spirv.Select %arg0, %arg1, %arg2 : i1, !spirv.array<4 x i32>
+  return
+}
+
+func.func @select_op_struct(%arg0: i1, %arg1: !spirv.struct<(i32, i32)>, %arg2: !spirv.struct<(i32, i32)>) -> () {
+  // CHECK: spirv.Select {{%.*}}, {{%.*}}, {{%.*}} : i1, !spirv.struct<(i32, i32)>
+  %0 = spirv.Select %arg0, %arg1, %arg2 : i1, !spirv.struct<(i32, i32)>
+  return
+}
+
+func.func @select_op_matrix(%arg0: i1, %arg1: !spirv.matrix<4 x vector<3xf32>>, %arg2: !spirv.matrix<4 x vector<3xf32>>) -> () {
+  // CHECK: spirv.Select {{%.*}}, {{%.*}}, {{%.*}} : i1, !spirv.matrix<4 x vector<3xf32>>
+  %0 = spirv.Select %arg0, %arg1, %arg2 : i1, !spirv.matrix<4 x vector<3xf32>>
+  return
+}
+
 // -----
 
 func.func @select_op(%arg0: i1) -> () {
diff --git a/mlir/test/Target/SPIRV/logical-ops.mlir b/mlir/test/Target/SPIRV/logical-ops.mlir
index d570815448b7c..39381461e4a79 100644
--- a/mlir/test/Target/SPIRV/logical-ops.mlir
+++ b/mlir/test/Target/SPIRV/logical-ops.mlir
@@ -138,3 +138,28 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.4, [Shader, Linkage, BFloat1
     spirv.Return
   }
 }
+
+// -----
+
+// Test select works with composite types
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.6, [Shader, Linkage], []> {
+  spirv.func @select_op_array(%arg0: i1, %arg1: !spirv.array<4 x i32>, %arg2: !spirv.array<4 x i32>) -> () "None" {
+    // CHECK: spirv.Select {{%.*}}, {{%.*}}, {{%.*}} : i1, !spirv.array<4 x i32>
+    %0 = spirv.Select %arg0, %arg1, %arg2 : i1, !spirv.array<4 x i32>
+    spirv.Return
+  }
+
+  spirv.func @select_op_struct(%arg0: i1, %arg1: !spirv.struct<(i32, i32)>, %arg2: !spirv.struct<(i32, i32)>) -> () "None" {
+    // CHECK: spirv.Select {{%.*}}, {{%.*}}, {{%.*}} : i1, !spirv.struct<(i32, i32)>
+    %0 = spirv.Select %arg0, %arg1, %arg2 : i1, !spirv.struct<(i32, i32)>
+    spirv.Return
+  }
+
+  spirv.func @select_op_matrix(%arg0: i1, %arg1: !spirv.matrix<4 x vector<3xf32>>, %arg2: !spirv.matrix<4 x vector<3xf32>>) -> () "None" {
+    // CHECK: spirv.Select {{%.*}}, {{%.*}}, {{%.*}} : i1, !spirv.matrix<4 x vector<3xf32>>
+    %0 = spirv.Select %arg0, %arg1, %arg2 : i1, !spirv.matrix<4 x vector<3xf32>>
+    spirv.Return
+  }
+}
+

``````````

</details>


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