[mlir][spirv] Improve vector extract/insert element conversion

[antiagainst]

Summary

• Fix type conversions around positions--we need to use the converted value from the adaptor.
• Convert constant position cases to composite extract/insert.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

 
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include <numeric>
 
 using namespace mlir;
 
 /// Gets the first integer value from `attr`, assuming it is an integer array

 
 struct VectorExtractElementOpConvert final
     : public OpConversionPattern<vector::ExtractElementOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::ExtractElementOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    Type vectorType =
-        getTypeConverter()->convertType(adaptor.getVector().getType());
-    if (!vectorType)
+    Type resultType = getTypeConverter()->convertType(extractOp.getType());
+    if (!resultType)
       return failure();
 
-    if (vectorType.isa<spirv::ScalarType>()) {
+    if (adaptor.getVector().getType().isa<spirv::ScalarType>()) {
       rewriter.replaceOp(extractOp, adaptor.getVector());
       return success();
     }
 
+    APInt cstPos;
+    if (matchPattern(adaptor.getPosition(), m_ConstantInt(&cstPos)))
+      rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
+          extractOp, resultType, adaptor.getVector(),
+          rewriter.getI32ArrayAttr({static_cast<int>(cstPos.getSExtValue())}));
+    else
       rewriter.replaceOpWithNewOp<spirv::VectorExtractDynamicOp>(
-        extractOp, extractOp.getType(), adaptor.getVector(),
-        extractOp.getPosition());
+          extractOp, resultType, adaptor.getVector(), adaptor.getPosition());
     return success();
   }
 };
 
 struct VectorInsertElementOpConvert final
     : public OpConversionPattern<vector::InsertElementOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::InsertElementOp insertOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type vectorType = getTypeConverter()->convertType(insertOp.getType());
     if (!vectorType)
       return failure();
 
     if (vectorType.isa<spirv::ScalarType>()) {
       rewriter.replaceOp(insertOp, adaptor.getSource());
       return success();
     }
 
+    APInt cstPos;
+    if (matchPattern(adaptor.getPosition(), m_ConstantInt(&cstPos)))
+      rewriter.replaceOpWithNewOp<spirv::CompositeInsertOp>(
+          insertOp, adaptor.getSource(), adaptor.getDest(),
+          cstPos.getSExtValue());
+    else
       rewriter.replaceOpWithNewOp<spirv::VectorInsertDynamicOp>(
           insertOp, vectorType, insertOp.getDest(), adaptor.getSource(),
-        insertOp.getPosition());
+          adaptor.getPosition());
     return success();
   }
 };
 
 struct VectorInsertStridedSliceOpConvert final
     : public OpConversionPattern<vector::InsertStridedSliceOp> {
   using OpConversionPattern::OpConversionPattern;
 

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

 //       CHECK:   spirv.VectorExtractDynamic %[[V]][%[[ID]]] : vector<4xf32>, i32
 func.func @extract_element(%arg0 : vector<4xf32>, %id : i32) -> f32 {
   %0 = vector.extractelement %arg0[%id : i32] : vector<4xf32>
   return %0: f32
 }
 
 // -----
 
+// CHECK-LABEL: @extract_element_cst
+//  CHECK-SAME: %[[V:.*]]: vector<4xf32>
+//       CHECK:   spirv.CompositeExtract %[[V]][1 : i32] : vector<4xf32>
+func.func @extract_element_cst(%arg0 : vector<4xf32>) -> f32 {
+  %idx = arith.constant 1 : i32
+  %0 = vector.extractelement %arg0[%idx : i32] : vector<4xf32>
+  return %0: f32
+}
+
+// -----
+
 // CHECK-LABEL: @extract_element_index
 func.func @extract_element_index(%arg0 : vector<4xf32>, %id : index) -> f32 {
-  // CHECK: vector.extractelement
+  // CHECK: spirv.VectorExtractDynamic
   %0 = vector.extractelement %arg0[%id : index] : vector<4xf32>
   return %0: f32
 }
 
 // -----
 
 // CHECK-LABEL: @extract_element_size5_vector
 func.func @extract_element_size5_vector(%arg0 : vector<5xf32>, %id : i32) -> f32 {

 //       CHECK:   spirv.VectorInsertDynamic %[[VAL]], %[[V]][%[[ID]]] : vector<4xf32>, i32
 func.func @insert_element(%val: f32, %arg0 : vector<4xf32>, %id : i32) -> vector<4xf32> {
   %0 = vector.insertelement %val, %arg0[%id : i32] : vector<4xf32>
   return %0: vector<4xf32>
 }
 
 // -----
 
+// CHECK-LABEL: @insert_element_cst
+//  CHECK-SAME: %[[VAL:.*]]: f32, %[[V:.*]]: vector<4xf32>
+//       CHECK:   spirv.CompositeInsert %[[VAL]], %[[V]][2 : i32] : f32 into vector<4xf32>
+func.func @insert_element_cst(%val: f32, %arg0 : vector<4xf32>) -> vector<4xf32> {
+  %idx = arith.constant 2 : i32
+  %0 = vector.insertelement %val, %arg0[%idx : i32] : vector<4xf32>
+  return %0: vector<4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @insert_element_index
 func.func @insert_element_index(%val: f32, %arg0 : vector<4xf32>, %id : index) -> vector<4xf32> {
-  // CHECK: vector.insertelement
+  // CHECK: spirv.VectorInsertDynamic
   %0 = vector.insertelement %val, %arg0[%id : index] : vector<4xf32>
   return %0: vector<4xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @insert_element_size5_vector
 func.func @insert_element_size5_vector(%val: f32, %arg0 : vector<5xf32>, %id : i32) -> vector<5xf32> {