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

[Jakub Kuderski]

================
@@ -133,6 +134,21 @@ class MatchBroadcastableShapes<string input1, string input2, string output>:
     "})">]>
   >;
 
+class TernaryMatchBroadcastableShapes<string input1, string input2, string input3, string output>:
+  PredOpTrait<"the shape of " # input1 # ", " # input2 # ", and " # input3 # " are compatible for broadcasting and the broadcast shape is equal to the " # output # " shape",
+  Implies<And<[CPred<HasRank<input1>.result>, CPred<HasRank<input2>.result>, CPred<HasRank<input3>.result>, CPred<HasRank<output>.result>,
+               CPred<Rank<input1>.result # " == " # Rank<input2>.result # " && " # Rank<input1>.result # " == " # Rank<input3>.result # " && " # Rank<input1>.result # " == " # Rank<output>.result>]>,
+    [CPred<"llvm::all_of_zip(" # Shape<input1>.result # ", " # Shape<input2>.result # ", " # Shape<input3>.result # ", " # Shape<output>.result # ", " #
+    "[](int64_t input1Dim, int64_t input2Dim, int64_t input3Dim, int64_t outputDim) { " #
+    "   bool dynamic = ShapedType::isDynamic(input1Dim) || ShapedType::isDynamic(input2Dim) || ShapedType::isDynamic(input3Dim) || ShapedType::isDynamic(outputDim);"
+    "   bool broadcastableInputs = (input1Dim == input2Dim || input1Dim == 1 || input2Dim == 1) && "
+    "                              (input1Dim == input3Dim || input1Dim == 1 || input3Dim == 1) && "
+    "                              (input2Dim == input3Dim || input2Dim == 1 || input3Dim == 1);"
+    "   bool broacastDimMatchesOutputDim = std::max(input1Dim, std::max(input2Dim, input3Dim)) == outputDim;"
+    "   return dynamic || (broadcastableInputs && broacastDimMatchesOutputDim);" #
+    "})">]>
+  >;
----------------
kuhar wrote:

If this is only going to be used for the select op, I'd implement this in C++ instead. Complicated logic in `CPred`s is difficult to maintain

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