[mlir][spirv] Add VectorExtractDynamicOp and vector.extractelement lowering

[abialas]

Summary

• VectorExtractDynamicOp in SPIRV dialect
• conversion from vector.extractelement to spirv VectorExtractDynamicOp

[ThomasRaoux]

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

89

I think we need to check that the original vector type is converted to a spv.vector. Large vector would be converted to an array. You can use !spirv::CompositeType::isValid(insertOp.getDestVectorType() as i the pattern above. You may also want a negative test for this case.

[antiagainst]

mlir/include/mlir/Dialect/SPIRV/SPIRVOps.td

532 ↗

(On Diff #302549)

Can this be put in SPIRVCompositeOps.td? We are basically following the spec's group regarding file organization. (This file is a bit different but it should be cleaned up..)

574 ↗

(On Diff #302549)

This cannot be aritary SPIR-V types? It should be SPV_Scalar?

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

89

+1. BTW, if you want to test out large vectors that are enabled by the Vector16 extension, you'll need to have proper spv.target_env that enables Vector16 extension. Otherwise it won't convert. You can find examples of how to specify such target env in various -> spirv conversion mlir files.

[abialas]

Thank You for the comments! Added requested changes. I wanted to keep this change not dependent on the Vector16 change, so did not add the unit test for VectorExtractDynamicOp with vector sizes 8 or 16, but verified that it is working in our use cases.

[antiagainst]

LGTM. Just one final nit: the tests for op and serialization should be added into composite-ops.mlir now given the op is defined in SPIRVCompositeOps.td.

[abialas]

moved tests to composite-ops

mlir/include/mlir/Dialect/SPIRV/SPIRVBase.td

 def SPV_OC_OpFunctionCall              : I32EnumAttrCase<"OpFunctionCall", 57>;
 def SPV_OC_OpVariable                  : I32EnumAttrCase<"OpVariable", 59>;
 def SPV_OC_OpLoad                      : I32EnumAttrCase<"OpLoad", 61>;
 def SPV_OC_OpStore                     : I32EnumAttrCase<"OpStore", 62>;
 def SPV_OC_OpCopyMemory                : I32EnumAttrCase<"OpCopyMemory", 63>;
 def SPV_OC_OpAccessChain               : I32EnumAttrCase<"OpAccessChain", 65>;
 def SPV_OC_OpDecorate                  : I32EnumAttrCase<"OpDecorate", 71>;
 def SPV_OC_OpMemberDecorate            : I32EnumAttrCase<"OpMemberDecorate", 72>;
+def SPV_OC_OpVectorExtractDynamic      : I32EnumAttrCase<"OpVectorExtractDynamic", 77>;
 def SPV_OC_OpCompositeConstruct        : I32EnumAttrCase<"OpCompositeConstruct", 80>;
 def SPV_OC_OpCompositeExtract          : I32EnumAttrCase<"OpCompositeExtract", 81>;
 def SPV_OC_OpCompositeInsert           : I32EnumAttrCase<"OpCompositeInsert", 82>;
 def SPV_OC_OpTranspose                 : I32EnumAttrCase<"OpTranspose", 84>;
 def SPV_OC_OpConvertFToU               : I32EnumAttrCase<"OpConvertFToU", 109>;
 def SPV_OC_OpConvertFToS               : I32EnumAttrCase<"OpConvertFToS", 110>;
 def SPV_OC_OpConvertSToF               : I32EnumAttrCase<"OpConvertSToF", 111>;
 def SPV_OC_OpConvertUToF               : I32EnumAttrCase<"OpConvertUToF", 112>;
SPV_I32EnumAttr<"Opcode", "valid SPIR-V instructions", [
       SPV_OC_OpTypeFloat, SPV_OC_OpTypeVector, SPV_OC_OpTypeMatrix,
       SPV_OC_OpTypeArray, SPV_OC_OpTypeRuntimeArray, SPV_OC_OpTypeStruct,
       SPV_OC_OpTypePointer, SPV_OC_OpTypeFunction, SPV_OC_OpTypeForwardPointer,
       SPV_OC_OpConstantTrue, SPV_OC_OpConstantFalse, SPV_OC_OpConstant,
       SPV_OC_OpConstantComposite, SPV_OC_OpConstantNull, SPV_OC_OpSpecConstantTrue,
       SPV_OC_OpSpecConstantFalse, SPV_OC_OpSpecConstant,
       SPV_OC_OpSpecConstantComposite, SPV_OC_OpFunction, SPV_OC_OpFunctionParameter,
       SPV_OC_OpFunctionEnd, SPV_OC_OpFunctionCall, SPV_OC_OpVariable, SPV_OC_OpLoad,
-      SPV_OC_OpStore, SPV_OC_OpCopyMemory, SPV_OC_OpAccessChain, SPV_OC_OpDecorate,
-      SPV_OC_OpMemberDecorate, SPV_OC_OpCompositeConstruct,
+      SPV_OC_OpStore, SPV_OC_OpCopyMemory, SPV_OC_OpAccessChain,
+      SPV_OC_OpDecorate, SPV_OC_OpMemberDecorate,
+      SPV_OC_OpVectorExtractDynamic, SPV_OC_OpCompositeConstruct,
       SPV_OC_OpCompositeExtract, SPV_OC_OpCompositeInsert, SPV_OC_OpTranspose,
       SPV_OC_OpConvertFToU, SPV_OC_OpConvertFToS, SPV_OC_OpConvertSToF,
       SPV_OC_OpConvertUToF, SPV_OC_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert,
       SPV_OC_OpBitcast, SPV_OC_OpSNegate, SPV_OC_OpFNegate, SPV_OC_OpIAdd,
       SPV_OC_OpFAdd, SPV_OC_OpISub, SPV_OC_OpFSub, SPV_OC_OpIMul, SPV_OC_OpFMul,
       SPV_OC_OpUDiv, SPV_OC_OpSDiv, SPV_OC_OpFDiv, SPV_OC_OpUMod, SPV_OC_OpSRem,
       SPV_OC_OpSMod, SPV_OC_OpFRem, SPV_OC_OpFMod, SPV_OC_OpMatrixTimesScalar,
       SPV_OC_OpMatrixTimesMatrix, SPV_OC_OpLogicalEqual, SPV_OC_OpLogicalNotEqual,

mlir/include/mlir/Dialect/SPIRV/SPIRVCompositeOps.td

def SPV_CompositeInsertOp : SPV_Op<"CompositeInsert", [NoSideEffect]> {
   );
 
   let builders = [
     OpBuilderDAG<(ins "Value":$object, "Value":$composite,
       "ArrayRef<int32_t>":$indices)>
   ];
 }
 
+// -----
+
+def SPV_VectorExtractDynamicOp : SPV_Op<"VectorExtractDynamic",
+            [NoSideEffect, TypesMatchWith<"type of 'value' matches element type of 'vector'",
+                     "vector", "result",
+                     "$_self.cast<mlir::VectorType>().getElementType()">]> {
+  let summary = [{
+    Extract a single, dynamically selected, component of a vector.
+  }];
+
+  let description = [{
+    Result Type must be a scalar type.
+
+    Vector must have a type OpTypeVector whose Component Type is Result
+    Type.
+
+    Index must be a scalar integer. It is interpreted as a 0-based index of
+    which component of Vector to extract.
+
+    Behavior is undefined if Index's value is less than zero or greater than
+    or equal to the number of components in Vector.
+
+    <!-- End of AutoGen section -->
+
+    ```
+  scalar-type ::= integer-type | float-type | boolean-type 
+    vector-extract-dynamic-op ::= `spv.VectorExtractDynamic ` ssa-use `[` ssa-use `]`
+                                    `:` `vector<` integer-literal `x` scalar-type `>` `,`
+                                    integer-type
+    ```mlir
+
+    #### Example:
+
+    ```
+    %2 = spv.VectorExtractDynamic %0[%1] : vector<8xf32>, i32
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_Vector:$vector,
+    SPV_Integer:$index
+  );
+
+  let results = (outs
+    SPV_Scalar:$result
+  );
+  
+  let verifier = [{ return success(); }];
+  
+  let assemblyFormat = [{
+    $vector `[` $index `]` attr-dict `:` type($vector) `,` type($index) 
+  }];
+  
+}
+
+// -----
+
 #endif // SPIRV_COMPOSITE_OPS

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

matchAndRewrite(vector::InsertOp insertOp, ArrayRef<Value> operands,
     vector::InsertOp::Adaptor adaptor(operands);
     int32_t id = insertOp.position().begin()->cast<IntegerAttr>().getInt();
     Value newInsert = rewriter.create<spirv::CompositeInsertOp>(
         insertOp.getLoc(), adaptor.source(), adaptor.dest(), id);
     rewriter.replaceOp(insertOp, newInsert);
     return success();
   }
 };
+
+struct VectorExtractElementOpConvert final
+    : public SPIRVOpLowering<vector::ExtractElementOp> {
+  using SPIRVOpLowering<vector::ExtractElementOp>::SPIRVOpLowering;
+  LogicalResult
+  matchAndRewrite(vector::ExtractElementOp extractElementOp,
+                  ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (!spirv::CompositeType::isValid(extractElementOp.getVectorType()))
+      return failure();
+    vector::ExtractElementOp::Adaptor adaptor(operands);
+    Value newExtractElement = rewriter.create<spirv::VectorExtractDynamicOp>(
+        extractElementOp.getLoc(), extractElementOp.getType(), adaptor.vector(),
+        extractElementOp.position());
+    rewriter.replaceOp(extractElementOp, newExtractElement);
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::populateVectorToSPIRVPatterns(MLIRContext *context,
                                          SPIRVTypeConverter &typeConverter,
                                          OwningRewritePatternList &patterns) {
   patterns.insert<VectorBroadcastConvert, VectorExtractOpConvert,
-                  VectorInsertOpConvert>(context, typeConverter);
+                  VectorInsertOpConvert, VectorExtractElementOpConvert>(
+      context, typeConverter);
 }
 
 namespace {
 struct LowerVectorToSPIRVPass
     : public ConvertVectorToSPIRVBase<LowerVectorToSPIRVPass> {
   void runOnOperation() override;
 };
 } // namespace

mlir/test/Conversion/VectorToSPIRV/simple.mlir

-// RUN: mlir-opt -split-input-file -convert-vector-to-spirv %s -o - | FileCheck %s
+// RUN: mlir-opt -split-input-file -convert-vector-to-spirv -verify-diagnostics %s -o - | FileCheck %s
 
 // CHECK-LABEL: broadcast
 //  CHECK-SAME: %[[A:.*]]: f32
 //       CHECK:   spv.CompositeConstruct %[[A]], %[[A]], %[[A]], %[[A]] : vector<4xf32>
 //       CHECK:   spv.CompositeConstruct %[[A]], %[[A]] : vector<2xf32>
 func @broadcast(%arg0 : f32) {
   %0 = vector.broadcast %arg0 : f32 to vector<4xf32>
   %1 = vector.broadcast %arg0 : f32 to vector<2xf32>
   spv.Return
 }
 
 // -----
 
 // CHECK-LABEL: extract_insert
 //  CHECK-SAME: %[[V:.*]]: vector<4xf32>
 //       CHECK:   %[[S:.*]] = spv.CompositeExtract %[[V]][1 : i32] : vector<4xf32>
 //       CHECK:   spv.CompositeInsert %[[S]], %[[V]][0 : i32] : f32 into vector<4xf32>
 func @extract_insert(%arg0 : vector<4xf32>) {
   %0 = vector.extract %arg0[1] : vector<4xf32>
   %1 = vector.insert %0, %arg0[0] : f32 into vector<4xf32>
   spv.Return
 }
+
+// -----
+
+// 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) {
+  %0 = vector.extractelement %arg0[%id : i32] : vector<4xf32>
+  spv.ReturnValue %0: f32
+}
+
+// -----
+
+func @extract_element_negative(%arg0 : vector<5xf32>, %id : i32) {
+// expected-error @+1 {{failed to legalize operation 'vector.extractelement'}}
+  %0 = vector.extractelement %arg0[%id : i32] : vector<5xf32>
+  spv.ReturnValue %0: f32
+}

mlir/test/Dialect/SPIRV/Serialization/composite-op.mlir

 // RUN: mlir-translate -split-input-file -test-spirv-roundtrip %s | FileCheck %s
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   spv.func @composite_insert(%arg0 : !spv.struct<(f32, !spv.struct<(!spv.array<4xf32>, f32)>)>, %arg1: !spv.array<4xf32>) -> !spv.struct<(f32, !spv.struct<(!spv.array<4xf32>, f32)>)> "None" {
     // CHECK: spv.CompositeInsert {{%.*}}, {{%.*}}[1 : i32, 0 : i32] : !spv.array<4 x f32> into !spv.struct<(f32, !spv.struct<(!spv.array<4 x f32>, f32)>)>
     %0 = spv.CompositeInsert %arg1, %arg0[1 : i32, 0 : i32] : !spv.array<4xf32> into !spv.struct<(f32, !spv.struct<(!spv.array<4xf32>, f32)>)>
     spv.ReturnValue %0: !spv.struct<(f32, !spv.struct<(!spv.array<4xf32>, f32)>)>
   }
   spv.func @composite_construct_vector(%arg0: f32, %arg1: f32, %arg2 : f32) -> vector<3xf32> "None" {
     // CHECK: spv.CompositeConstruct {{%.*}}, {{%.*}}, {{%.*}} : vector<3xf32>
     %0 = spv.CompositeConstruct %arg0, %arg1, %arg2 : vector<3xf32>
     spv.ReturnValue %0: vector<3xf32>
   }
+  spv.func @vector_dynamic_extract(%vec: vector<4xf32>, %id : i32) -> f32 "None" {
+    // CHECK: spv.VectorExtractDynamic %{{.*}}[%{{.*}}] : vector<4xf32>, i32
+    %0 = spv.VectorExtractDynamic %vec[%id] : vector<4xf32>, i32
+    spv.ReturnValue %0: f32
+  }
 }

mlir/test/Dialect/SPIRV/composite-ops.mlir

 
 // -----
 
 func @composite_insert_invalid_result_type(%arg0: !spv.array<4xf32>, %arg1 : f32) -> !spv.array<4xf64> {
   // expected-error @+1 {{result type should be the same as the composite type, but found '!spv.array<4 x f32>' vs '!spv.array<4 x f64>'}}
   %0 = "spv.CompositeInsert"(%arg1, %arg0) {indices = [0: i32]} : (f32, !spv.array<4xf32>) -> !spv.array<4xf64>
   return %0: !spv.array<4xf64>
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spv.VectorExtractDynamic
+//===----------------------------------------------------------------------===//
+
+func @vector_dynamic_extract(%vec: vector<4xf32>, %id : i32) -> f32 {
+  // CHECK: spv.VectorExtractDynamic %{{.*}}[%{{.*}}] : vector<4xf32>, i32
+  %0 = spv.VectorExtractDynamic %vec[%id] : vector<4xf32>, i32
+  return %0 : f32
+}