[mlir][spirv] Add VectorInsertDynamicOp and vector.insertelement lowering

[abialas]

Summary

VectorInsertDynamicOp in SPIRV dialect
conversion from vector.insertelement to spirv VectorInsertDynamicOp

[ThomasRaoux]

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

257

nit: alignment

285

Do we need $component type? This looks redundant as it needs to match vector element type.

[abialas]

Added changes requested by Thomas

[ThomasRaoux]

Looks good!

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

230

nit: Can you break this line so that it doesn't go beyond 80 columns

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

 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_OpVectorInsertDynamic       : I32EnumAttrCase<"OpVectorInsertDynamic", 78>;
 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_OpVectorExtractDynamic, SPV_OC_OpCompositeConstruct,
+      SPV_OC_OpStore, SPV_OC_OpCopyMemory, SPV_OC_OpAccessChain, SPV_OC_OpDecorate,
+      SPV_OC_OpMemberDecorate, SPV_OC_OpVectorExtractDynamic,
+      SPV_OC_OpVectorInsertDynamic, 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

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'",
+            [NoSideEffect,
+            TypesMatchWith<"type of 'result' 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.
def SPV_VectorExtractDynamicOp : SPV_Op<"VectorExtractDynamic",
   let assemblyFormat = [{
     $vector `[` $index `]` attr-dict `:` type($vector) `,` type($index) 
   }];
   
 }
 
 // -----
 
+def SPV_VectorInsertDynamicOp : SPV_Op<"VectorInsertDynamic",
+        [NoSideEffect,
+        TypesMatchWith<"type of 'component' matches element type of 'vector'",
+                "vector", "component",
+                "$_self.cast<mlir::VectorType>().getElementType()">,
+                AllTypesMatch<["vector", "result"]>]> {
+  let summary = [{
+    Make a copy of a vector, with a single, variably selected, component
+    modified.
+  }];
+
+  let description = [{
+    Result Type must be an OpTypeVector.
+
+    Vector must have the same type as Result Type and is the vector that the
+    non-written components are copied from.
+
+    Component is the value supplied for the component selected by Index. It
+    must have the same type as the type of components in Result Type.
+
+    Index must be a scalar integer. It is interpreted as a 0-based index of
+    which component to modify.
+
+    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-insert-dynamic-op ::= `spv.VectorInsertDynamic ` ssa-use `,`
+                                  ssa-use `[` ssa-use `]`
+                                  `:` `vector<` integer-literal `x` scalar-type `>` `,`
+                                  integer-type
+    ```mlir
+
+    #### Example:
+
+    ```
+    %scalar = ... : f32
+    %2 = spv.VectorInsertDynamic %scalar %0[%1] : f32, vector<8xf32>, i32
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_Vector:$vector,
+    SPV_Scalar:$component,
+    SPV_Integer:$index
+  );
+
+  let results = (outs
+    SPV_Vector:$result
+  );
+  
+  let verifier = [{ return success(); }];
+  
+  let assemblyFormat = [{
+    $component `,` $vector `[` $index `]` attr-dict `:` type($vector) `,` type($index) 
+  }];
+}
+
+// -----
+
 #endif // SPIRV_COMPOSITE_OPS

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

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

mlir/test/Conversion/VectorToSPIRV/simple.mlir

 
 // -----
 
 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
 }
+
+// -----
+
+// CHECK-LABEL: insert_element
+//  CHECK-SAME: %[[VAL:.*]]: f32, %[[V:.*]]: vector<4xf32>, %[[ID:.*]]: i32
+//       CHECK:   spv.VectorInsertDynamic %[[VAL]], %[[V]][%[[ID]]] : vector<4xf32>, i32
+func @insert_element(%val: f32, %arg0 : vector<4xf32>, %id : i32) {
+  %0 = vector.insertelement %val, %arg0[%id : i32] : vector<4xf32>
+  spv.ReturnValue %0: vector<4xf32>
+}
+
+// -----
+
+func @insert_element_negative(%val: f32, %arg0 : vector<5xf32>, %id : i32) {
+// expected-error @+1 {{failed to legalize operation 'vector.insertelement'}}
+  %0 = vector.insertelement %val, %arg0[%id : i32] : vector<5xf32>
+  spv.Return
+}

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

spv.func @composite_construct_vector(%arg0: f32, %arg1: f32, %arg2 : f32) -> vector<3xf32> "None" {
     %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
   }
+  spv.func @vector_dynamic_insert(%val: f32, %vec: vector<4xf32>, %id : i32) -> vector<4xf32> "None" {
+    // CHECK: spv.VectorInsertDynamic %{{.*}}, %{{.*}}[%{{.*}}] : vector<4xf32>, i32
+    %0 = spv.VectorInsertDynamic %val, %vec[%id] : vector<4xf32>, i32
+    spv.ReturnValue %0: vector<4xf32>
+  }
 }

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

 // 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
 }
+
+//===----------------------------------------------------------------------===//
+// spv.VectorInsertDynamic
+//===----------------------------------------------------------------------===//
+
+func @vector_dynamic_insert(%val: f32, %vec: vector<4xf32>, %id : i32) -> vector<4xf32> {
+  // CHECK: spv.VectorInsertDynamic %{{.*}}, %{{.*}}[%{{.*}}] : vector<4xf32>, i32
+  %0 = spv.VectorInsertDynamic %val, %vec[%id] : vector<4xf32>, i32
+  return %0 : vector<4xf32>
+}