[mlir][spirv] Support size-1 vector.insert during conversion

[antiagainst]

Summary

Size-1 vectors are converted into scalars in SPIR-V so these ops
go away entirely.

Along the way, update the pass to perform partial conversion
so that we can have more focused tests.

[mravishankar]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

161

This looks like a canonicalization at vector dialect then a something to do in SPIR-V. Or do the normal thing in SPIR-V and canonicalize in SPIR-V dialect. Seems off to do this during conversion.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp

41 ↗

(On Diff #393467)

Is this change related somehow. I dont see the correlation

[antiagainst]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

161

Crossing the vector-scalar boundary doesn't seem like a generic transformation to me that we'd want to put in vector dialect as a canonicalization pattern. Vector dialect should be working on vectors, even they are of size one. Going to scalar is not always wanted at that level, given that they won't work anymore with other vector-level ops.

I consider going to scalars as lowering so it makes sense to happen during vector to SPIR-V conversion time to me. (Note that size-1 vector is valid vector and valid input to SPIR-V conversion; the fact that we are lowering them to scalars is SPIR-V's particular requirements.)

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp

41 ↗

(On Diff #393467)

This is needed for cleaning up the tests. I updated the patch comment to be clear. :)

[antiagainst]

Move pass and test change to its own patch: https://reviews.llvm.org/D115756

[antiagainst]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRVPass.cpp

41 ↗

(On Diff #393467)

Actually, I split this out as its own patch: https://reviews.llvm.org/D115756.

[mravishankar]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

160

Nit: type s/Speical/Special

[mravishankar]

I wonder if this will lead to unrealized conversion casts not being able to be eliminated since the uses might actually need a vector.

[antiagainst]

In D115517#3258606, @mravishankar wrote:

I wonder if this will lead to unrealized conversion casts not being able to be eliminated since the uses might actually need a vector.

Size-1 vector are converted to scalars by the SPIR-V type converter globally. So shouldn't be a big problem. If it happens, it's an indication a missing op conversion.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

 
 struct VectorInsertOpConvert final
     : public OpConversionPattern<vector::InsertOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::InsertOp insertOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
+    // Special case for inserting scalar values into size-1 vectors.
+    if (insertOp.getSourceType().isIntOrFloat() &&
+        insertOp.getDestVectorType().getNumElements() == 1) {
+      rewriter.replaceOp(insertOp, adaptor.source());
+      return success();
+    }
+
     if (insertOp.getSourceType().isa<VectorType>() ||
         !spirv::CompositeType::isValid(insertOp.getDestVectorType()))
       return failure();
     int32_t id = getFirstIntValue(insertOp.position());
     rewriter.replaceOpWithNewOp<spirv::CompositeInsertOp>(
         insertOp, adaptor.source(), adaptor.dest(), id);
     return success();
   }
struct VectorInsertStridedSliceOpConvert final
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::InsertStridedSliceOp insertOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Value srcVector = adaptor.getOperands().front();
     Value dstVector = adaptor.getOperands().back();
 
-    // Insert scalar values not supported yet.
-    if (srcVector.getType().isa<spirv::ScalarType>() ||
-        dstVector.getType().isa<spirv::ScalarType>())
-      return failure();
-
     uint64_t stride = getFirstIntValue(insertOp.strides());
     if (stride != 1)
       return failure();
+    uint64_t offset = getFirstIntValue(insertOp.offsets());
+
+    if (srcVector.getType().isa<spirv::ScalarType>()) {
+      assert(!dstVector.getType().isa<spirv::ScalarType>());
+      rewriter.replaceOpWithNewOp<spirv::CompositeInsertOp>(
+          insertOp, dstVector.getType(), srcVector, dstVector,
+          rewriter.getI32ArrayAttr(offset));
+      return success();
+    }
 
     uint64_t totalSize =
         dstVector.getType().cast<VectorType>().getNumElements();
     uint64_t insertSize =
         srcVector.getType().cast<VectorType>().getNumElements();
-    uint64_t offset = getFirstIntValue(insertOp.offsets());
 
     SmallVector<int32_t, 2> indices(totalSize);
     std::iota(indices.begin(), indices.end(), 0);
     std::iota(indices.begin() + offset, indices.begin() + offset + insertSize,
               totalSize);
 
     rewriter.replaceOpWithNewOp<spirv::VectorShuffleOp>(
         insertOp, dstVector.getType(), dstVector, srcVector,

mlir/test/Conversion/VectorToSPIRV/simple.mlir

 //       CHECK:   spv.CompositeInsert %[[S]], %[[V]][2 : i32] : f32 into vector<4xf32>
 func @insert(%arg0 : vector<4xf32>, %arg1: f32) -> vector<4xf32> {
   %1 = vector.insert %arg1, %arg0[2] : f32 into vector<4xf32>
   return %1: vector<4xf32>
 }
 
 // -----
 
+// CHECK-LABEL: @insert_size1_vector
+//  CHECK-SAME: %[[V:.*]]: vector<1xf32>, %[[S:.*]]: f32
+//       CHECK:   %[[R:.+]] = builtin.unrealized_conversion_cast %[[S]]
+//       CHECK:   return %[[R]]
+func @insert_size1_vector(%arg0 : vector<1xf32>, %arg1: f32) -> vector<1xf32> {
+  %1 = vector.insert %arg1, %arg0[0] : f32 into vector<1xf32>
+  return %1 : vector<1xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @extract_element
 //  CHECK-SAME: %[[V:.*]]: vector<4xf32>, %[[ID:.*]]: i32
 //       CHECK:   spv.VectorExtractDynamic %[[V]][%[[ID]]] : vector<4xf32>, i32
 func @extract_element(%arg0 : vector<4xf32>, %id : i32) -> f32 {
   %0 = vector.extractelement %arg0[%id : i32] : vector<4xf32>
   return %0: f32
 }
 

 //       CHECK:   spv.VectorShuffle [0 : i32, 4 : i32, 5 : i32, 3 : i32] %[[ALL]] : vector<4xf32>, %[[PART]] : vector<2xf32> -> vector<4xf32>
 func @insert_strided_slice(%arg0: vector<2xf32>, %arg1: vector<4xf32>) -> vector<4xf32> {
   %0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [1], strides = [1]} : vector<2xf32> into vector<4xf32>
   return %0 : vector<4xf32>
 }
 
 // -----
 
+// CHECK-LABEL: @insert_size1_vector
+//  CHECK-SAME: %[[SUB:.*]]: vector<1xf32>, %[[FULL:.*]]: vector<3xf32>
+//       CHECK:   %[[S:.+]] = builtin.unrealized_conversion_cast %[[SUB]]
+//       CHECK:   spv.CompositeInsert %[[S]], %[[FULL]][2 : i32] : f32 into vector<3xf32>
+func @insert_size1_vector(%arg0 : vector<1xf32>, %arg1: vector<3xf32>) -> vector<3xf32> {
+  %1 = vector.insert_strided_slice %arg0, %arg1 {offsets = [2], strides = [1]} : vector<1xf32> into vector<3xf32>
+  return %1 : vector<3xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @fma
 //  CHECK-SAME: %[[A:.*]]: vector<4xf32>, %[[B:.*]]: vector<4xf32>, %[[C:.*]]: vector<4xf32>
 //       CHECK:   spv.GLSL.Fma %[[A]], %[[B]], %[[C]] : vector<4xf32>
 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>
 }