[mlir][Vector] Add support for 0-D vectors to vector.insert/extract

[dcaballe]

Summary

This is part of the process to remove vector.insertelement/extractelement
from the Vector dialect.

RFC: https://discourse.llvm.org/t/rfc-psa-remove-vector-extractelement-and-vector-insertelement-ops-in-favor-of-vector-extract-and-vector-insert-ops

[dcaballe]

Add more tets, doc examples and fix bug

[nicolasvasilache]

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

704

Since your objective is to ultimately remove insertelement/extractelement, can you comment on the symmetry breaking property of this PR ?

vector.insert can now insert either a scalar or a 0-d vector into a larger vector but vector.extract must extract a scalar and cannot extract a 0-d vector.

I don't know yet how this plays in your full design but potential lack of symmetry between insert and extract ops has always been something to think about.

Thoughts?

[dcaballe]

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

704

That's a good point! The main problem with vector.extract is that the extracted type is inferred from the source type, to the point that is not even represented in the assembly format of the instruction (e.g., vector.extract %0[3, 3, 3]: vector<4x8x16xf32>). We could change its assembly format but that would be a massive change so I would prefer not to go there unless it's strictly necessary. Perhaps we could restrict both vector.extract/insert to only scalars and use vector.broadcast, vector.shape_cast and other shape transforming ops to convert 0-D from/to scalars?

[awarzynski]

mlir/test/Dialect/Vector/invalid.mlir

244

Looks like the diagnostic in VectorOps.cpp is off? I think that it's this one: https://github.com/llvm/llvm-project/blob/678360fd9d8517a978c5f46364a56412d70e15c7/mlir/lib/Dialect/Vector/IR/VectorOps.cpp#L2324-L2326

[dcaballe]

Addressed feedback.

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

704

Let's keep it as is right now and revisit if there is any problem. The broadcast/shape_cast approach described above seems reasonable to me so it's not that we won't be able to represent all the cases in one way or another.

[jsetoain]

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

698–699

Am I right to assume that support for values as $position shall come in a different patch?

[dcaballe]

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

698–699

Yes, I have that almost done :)

[dcaballe]

Please, let me know if there are any other comments. Otherwise, this should be ready to go.

[awarzynski]

mlir/test/Dialect/Vector/invalid.mlir

244

[nit] Could you add "expected"? For consistency with other diagnostics and also because that makes it easier for my brain to parse :)

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

def Vector_ExtractElementOp :
   let hasFolder = 1;
 }
 
 def Vector_ExtractOp :
   Vector_Op<"extract", [Pure,
      PredOpTrait<"operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      DeclareOpInterfaceMethods<InferTypeOpInterface>]>,
-    Arguments<(ins AnyVector:$vector, I64ArrayAttr:$position)>,
+    Arguments<(ins AnyVectorOfAnyRank:$vector, I64ArrayAttr:$position)>,
     Results<(outs AnyType)> {
   let summary = "extract operation";
   let description = [{
     Takes an n-D vector and a k-D position and extracts the (n-k)-D vector at
-    the proper position. Degenerates to an element type in the 0-D case.
+    the proper position. Degenerates to an element type if n-k is zero.
 
     Example:
 
     ```mlir
     %1 = vector.extract %0[3]: vector<4x8x16xf32>
     %2 = vector.extract %0[3, 3, 3]: vector<4x8x16xf32>
+    %3 = vector.extract %1[]: vector<f32>
     ```
   }];
   let builders = [
     OpBuilder<(ins "Value":$source, "ArrayRef<int64_t>":$position)>,
     // Convenience builder which assumes the values in `position` are defined by
     // ConstantIndexOp.
     OpBuilder<(ins "Value":$source, "ValueRange":$position)>
   ];
def Vector_InsertElementOp :
   let hasFolder = 1;
 }
 
 def Vector_InsertOp :
   Vector_Op<"insert", [Pure,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      AllTypesMatch<["dest", "res"]>]>,
-     Arguments<(ins AnyType:$source, AnyVector:$dest, I64ArrayAttr:$position)>,
-     Results<(outs AnyVector:$res)> {
+     Arguments<(ins AnyType:$source, AnyVectorOfAnyRank:$dest, I64ArrayAttr:$position)>,
+     Results<(outs AnyVectorOfAnyRank:$res)> {
   let summary = "insert operation";
   let description = [{
     Takes an n-D source vector, an (n+k)-D destination vector and a k-D position
     and inserts the n-D source into the (n+k)-D destination at the proper
-    position. Degenerates to a scalar source type when n = 0.
+    position. Degenerates to a scalar or a 0-d vector source type when n = 0.
 
     Example:
 
     ```mlir
     %2 = vector.insert %0, %1[3] : vector<8x16xf32> into vector<4x8x16xf32>
     %5 = vector.insert %3, %4[3, 3, 3] : f32 into vector<4x8x16xf32>
+    %8 = vector.insert %6, %7[] : f32 into vector<f32>
+    %11 = vector.insert %9, %10[3, 3, 3] : vector<f32> into vector<4x8x16xf32>
     ```
   }];
   let assemblyFormat = [{
     $source `,` $dest $position attr-dict `:` type($source) `into` type($dest)
   }];
 
   let builders = [
     OpBuilder<(ins "Value":$source, "Value":$dest,

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

ExtractOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
                             ValueRange operands, DictionaryAttr attributes,
                             OpaqueProperties properties, RegionRange,
                             SmallVectorImpl<Type> &inferredReturnTypes) {
   ExtractOp::Adaptor op(operands, attributes, properties);
   auto vectorType = llvm::cast<VectorType>(op.getVector().getType());
   if (static_cast<int64_t>(op.getPosition().size()) == vectorType.getRank()) {
     inferredReturnTypes.push_back(vectorType.getElementType());
   } else {
-    auto n =
-        std::min<size_t>(op.getPosition().size(), vectorType.getRank() - 1);
+    auto n = std::min<size_t>(op.getPosition().size(), vectorType.getRank());
     inferredReturnTypes.push_back(VectorType::get(
         vectorType.getShape().drop_front(n), vectorType.getElementType()));
   }
   return success();
 }
 
 bool ExtractOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   // Allow extracting 1-element vectors instead of scalars.
void InsertOp::build(OpBuilder &builder, OperationState &result, Value source,
   build(builder, result, source, dest, positionConstants);
 }
 
 LogicalResult InsertOp::verify() {
   auto positionAttr = getPosition().getValue();
   auto destVectorType = getDestVectorType();
   if (positionAttr.size() > static_cast<unsigned>(destVectorType.getRank()))
     return emitOpError(
-        "expected position attribute of rank smaller than dest vector rank");
+        "expected position attribute of rank no greater than dest vector rank");
   auto srcVectorType = llvm::dyn_cast<VectorType>(getSourceType());
   if (srcVectorType &&
       (static_cast<unsigned>(srcVectorType.getRank()) + positionAttr.size() !=
        static_cast<unsigned>(destVectorType.getRank())))
     return emitOpError("expected position attribute rank + source rank to "
                        "match dest vector rank");
   if (!srcVectorType &&
       (positionAttr.size() != static_cast<unsigned>(destVectorType.getRank())))

mlir/test/Dialect/Vector/invalid.mlir

 
 func.func @extract_precise_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[3, 7, 16] : vector<4x8x16xf32>
 }
 
 // -----
 
+func.func @extract_0d(%arg0: vector<f32>) {
+  // expected-error@+1 {{expected position attribute of rank smaller than vector rank}}
+  %1 = vector.extract %arg0[0] : vector<f32>
+}
+
+// -----
+
 func.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.func @insert_element(%arg0: f32, %arg1: vector<f32>) {
func.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}}
   %0 = "vector.insertelement" (%arg0, %arg1, %c) : (i32, vector<4xf32>, i32) -> (vector<4xf32>)
 }
 
 // -----
 
 func.func @insert_vector_type(%a: f32, %b: vector<4x8x16xf32>) {
-  // expected-error@+1 {{expected position attribute of rank smaller than dest vector rank}}
+  // expected-error@+1 {{expected position attribute of rank no greater than dest vector rank}}
   %1 = vector.insert %a, %b[3, 3, 3, 3, 3, 3] : f32 into vector<4x8x16xf32>
 }
 
 // -----
 
 func.func @insert_vector_type(%a: vector<4xf32>, %b: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute rank + source rank to match dest vector rank}}
   %1 = vector.insert %a, %b[3] : vector<4xf32> into vector<4x8x16xf32>

 
 func.func @insert_precise_position_overflow(%a: f32, %b: vector<4x8x16xf32>) {
   // expected-error@+1 {{expected position attribute #1 to be a non-negative integer smaller than the corresponding dest vector dimension}}
   %1 = vector.insert %a, %b[4, 7, 15] : f32 into vector<4x8x16xf32>
 }
 
 // -----
 
+func.func @insert_0d(%a: vector<f32>, %b: vector<4x8x16xf32>) {
+  // expected-error@+1 {{expected position attribute rank + source rank to match dest vector rank}}
+  %1 = vector.insert %a, %b[2, 6] : vector<f32> into vector<4x8x16xf32>
+}
+
+// -----
+
+func.func @insert_0d(%a: f32, %b: vector<f32>) {
+  // expected-error@+1 {{expected position attribute of rank no greater than dest vector rank}}
+  %1 = vector.insert %a, %b[0] : f32 into vector<f32>
+}
+
+// -----
+
 func.func @outerproduct_num_operands(%arg0: f32) {
   // expected-error@+1 {{expected at least 2 operands}}
   %1 = vector.outerproduct %arg0 : f32, f32
 }
 // -----
 
 func.func @outerproduct_non_vector_operand(%arg0: f32) {
   // expected-error@+1 {{expected vector type for operand #1}}

mlir/test/Dialect/Vector/ops.mlir

func.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: @extract_0d
+func.func @extract_0d(%a: vector<f32>) -> f32 {
+  // CHECK-NEXT: vector.extract %{{.*}}[] : vector<f32>
+  %0 = vector.extract %a[] : vector<f32>
+  return %0 : f32
+}
+
 // CHECK-LABEL: @insert_element_0d
 func.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.func @insert(%a: f32, %b: vector<16xf32>, %c: vector<8x16xf32>, %res: vector<4x8x16xf32>) -> vector<4x8x16xf32> {
   %2 = vector.insert %b, %res[3, 3] : vector<16xf32> into vector<4x8x16xf32>
   // CHECK: vector.insert %{{.*}}, %{{.*}}[3, 3, 3] : f32 into vector<4x8x16xf32>
   %3 = vector.insert %a, %res[3, 3, 3] : f32 into vector<4x8x16xf32>
   // CHECK: vector.insert %{{.*}}, %{{.*}}[] : vector<4x8x16xf32> into vector<4x8x16xf32>
   %4 = vector.insert %3, %3[] : vector<4x8x16xf32> into vector<4x8x16xf32>
   return %4 : vector<4x8x16xf32>
 }
 
+// CHECK-LABEL: @insert_0d
+func.func @insert_0d(%a: f32, %b: vector<f32>, %c: vector<2x3xf32>) -> (vector<f32>, vector<2x3xf32>) {
+  // CHECK-NEXT: vector.insert %{{.*}}, %{{.*}}[] : f32 into vector<f32>
+  %1 = vector.insert %a,  %b[] : f32 into vector<f32>
+  // CHECK-NEXT: vector.insert %{{.*}}, %{{.*}}[0, 1] : vector<f32> into vector<2x3xf32>
+  %2 = vector.insert %b,  %c[0, 1] : vector<f32> into vector<2x3xf32>
+  return %1, %2 : vector<f32>, vector<2x3xf32>
+}
+
 // CHECK-LABEL: @outerproduct
 func.func @outerproduct(%arg0: vector<4xf32>, %arg1: vector<8xf32>, %arg2: vector<4x8xf32>) -> vector<4x8xf32> {
   // CHECK: vector.outerproduct {{.*}} : vector<4xf32>, vector<8xf32>
   %0 = vector.outerproduct %arg0, %arg1 : vector<4xf32>, vector<8xf32>
   // CHECK: vector.outerproduct {{.*}}, {{.*}}, {{.*}} : vector<4xf32>, vector<8xf32>
   %1 = vector.outerproduct %arg0, %arg1, %arg2 : vector<4xf32>, vector<8xf32>
   return %1 : vector<4x8xf32>
 }