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

[nicolasvasilache]

Summary

This revision starts making concrete use of 0-d vectors to extend the semantics of
ExtractElementOp.
In the process a new VectorOfAnyRank Tablegen OpBase.td is added to allow progressive transition to supporting 0-d vectors by gradually opting in.

[nicolasvasilache]

Rebase past windows build error.

[dcaballe]

mlir/include/mlir/IR/OpBase.td

216

For how long do you think we would need this temporary definitions?

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

424

I'm probably missing something obvious but why do we extract position 1 instead of 0?

[nicolasvasilache]

Bugfixes.

mlir/include/mlir/IR/OpBase.td

216

this can go away once ops that want to use 0-d are updated and the footguns are removed.

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

424

err.. because, you know .. bugs. Thanks!

[gysit]

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

424

I am surprised the index is not zero? From the tests above it seems like LLVM uses zero based indexing?

The example:
<result> = extractelement <4 x i32> %vec, i32 0 ; yields i32
from https://llvm.org/docs/LangRef.html#extractelement-instruction also uses index 0.

[nicolasvasilache]

@gysit @dcaballe thanks for spotting the bug, it is now fixed.

I find it quite scary that LLVM dit not complain a second about me accessing element 1 of a vector<1xf32> and produced the result I was expecting ...

[gysit]

In D114387#3148656, @nicolasvasilache wrote:

I find it quite scary that LLVM dit not complain a second about me accessing element 1 of a vector<1xf32> and produced the result I was expecting ...

True! I assume it is lowered to n-d hardware vector and the scalar is broadcasted to all lanes?

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

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

mlir/include/mlir/IR/OpBase.td

 // * Use uncategorized constraint to specify
 //   * Multi-entity constraints in rewrite rules
 
 //===----------------------------------------------------------------------===//
 // Common predicates
 //===----------------------------------------------------------------------===//
 
 // Whether a type is a VectorType.
-def IsVectorTypePred : CPred<"$_self.isa<::mlir::VectorType>()">;
+// Explicitly disallow 0-D vectors for now until we have good enough coverage.
+def IsVectorTypePred : And<[CPred<"$_self.isa<::mlir::VectorType>()">,
+                            CPred<"$_self.cast<::mlir::VectorType>().getRank() > 0">]>;
+
+// Temporary vector type clone that allows gradual transition to 0-D vectors.
+def IsVectorOfAnyRankTypePred : CPred<"$_self.isa<::mlir::VectorType>()">;
 
 // Whether a type is a TensorType.
 def IsTensorTypePred : CPred<"$_self.isa<::mlir::TensorType>()">;
 
 // Whether a type is a MemRefType.
 def IsMemRefTypePred : CPred<"$_self.isa<::mlir::MemRefType>()">;
 
 // Whether a type is an UnrankedMemRefType
Or<!foreach(rank, ranks,
                          == }]
                       # rank>)>]>;
 
 // Vector types.
 
 class VectorOf<list<Type> allowedTypes> :
   ShapedContainerType<allowedTypes, IsVectorTypePred, "vector",
                       "::mlir::VectorType">;
+// Temporary vector type clone that allows gradual transition to 0-D vectors.
+class VectorOfAnyRankOf<list<Type> allowedTypes> :
+  ShapedContainerType<allowedTypes, IsVectorOfAnyRankTypePred, "vector",
+                      "::mlir::VectorType">;
 
 // Whether the number of elements of a vector is from the given
 // `allowedRanks` list
 class IsVectorOfRankPred<list<int> allowedRanks> :
   And<[IsVectorTypePred,
        Or<!foreach(allowedlength, allowedRanks,
                    CPred<[{$_self.cast<::mlir::VectorType>().getRank()
                            == }]

 class VectorOfLengthAndType<list<int> allowedLengths,
                             list<Type> allowedTypes> : Type<
   And<[VectorOf<allowedTypes>.predicate,
        VectorOfLength<allowedLengths>.predicate]>,
   VectorOf<allowedTypes>.summary # VectorOfLength<allowedLengths>.summary,
   "::mlir::VectorType">;
 
 def AnyVector : VectorOf<[AnyType]>;
+// Temporary vector type clone that allows gradual transition to 0-D vectors.
+def AnyVectorOfAnyRank : VectorOfAnyRankOf<[AnyType]>;
 
 // Shaped types.
 
 def AnyShaped: ShapedContainerType<[AnyType], IsShapedTypePred, "shaped",
                                    "::mlir::ShapedType">;
 
 // Tensor types.
 

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

if (ShapedType::isDynamicStrideOrOffset(offset))
     return {};
 
   Type elementType = convertType(type.getElementType());
   if (!elementType)
     return {};
   return LLVM::LLVMPointerType::get(elementType, type.getMemorySpaceAsInt());
 }
 
-/// Convert an n-D vector type to an LLVM vector type via (n-1)-D array type
-/// when n > 1. For example, `vector<4 x f32>` remains as is while,
-/// `vector<4x8x16xf32>` converts to `!llvm.array<4xarray<8 x vector<16xf32>>>`.
+/// Convert an n-D vector type to an LLVM vector type:
+///  * 0-D `vector<T>` are converted to vector<1xT>
+///  * 1-D `vector<axT>` remains as is while,
+///  * n>1 `vector<ax...xkxT>` convert via an (n-1)-D array type to
+///    `!llvm.array<ax...array<jxvector<kxT>>>`.
 Type LLVMTypeConverter::convertVectorType(VectorType type) {
   auto elementType = convertType(type.getElementType());
   if (!elementType)
     return {};
+  if (type.getShape().empty())
+    return VectorType::get({1}, elementType);
   Type vectorType = VectorType::get(type.getShape().back(), elementType);
   assert(LLVM::isCompatibleVectorType(vectorType) &&
          "expected vector type compatible with the LLVM dialect");
   auto shape = type.getShape();
   for (int i = shape.size() - 2; i >= 0; --i)
     vectorType = LLVM::LLVMArrayType::get(vectorType, shape[i]);
   return vectorType;
 }

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

static VectorType reducedVectorTypeBack(VectorType tp) {
   return VectorType::get(tp.getShape().take_back(), tp.getElementType());
 }
 
 // Helper that picks the proper sequence for inserting.
 static Value insertOne(ConversionPatternRewriter &rewriter,
                        LLVMTypeConverter &typeConverter, Location loc,
                        Value val1, Value val2, Type llvmType, int64_t rank,
                        int64_t pos) {
+  assert(rank > 0 && "0-D vector corner case should have been handled already");
   if (rank == 1) {
     auto idxType = rewriter.getIndexType();
     auto constant = rewriter.create<LLVM::ConstantOp>(
         loc, typeConverter.convertType(idxType),
         rewriter.getIntegerAttr(idxType, pos));
     return rewriter.create<LLVM::InsertElementOp>(loc, llvmType, val1, val2,
                                                   constant);
   }
   return rewriter.create<LLVM::InsertValueOp>(loc, llvmType, val1, val2,
                                               rewriter.getI64ArrayAttr(pos));
 }
 
 // Helper that picks the proper sequence for extracting.
 static Value extractOne(ConversionPatternRewriter &rewriter,
                         LLVMTypeConverter &typeConverter, Location loc,
                         Value val, Type llvmType, int64_t rank, int64_t pos) {
+  assert(rank > 0 && "0-D vector corner case should have been handled already");
   if (rank == 1) {
     auto idxType = rewriter.getIndexType();
     auto constant = rewriter.create<LLVM::ConstantOp>(
         loc, typeConverter.convertType(idxType),
         rewriter.getIntegerAttr(idxType, pos));
     return rewriter.create<LLVM::ExtractElementOp>(loc, llvmType, val,
                                                    constant);
   }
matchAndRewrite(vector::ExtractElementOp extractEltOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto vectorType = extractEltOp.getVectorType();
     auto llvmType = typeConverter->convertType(vectorType.getElementType());
 
     // Bail if result type cannot be lowered.
     if (!llvmType)
       return failure();
 
+    if (vectorType.getRank() == 0) {
+      Location loc = extractEltOp.getLoc();
+      auto idxType = rewriter.getIndexType();
+      auto zero = rewriter.create<LLVM::ConstantOp>(
+          loc, typeConverter->convertType(idxType),
+          rewriter.getIntegerAttr(idxType, 0));
+      rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>(
+          extractEltOp, llvmType, adaptor.vector(), zero);
+      return success();
+    }
+
     rewriter.replaceOpWithNewOp<LLVM::ExtractElementOp>(
         extractEltOp, llvmType, adaptor.vector(), adaptor.position());
     return success();
   }
 };
 
 class VectorExtractOpConversion
     : public ConvertOpToLLVMPattern<vector::ExtractOp> {

mlir/lib/Dialect/Vector/VectorOps.cpp

results.add<CanonicalizeContractAdd<arith::AddIOp>,
               CanonicalizeContractAdd<arith::AddFOp>>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
 void vector::ExtractElementOp::build(OpBuilder &builder, OperationState &result,
+                                     Value source) {
+  result.addOperands({source});
+  result.addTypes(source.getType().cast<VectorType>().getElementType());
+}
+
+void vector::ExtractElementOp::build(OpBuilder &builder, OperationState &result,
                                      Value source, Value position) {
   result.addOperands({source, position});
   result.addTypes(source.getType().cast<VectorType>().getElementType());
 }
 
 static LogicalResult verify(vector::ExtractElementOp op) {
   VectorType vectorType = op.getVectorType();
+  if (vectorType.getRank() == 0) {
+    if (op.position())
+      return op.emitOpError("expected position to be empty with 0-D vector");
+    return success();
+  }
   if (vectorType.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();
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractOp
 //===----------------------------------------------------------------------===//
 
 static Type inferExtractOpResultType(VectorType vectorType,

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

 //       CHECK:   %[[i2:.*]] = llvm.insertvalue %[[e2]], %[[i1]][1] : !llvm.array<3 x vector<4xf32>>
 //       CHECK:   %[[e3:.*]] = llvm.extractvalue %[[VAL_1]][1] : !llvm.array<2 x vector<4xf32>>
 //       CHECK:   %[[i3:.*]] = llvm.insertvalue %[[e3]], %[[i2]][2] : !llvm.array<3 x vector<4xf32>>
 //       CHECK:   %[[VAL_3:.*]] = builtin.unrealized_conversion_cast %[[i3]] : !llvm.array<3 x vector<4xf32>> to vector<3x4xf32>
 //       CHECK:   return %[[VAL_3]] : vector<3x4xf32>
 
 // -----
 
+// CHECK-LABEL: @extract_element_0d
+func @extract_element_0d(%a: vector<f32>) -> f32 {
+  // CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : index) : i64
+  // CHECK: llvm.extractelement %{{.*}}[%[[C0]] : {{.*}}] : vector<1xf32>
+  %1 = vector.extractelement %a[] : vector<f32>
+  return %1 : f32
+}
+
+// -----
+
 func @extract_element(%arg0: vector<16xf32>) -> f32 {
   %0 = arith.constant 15 : i32
   %1 = vector.extractelement %arg0[%0 : i32]: vector<16xf32>
   return %1 : f32
 }
 // CHECK-LABEL: @extract_element(
 // CHECK-SAME: %[[A:.*]]: vector<16xf32>)
 //       CHECK:   %[[c:.*]] = arith.constant 15 : i32

mlir/test/Dialect/Vector/invalid.mlir

 
 func @shuffle_empty_mask(%arg0: vector<2xf32>, %arg1: vector<2xf32>) {
   // expected-error@+1 {{'vector.shuffle' invalid mask length}}
   %1 = vector.shuffle %arg0, %arg1 [] : vector<2xf32>, vector<2xf32>
 }
 
 // -----
 
+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_element(%arg0: vector<4x4xf32>) {
   %c = arith.constant 3 : i32
-  // expected-error@+1 {{'vector.extractelement' op expected 1-D vector}}
+  // expected-error@+1 {{unexpected >1 vector rank}}
   %1 = vector.extractelement %arg0[%c : i32] : vector<4x4xf32>
 }
 
 // -----
 
 func @extract_vector_type(%arg0: index) {
   // expected-error@+1 {{expected vector type}}
   %1 = vector.extract %arg0[] : index

mlir/test/Dialect/Vector/ops.mlir

 
 // CHECK-LABEL: @shuffle2D
 func @shuffle2D(%a: vector<1x4xf32>, %b: vector<2x4xf32>) -> vector<3x4xf32> {
   // CHECK: vector.shuffle %{{.*}}, %{{.*}}[0, 1, 2] : vector<1x4xf32>, vector<2x4xf32>
   %1 = vector.shuffle %a, %b[0, 1, 2] : vector<1x4xf32>, vector<2x4xf32>
   return %1 : vector<3x4xf32>
 }
 
+// CHECK-LABEL: @extract_element_0d
+func @extract_element_0d(%a: vector<f32>) -> f32 {
+  // CHECK-NEXT: vector.extractelement %{{.*}}[] : vector<f32>
+  %1 = vector.extractelement %a[] : vector<f32>
+  return %1 : f32
+}
+
 // CHECK-LABEL: @extract_element
 func @extract_element(%a: vector<16xf32>) -> f32 {
   // CHECK:      %[[C15:.*]] = arith.constant 15 : i32
   %c = arith.constant 15 : i32
   // CHECK-NEXT: vector.extractelement %{{.*}}[%[[C15]] : i32] : vector<16xf32>
   %1 = vector.extractelement %a[%c : i32] : vector<16xf32>
   return %1 : f32
 }

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 @entry() {
+  %1 = arith.constant dense<42.0> : vector<f32>
+  call  @extract_element_0d(%1) : (vector<f32>) -> ()
+  return
+}