[mlir][Vector] Thread 0-d vectors through InsertElementOp.

[nicolasvasilache]

Summary

This revision makes concrete use of 0-d vectors to extend the semantics of
InsertElementOp.

[pifon2a]

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

518

nit: wrap the line

[dcaballe]

LGTM. Just one minor question. Thanks!

mlir/test/Dialect/Vector/ops.mlir

198

Would it be ok to use %b[] with a 1D vector? Could we add a test for that either to invalid or here depending on the answer?

[nicolasvasilache]

Rebase past windows build error.

[nicolasvasilache]

Address comments and fix bugs.

mlir/include/mlir/Dialect/Vector/VectorOps.td

 }
 
 def Vector_InsertElementOp :
   Vector_Op<"insertelement", [NoSideEffect,
      TypesMatchWith<"source operand type matches element type of result",
                     "result", "source",
                     "$_self.cast<ShapedType>().getElementType()">,
      AllTypesMatch<["dest", "result"]>]>,
-     Arguments<(ins AnyType:$source, AnyVector:$dest,
-                    AnySignlessIntegerOrIndex:$position)>,
-     Results<(outs AnyVector:$result)> {
+     Arguments<(ins AnyType:$source, AnyVectorOfAnyRank:$dest,
+                    Optional<AnySignlessIntegerOrIndex>:$position)>,
+     Results<(outs AnyVectorOfAnyRank:$result)> {
   let summary = "insertelement operation";
   let description = [{
-    Takes a scalar source, an 1-D destination vector and a dynamic index
-    position and inserts the source into the destination at the proper
-    position.  Note that this instruction resembles vector.insert, but
-    is restricted to 1-D vectors and relaxed to dynamic indices. It is
-    meant to be closer to LLVM's version:
+    Takes a scalar source, a 0-D or 1-D destination vector and a dynamic index
+    position and inserts the source into the destination at the proper position.
+
+    Note that this instruction resembles vector.insert, but is restricted to 0-D
+    and 1-D vectors and relaxed to dynamic indices. 
+
+    It is meant to be closer to LLVM's version:
     https://llvm.org/docs/LangRef.html#insertelement-instruction
 
     Example:
 
     ```mlir
     %c = arith.constant 15 : i32
     %f = arith.constant 0.0f : f32
     %1 = vector.insertelement %f, %0[%c : i32]: vector<16xf32>
+    %2 = vector.insertelement %f, %z[]: vector<f32>
     ```
   }];
   let assemblyFormat = [{
-    $source `,` $dest `[` $position `:` type($position) `]` attr-dict `:`
+    $source `,` $dest `[` ($position^ `:` type($position))? `]`  attr-dict `:`
     type($result)
   }];
 
   let builders = [
+    // 0-D builder.
+    OpBuilder<(ins "Value":$source, "Value":$dest)>,
+    // 1-D + position builder.
     OpBuilder<(ins "Value":$source, "Value":$dest, "Value":$position)>
   ];
   let extraClassDeclaration = [{
     Type getSourceType() { return source().getType(); }
     VectorType getDestVectorType() {
       return dest().getType().cast<VectorType>();
     }
   }];

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

matchAndRewrite(vector::InsertElementOp insertEltOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto vectorType = insertEltOp.getDestVectorType();
     auto llvmType = typeConverter->convertType(vectorType);
 
     // Bail if result type cannot be lowered.
     if (!llvmType)
       return failure();
 
+    if (vectorType.getRank() == 0) {
+      Location loc = insertEltOp.getLoc();
+      auto idxType = rewriter.getIndexType();
+      auto zero = rewriter.create<LLVM::ConstantOp>(
+          loc, typeConverter->convertType(idxType),
+          rewriter.getIntegerAttr(idxType, 0));
+      rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>(
+          insertEltOp, llvmType, adaptor.dest(), adaptor.source(), zero);
+      return success();
+    }
+
     rewriter.replaceOpWithNewOp<LLVM::InsertElementOp>(
         insertEltOp, llvmType, adaptor.dest(), adaptor.source(),
         adaptor.position());
     return success();
   }
 };
 
 class VectorInsertOpConversion

mlir/lib/Dialect/Vector/VectorOps.cpp

static ParseResult parseShuffleOp(OpAsmParser &parser, OperationState &result) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // InsertElementOp
 //===----------------------------------------------------------------------===//
 
 void InsertElementOp::build(OpBuilder &builder, OperationState &result,
+                            Value source, Value dest) {
+  result.addOperands({source, dest});
+  result.addTypes(dest.getType());
+}
+
+void InsertElementOp::build(OpBuilder &builder, OperationState &result,
                             Value source, Value dest, Value position) {
   result.addOperands({source, dest, position});
   result.addTypes(dest.getType());
 }
 
 static LogicalResult verify(InsertElementOp op) {
   auto dstVectorType = op.getDestVectorType();
+  if (dstVectorType.getRank() == 0) {
+    if (op.position())
+      return op.emitOpError("expected position to be empty with 0-D vector");
+    return success();
+  }
   if (dstVectorType.getRank() != 1)
-    return op.emitOpError("expected 1-D vector");
+    return op.emitOpError("unexpected >1 vector rank");
+  if (!op.position())
+    return op.emitOpError("expected position for 1-D vector");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // InsertOp
 //===----------------------------------------------------------------------===//
 
 void InsertOp::build(OpBuilder &builder, OperationState &result, Value source,

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 // CHECK-LABEL: @extract_element_from_vec_3d
 //       CHECK:   llvm.extractvalue {{.*}}[0, 0] : !llvm.array<4 x array<3 x vector<16xf32>>>
 //       CHECK:   llvm.mlir.constant(0 : i64) : i64
 //       CHECK:   llvm.extractelement {{.*}}[{{.*}} : i64] : vector<16xf32>
 //       CHECK:   return {{.*}} : f32
 
 // -----
 
+// CHECK-LABEL: @insert_element_0d
+// CHECK-SAME: %[[A:.*]]: f32,
+func @insert_element_0d(%a: f32, %b: vector<f32>) -> vector<f32> {
+  // CHECK: %[[B:.*]] =  builtin.unrealized_conversion_cast %{{.*}} :
+  // CHECK:   vector<f32> to vector<1xf32>
+  // CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : index) : i64
+  // CHECK: %[[x:.*]] = llvm.insertelement %[[A]], %[[B]][%[[C0]] : {{.*}}] : vector<1xf32>
+  %1 = vector.insertelement %a, %b[] : vector<f32>
+  return %1 : vector<f32>
+}
+
+// -----
+
 func @insert_element(%arg0: f32, %arg1: vector<4xf32>) -> vector<4xf32> {
   %0 = arith.constant 3 : i32
   %1 = vector.insertelement %arg0, %arg1[%0 : i32] : vector<4xf32>
   return %1 : vector<4xf32>
 }
 // CHECK-LABEL: @insert_element(
 // CHECK-SAME: %[[A:.*]]: f32,
 // CHECK-SAME: %[[B:.*]]: vector<4xf32>)

mlir/test/Dialect/Vector/invalid.mlir

 
 func @extract_element(%arg0: vector<f32>) {
   %c = arith.constant 3 : i32
   // expected-error@+1 {{expected position to be empty with 0-D vector}}
   %1 = vector.extractelement %arg0[%c : i32] : vector<f32>
 }
 
 // -----
  
 func @extract_element(%arg0: vector<4xf32>) {
   %c = arith.constant 3 : i32
   // expected-error@+1 {{expected position for 1-D vector}}
   %1 = vector.extractelement %arg0[] : vector<4xf32>
 }
 
 // -----
 

 
 func @extract_position_overflow(%arg0: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute #3 to be a non-negative integer smaller than the corresponding vector dimension}}
   %1 = vector.extract %arg0[0, 0, -1] : vector<4x8x16xf32>
 }
 
 // -----
 
+func @insert_element(%arg0: f32, %arg1: vector<f32>) {
+  %c = arith.constant 3 : i32
+  // expected-error@+1 {{expected position to be empty with 0-D vector}}
+  %0 = vector.insertelement %arg0, %arg1[%c : i32] : vector<f32>
+}
+
+// -----
+
+func @insert_element(%arg0: f32, %arg1: vector<4xf32>) {
+  %c = arith.constant 3 : i32
+  // expected-error@+1 {{expected position for 1-D vector}}
+  %0 = vector.insertelement %arg0, %arg1[] : vector<4xf32>
+}
+
+// -----
+
 func @insert_element(%arg0: f32, %arg1: vector<4x4xf32>) {
   %c = arith.constant 3 : i32
-  // expected-error@+1 {{'vector.insertelement' op expected 1-D vector}}
+  // expected-error@+1 {{unexpected >1 vector rank}}
   %0 = vector.insertelement %arg0, %arg1[%c : i32] : vector<4x4xf32>
 }
 
 // -----
 
 func @insert_element_wrong_type(%arg0: i32, %arg1: vector<4xf32>) {
   %c = arith.constant 3 : i32
   // expected-error@+1 {{'vector.insertelement' op failed to verify that source operand type matches element type of result}}

mlir/test/Dialect/Vector/ops.mlir

func @extract(%arg0: vector<4x8x16xf32>) -> (vector<4x8x16xf32>, vector<8x16xf32>, vector<16xf32>, f32) {
   %1 = vector.extract %arg0[3] : vector<4x8x16xf32>
   // CHECK-NEXT: vector.extract {{.*}}[3, 3] : vector<4x8x16xf32>
   %2 = vector.extract %arg0[3, 3] : vector<4x8x16xf32>
   // CHECK-NEXT: vector.extract {{.*}}[3, 3, 3] : vector<4x8x16xf32>
   %3 = vector.extract %arg0[3, 3, 3] : vector<4x8x16xf32>
   return %0, %1, %2, %3 : vector<4x8x16xf32>, vector<8x16xf32>, vector<16xf32>, f32
 }
 
+// CHECK-LABEL: @insert_element_0d
+func @insert_element_0d(%a: f32, %b: vector<f32>) -> vector<f32> {
+  // CHECK-NEXT: vector.insertelement %{{.*}}, %{{.*}}[] : vector<f32>
+  %1 = vector.insertelement %a, %b[] : vector<f32>
+  return %1 : vector<f32>
+}
+
 // CHECK-LABEL: @insert_element
 func @insert_element(%a: f32, %b: vector<16xf32>) -> vector<16xf32> {
   // CHECK:      %[[C15:.*]] = arith.constant 15 : i32
   %c = arith.constant 15 : i32
   // CHECK-NEXT: vector.insertelement %{{.*}}, %{{.*}}[%[[C15]] : i32] : vector<16xf32>
   %1 = vector.insertelement %a, %b[%c : i32] : vector<16xf32>
   return %1 : vector<16xf32>
 }

mlir/test/Integration/Dialect/Vector/CPU/test-0-d-vectors.mlir

 // RUN: mlir-opt %s -convert-scf-to-std -convert-vector-to-llvm -convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 func @extract_element_0d(%a: vector<f32>) {
   %1 = vector.extractelement %a[] : vector<f32>
   // CHECK: 42
   vector.print %1: f32
   return
 }
 
+func @insert_element_0d(%a: f32, %b: vector<f32>) -> (vector<f32>) {
+  %1 = vector.insertelement %a, %b[] : vector<f32>
+  return %1: vector<f32>
+}
+
 func @entry() {
-  %1 = arith.constant dense<42.0> : vector<f32>
-  call  @extract_element_0d(%1) : (vector<f32>) -> ()
+  %0 = arith.constant 42.0 : f32
+  %1 = arith.constant dense<0.0> : vector<f32>
+  %2 = call  @insert_element_0d(%0, %1) : (f32, vector<f32>) -> (vector<f32>)
+  call  @extract_element_0d(%2) : (vector<f32>) -> ()
   return
 }