[mlir][StandardToSPIRV] Add support for lowering uitofp to SPIR-V

[hanchung]

Summary

• Extend spirv::ConstantOp::getZero/One to handle float, vector of int, and vector of float.
• Refactor ZeroExtendI1Pattern to use getZero/One methods.
• Add one more test for lowering std.zexti which extends vector<4xi1> to vector<4xi64>.

[antiagainst]

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

543

We should generate spv.constant here?

Actually for spirv::ConstantOp we have utility methods getZero and getOne, but now it only supports scalar integer. Could you extend those methods to support float and vectors? So we don't need to have this logic duplicated in all patterns needing zero/one (which is quite common).

[hanchung]

Extend getZero and getOne to support floating-point, vector of int/float.

[hanchung]

I found that we have special handling for boolean, but I couldn't test it with vector<i1>. I think I can add another pattern to truncate a vector of integer to vector<*xi1> (which is missing) and cover the test there.

[hanchung]

do not use 'else' after 'return'

mlir/lib/Conversion/StandardToSPIRV/StandardToSPIRV.cpp

matchAndRewrite(ZeroExtendIOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto srcType = operands.front().getType();
     if (!isBoolScalarOrVector(srcType))
       return failure();
 
     auto dstType =
         this->getTypeConverter()->convertType(op.getResult().getType());
     Location loc = op.getLoc();
-    Attribute zeroAttr, oneAttr;
-    if (auto vectorType = dstType.dyn_cast<VectorType>()) {
-      zeroAttr = DenseElementsAttr::get(vectorType, 0);
-      oneAttr = DenseElementsAttr::get(vectorType, 1);
-    } else {
-      zeroAttr = IntegerAttr::get(dstType, 0);
-      oneAttr = IntegerAttr::get(dstType, 1);
-    }
-    Value zero = rewriter.create<ConstantOp>(loc, zeroAttr);
-    Value one = rewriter.create<ConstantOp>(loc, oneAttr);
+    Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
+    Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
+    rewriter.template replaceOpWithNewOp<spirv::SelectOp>(
+        op, dstType, operands.front(), one, zero);
+    return success();
+};
+
+/// Converts std.uitofp to spv.Select if the type of source is i1 or vector of
+/// i1.
+class UIToFPI1Pattern final : public OpConversionPattern<UIToFPOp> {
+public:
+  using OpConversionPattern<UIToFPOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(UIToFPOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto srcType = operands.front().getType();
+    if (!isBoolScalarOrVector(srcType))
+      return failure();
+
+    auto dstType =
+        this->getTypeConverter()->convertType(op.getResult().getType());
+    Location loc = op.getLoc();
+    Value zero = spirv::ConstantOp::getZero(dstType, loc, rewriter);
+    Value one = spirv::ConstantOp::getOne(dstType, loc, rewriter);
     rewriter.template replaceOpWithNewOp<spirv::SelectOp>(
         op, dstType, operands.front(), one, zero);
     return success();
   }
 };
 
 /// Converts type-casting standard operations to SPIR-V operations.
 template <typename StdOp, typename SPIRVOp>
matchAndRewrite(StdOp operation, ArrayRef<Value> operands,
     } else {
       rewriter.template replaceOpWithNewOp<SPIRVOp>(operation, dstType,
                                                     operands);
     }
     return success();
   }
 };
 
 /// Converts std.xor to SPIR-V operations.
 class XOrOpPattern final : public OpConversionPattern<XOrOp> {
 public:
   using OpConversionPattern<XOrOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(XOrOp xorOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
patterns.insert<
 
       // Memory patterns
       AllocOpPattern, DeallocOpPattern, IntLoadOpPattern, IntStoreOpPattern,
       LoadOpPattern, StoreOpPattern,
 
       ReturnOpPattern, SelectOpPattern,
 
       // Type cast patterns
-      ZeroExtendI1Pattern, TypeCastingOpPattern<IndexCastOp, spirv::SConvertOp>,
+      UIToFPI1Pattern, ZeroExtendI1Pattern,
+      TypeCastingOpPattern<IndexCastOp, spirv::SConvertOp>,
       TypeCastingOpPattern<SIToFPOp, spirv::ConvertSToFOp>,
+      TypeCastingOpPattern<UIToFPOp, spirv::ConvertUToFOp>,
       TypeCastingOpPattern<ZeroExtendIOp, spirv::UConvertOp>,
       TypeCastingOpPattern<TruncateIOp, spirv::SConvertOp>,
       TypeCastingOpPattern<FPToSIOp, spirv::ConvertFToSOp>,
       TypeCastingOpPattern<FPExtOp, spirv::FConvertOp>,
       TypeCastingOpPattern<FPTruncOp, spirv::FConvertOp>>(typeConverter,
                                                           context);
 }
 } // namespace mlir

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

 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/CallInterfaces.h"
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/bit.h"
 
 using namespace mlir;
 
 // TODO: generate these strings using ODS.
 static constexpr const char kMemoryAccessAttrName[] = "memory_access";
 static constexpr const char kSourceMemoryAccessAttrName[] =
spirv::ConstantOp spirv::ConstantOp::getZero(Type type, Location loc,
   if (auto intType = type.dyn_cast<IntegerType>()) {
     unsigned width = intType.getWidth();
     if (width == 1)
       return builder.create<spirv::ConstantOp>(loc, type,
                                                builder.getBoolAttr(false));
     return builder.create<spirv::ConstantOp>(
         loc, type, builder.getIntegerAttr(type, APInt(width, 0)));
   }
+  if (auto floatType = type.dyn_cast<FloatType>()) {
+    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()));
+    }
+    if (elemType.isa<FloatType>()) {
+      return builder.create<spirv::ConstantOp>(
+          loc, type,
+          DenseFPElementsAttr::get(vectorType,
+                                   FloatAttr::get(elemType, 0.0).getValue()));
+    }
+  }
 
   llvm_unreachable("unimplemented types for ConstantOp::getZero()");
 }
 
 spirv::ConstantOp spirv::ConstantOp::getOne(Type type, Location loc,
                                             OpBuilder &builder) {
   if (auto intType = type.dyn_cast<IntegerType>()) {
     unsigned width = intType.getWidth();
     if (width == 1)
       return builder.create<spirv::ConstantOp>(loc, type,
                                                builder.getBoolAttr(true));
     return builder.create<spirv::ConstantOp>(
         loc, type, builder.getIntegerAttr(type, APInt(width, 1)));
   }
+  if (auto floatType = type.dyn_cast<FloatType>()) {
+    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()));
+    }
+    if (elemType.isa<FloatType>()) {
+      return builder.create<spirv::ConstantOp>(
+          loc, type,
+          DenseFPElementsAttr::get(vectorType,
+                                   FloatAttr::get(elemType, 1.0).getValue()));
+    }
+  }
 
   llvm_unreachable("unimplemented types for ConstantOp::getOne()");
 }
 
 //===----------------------------------------------------------------------===//
 // spv.EntryPoint
 //===----------------------------------------------------------------------===//
 

mlir/test/Conversion/StandardToSPIRV/std-ops-to-spirv.mlir

 
 // CHECK-LABEL: @sitofp2
 func @sitofp2(%arg0 : i64) -> f64 {
   // CHECK: spv.ConvertSToF %{{.*}} : i64 to f64
   %0 = std.sitofp %arg0 : i64 to f64
   return %0 : f64
 }
 
+// CHECK-LABEL: @uitofp_i16_f32
+func @uitofp_i16_f32(%arg0: i16) -> f32 {
+  // CHECK: spv.ConvertUToF %{{.*}} : i16 to f32
+  %0 = std.uitofp %arg0 : i16 to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: @uitofp_i32_f32
+func @uitofp_i32_f32(%arg0 : i32) -> f32 {
+  // CHECK: spv.ConvertUToF %{{.*}} : i32 to f32
+  %0 = std.uitofp %arg0 : i32 to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: @uitofp_i1_f32
+func @uitofp_i1_f32(%arg0 : i1) -> f32 {
+  // CHECK: %[[ZERO:.+]] = spv.constant 0.000000e+00 : f32
+  // CHECK: %[[ONE:.+]] = spv.constant 1.000000e+00 : f32
+  // CHECK: spv.Select %{{.*}}, %[[ONE]], %[[ZERO]] : i1, f32
+  %0 = std.uitofp %arg0 : i1 to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: @uitofp_i1_f64
+func @uitofp_i1_f64(%arg0 : i1) -> f64 {
+  // CHECK: %[[ZERO:.+]] = spv.constant 0.000000e+00 : f64
+  // CHECK: %[[ONE:.+]] = spv.constant 1.000000e+00 : f64
+  // CHECK: spv.Select %{{.*}}, %[[ONE]], %[[ZERO]] : i1, f64
+  %0 = std.uitofp %arg0 : i1 to f64
+  return %0 : f64
+}
+
+// CHECK-LABEL: @uitofp_vec_i1_f32
+func @uitofp_vec_i1_f32(%arg0 : vector<4xi1>) -> vector<4xf32> {
+  // CHECK: %[[ZERO:.+]] = spv.constant dense<0.000000e+00> : vector<4xf32>
+  // CHECK: %[[ONE:.+]] = spv.constant dense<1.000000e+00> : vector<4xf32>
+  // CHECK: spv.Select %{{.*}}, %[[ONE]], %[[ZERO]] : vector<4xi1>, vector<4xf32>
+  %0 = std.uitofp %arg0 : vector<4xi1> to vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
+// CHECK-LABEL: @uitofp_vec_i1_f64
+spv.func @uitofp_vec_i1_f64(%arg0: vector<4xi1>) -> vector<4xf64> "None" {
+  // CHECK: %[[ZERO:.+]] = spv.constant dense<0.000000e+00> : vector<4xf64>
+  // CHECK: %[[ONE:.+]] = spv.constant dense<1.000000e+00> : vector<4xf64>
+  // CHECK: spv.Select %{{.*}}, %[[ONE]], %[[ZERO]] : vector<4xi1>, vector<4xf64>
+  %0 = spv.constant dense<0.000000e+00> : vector<4xf64>
+  %1 = spv.constant dense<1.000000e+00> : vector<4xf64>
+  %2 = spv.Select %arg0, %1, %0 : vector<4xi1>, vector<4xf64>
+  spv.ReturnValue %2 : vector<4xf64>
+}
+
 // CHECK-LABEL: @zexti1
 func @zexti1(%arg0: i16) -> i64 {
   // CHECK: spv.UConvert %{{.*}} : i16 to i64
   %0 = std.zexti %arg0 : i16 to i64
   return %0 : i64
 }
 
 // CHECK-LABEL: @zexti2

 func @zexti4(%arg0 : vector<4xi1>) -> vector<4xi32> {
   // CHECK: %[[ZERO:.+]] = spv.constant dense<0> : vector<4xi32>
   // CHECK: %[[ONE:.+]] = spv.constant dense<1> : vector<4xi32>
   // CHECK: spv.Select %{{.*}}, %[[ONE]], %[[ZERO]] : vector<4xi1>, vector<4xi32>
   %0 = std.zexti %arg0 : vector<4xi1> to vector<4xi32>
   return %0 : vector<4xi32>
 }
 
+// CHECK-LABEL: @zexti5
+func @zexti5(%arg0 : vector<4xi1>) -> vector<4xi64> {
+  // CHECK: %[[ZERO:.+]] = spv.constant dense<0> : vector<4xi64>
+  // CHECK: %[[ONE:.+]] = spv.constant dense<1> : vector<4xi64>
+  // CHECK: spv.Select %{{.*}}, %[[ONE]], %[[ZERO]] : vector<4xi1>, vector<4xi64>
+  %0 = std.zexti %arg0 : vector<4xi1> to vector<4xi64>
+  return %0 : vector<4xi64>
+}
+
 // CHECK-LABEL: @trunci1
 func @trunci1(%arg0 : i64) -> i16 {
   // CHECK: spv.SConvert %{{.*}} : i64 to i16
   %0 = std.trunci %arg0 : i64 to i16
   return %0 : i16
 }
 
 // CHECK-LABEL: @trunci2