[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

703

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

703

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 ↗

(On Diff #530289)

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

703

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

697–698

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

697–698

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 ↗

(On Diff #530289)

[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>
     ```
   }];
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. Can degenerate to a scalar 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>
     ```
   }];

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>) -> vector<f32> {
+  // CHECK-NEXT: vector.insert %{{.*}}, %{{.*}}[] : f32 into vector<f32>
+  %1 = vector.insert %a,  %b[] : f32 into vector<f32>
+  return %1 : vector<f32>
+}
+
 // 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>
 }