[Mlir-commits] [mlir] [mlir][spirv] Add comparison and elementwise ternary ops in TOSA Ext Inst Set (PR #186356)

[Jakub Kuderski]

================
@@ -1742,4 +1771,134 @@ def SPIRV_TosaSinOp : SPIRV_TosaFloatElementwiseUnaryOp<"Sin", 43> {
 }
 
 
+def SPIRV_TosaSelectOp : SPIRV_TosaOpWithResult<"Select", 44, [Pure,
+  AllElementTypesMatch<["true_value", "false_value", "output"]>,
+  AllRanksMatch<["condition", "true_value", "false_value", "output"]>,
+  TernaryMatchBroadcastableShapes<"condition", "true_value", "false_value", "output">,
+  DeclareOpInterfaceMethods<SelectLikeOpInterface>]> {
+  let summary = "Select operator.";
+
+  let description = [{
+    Elementwise Select of the output based on a condition.
+
+    References:
+      * https://github.khronos.org/SPIRV-Registry/extended/TOSA.001000.1.html#_select
+      * https://www.mlplatform.org/tosa/tosa_spec_1_0_1.html#_select
+
+    #### Example:
+    ```mlir
+    %0 = spirv.Tosa.Select %arg0, %arg1, %arg2 : !spirv.arm.tensor<4x1x4x5xi1>, !spirv.arm.tensor<4x6x4x5xi8>, !spirv.arm.tensor<4x6x4x5xi8> -> !spirv.arm.tensor<4x6x4x5xi8>
+    %0 = spirv.Tosa.Select %arg0, %arg1, %arg2 : !spirv.arm.tensor<9x2x15x8xi1>, !spirv.arm.tensor<9x2x15x8xf16>, !spirv.arm.tensor<9x1x15x8xf16> -> !spirv.arm.tensor<9x2x15x8xf16>
----------------
kuhar wrote:

maybe also call these con / trueVal / falseVal here to make it self-documenting?

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