[Mlir-commits] [mlir] [mlir][spirv] Add pattern matching for arith.index_cast i1 to index (PR #155729)

[Ian Li]

https://github.com/ianayl created https://github.com/llvm/llvm-project/pull/155729

Currently, `arith.index_cast` gets converted to `OpSConvert`: https://github.com/llvm/llvm-project/blob/9bf5bf3baf3c7aec82cdd235c6a2fd57b4dd55ab/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp#L1331 [OpSConvert requires its operands to be of integer type](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpSConvert), which poses an issue for `i1` since SPIRV distinguishes between booleans and integers. As a result, the following example would not be converted, leaving behind illegal ops:
```
%0 = arith.index_cast %arg0 : i1 to index
```
This PR adds additional logic to convert index_casts on `i1` into SPIRV index equivalents.

>From f446699bc016509b7a7c6c0a2170b61d0b8709c8 Mon Sep 17 00:00:00 2001
From: Ian Li <ian.li at intel.com>
Date: Wed, 27 Aug 2025 16:51:17 -0700
Subject: [PATCH] [mlir][spirv] Add pattern matching for arith.index_cast i1 to
 index

---
 .../Conversion/ArithToSPIRV/ArithToSPIRV.cpp  | 37 ++++++++++++++++++-
 .../ArithToSPIRV/arith-to-spirv.mlir          |  7 ++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp b/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp
index 265293b83f84c..172f322a12fd8 100644
--- a/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp
+++ b/mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp
@@ -607,6 +607,41 @@ struct UIToFPI1Pattern final : public OpConversionPattern<arith::UIToFPOp> {
   }
 };
 
+//===----------------------------------------------------------------------===//
+// IndexCastOp
+//===----------------------------------------------------------------------===//
+
+/// Converts arith.index_cast to spirv.Select if the type of source is i1 or
+/// vector of i1.
+struct IndexCastI1Pattern final : public OpConversionPattern<arith::IndexCastOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::IndexCastOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type srcType = adaptor.getOperands().front().getType();
+    if (!srcType.isInteger(1))
+      return failure();
+
+    Type dstType = getTypeConverter()->convertType(op.getType());
+    if (!dstType)
+      return getTypeConversionFailure(rewriter, op);
+    // if (!dstType.isIndex()) {
+    //   llvm::errs() << "why doesnt this work?\n";
+    //   return failure();
+    // }
+
+    auto *converter = this->template getTypeConverter<SPIRVTypeConverter>();
+    Location loc = op.getLoc();
+    Type spirvI32T = converter->getIndexType();
+    Value zero = spirv::ConstantOp::getZero(spirvI32T, loc, rewriter);
+    Value one = spirv::ConstantOp::getOne(spirvI32T, loc, rewriter);
+    auto newOp = rewriter.replaceOpWithNewOp<spirv::SelectOp>(
+        op, dstType, adaptor.getOperands().front(), one, zero);
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // ExtSIOp
 //===----------------------------------------------------------------------===//
@@ -1328,7 +1363,7 @@ void mlir::arith::populateArithToSPIRVPatterns(
     TypeCastingOpPattern<arith::SIToFPOp, spirv::ConvertSToFOp>,
     TypeCastingOpPattern<arith::FPToUIOp, spirv::ConvertFToUOp>,
     TypeCastingOpPattern<arith::FPToSIOp, spirv::ConvertFToSOp>,
-    TypeCastingOpPattern<arith::IndexCastOp, spirv::SConvertOp>,
+    TypeCastingOpPattern<arith::IndexCastOp, spirv::SConvertOp>, IndexCastI1Pattern,
     TypeCastingOpPattern<arith::IndexCastUIOp, spirv::UConvertOp>,
     TypeCastingOpPattern<arith::BitcastOp, spirv::BitcastOp>,
     CmpIOpBooleanPattern, CmpIOpPattern,
diff --git a/mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir b/mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir
index 6e2352e706acc..8bb63fff861ce 100644
--- a/mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir
+++ b/mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir
@@ -734,6 +734,13 @@ func.func @index_castui4(%arg0: index) {
   return
 }
 
+// CHECK-LABEL: index_casti1_1
+func.func @index_casti1_1(%arg0 : i1) -> index {
+  // CHECK: spirv.Select %{{.+}}, %{{.+}}, %{{.+}} : i1, i32
+  %0 = arith.index_cast %arg0 : i1 to index
+  return %0 : index
+}
+
 // CHECK-LABEL: @bit_cast
 func.func @bit_cast(%arg0: vector<2xf32>, %arg1: i64) {
   // CHECK: spirv.Bitcast %{{.+}} : vector<2xf32> to vector<2xi32>