[Mlir-commits] [mlir] [mlir][LLVM] `ArithToLLVM`: Add 1:N support for `arith.select` lowering (PR #153944)

[Markus Böck]

================
@@ -479,6 +489,36 @@ CmpFOpLowering::matchAndRewrite(arith::CmpFOp op, OpAdaptor adaptor,
       rewriter);
 }
 
+//===----------------------------------------------------------------------===//
+// SelectOpOneToNLowering
+//===----------------------------------------------------------------------===//
+
+/// Pattern for arith.select where the true/false values lower to multiple
+/// SSA values (1:N conversion). This pattern generates multiple arith.select
+/// than can be lowered by the 1:1 arith.select pattern.
+LogicalResult SelectOpOneToNLowering::matchAndRewrite(
+    arith::SelectOp op, Adaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  // In case of a 1:1 conversion, the 1:1 pattern will match.
+  if (llvm::hasSingleElement(adaptor.getTrueValue()))
+    return rewriter.notifyMatchFailure(
+        op, "not a 1:N conversion, 1:1 pattern will match");
+  if (!op.getCondition().getType().isInteger(1))
+    return rewriter.notifyMatchFailure(op,
+                                       "non-i1 conditions are not supported");
+  if (!llvm::hasSingleElement(adaptor.getCondition()))
+    return rewriter.notifyMatchFailure(
+        op, "1:N condition conversion is not supported");
+  SmallVector<Value> results;
+  for (auto [trueValue, falseValue] :
+       llvm::zip_equal(adaptor.getTrueValue(), adaptor.getFalseValue()))
+    results.push_back(arith::SelectOp::create(
+        rewriter, op.getLoc(), llvm::getSingleElement(adaptor.getCondition()),
----------------
zero9178 wrote:

Since we are only creating a new arith.select, would it make sense to use `op.getCondition` here to use the pre-type conversion (i.e. `i1`) value here?

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