add support for single element vectors

Author: ianayl

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

@@ -622,10 +622,14 @@ struct IndexCastIndexI1Pattern final
     if (!isBoolScalarOrVector(op.getType()))
       return failure();
 
+    Type dstType = getTypeConverter()->convertType(op.getType());
+    if (!dstType)
+      return getTypeConversionFailure(rewriter, op);
+
     Location loc = op.getLoc();
     Value zeroIdx =
-        spirv::ConstantOp::getZero(adaptor.getIn().getType(), loc, rewriter);
-    rewriter.replaceOpWithNewOp<spirv::INotEqualOp>(op, op.getType(), zeroIdx,
+       spirv::ConstantOp::getZero(adaptor.getIn().getType(), loc, rewriter);
+    rewriter.replaceOpWithNewOp<spirv::INotEqualOp>(op, dstType, zeroIdx,
                                                     adaptor.getIn());
     return success();
   }

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

@@ -743,7 +743,16 @@ func.func @index_castindexi1_1(%arg0 : index) {
 }
 
 // CHECK-LABEL: index_castindexi1_2
-func.func @index_castindexi1_2(%arg0 : vector<3xindex>) {
+func.func @index_castindexi1_2(%arg0 : vector<1xindex>) -> vector<1xi1> {
+  // Single-element vectors do not exist in SPIRV.
+  // CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i32
+  // CHECK: spirv.INotEqual %[[ZERO]], %{{.+}} : i32
+  %0 = arith.index_cast %arg0 : vector<1xindex> to vector<1xi1>
+  return %0 : vector<1xi1>
+}
+
+// CHECK-LABEL: index_castindexi1_3
+func.func @index_castindexi1_3(%arg0 : vector<3xindex>) {
   // CHECK: %[[ZERO:.+]] = spirv.Constant dense<0> : vector<3xi32>
   // CHECK: spirv.INotEqual %[[ZERO]], %{{.+}} : vector<3xi32>
   %0 = arith.index_cast %arg0 : vector<3xindex> to vector<3xi1>