[mlir][spirv] Fix integer type emulation with extension/truncation

[antiagainst]

Summary

For integer extension or truncation with type emulation, we need
to make sure we perform masking or shifting to discard unwanted
bits to avoid pollute consumer ops.

[antiagainst]

Fix

[kuhar]

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

112–113

nit: I implemented a similar helper in wide int emulation: https://github.com/llvm/llvm-project/blob/da56750f82a60b5872679424fc6ae78a42c786a0/mlir/lib/Dialect/Arith/Transforms/EmulateWideInt.cpp#L62-L74.

I think that implementation is a little bit more concise and uses SplatElementsAttr::get without a temporary vector.

565–568

Should we assert that the shift amount is as expected? ([0, BW-1])

579

the right *sign* (i.e., the MSB) bit?

I think the sign bit may not be the leading bit depending on the endiannes, no?

620

Outdated comment? This pattern doesn't generate any selects.

643

Should we do uint64_t(1) << bw if the helper expects u64 anyway? Or even better, use one of the helpers in MathExtras to create a mask: https://llvm.org/doxygen/namespacellvm.html#a79dff91526e31b1a05f59eed6c189eb4?

689

Same: the comment doesn't match the code

711

Also here

[antiagainst]

Rebase

[antiagainst]

Address comments

[kuhar]

LGTM

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

if (type.isInteger(1))
     return true;
 
   if (auto vecType = dyn_cast<VectorType>(type))
     return vecType.getElementType().isInteger(1);
 
   return false;
 }
 
+/// Creates a scalar/vector integer constant.
+static Value getScalarOrVectorConstInt(Type type, uint64_t value,
+                                       OpBuilder &builder, Location loc) {
+  if (auto vectorType = dyn_cast<VectorType>(type)) {
+    Attribute element = IntegerAttr::get(vectorType.getElementType(), value);
+    SmallVector<Attribute> values(vectorType.getNumElements(), element);
+    auto attr = DenseElementsAttr::get(vectorType, values);
+    return builder.create<spirv::ConstantOp>(loc, vectorType, attr);
+  }
+
+  if (auto intType = dyn_cast<IntegerType>(type))
+    return builder.create<spirv::ConstantOp>(
+        loc, type, builder.getIntegerAttr(type, value));
+
+  return nullptr;
+}
+
 /// Returns true if scalar/vector type `a` and `b` have the same number of
 /// bitwidth.
 static bool hasSameBitwidth(Type a, Type b) {
   auto getNumBitwidth = [](Type type) {
     unsigned bw = 0;
     if (type.isIntOrFloat())
       bw = type.getIntOrFloatBitWidth();
     else if (auto vecType = dyn_cast<VectorType>(type))
matchAndRewrite(arith::ExtSIOp op, OpAdaptor adaptor,
 
     Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
     rewriter.replaceOpWithNewOp<spirv::SelectOp>(op, dstType, operand, allOnes,
                                                  zero);
     return success();
   }
 };
 
+/// Converts arith.extsi to spirv.Select if the type of source is neither i1 nor
+/// vector of i1.
+struct ExtSIPattern final : public OpConversionPattern<arith::ExtSIOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::ExtSIOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type srcType = adaptor.getIn().getType();
+    if (isBoolScalarOrVector(srcType))
+      return failure();
+
+    Type dstType = getTypeConverter()->convertType(op.getType());
+    if (!dstType)
+      return getTypeConversionFailure(rewriter, op);
+
+    if (dstType == srcType) {
+      // We can have the same source and destination type due to type emulation.
+      // Perform bit shifting to make sure we have the proper leading set bits.
+
+      unsigned srcBW =
+          getElementTypeOrSelf(op.getIn().getType()).getIntOrFloatBitWidth();
+      unsigned dstBW =
+          getElementTypeOrSelf(op.getType()).getIntOrFloatBitWidth();
+      Value shiftSize = getScalarOrVectorConstInt(dstType, dstBW - srcBW,
+                                                  rewriter, op.getLoc());
+
+      // First shift left to sequeeze out all leading bits beyond the original
+      // bitwidth. Here we need to use the original source and result type's
+      // bitwidth.
+      auto shiftLOp = rewriter.create<spirv::ShiftLeftLogicalOp>(
+          op.getLoc(), dstType, adaptor.getIn(), shiftSize);
+
+      // Then we perform arithmetic right shift to make sure we have the right
+      // leading set bits for negative values.
+      rewriter.replaceOpWithNewOp<spirv::ShiftRightArithmeticOp>(
+          op, dstType, shiftLOp, shiftSize);
+    } else {
+      rewriter.replaceOpWithNewOp<spirv::SConvertOp>(op, dstType,
+                                                     adaptor.getOperands());
+    }
+
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // ExtUIOp
 //===----------------------------------------------------------------------===//
 
 /// Converts arith.extui to spirv.Select if the type of source is i1 or vector
 /// of i1.
 struct ExtUII1Pattern final : public OpConversionPattern<arith::ExtUIOp> {
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(arith::ExtUIOp op, OpAdaptor adaptor,
     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);
     return success();
   }
 };
 
+/// Converts arith.extui to spirv.Select if the type of source is neither i1 nor
+/// vector of i1.
+struct ExtUIPattern final : public OpConversionPattern<arith::ExtUIOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::ExtUIOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type srcType = adaptor.getIn().getType();
+    if (isBoolScalarOrVector(srcType))
+      return failure();
+
+    Type dstType = getTypeConverter()->convertType(op.getType());
+    if (!dstType)
+      return getTypeConversionFailure(rewriter, op);
+
+    if (dstType == srcType) {
+      // We can have the same source and destination type due to type emulation.
+      // Perform bit masking to make sure we don't pollute downstream consumers
+      // with unwanted bits. Here we need to use the original source type's
+      // bitwidth.
+      unsigned bitwidth =
+          getElementTypeOrSelf(op.getIn().getType()).getIntOrFloatBitWidth();
+      Value mask = getScalarOrVectorConstInt(dstType, (1u << bitwidth) - 1,
+                                             rewriter, op.getLoc());
+      rewriter.replaceOpWithNewOp<spirv::BitwiseAndOp>(op, dstType,
+                                                       adaptor.getIn(), mask);
+    } else {
+      rewriter.replaceOpWithNewOp<spirv::UConvertOp>(op, dstType,
+                                                     adaptor.getOperands());
+    }
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // TruncIOp
 //===----------------------------------------------------------------------===//
 
 /// Converts arith.trunci to spirv.Select if the type of result is i1 or vector
 /// of i1.
 struct TruncII1Pattern final : public OpConversionPattern<arith::TruncIOp> {
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(arith::TruncIOp op, OpAdaptor adaptor,
 
     Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
     Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
     rewriter.replaceOpWithNewOp<spirv::SelectOp>(op, dstType, isOne, one, zero);
     return success();
   }
 };
 
+/// Converts arith.trunci to spirv.Select if the type of result is neither i1
+/// nor vector of i1.
+struct TruncIPattern final : public OpConversionPattern<arith::TruncIOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::TruncIOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type srcType = adaptor.getIn().getType();
+    Type dstType = getTypeConverter()->convertType(op.getType());
+    if (!dstType)
+      return getTypeConversionFailure(rewriter, op);
+
+    if (isBoolScalarOrVector(dstType))
+      return failure();
+
+    if (dstType == srcType) {
+      // We can have the same source and destination type due to type emulation.
+      // Perform bit masking to make sure we don't pollute downstream consumers
+      // with unwanted bits. Here we need to use the original result type's
+      // bitwidth.
+      unsigned bw = getElementTypeOrSelf(op.getType()).getIntOrFloatBitWidth();
+      Value mask = getScalarOrVectorConstInt(dstType, (1u << bw) - 1, rewriter,
+                                             op.getLoc());
+      rewriter.replaceOpWithNewOp<spirv::BitwiseAndOp>(op, dstType,
+                                                       adaptor.getIn(), mask);
+    } else {
+      // Given this is truncation, either SConvertOp or UConvertOp works.
+      rewriter.replaceOpWithNewOp<spirv::SConvertOp>(op, dstType,
+                                                     adaptor.getOperands());
+    }
+    return success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // TypeCastingOp
 //===----------------------------------------------------------------------===//
 
 /// Converts type-casting standard operations to SPIR-V operations.
 template <typename Op, typename SPIRVOp>
 struct TypeCastingOpPattern final : public OpConversionPattern<Op> {
   using OpConversionPattern<Op>::OpConversionPattern;
patterns.add<
     spirv::ElementwiseOpPattern<arith::NegFOp, spirv::FNegateOp>,
     spirv::ElementwiseOpPattern<arith::AddFOp, spirv::FAddOp>,
     spirv::ElementwiseOpPattern<arith::SubFOp, spirv::FSubOp>,
     spirv::ElementwiseOpPattern<arith::MulFOp, spirv::FMulOp>,
     spirv::ElementwiseOpPattern<arith::DivFOp, spirv::FDivOp>,
     spirv::ElementwiseOpPattern<arith::RemFOp, spirv::FRemOp>,
     spirv::ElementwiseOpPattern<arith::IsNanOp, spirv::IsNanOp>,
     spirv::ElementwiseOpPattern<arith::IsInfOp, spirv::IsInfOp>,
-    TypeCastingOpPattern<arith::ExtUIOp, spirv::UConvertOp>, ExtUII1Pattern,
-    TypeCastingOpPattern<arith::ExtSIOp, spirv::SConvertOp>, ExtSII1Pattern,
+    ExtUIPattern, ExtUII1Pattern,
+    ExtSIPattern, ExtSII1Pattern,
     TypeCastingOpPattern<arith::ExtFOp, spirv::FConvertOp>,
-    TypeCastingOpPattern<arith::TruncIOp, spirv::SConvertOp>, TruncII1Pattern,
+    TruncIPattern, TruncII1Pattern,
     TypeCastingOpPattern<arith::TruncFOp, spirv::FConvertOp>,
     TypeCastingOpPattern<arith::UIToFPOp, spirv::ConvertUToFOp>, UIToFPI1Pattern,
     TypeCastingOpPattern<arith::SIToFPOp, spirv::ConvertSToFOp>,
     TypeCastingOpPattern<arith::FPToUIOp, spirv::ConvertFToUOp>,
     TypeCastingOpPattern<arith::FPToSIOp, spirv::ConvertFToSOp>,
     TypeCastingOpPattern<arith::IndexCastOp, spirv::SConvertOp>,
     TypeCastingOpPattern<arith::IndexCastUIOp, spirv::UConvertOp>,
     TypeCastingOpPattern<arith::BitcastOp, spirv::BitcastOp>,

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

 // CHECK-LABEL: @fpext2
 // CHECK-SAME: %[[ARG:.*]]: f32
 func.func @fpext2(%arg0 : f32) -> f64 {
   // CHECK-NEXT: spirv.FConvert %[[ARG]] : f32 to f64
   %0 = arith.extf %arg0 : f32 to f64
   return %0: f64
 }
 
+// CHECK-LABEL: @trunci4
+//  CHECK-SAME: %[[ARG:.*]]: i32
+func.func @trunci4(%arg0 : i32) -> i4 {
+  // CHECK: %[[MASK:.+]] = spirv.Constant 15 : i32
+  // CHECK: %[[AND:.+]] = spirv.BitwiseAnd %[[ARG]], %[[MASK]] : i32
+  %0 = arith.trunci %arg0 : i32 to i4
+  // CHECK: %[[RET:.+]] = builtin.unrealized_conversion_cast %[[AND]] : i32 to i4
+  // CHECK: return %[[RET]] : i4
+  return %0 : i4
+}
+
+// CHECK-LABEL: @zexti4
+func.func @zexti4(%arg0: i4) -> i32 {
+  // CHECK: %[[INPUT:.+]] = builtin.unrealized_conversion_cast %{{.+}} : i4 to i32
+  // CHECK: %[[MASK:.+]] = spirv.Constant 15 : i32
+  // CHECK: %[[AND:.+]] = spirv.BitwiseAnd %[[INPUT]], %[[MASK]] : i32
+  %0 = arith.extui %arg0 : i4 to i32
+  // CHECK: return %[[AND]] : i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: @sexti4
+func.func @sexti4(%arg0: i4) -> i32 {
+  // CHECK: %[[INPUT:.+]] = builtin.unrealized_conversion_cast %arg0 : i4 to i32
+  // CHECK: %[[SIZE:.+]] = spirv.Constant 28 : i32
+  // CHECK: %[[SL:.+]] = spirv.ShiftLeftLogical %[[INPUT]], %[[SIZE]] : i32, i32
+  // CHECK: %[[SR:.+]] = spirv.ShiftRightArithmetic %[[SL]], %[[SIZE]] : i32, i32
+  %0 = arith.extsi %arg0 : i4 to i32
+  // CHECK: return %[[SR]] : i32
+  return %0 : i32
+}
+
 } // end module
 
 // -----
 
 // Checks that cast types will be adjusted when missing special capabilities for
 // certain non-32-bit scalar types.
 module attributes {
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Float16], []>, #spirv.resource_limits<>>