[mlir][spirv] Convert single element vector.splat/fma

[antiagainst]

Update

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

 };
 
 struct VectorFmaOpConvert final : public OpConversionPattern<vector::FMAOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::FMAOp fmaOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (!spirv::CompositeType::isValid(fmaOp.getVectorType()))
+    Type dstType = getTypeConverter()->convertType(fmaOp.getType());
+    if (!dstType)
       return failure();
     rewriter.replaceOpWithNewOp<spirv::GLSLFmaOp>(
-        fmaOp, fmaOp.getType(), adaptor.getLhs(), adaptor.getRhs(),
-        adaptor.getAcc());
+        fmaOp, dstType, adaptor.getLhs(), adaptor.getRhs(), adaptor.getAcc());
     return success();
   }
 };
 
 struct VectorInsertOpConvert final
     : public OpConversionPattern<vector::InsertOp> {
   using OpConversionPattern::OpConversionPattern;
 

 
 class VectorSplatPattern final : public OpConversionPattern<vector::SplatOp> {
 public:
   using OpConversionPattern<vector::SplatOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::SplatOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    VectorType dstVecType = op.getType();
-    if (!spirv::CompositeType::isValid(dstVecType))
+    Type dstType = getTypeConverter()->convertType(op.getType());
+    if (!dstType)
       return failure();
+    if (dstType.isa<spirv::ScalarType>()) {
+      rewriter.replaceOp(op, adaptor.getInput());
+    } else {
+      auto dstVecType = dstType.cast<VectorType>();
       SmallVector<Value, 4> source(dstVecType.getNumElements(),
                                    adaptor.getInput());
-    rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(op, dstVecType,
+      rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(op, dstType,
                                                                source);
+    }
     return success();
   }
 };
 
 struct VectorShuffleOpConvert final
     : public OpConversionPattern<vector::ShuffleOp> {
   using OpConversionPattern::OpConversionPattern;
 

mlir/test/Conversion/VectorToSPIRV/vector-to-spirv.mlir

 //       CHECK:   spv.GLSL.Fma %[[A]], %[[B]], %[[C]] : vector<4xf32>
 func.func @fma(%a: vector<4xf32>, %b: vector<4xf32>, %c: vector<4xf32>) -> vector<4xf32> {
   %0 = vector.fma %a, %b, %c: vector<4xf32>
   return %0 : vector<4xf32>
 }
 
 // -----
 
+// CHECK-LABEL: @fma_size1_vector
+//       CHECK:   spv.GLSL.Fma %{{.+}} : f32
+func.func @fma_size1_vector(%a: vector<1xf32>, %b: vector<1xf32>, %c: vector<1xf32>) -> vector<1xf32> {
+  %0 = vector.fma %a, %b, %c: vector<1xf32>
+  return %0 : vector<1xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @splat
 //  CHECK-SAME: (%[[A:.+]]: f32)
 //       CHECK:   %[[VAL:.+]] = spv.CompositeConstruct %[[A]], %[[A]], %[[A]], %[[A]] : vector<4xf32>
 //       CHECK:   return %[[VAL]]
 func.func @splat(%f : f32) -> vector<4xf32> {
   %splat = vector.splat %f : vector<4xf32>
   return %splat : vector<4xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @splat_size1_vector
+//  CHECK-SAME: (%[[A:.+]]: f32)
+//       CHECK:   %[[VAL:.+]] = builtin.unrealized_conversion_cast %[[A]]
+//       CHECK:   return %[[VAL]]
+func.func @splat_size1_vector(%f : f32) -> vector<1xf32> {
+  %splat = vector.splat %f : vector<1xf32>
+  return %splat : vector<1xf32>
+}
+
+// -----
+
 // CHECK-LABEL:  func @shuffle
 //  CHECK-SAME:  %[[ARG0:.+]]: vector<1xf32>, %[[ARG1:.+]]: vector<1xf32>
 //       CHECK:    %[[V0:.+]] = builtin.unrealized_conversion_cast %[[ARG0]]
 //       CHECK:    %[[V1:.+]] = builtin.unrealized_conversion_cast %[[ARG1]]
 //       CHECK:    spv.CompositeConstruct %[[V0]], %[[V1]], %[[V1]], %[[V0]] : vector<4xf32>
 func.func @shuffle(%v0 : vector<1xf32>, %v1: vector<1xf32>) -> vector<4xf32> {
   %shuffle = vector.shuffle %v0, %v1 [0, 1, 1, 0] : vector<1xf32>, vector<1xf32>
   return %shuffle : vector<4xf32>