[mlir][spirv] Support size-1 vector/tensor constant during conversion

[antiagainst]

Summary

Depends On D115517

[mravishankar]

Minor comment. LGTM otherwise.

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

370

IIUC its part of the constant op verification that the type of attribute and type of the result are the same?

[antiagainst]

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

370

Yup I think so for arith::ConstantOp.

mlir/lib/Conversion/ArithmeticToSPIRV/ArithmeticToSPIRV.cpp

 //===----------------------------------------------------------------------===//
 // ConstantOp with composite type
 //===----------------------------------------------------------------------===//
 
 LogicalResult ConstantCompositeOpPattern::matchAndRewrite(
     arith::ConstantOp constOp, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   auto srcType = constOp.getType().dyn_cast<ShapedType>();
-  if (!srcType)
+  if (!srcType || srcType.getNumElements() == 1)
     return failure();
 
   // arith.constant should only have vector or tenor types.
   assert((srcType.isa<VectorType, RankedTensorType>()));
 
   auto dstType = getTypeConverter()->convertType(srcType);
   if (!dstType)
     return failure();

 //===----------------------------------------------------------------------===//
 // ConstantOp with scalar type
 //===----------------------------------------------------------------------===//
 
 LogicalResult ConstantScalarOpPattern::matchAndRewrite(
     arith::ConstantOp constOp, OpAdaptor adaptor,
     ConversionPatternRewriter &rewriter) const {
   Type srcType = constOp.getType();
+  if (auto shapedType = srcType.dyn_cast<ShapedType>()) {
+    if (shapedType.getNumElements() != 1)
+      return failure();
+    srcType = shapedType.getElementType();
+  }
   if (!srcType.isIntOrIndexOrFloat())
     return failure();
 
+  Attribute cstAttr = constOp.getValue();
+  if (cstAttr.getType().isa<ShapedType>())
+    cstAttr = cstAttr.cast<DenseElementsAttr>().getSplatValue<Attribute>();
+
   Type dstType = getTypeConverter()->convertType(srcType);
   if (!dstType)
     return failure();
 
   // Floating-point types.
   if (srcType.isa<FloatType>()) {
-    auto srcAttr = constOp.getValue().cast<FloatAttr>();
+    auto srcAttr = cstAttr.cast<FloatAttr>();
     auto dstAttr = srcAttr;
 
     // Floating-point types not supported in the target environment are all
     // converted to float type.
     if (srcType != dstType) {
       dstAttr = convertFloatAttr(srcAttr, dstType.cast<FloatType>(), rewriter);
       if (!dstAttr)
         return failure();
     }
 
     rewriter.replaceOpWithNewOp<spirv::ConstantOp>(constOp, dstType, dstAttr);
     return success();
   }
 
   // Bool type.
   if (srcType.isInteger(1)) {
     // arith.constant can use 0/1 instead of true/false for i1 values. We need
     // to handle that here.
-    auto dstAttr = convertBoolAttr(constOp.getValue(), rewriter);
+    auto dstAttr = convertBoolAttr(cstAttr, rewriter);
     if (!dstAttr)
       return failure();
     rewriter.replaceOpWithNewOp<spirv::ConstantOp>(constOp, dstType, dstAttr);
     return success();
   }
 
   // IndexType or IntegerType. Index values are converted to 32-bit integer
   // values when converting to SPIR-V.
-  auto srcAttr = constOp.getValue().cast<IntegerAttr>();
+  auto srcAttr = cstAttr.cast<IntegerAttr>();
   auto dstAttr =
       convertIntegerAttr(srcAttr, dstType.cast<IntegerType>(), rewriter);
   if (!dstAttr)
     return failure();
   rewriter.replaceOpWithNewOp<spirv::ConstantOp>(constOp, dstType, dstAttr);
   return success();
 }
 

mlir/test/Conversion/ArithmeticToSPIRV/arithmetic-to-spirv.mlir

func @constant_64bit() {
   %1 = arith.constant 5.0 : f64
   // CHECK: spv.Constant dense<[2, 3]> : vector<2xi64>
   %2 = arith.constant dense<[2, 3]> : vector<2xi64>
   // CHECK: spv.Constant dense<4.000000e+00> : tensor<5xf64> : !spv.array<5 x f64, stride=8>
   %3 = arith.constant dense<4.0> : tensor<5xf64>
   return
 }
 
+// CHECK-LABEL: @constant_size1
+func @constant_size1() {
+  // CHECK: spv.Constant true
+  %0 = arith.constant dense<true> : tensor<1xi1>
+  // CHECK: spv.Constant 4 : i64
+  %1 = arith.constant dense<4> : vector<1xi64>
+  // CHECK: spv.Constant 5.000000e+00 : f64
+  %2 = arith.constant dense<5.0> : tensor<1xf64>
+  return
+}
+
 } // end module
 
 // -----
 
 // Check that constants are converted to 32-bit when no special capability.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, {}>
 } {
func @constant_64bit() {
   %2 = arith.constant dense<[2, 3]> : vector<2xi64>
   // CHECK: spv.Constant dense<4.000000e+00> : tensor<5xf32> : !spv.array<5 x f32, stride=4>
   %3 = arith.constant dense<4.0> : tensor<5xf64>
   // CHECK: spv.Constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : tensor<4xf32> : !spv.array<4 x f32, stride=4>
   %4 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf16>
   return
 }
 
+// CHECK-LABEL: @constant_size1
+func @constant_size1() {
+  // CHECK: spv.Constant 4 : i32
+  %0 = arith.constant dense<4> : vector<1xi64>
+  // CHECK: spv.Constant 5.000000e+00 : f32
+  %1 = arith.constant dense<5.0> : tensor<1xf64>
+  return
+}
+
 // CHECK-LABEL: @corner_cases
 func @corner_cases() {
   // CHECK: %{{.*}} = spv.Constant -1 : i32
   %0 = arith.constant 4294967295  : i64 // 2^32 - 1
   // CHECK: %{{.*}} = spv.Constant 2147483647 : i32
   %1 = arith.constant 2147483647  : i64 // 2^31 - 1
   // CHECK: %{{.*}} = spv.Constant -2147483648 : i32
   %2 = arith.constant 2147483648  : i64 // 2^31