Amend reviewer comments, add vector support

Author: ianayl

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

@@ -611,16 +611,16 @@ struct UIToFPI1Pattern final : public OpConversionPattern<arith::UIToFPOp> {
 // IndexCastOp
 //===----------------------------------------------------------------------===//
 
-/// Converts arith.index_cast to spirv.Select if the source type is i1
+/// Converts arith.index_cast to spirv.Select if the source type is i1.
 struct IndexCastI1IndexPattern 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))
+    Type srcType = adaptor.getIn().getType();
+    if (!isBoolScalarOrVector(srcType))
       return failure();
 
     Type dstType = getTypeConverter()->convertType(op.getType());
@@ -631,7 +631,7 @@ struct IndexCastI1IndexPattern final
     Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
     Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
     rewriter.replaceOpWithNewOp<spirv::SelectOp>(
-        op, dstType, adaptor.getOperands().front(), one, zero);
+        op, dstType, adaptor.getIn(), one, zero);
     return success();
   }
 };

mlir/test/Conversion/ArithToSPIRV/arith-to-spirv.mlir

@@ -734,15 +734,24 @@ func.func @index_castui4(%arg0: index) {
   return
 }
 
-// CHECK-LABEL: index_casti1index
-func.func @index_casti1index(%arg0 : i1) {
+// CHECK-LABEL: index_casti1index_1
+func.func @index_casti1index_1(%arg0 : i1) {
   // CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
   // CHECK: %[[ONE:.+]] = spirv.Constant 1 : i32
   // CHECK: spirv.Select %{{.+}}, %[[ONE]], %[[ZERO]] : i1, i32
   %0 = arith.index_cast %arg0 : i1 to index
   return
 }
 
+// CHECK-LABEL: index_casti1index_2
+func.func @index_casti1index_2(%arg0 : vector<3xi1>) {
+  // CHECK: %[[ZERO:.+]] = spirv.Constant dense<0> : vector<3xi32>
+  // CHECK: %[[ONE:.+]] = spirv.Constant dense<1> : vector<3xi32>
+  // CHECK: spirv.Select %{{.+}}, %[[ONE]], %[[ZERO]] : vector<3xi1>, vector<3xi32>
+  %0 = arith.index_cast %arg0 : vector<3xi1> to vector<3xindex>
+  return
+}
+
 // CHECK-LABEL: @bit_cast
 func.func @bit_cast(%arg0: vector<2xf32>, %arg1: i64) {
   // CHECK: spirv.Bitcast %{{.+}} : vector<2xf32> to vector<2xi32>