SPIRV vector.mask lowering: use 64-bit type. (#1093)

So far, most of the kernels use 16 or 32 vector sizes. But there is a
kernel that returns an i8 tensor and the vector size = 64. So, there
was an overflow when computing the mask (1 << valid_elements). This PR
updates the data type from 32-bit to 64-bit when computing the mask.
It also adds an assert to catch more easily this type of issue.

Author: antonio-cortes-perez

lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp

@@ -93,31 +93,31 @@ class VectorMaskConversionPattern final
       return mlir::failure();
 
     auto vWidth = vTy.getNumElements();
+    assert(vWidth <= 64 && "vector.create_mask supports vector widths <= 64");
     auto vWidthConst = rewriter.create<mlir::arith::ConstantOp>(
-        vMaskOp.getLoc(), rewriter.getI32IntegerAttr(vWidth));
+        vMaskOp.getLoc(), rewriter.getI64IntegerAttr(vWidth));
     auto maskVal = adaptor.getOperands()[0];
     maskVal = rewriter.create<mlir::arith::TruncIOp>(
-        vMaskOp.getLoc(), rewriter.getI32Type(), maskVal);
+        vMaskOp.getLoc(), rewriter.getI64Type(), maskVal);
 
     // maskVal < vWidth
     auto cmp = rewriter.create<mlir::arith::CmpIOp>(
         vMaskOp.getLoc(), mlir::arith::CmpIPredicate::slt, maskVal,
         vWidthConst);
     auto one = rewriter.create<mlir::arith::ConstantOp>(
-        vMaskOp.getLoc(), rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
+        vMaskOp.getLoc(), rewriter.getI64IntegerAttr(1));
     auto shift = rewriter.create<mlir::spirv::ShiftLeftLogicalOp>(
         vMaskOp.getLoc(), one, maskVal);
     auto mask1 =
         rewriter.create<mlir::arith::SubIOp>(vMaskOp.getLoc(), shift, one);
     auto mask2 = rewriter.create<mlir::arith::ConstantOp>(
-        vMaskOp.getLoc(), rewriter.getI32Type(),
-        rewriter.getI32IntegerAttr(0xFFFFFFFF));
+        vMaskOp.getLoc(), rewriter.getI64IntegerAttr(-1)); // all ones
     mlir::Value sel = rewriter.create<mlir::arith::SelectOp>(vMaskOp.getLoc(),
                                                              cmp, mask1, mask2);
 
     // maskVal < 0
     auto zero = rewriter.create<mlir::arith::ConstantOp>(
-        vMaskOp.getLoc(), rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
+        vMaskOp.getLoc(), rewriter.getI64IntegerAttr(0));
     auto cmp2 = rewriter.create<mlir::arith::CmpIOp>(
         vMaskOp.getLoc(), mlir::arith::CmpIPredicate::slt, maskVal, zero);
     sel = rewriter.create<mlir::arith::SelectOp>(vMaskOp.getLoc(), cmp2, zero,

test/Conversion/GPUToSPIRV/create_mask.mlir

@@ -16,18 +16,17 @@ module attributes {
 
 // CHECK-LABEL: spirv.func @create_mask
 // CHECK-SAME: %[[MASK_VAL:[[:alnum:]]+]]: i64
-// CHECK-NEXT: %[[VECTOR_WIDTH:.*]] = spirv.Constant 16 : i32
-// CHECK-NEXT: %[[MASK_VAL_I32:.*]] = spirv.SConvert %[[MASK_VAL]] : i64 to i32
-// CHECK-NEXT: %[[CMP1:.*]] = spirv.SLessThan %[[MASK_VAL_I32]], %[[VECTOR_WIDTH]] : i32
-// CHECK-NEXT: %[[ONE:.*]] = spirv.Constant 1 : i32
-// CHECK-NEXT: %[[SHIFT:.*]] = spirv.ShiftLeftLogical %[[ONE]], %[[MASK_VAL_I32]] : i32, i32
-// CHECK-NEXT: %[[MASK:.*]] = spirv.ISub %[[SHIFT]], %[[ONE]] : i32
-// CHECK-NEXT: %[[MASK_ONES:.*]] = spirv.Constant -1 : i32
-// CHECK-NEXT: %[[SELECT1:.*]] = spirv.Select %[[CMP1]], %[[MASK]], %[[MASK_ONES]] : i1, i32
-// CHECK-NEXT: %[[ZERO:.*]] = spirv.Constant 0 : i32
-// CHECK-NEXT: %[[CMP2:.*]] = spirv.SLessThan %[[MASK_VAL_I32]], %[[ZERO]] : i32
-// CHECK-NEXT: %[[SELECT2:.*]] = spirv.Select %[[CMP2]], %[[ZERO]], %[[SELECT1]] : i1, i32
-// CHECK-NEXT: %[[CAST:.*]] = spirv.SConvert %[[SELECT2]] : i32 to i16
+// CHECK-NEXT: %[[VECTOR_WIDTH:.*]] = spirv.Constant 16 : i64
+// CHECK-NEXT: %[[CMP1:.*]] = spirv.SLessThan %[[MASK_VAL]], %[[VECTOR_WIDTH]] : i64
+// CHECK-NEXT: %[[ONE:.*]] = spirv.Constant 1 : i64
+// CHECK-NEXT: %[[SHIFT:.*]] = spirv.ShiftLeftLogical %[[ONE]], %[[MASK_VAL]] : i64, i64
+// CHECK-NEXT: %[[MASK:.*]] = spirv.ISub %[[SHIFT]], %[[ONE]] : i64
+// CHECK-NEXT: %[[MASK_ONES:.*]] = spirv.Constant -1 : i64
+// CHECK-NEXT: %[[SELECT1:.*]] = spirv.Select %[[CMP1]], %[[MASK]], %[[MASK_ONES]] : i1, i64
+// CHECK-NEXT: %[[ZERO:.*]] = spirv.Constant 0 : i64
+// CHECK-NEXT: %[[CMP2:.*]] = spirv.SLessThan %[[MASK_VAL]], %[[ZERO]] : i64
+// CHECK-NEXT: %[[SELECT2:.*]] = spirv.Select %[[CMP2]], %[[ZERO]], %[[SELECT1]] : i1, i64
+// CHECK-NEXT: %[[CAST:.*]] = spirv.SConvert %[[SELECT2]] : i64 to i16
 // CHECK-NEXT: spirv.Bitcast %[[CAST]] : i16 to vector<16xi1>
 // CHECK-NEXT: spirv.Return