[mlir][spirv] Convert some 0-D vector extract/insertelement ops

[antiagainst]

reorder

[kuhar]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

65–66

nit: Type resultType = ? In D132613, @Mogball suggested using Type instead of auto in this case, but I'm not sure how much consensus there is on this

Also in the other matchAndRewrite below.

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

125

why not .+?

[antiagainst]

Address comments

[antiagainst]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

65–66

Yeah, I think from the code it's clearly some Type. So not sure how much additional benefits it brings. Though given auto and Type are both 4 chars, I'm fine switching to Type since it's not longer.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

 
 struct VectorBitcastConvert final
     : public OpConversionPattern<vector::BitCastOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::BitCastOp bitcastOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto dstType = getTypeConverter()->convertType(bitcastOp.getType());
+    Type dstType = getTypeConverter()->convertType(bitcastOp.getType());
     if (!dstType)
       return failure();
 
     if (dstType == adaptor.getSource().getType())
       rewriter.replaceOp(bitcastOp, adaptor.getSource());
     else
       rewriter.replaceOpWithNewOp<spirv::BitcastOp>(bitcastOp, dstType,
                                                     adaptor.getSource());
 
     return success();
   }
 };
 
 struct VectorBroadcastConvert final
     : public OpConversionPattern<vector::BroadcastOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(vector::BroadcastOp broadcastOp, OpAdaptor adaptor,
+  matchAndRewrite(vector::BroadcastOp castOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (broadcastOp.getSource().getType().isa<VectorType>() ||
-        !spirv::CompositeType::isValid(broadcastOp.getVectorType()))
+    Type resultType = getTypeConverter()->convertType(castOp.getVectorType());
+    if (!resultType)
       return failure();
-    SmallVector<Value, 4> source(broadcastOp.getVectorType().getNumElements(),
+
+    if (resultType.isa<spirv::ScalarType>()) {
+      rewriter.replaceOp(castOp, adaptor.getSource());
+      return success();
+    }
+
+    SmallVector<Value, 4> source(castOp.getVectorType().getNumElements(),
                                  adaptor.getSource());
     rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
-        broadcastOp, broadcastOp.getVectorType(), source);
+        castOp, castOp.getVectorType(), source);
     return success();
   }
 };
 
 struct VectorExtractOpConvert final
     : public OpConversionPattern<vector::ExtractOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Only support extracting a scalar value now.
     VectorType resultVectorType = extractOp.getType().dyn_cast<VectorType>();
     if (resultVectorType && resultVectorType.getNumElements() > 1)
       return failure();
 
-    auto dstType = getTypeConverter()->convertType(extractOp.getType());
+    Type dstType = getTypeConverter()->convertType(extractOp.getType());
     if (!dstType)
       return failure();
 
     if (adaptor.getVector().getType().isa<spirv::ScalarType>()) {
       rewriter.replaceOp(extractOp, adaptor.getVector());
       return success();
     }
 
     int32_t id = getFirstIntValue(extractOp.getPosition());
     rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
         extractOp, adaptor.getVector(), id);
     return success();
   }
 };
 
 struct VectorExtractStridedSliceOpConvert final
     : public OpConversionPattern<vector::ExtractStridedSliceOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::ExtractStridedSliceOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto dstType = getTypeConverter()->convertType(extractOp.getType());
+    Type dstType = getTypeConverter()->convertType(extractOp.getType());
     if (!dstType)
       return failure();
 
     uint64_t offset = getFirstIntValue(extractOp.getOffsets());
     uint64_t size = getFirstIntValue(extractOp.getSizes());
     uint64_t stride = getFirstIntValue(extractOp.getStrides());
     if (stride != 1)
       return failure();
struct VectorInsertOpConvert final
   }
 };
 
 struct VectorExtractElementOpConvert final
     : public OpConversionPattern<vector::ExtractElementOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(vector::ExtractElementOp extractElementOp, OpAdaptor adaptor,
+  matchAndRewrite(vector::ExtractElementOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (!spirv::CompositeType::isValid(extractElementOp.getVectorType()))
+    Type vectorType =
+        getTypeConverter()->convertType(adaptor.getVector().getType());
+    if (!vectorType)
       return failure();
+
+    if (vectorType.isa<spirv::ScalarType>()) {
+      rewriter.replaceOp(extractOp, adaptor.getVector());
+      return success();
+    }
+
     rewriter.replaceOpWithNewOp<spirv::VectorExtractDynamicOp>(
-        extractElementOp, extractElementOp.getType(), adaptor.getVector(),
-        extractElementOp.getPosition());
+        extractOp, extractOp.getType(), adaptor.getVector(),
+        extractOp.getPosition());
     return success();
   }
 };
 
 struct VectorInsertElementOpConvert final
     : public OpConversionPattern<vector::InsertElementOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(vector::InsertElementOp insertElementOp, OpAdaptor adaptor,
+  matchAndRewrite(vector::InsertElementOp insertOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (!spirv::CompositeType::isValid(insertElementOp.getDestVectorType()))
+    Type vectorType = getTypeConverter()->convertType(insertOp.getType());
+    if (!vectorType)
       return failure();
+
+    if (vectorType.isa<spirv::ScalarType>()) {
+      rewriter.replaceOp(insertOp, adaptor.getSource());
+      return success();
+    }
+
     rewriter.replaceOpWithNewOp<spirv::VectorInsertDynamicOp>(
-        insertElementOp, insertElementOp.getType(), insertElementOp.getDest(),
-        adaptor.getSource(), insertElementOp.getPosition());
+        insertOp, vectorType, insertOp.getDest(), adaptor.getSource(),
+        insertOp.getPosition());
     return success();
   }
 };
 
 struct VectorInsertStridedSliceOpConvert final
     : public OpConversionPattern<vector::InsertStridedSliceOp> {
   using OpConversionPattern::OpConversionPattern;
 
struct VectorShuffleOpConvert final
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::ShuffleOp shuffleOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto oldResultType = shuffleOp.getVectorType();
     if (!spirv::CompositeType::isValid(oldResultType))
       return failure();
-    auto newResultType = getTypeConverter()->convertType(oldResultType);
+    Type newResultType = getTypeConverter()->convertType(oldResultType);
 
     auto oldSourceType = shuffleOp.getV1VectorType();
     if (oldSourceType.getNumElements() > 1) {
       SmallVector<int32_t, 4> components = llvm::to_vector<4>(
           llvm::map_range(shuffleOp.getMask(), [](Attribute attr) -> int32_t {
             return attr.cast<IntegerAttr>().getValue().getZExtValue();
           }));
       rewriter.replaceOpWithNewOp<spirv::VectorShuffleOp>(

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

return IntegerType::get(targetEnv.getContext(), /*width=*/32,
                           intType.getSignedness());
 }
 
 /// Converts a vector `type` to a suitable type under the given `targetEnv`.
 static Type convertVectorType(const spirv::TargetEnv &targetEnv,
                               const SPIRVTypeConverter::Options &options,
                               VectorType type,
                               Optional<spirv::StorageClass> storageClass = {}) {
-  if (type.getRank() == 1 && type.getNumElements() == 1)
+  if (type.getRank() <= 1 && type.getNumElements() == 1)
     return type.getElementType();
 
   if (!spirv::CompositeType::isValid(type)) {
     // TODO: Vector types with more than four elements can be translated into
     // array types.
     LLVM_DEBUG(llvm::dbgs() << type << " illegal: > 4-element unimplemented\n");
     return nullptr;
   }

mlir/test/Conversion/FuncToSPIRV/types-to-spirv.mlir

func.func @float_vector(
   %arg0: vector<2xf16>,
   %arg1: vector<3xf64>
 ) { return }
 
 // CHECK-LABEL: spv.func @one_element_vector
 // CHECK-SAME: %{{.+}}: i32
 func.func @one_element_vector(%arg0: vector<1xi32>) { return }
 
+// CHECK-LABEL: spv.func @zerod_vector
+//  CHECK-SAME: %{{.+}}: f32
+func.func @zerod_vector(%arg0: vector<f32>) { return }
+
 } // end module
 
 // -----
 
 // Check that > 4-element vectors are not supported.
 module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, #spv.resource_limits<>>
 } {

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

 func.func @extract_element_size5_vector(%arg0 : vector<5xf32>, %id : i32) -> f32 {
   // CHECK: vector.extractelement
   %0 = vector.extractelement %arg0[%id : i32] : vector<5xf32>
   return %0: f32
 }
 
 // -----
 
+// CHECK-LABEL: @extract_element_size1_vector
+//  CHECK-SAME: (%[[S:.+]]: f32
+func.func @extract_element_size1_vector(%arg0 : f32, %i: index) -> f32 {
+  %bcast = vector.broadcast %arg0 : f32 to vector<1xf32>
+  %0 = vector.extractelement %bcast[%i : index] : vector<1xf32>
+  // CHECK: return %[[S]]
+  return %0: f32
+}
+
+// -----
+
+// CHECK-LABEL: @extract_element_0d_vector
+//  CHECK-SAME: (%[[S:.+]]: f32)
+func.func @extract_element_0d_vector(%arg0 : f32) -> f32 {
+  %bcast = vector.broadcast %arg0 : f32 to vector<f32>
+  %0 = vector.extractelement %bcast[] : vector<f32>
+  // CHECK: return %[[S]]
+  return %0: f32
+}
+
+// -----
+
 // CHECK-LABEL: @extract_strided_slice
 //  CHECK-SAME: %[[ARG:.+]]: vector<4xf32>
 //       CHECK:   spv.VectorShuffle [1 : i32, 2 : i32] %[[ARG]] : vector<4xf32>, %[[ARG]] : vector<4xf32> -> vector<2xf32>
 //       CHECK:   spv.CompositeExtract %[[ARG]][1 : i32] : vector<4xf32>
 func.func @extract_strided_slice(%arg0: vector<4xf32>) -> (vector<2xf32>, vector<1xf32>) {
   %0 = vector.extract_strided_slice %arg0 {offsets = [1], sizes = [2], strides = [1]} : vector<4xf32> to vector<2xf32>
   %1 = vector.extract_strided_slice %arg0 {offsets = [1], sizes = [1], strides = [1]} : vector<4xf32> to vector<1xf32>
   return %0, %1 : vector<2xf32>, vector<1xf32>

 func.func @insert_element_size5_vector(%val: f32, %arg0 : vector<5xf32>, %id : i32) -> vector<5xf32> {
   // CHECK: vector.insertelement
   %0 = vector.insertelement %val, %arg0[%id : i32] : vector<5xf32>
   return %0 : vector<5xf32>
 }
 
 // -----
 
+// CHECK-LABEL: @insert_element_size1_vector
+//  CHECK-SAME: (%[[S:[a-z0-9]+]]: f32
+func.func @insert_element_size1_vector(%scalar: f32, %vector : vector<1xf32>, %i: index) -> vector<1xf32> {
+  %0 = vector.insertelement %scalar, %vector[%i : index] : vector<1xf32>
+  // CHECK: %[[V:.+]] = builtin.unrealized_conversion_cast %arg0 : f32 to vector<1xf32>
+  // CHECK: return %[[V]]
+  return %0: vector<1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @insert_element_0d_vector
+//  CHECK-SAME: (%[[S:[a-z0-9]+]]: f32
+func.func @insert_element_0d_vector(%scalar: f32, %vector : vector<f32>) -> vector<f32> {
+  %0 = vector.insertelement %scalar, %vector[] : vector<f32>
+  // CHECK: %[[V:.+]] = builtin.unrealized_conversion_cast %arg0 : f32 to vector<f32>
+  // CHECK: return %[[V]]
+  return %0: vector<f32>
+}
+
+// -----
+
 // CHECK-LABEL: @insert_strided_slice
 //  CHECK-SAME: %[[PART:.+]]: vector<2xf32>, %[[ALL:.+]]: vector<4xf32>
 //       CHECK:   spv.VectorShuffle [0 : i32, 4 : i32, 5 : i32, 3 : i32] %[[ALL]] : vector<4xf32>, %[[PART]] : vector<2xf32> -> vector<4xf32>
 func.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>
 }