[mlir][spirv] Check type legality using converter for vectors

[antiagainst]

Summary

This allows index vectors to be converted to SPIR-V.

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 {
+    if (isa<VectorType>(insertOp.getSourceType()))
+      return rewriter.notifyMatchFailure(insertOp, "unsupported vector source");
+    if (!getTypeConverter()->convertType(insertOp.getDestVectorType()))
+      return rewriter.notifyMatchFailure(insertOp,
+                                         "unsupported dest vector type");
+
     // Special case for inserting scalar values into size-1 vectors.
     if (insertOp.getSourceType().isIntOrFloat() &&
         insertOp.getDestVectorType().getNumElements() == 1) {
       rewriter.replaceOp(insertOp, adaptor.getSource());
       return success();
     }
 
-    if (isa<VectorType>(insertOp.getSourceType()) ||
-        !spirv::CompositeType::isValid(insertOp.getDestVectorType()))
-      return failure();
     int32_t id = getFirstIntValue(insertOp.getPosition());
     rewriter.replaceOpWithNewOp<spirv::CompositeInsertOp>(
         insertOp, adaptor.getSource(), adaptor.getDest(), id);
     return success();
   }
 };
 
 struct VectorExtractElementOpConvert final

 struct VectorShuffleOpConvert final
     : public OpConversionPattern<vector::ShuffleOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::ShuffleOp shuffleOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto oldResultType = shuffleOp.getResultVectorType();
-    if (!spirv::CompositeType::isValid(oldResultType))
-      return failure();
     Type newResultType = getTypeConverter()->convertType(oldResultType);
+    if (!newResultType)
+      return rewriter.notifyMatchFailure(shuffleOp,
+                                         "unsupported result vector type");
 
     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 cast<IntegerAttr>(attr).getValue().getZExtValue();
           }));
       rewriter.replaceOpWithNewOp<spirv::VectorShuffleOp>(

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

 //       CHECK:   spirv.CompositeInsert %[[S]], %[[V]][2 : i32] : f32 into vector<4xf32>
 func.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_index_vector
+//       CHECK:   spirv.CompositeInsert %{{.+}}, %{{.+}}[2 : i32] : i32 into vector<4xi32>
+func.func @insert_index_vector(%arg0 : vector<4xindex>, %arg1: index) -> vector<4xindex> {
+  %1 = vector.insert %arg1, %arg0[2] : index into vector<4xindex>
+  return %1: vector<4xindex>
+}
+
+// -----
+
 // CHECK-LABEL: @insert_size1_vector
 //  CHECK-SAME: %[[V:.*]]: vector<1xf32>, %[[S:.*]]: f32
 //       CHECK:   %[[R:.+]] = builtin.unrealized_conversion_cast %[[S]]
 //       CHECK:   return %[[R]]
 func.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:    spirv.CompositeConstruct %[[V0]], %[[V1]], %[[V1]], %[[V0]] : (f32, f32, f32, f32) -> 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>
 }
 
 // -----
 
+// CHECK-LABEL:  func @shuffle_index_vector
+//  CHECK-SAME:  %[[ARG0:.+]]: vector<1xindex>, %[[ARG1:.+]]: vector<1xindex>
+//       CHECK:    %[[V0:.+]] = builtin.unrealized_conversion_cast %[[ARG0]]
+//       CHECK:    %[[V1:.+]] = builtin.unrealized_conversion_cast %[[ARG1]]
+//       CHECK:    spirv.CompositeConstruct %[[V0]], %[[V1]], %[[V1]], %[[V0]] : (i32, i32, i32, i32) -> vector<4xi32>
+func.func @shuffle_index_vector(%v0 : vector<1xindex>, %v1: vector<1xindex>) -> vector<4xindex> {
+  %shuffle = vector.shuffle %v0, %v1 [0, 1, 1, 0] : vector<1xindex>, vector<1xindex>
+  return %shuffle : vector<4xindex>
+}
+
+// -----
+
 // CHECK-LABEL:  func @shuffle
 //  CHECK-SAME:  %[[V0:.+]]: vector<3xf32>, %[[V1:.+]]: vector<3xf32>
 //       CHECK:    spirv.VectorShuffle [3 : i32, 2 : i32, 5 : i32, 1 : i32] %[[V0]] : vector<3xf32>, %[[V1]] : vector<3xf32> -> vector<4xf32>
 func.func @shuffle(%v0 : vector<3xf32>, %v1: vector<3xf32>) -> vector<4xf32> {
   %shuffle = vector.shuffle %v0, %v1 [3, 2, 5, 1] : vector<3xf32>, vector<3xf32>
   return %shuffle : vector<4xf32>
 }