[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

527

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,
     Value zero = rewriter.create<ConstantOp>(loc, zeroAttr);
     Value one = rewriter.create<ConstantOp>(loc, oneAttr);
     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();
+    Attribute zeroAttr, oneAttr;
+    if (auto vectorType = dstType.dyn_cast<VectorType>()) {
+      Type elemType = vectorType.getElementType();
+      zeroAttr = DenseFPElementsAttr::get(
+          vectorType, FloatAttr::get(elemType, 0.0).getValue());
+      oneAttr = DenseFPElementsAttr::get(
+          vectorType, FloatAttr::get(elemType, 1.0).getValue());
+    } else {
+      zeroAttr = FloatAttr::get(dstType, 0.0);
+      oneAttr = FloatAttr::get(dstType, 1.0);
+    }
+    Value zero = rewriter.create<ConstantOp>(loc, zeroAttr);
+    Value one = rewriter.create<ConstantOp>(loc, oneAttr);
+    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>
 class TypeCastingOpPattern final : public OpConversionPattern<StdOp> {
 public:
   using OpConversionPattern<StdOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(StdOp operation, ArrayRef<Value> operands,
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/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: @uitofp1
+func @uitofp1(%arg0: i16) -> f32 {
+  // CHECK: spv.ConvertUToF %{{.*}} : i16 to f32
+  %0 = std.uitofp %arg0 : i16 to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: @uitofp2
+func @uitofp2(%arg0 : i32) -> f32 {
+  // CHECK: spv.ConvertUToF %{{.*}} : i32 to f32
+  %0 = std.uitofp %arg0 : i32 to f32
+  return %0 : f32
+}
+
+// CHECK-LABEL: @uitofp3
+func @uitofp3(%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: @uitofp4
+func @uitofp4(%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: @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