[mlir][arith][spirv] Handle i1 sign extension in arith-to-spirv

[kuhar]

Summary

Also fix some surrounding nits.

mlir/lib/Conversion/ArithToSPIRV/ArithToSPIRV.cpp

 #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/Support/Debug.h"
+#include <cassert>
 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTARITHTOSPIRV
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
 #define DEBUG_TYPE "arith-to-spirv-pattern"
 

 struct UIToFPI1Pattern final : public OpConversionPattern<arith::UIToFPOp> {
   using OpConversionPattern<arith::UIToFPOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::UIToFPOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Converts arith.extsi to spirv.Select if the type of source is i1 or vector
+/// of i1.
+struct ExtSII1Pattern final : public OpConversionPattern<arith::ExtSIOp> {
+  using OpConversionPattern<arith::ExtSIOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::ExtSIOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// 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<arith::ExtUIOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(arith::ExtUIOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
UIToFPI1Pattern::matchAndRewrite(arith::UIToFPOp op, OpAdaptor adaptor,
   Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
   Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
   rewriter.template replaceOpWithNewOp<spirv::SelectOp>(
       op, dstType, adaptor.getOperands().front(), one, zero);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// ExtSII1Pattern
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+ExtSII1Pattern::matchAndRewrite(arith::ExtSIOp op, OpAdaptor adaptor,
+                                ConversionPatternRewriter &rewriter) const {
+  Value operand = adaptor.getIn();
+  if (!isBoolScalarOrVector(operand.getType()))
+    return failure();
+
+  Location loc = op.getLoc();
+  Type dstType = getTypeConverter()->convertType(op.getResult().getType());
+
+  Value allOnes;
+  if (auto intTy = dstType.dyn_cast<IntegerType>()) {
+    unsigned componentBitwidth = intTy.getWidth();
+    allOnes = rewriter.create<spirv::ConstantOp>(
+        loc, intTy,
+        rewriter.getIntegerAttr(intTy, APInt::getAllOnes(componentBitwidth)));
+  } else if (auto vectorTy = dstType.dyn_cast<VectorType>()) {
+    unsigned componentBitwidth = vectorTy.getElementTypeBitWidth();
+    allOnes = rewriter.create<spirv::ConstantOp>(
+        loc, vectorTy,
+        SplatElementsAttr::get(vectorTy, APInt::getAllOnes(componentBitwidth)));
+  } else {
+    return rewriter.notifyMatchFailure(
+        loc, llvm::formatv("unhandled type: {0}", dstType));
+  }
+
+  Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
+  rewriter.replaceOpWithNewOp<spirv::SelectOp>(op, dstType, operand, allOnes,
+                                               zero);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // ExtUII1Pattern
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ExtUII1Pattern::matchAndRewrite(arith::ExtUIOp op, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
   auto srcType = adaptor.getOperands().front().getType();
   if (!isBoolScalarOrVector(srcType))
patterns.add<
     spirv::ElementwiseOpPattern<arith::ShRSIOp, spirv::ShiftRightArithmeticOp>,
     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>,
     TypeCastingOpPattern<arith::ExtUIOp, spirv::UConvertOp>, ExtUII1Pattern,
-    TypeCastingOpPattern<arith::ExtSIOp, spirv::SConvertOp>,
+    TypeCastingOpPattern<arith::ExtSIOp, spirv::SConvertOp>, ExtSII1Pattern,
     TypeCastingOpPattern<arith::ExtFOp, spirv::FConvertOp>,
     TypeCastingOpPattern<arith::TruncIOp, spirv::SConvertOp>, TruncII1Pattern,
     TypeCastingOpPattern<arith::TruncFOp, spirv::FConvertOp>,
     TypeCastingOpPattern<arith::UIToFPOp, spirv::ConvertUToFOp>, UIToFPI1Pattern,
     TypeCastingOpPattern<arith::SIToFPOp, spirv::ConvertSToFOp>,
     TypeCastingOpPattern<arith::FPToSIOp, spirv::ConvertFToSOp>,
     TypeCastingOpPattern<arith::IndexCastOp, spirv::SConvertOp>,
     TypeCastingOpPattern<arith::IndexCastUIOp, spirv::UConvertOp>,

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

return builder.create<spirv::ConstantOp>(
         loc, type, builder.getFloatAttr(floatType, 0.0));
   }
   if (auto vectorType = type.dyn_cast<VectorType>()) {
     Type elemType = vectorType.getElementType();
     if (elemType.isa<IntegerType>()) {
       return builder.create<spirv::ConstantOp>(
           loc, type,
           DenseElementsAttr::get(vectorType,
-                                 IntegerAttr::get(elemType, 0.0).getValue()));
+                                 IntegerAttr::get(elemType, 0).getValue()));
     }
     if (elemType.isa<FloatType>()) {
       return builder.create<spirv::ConstantOp>(
           loc, type,
           DenseFPElementsAttr::get(vectorType,
                                    FloatAttr::get(elemType, 0.0).getValue()));
     }
   }
return builder.create<spirv::ConstantOp>(
         loc, type, builder.getFloatAttr(floatType, 1.0));
   }
   if (auto vectorType = type.dyn_cast<VectorType>()) {
     Type elemType = vectorType.getElementType();
     if (elemType.isa<IntegerType>()) {
       return builder.create<spirv::ConstantOp>(
           loc, type,
           DenseElementsAttr::get(vectorType,
-                                 IntegerAttr::get(elemType, 1.0).getValue()));
+                                 IntegerAttr::get(elemType, 1).getValue()));
     }
     if (elemType.isa<FloatType>()) {
       return builder.create<spirv::ConstantOp>(
           loc, type,
           DenseFPElementsAttr::get(vectorType,
                                    FloatAttr::get(elemType, 1.0).getValue()));
     }
   }

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

 
 // CHECK-LABEL: @sexti2
 func.func @sexti2(%arg0 : i32) -> i64 {
   // CHECK: spirv.SConvert %{{.*}} : i32 to i64
   %0 = arith.extsi %arg0 : i32 to i64
   return %0 : i64
 }
 
+// CHECK-LABEL: @sext_bool_scalar
+// CHECK-SAME:  ([[ARG:%.+]]: i1) -> i32
+func.func @sext_bool_scalar(%arg0 : i1) -> i32 {
+  // CHECK-DAG:  [[ONES:%.+]] = spirv.Constant -1 : i32
+  // CHECK-DAG:  [[ZERO:%.+]] = spirv.Constant 0 : i32
+  // CHECK:      [[SEL:%.+]]  = spirv.Select [[ARG]], [[ONES]], [[ZERO]] : i1, i32
+  // CHECK-NEXT: return [[SEL]] : i32
+  %0 = arith.extsi %arg0 : i1 to i32
+  return %0 : i32
+}
+
+// CHECK-LABEL: @sext_bool_vector
+// CHECK-SAME:  ([[ARG:%.+]]: vector<3xi1>) -> vector<3xi32>
+func.func @sext_bool_vector(%arg0 : vector<3xi1>) -> vector<3xi32> {
+  // CHECK-DAG:  [[ONES:%.+]] = spirv.Constant dense<-1> : vector<3xi32>
+  // CHECK-DAG:  [[ZERO:%.+]] = spirv.Constant dense<0> : vector<3xi32>
+  // CHECK:      [[SEL:%.+]]  = spirv.Select [[ARG]], [[ONES]], [[ZERO]] : vector<3xi1>, vector<3xi32>
+  // CHECK-NEXT: return [[SEL]] : vector<3xi32>
+  %0 = arith.extsi %arg0 : vector<3xi1> to vector<3xi32>
+  return %0 : vector<3xi32>
+}
+
 // CHECK-LABEL: @zexti1
 func.func @zexti1(%arg0: i16) -> i64 {
   // CHECK: spirv.UConvert %{{.*}} : i16 to i64
   %0 = arith.extui %arg0 : i16 to i64
   return %0 : i64
 }
 
 // CHECK-LABEL: @zexti2