[mlir][VectorOps][EDSC] Add EDSC for VectorOps

[nicolasvasilache]

Summary

This revision adds EDSC support for VectorOps to enable the creation of a vector_matmul declaratively. The vector_matmul is a simple configuration
of the vector.contract op that follows the StructuredOps abstraction.

[nicolasvasilache]

Rebase.

[nicolasvasilache]

mlir/include/mlir/IR/AffineMap.h

72

I am really not happy about this duplication but I tried a bunch of different things and could not get it down to one single form.
@rriddle @mehdi_amini suggestions most welcome!

mlir/lib/Dialect/VectorOps/EDSC/Builders.cpp

31

One may think it is not great to materialize here but after spending too much time, I got to the following and still got compilation errors with issues converting to const StringRef *.
At this point I find it counterproductive and prefer the current form.

using IteratorTypeList = ArrayRef<IteratorType>;
class iterator_value_iterator final
    : llvm::mapped_iterator<IteratorTypeList::iterator,
                            StringRef (*)(IteratorType)> {
public:
  explicit iterator_value_iterator(IteratorTypeList::iterator it)
      : llvm::mapped_iterator<IteratorTypeList::iterator,
                              StringRef (*)(IteratorType)>(
            it, [](IteratorType iter) { return toString(iter); }) {}
  StringRef operator*() const { return toString(*this->I); }
};
 
llvm::iterator_range<iterator_value_iterator>
makeRange(ArrayRef<IteratorType> iteratorTypes) {
  return llvm::make_range(iterator_value_iterator(iteratorTypes.begin()),
                          iterator_value_iterator(iteratorTypes.end()));
}
 
Value mlir::edsc::ops::vector_contraction(
    StructuredIndexed A, StructuredIndexed B, StructuredIndexed C,
    ArrayRef<IteratorType> iteratorTypes) {
  using IndexingExprs = ArrayRef<ArrayRef<AffineExpr>>;
  auto range = makeRange(iteratorTypes);
  return vector_contract(
      A.getValue(), B.getValue(), C.getValue(),
      IndexingExprs{A.getExprs(), B.getExprs(), C.getExprs()},
      ArrayRef<StringRef>{range.begin(), range.end()});
}

[mehdi_amini]

mlir/include/mlir/IR/AffineMap.h

66

You wrote "Returns the AffineMap" but it returns a vector of AffineMap?

mlir/lib/Dialect/VectorOps/EDSC/Builders.cpp

31

What about materializing directly as an Array of StringAttr here instead of doing it in the builder?
Right now you materialize here first and then the builder will trigger another call to functional::map()

mlir/lib/IR/AffineMap.cpp

147

Using templates could avoid the temporary exprsVector, but I'm not sure it really matters...

[nicolasvasilache]

mlir/lib/Dialect/VectorOps/EDSC/Builders.cpp

31

That falls into the category of multiple attrs and having builders for multiple combinations of attrs vs underlying values.
Ideally I wouldn't fall in 2^n behavior here.
So either I'd go full attributes or full underlying values and reconstruct attributes, which is what I went for here.

Then orthogonally, the materialization is really a "proper API with ranges" issues on which I gave up in this revision.
If we were to care about optimization of this particular thing I'd prefer more effort into making ranges more usable.

mlir/lib/IR/AffineMap.cpp

147

Good enough, thanks!

[nicolasvasilache]

Address review comments.

[nicolasvasilache]

Flushing out some answers.

[mehdi_amini]

mlir/lib/Dialect/VectorOps/EDSC/Builders.cpp

31

I'm not sure what you mean: your builder API is encouraging a more costly model: you materialize two SmallVector instead of one.
Just *not* adding your builder seems to both reduce the number of builder and increasing efficiency.

mlir/include/mlir/Dialect/Linalg/EDSC/Builders.h

private:
   using BuilderType =
       typename std::conditional_t<std::is_same<LoopTy, loop::ParallelOp>::value,
                                   ParallelLoopNestBuilder,
                                   LoopOrAffineLoopBuilder>;
 
   std::unique_ptr<BuilderType> builder;
 };
 
-enum class IterType { Parallel, Reduction };
-
-inline StringRef toString(IterType t) {
-  switch (t) {
-  case IterType::Parallel:
-    return getParallelIteratorTypeName();
-  case IterType::Reduction:
-    return getReductionIteratorTypeName();
-  }
-  llvm_unreachable("Unsupported IterType");
-}
-
-/// A StructuredIndexed represents an indexable quantity that is either:
-/// 1. a captured value, which is suitable for buffer and tensor operands, or;
-/// 2. a captured type, which is suitable for tensor return values.
-///
-/// A StructuredIndexed itself is indexed and passed to `makeGenericLinalgOp`.
-/// It enable an idiomatic syntax for index expressions such as:
-///
-/// ```
-///      StructuredIndexed A(buffer_or_tensor_value), B(buffer_or_tensor_value),
-///        C(buffer_value_or_tensor_type);
-///      makeGenericLinalgOp({A({m, n}), B({k, n})}, {C({m, n})}, ... );
-/// ```
-struct StructuredIndexed : public ValueHandle {
-  StructuredIndexed(Type type) : ValueHandle(type) {}
-  StructuredIndexed(Value value) : ValueHandle(value) {}
-  StructuredIndexed(ValueHandle valueHandle) : ValueHandle(valueHandle) {}
-  StructuredIndexed operator()(ArrayRef<AffineExpr> indexings) {
-    return StructuredIndexed(*this, indexings);
-  }
-
-  ArrayRef<AffineExpr> getExprs() { return exprs; }
-
-private:
-  StructuredIndexed(Type t, ArrayRef<AffineExpr> indexings)
-      : ValueHandle(t), exprs(indexings.begin(), indexings.end()) {
-    assert(t.isa<RankedTensorType>() && "RankedTensor expected");
-  }
-  StructuredIndexed(Value v, ArrayRef<AffineExpr> indexings)
-      : ValueHandle(v), exprs(indexings.begin(), indexings.end()) {
-    assert((v.getType().isa<MemRefType>() ||
-            v.getType().isa<RankedTensorType>()) &&
-           "MemRef or RankedTensor expected");
-  }
-  StructuredIndexed(ValueHandle vh, ArrayRef<AffineExpr> indexings)
-      : ValueHandle(vh), exprs(indexings.begin(), indexings.end()) {}
-
-  SmallVector<AffineExpr, 4> exprs;
-};
-
 inline void defaultRegionBuilder(ArrayRef<BlockArgument> args) {}
 
 /// Build a `linalg.generic` op with the specified `inputs`, `outputs` and
 /// `region`.
 ///
 /// `otherValues` and `otherAttributes` may be passed and will be appended as
 /// operands and attributes respectively.
 ///
 /// Prerequisites:
 /// =============
 ///
 /// 1. `inputs` may contain StructuredIndexed that capture either buffer or
 /// tensor values.
 /// 2. `outputs` may contain StructuredIndexed that capture either buffer values
 /// or tensor types. If both buffer values and tensor types are present, then
 /// all buffer values must appear before any tensor type. Without this
 /// restriction output tensor results would need to be reordered, which would
 /// result in surprising behavior when combined with region definition.
 Operation *makeGenericLinalgOp(
-    ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
+    ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder =
         defaultRegionBuilder,
     ArrayRef<Value> otherValues = {}, ArrayRef<Attribute> otherAttributes = {});
 
 namespace ops {
 using edsc::StructuredIndexed;
 using edsc::ValueHandle;

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

 
 inline unsigned getNumIterators(ArrayAttr iteratorTypes) {
   unsigned res = 0;
   for (auto n : getAllIteratorTypeNames())
     res += getNumIterators(n, iteratorTypes);
   return res;
 }
 
+/// Typed representation for loop type strings.
+enum class IteratorType { Parallel, Reduction };
+
+inline StringRef toString(IteratorType t) {
+  switch (t) {
+  case IteratorType::Parallel:
+    return getParallelIteratorTypeName();
+  case IteratorType::Reduction:
+    return getReductionIteratorTypeName();
+  }
+  llvm_unreachable("Unsupported IteratorType");
+}
+
 } // end namespace mlir
 
 #endif // MLIR_UTILS_STRUCTUREDOPSUTILS_H

mlir/include/mlir/Dialect/VectorOps/EDSC/Builders.h

+//===- Builders.h - MLIR Declarative Vector Builders ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Provides intuitive composable interfaces for building structured MLIR
+// snippets in a declarative fashion.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_VECTOR_EDSC_BUILDERS_H_
+#define MLIR_DIALECT_VECTOR_EDSC_BUILDERS_H_
+
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/EDSC/Builders.h"
+#include "mlir/EDSC/Intrinsics.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/Builders.h"
+
+namespace mlir {
+namespace edsc {
+namespace ops {
+
+/// Build a generic vector contraction, that is a `vector.contract` op with
+/// specified `iteratorTypes`. The client is responsible for specifying proper
+/// indexings when creating the StructuredIndexed.
+/// The computation represents a notional (A * B + C) where indexings specify
+/// which dimensions are reduced and reordered.
+/// Return the result of the `vector.contract` op
+///
+/// Prerequisites:
+/// A, B and C capture values of proper vector types, and indexing expressions
+/// that match semantics of the `vector.contract` op.
+Value vector_contraction(StructuredIndexed A, StructuredIndexed B,
+                         StructuredIndexed C,
+                         ArrayRef<IteratorType> iteratorTypes);
+
+/// Build a generic vector contraction that computes a matmul on vectors.
+/// Return the result of C(i, j) + sum_k {A(i, k) * B(k, j)} on vectors.
+///
+/// Prerequisites:
+/// A, B and C capture values of proper vector types. For instance
+/// `A: vector<4x8xf32>`, `B: vector<8x16f32>` and `C: vector<4x16xf32>`.
+Value vector_matmul(Value A, Value B, Value C);
+
+} // namespace ops
+} // namespace edsc
+} // namespace mlir
+
+#endif // MLIR_DIALECT_VECTOR_EDSC_BUILDERS_H_

mlir/include/mlir/Dialect/VectorOps/EDSC/Intrinsics.h

+//===- Intrinsics.h - MLIR EDSC Intrinsics for VectorOps --------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_VECTOROPS_EDSC_INTRINSICS_H_
+#define MLIR_DIALECT_VECTOROPS_EDSC_INTRINSICS_H_
+
+#include "mlir/Dialect/VectorOps/VectorOps.h"
+#include "mlir/EDSC/Builders.h"
+#include "mlir/EDSC/Intrinsics.h"
+
+namespace mlir {
+namespace edsc {
+namespace intrinsics {
+
+using vector_contract = ValueBuilder<vector::ContractionOp>;
+
+} // namespace intrinsics
+} // namespace edsc
+} // namespace mlir
+
+#endif // MLIR_DIALECT_VECTOROPS_EDSC_INTRINSICS_H_

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

Examples:
       %lhs_mask = vector.constant_mask [7, 8, 16, 15] : vector<7x8x16x15xi1>
       %rhs_mask = vector.constant_mask [8, 16, 7, 5] : vector<8x16x7x5xi1>
 
       %5 = vector.contract #contraction_trait %0, %1, %2, %lhs_mask, %rhs_mask
          : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x8x5xf32>
   }];
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value lhs, Value rhs, "
-    "Value acc, ArrayAttr indexingMaps, ArrayAttr iteratorTypes">];
+    "Value acc, ArrayAttr indexingMaps, ArrayAttr iteratorTypes">,
+    OpBuilder<
+      "Builder *builder, OperationState &result, Value lhs, Value rhs, "
+      "Value acc, ArrayRef<ArrayRef<AffineExpr>> indexingExprs, "
+      "ArrayRef<StringRef> iteratorTypes">];
   let extraClassDeclaration = [{
     VectorType getLhsType() {
       return lhs().getType().cast<VectorType>();
     }
     VectorType getRhsType() {
       return rhs().getType().cast<VectorType>();
     }
     Type getAccType() { return acc().getType(); }

mlir/include/mlir/EDSC/Builders.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_EDSC_BUILDERS_H_
 #define MLIR_EDSC_BUILDERS_H_
 
 #include "mlir/Dialect/AffineOps/AffineOps.h"
 #include "mlir/Dialect/LoopOps/LoopOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 namespace mlir {
 
 namespace edsc {
 
 struct index_type {
public:
   operator bool() { return block != nullptr; }
   operator mlir::Block *() { return block; }
   mlir::Block *getBlock() { return block; }
 
 private:
   mlir::Block *block;
 };
 
+/// A StructuredIndexed represents an indexable quantity that is either:
+/// 1. a captured value, which is suitable for buffer and tensor operands, or;
+/// 2. a captured type, which is suitable for tensor return values.
+///
+/// A StructuredIndexed itself is indexed and passed to `makeGenericLinalgOp`.
+/// It enable an idiomatic syntax for index expressions such as:
+///
+/// ```
+///      StructuredIndexed A(buffer_or_tensor_value), B(buffer_or_tensor_value),
+///        C(buffer_value_or_tensor_type);
+///      makeGenericLinalgOp({A({m, n}), B({k, n})}, {C({m, n})}, ... );
+/// ```
+struct StructuredIndexed : public ValueHandle {
+  StructuredIndexed(Type type) : ValueHandle(type) {}
+  StructuredIndexed(Value value) : ValueHandle(value) {}
+  StructuredIndexed(ValueHandle valueHandle) : ValueHandle(valueHandle) {}
+  StructuredIndexed operator()(ArrayRef<AffineExpr> indexings) {
+    return this->hasValue() ? StructuredIndexed(this->getValue(), indexings)
+                            : StructuredIndexed(this->getType(), indexings);
+  }
+
+  StructuredIndexed(Type t, ArrayRef<AffineExpr> indexings)
+      : ValueHandle(t), exprs(indexings.begin(), indexings.end()) {
+    assert(t.isa<RankedTensorType>() && "RankedTensor expected");
+  }
+  StructuredIndexed(Value v, ArrayRef<AffineExpr> indexings)
+      : ValueHandle(v), exprs(indexings.begin(), indexings.end()) {
+    assert((v.getType().isa<MemRefType>() ||
+            v.getType().isa<RankedTensorType>() ||
+            v.getType().isa<VectorType>()) &&
+           "MemRef, RankedTensor or Vector expected");
+  }
+  StructuredIndexed(ValueHandle vh, ArrayRef<AffineExpr> indexings)
+      : ValueHandle(vh), exprs(indexings.begin(), indexings.end()) {}
+
+  ArrayRef<AffineExpr> getExprs() { return exprs; }
+
+private:
+  SmallVector<AffineExpr, 4> exprs;
+};
+
 template <typename Op, typename... Args>
 OperationHandle OperationHandle::create(Args... args) {
   return OperationHandle(ScopedContext::getBuilder()
                              .create<Op>(ScopedContext::getLocation(), args...)
                              .getOperation());
 }
 
 template <typename Op, typename... Args>

mlir/include/mlir/IR/AffineMap.h

public:
   /// The permutation is expressed as a non-empty vector of integers.
   /// E.g. the permutation `(i,j,k) -> (j,k,i)` will be expressed with
   /// `permutation = [1,2,0]`. All values in `permutation` must be
   /// integers, in the range 0..`permutation.size()-1` without duplications
   /// (i.e. `[1,1,2]` is an invalid permutation).
   static AffineMap getPermutationMap(ArrayRef<unsigned> permutation,
                                      MLIRContext *context);
 
+  /// Returns the AffineMap with as many results as `exprs.size()`, as many dims
+  /// as the largest dim in `exprs` and as many symbols as the largest symbol in
+  /// `exprs`.
+  static SmallVector<AffineMap, 4>
+  inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList);
+  static SmallVector<AffineMap, 4>
+  inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList);
+
   MLIRContext *getContext() const;
 
   explicit operator bool() { return map != nullptr; }
   bool operator==(AffineMap other) const { return other.map == map; }
   bool operator!=(AffineMap other) const { return !(other.map == map); }
 
   /// Returns true if this affine map is an identity affine map.
   /// An identity affine map corresponds to an identity affine function on the

mlir/lib/Dialect/Linalg/EDSC/Builders.cpp

 //===- Builders.cpp - MLIR Declarative Linalg Builders --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/EDSC/Builders.h"
 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/EDSC/Builders.h"
 #include "mlir/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Support/Functional.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
for (Value range : ranges) {
     steps.emplace_back(rangeOp.step());
   }
   builder = std::make_unique<ParallelLoopNestBuilder>(ivs, lbs, ubs, steps);
 }
 
 } // namespace edsc
 } // namespace mlir
 
-static void getMaxDimIndex(ArrayRef<StructuredIndexed> structuredIndices,
-                           unsigned &pos) {
-  for (auto sidx : structuredIndices) {
-    for (auto expr : sidx.getExprs()) {
-      expr.walk([&pos](AffineExpr e) {
-        if (auto d = e.dyn_cast<AffineDimExpr>())
-          pos = std::max(pos, d.getPosition());
-      });
-    }
-  }
-}
-
 Operation *mlir::edsc::makeGenericLinalgOp(
-    ArrayRef<IterType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
+    ArrayRef<IteratorType> iteratorTypes, ArrayRef<StructuredIndexed> inputs,
     ArrayRef<StructuredIndexed> outputs,
     function_ref<void(ArrayRef<BlockArgument>)> regionBuilder,
     ArrayRef<Value> otherValues, ArrayRef<Attribute> otherAttributes) {
   for (unsigned i = 0, e = outputs.size(); i + 1 < e; ++i)
     assert(!(outputs[i].getType().isa<RankedTensorType>() &&
              outputs[i + 1].getType().isa<MemRefType>()) &&
            "output tensors must be passed after output buffers");
   auto &builder = edsc::ScopedContext::getBuilder();
   auto *ctx = builder.getContext();
   unsigned nInputs = inputs.size();
   unsigned nOutputs = outputs.size();
-  unsigned maxPos = 0;
-  getMaxDimIndex(inputs, maxPos);
-  getMaxDimIndex(outputs, maxPos);
-  // maxPos is 0 indexed, need to turn this into a count (i.e. +1)
-  unsigned nDims = maxPos + 1;
-
-  SmallVector<AffineMap, 4> maps;
-  maps.reserve(nInputs + nOutputs);
-  for (auto in : inputs)
-    maps.push_back(
-        AffineMap::get(/*dimCount=*/nDims, /*symbolCount=*/0, in.getExprs()));
-  for (auto out : outputs)
-    maps.push_back(
-        AffineMap::get(/*dimCount=*/nDims, /*symbolCount=*/0, out.getExprs()));
+
+  SmallVector<SmallVector<AffineExpr, 4>, 4> exprsList;
+  exprsList.reserve(nInputs + nOutputs);
+  for (auto structuredIndexed : inputs)
+    exprsList.emplace_back(structuredIndexed.getExprs().begin(),
+                           structuredIndexed.getExprs().end());
+  for (auto structuredIndexed : outputs)
+    exprsList.emplace_back(structuredIndexed.getExprs().begin(),
+                           structuredIndexed.getExprs().end());
+  auto maps = AffineMap::inferFromExprList(exprsList);
 
   unsigned nViews = nInputs + nOutputs;
   SmallVector<Value, 4> values;
   values.reserve(nViews);
   values.append(inputs.begin(), inputs.end());
   std::copy_if(outputs.begin(), outputs.end(), std::back_inserter(values),
                [](StructuredIndexed s) { return s.hasValue(); });
   SmallVector<Type, 4> types;
void mlir::edsc::ops::macRegionBuilder(ArrayRef<BlockArgument> args) {
   assert(args.size() == 3 && "expected 3 block arguments");
   ValueHandle a(args[0]), b(args[1]), c(args[2]);
   linalg_yield((c + a * b).getValue());
 }
 
 Operation *mlir::edsc::ops::linalg_pointwise(UnaryPointwiseOpBuilder unaryOp,
                                              StructuredIndexed I,
                                              StructuredIndexed O) {
-  SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
-                                           edsc::IterType::Parallel);
+  SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
+                                         IteratorType::Parallel);
   if (O.getType().isa<RankedTensorType>()) {
     auto fun = [&unaryOp](ArrayRef<BlockArgument> args) {
       assert(args.size() == 1 && "expected 1 block arguments");
       ValueHandle a(args[0]);
       linalg_yield(unaryOp(a));
     };
     return makeGenericLinalgOp(iterTypes, {I}, {O}, fun);
   }
Operation *mlir::edsc::ops::linalg_pointwise_tanh(StructuredIndexed I,
   return linalg_pointwise(unOp, I, O);
 }
 
 /// Binary pointwise operation (with broadcast) entry point.
 Operation *mlir::edsc::ops::linalg_pointwise(BinaryPointwiseOpBuilder binaryOp,
                                              StructuredIndexed I1,
                                              StructuredIndexed I2,
                                              StructuredIndexed O) {
-  SmallVector<edsc::IterType, 4> iterTypes(O.getExprs().size(),
-                                           edsc::IterType::Parallel);
+  SmallVector<IteratorType, 4> iterTypes(O.getExprs().size(),
+                                         IteratorType::Parallel);
   if (O.getType().isa<RankedTensorType>()) {
     auto fun = [&binaryOp](ArrayRef<BlockArgument> args) {
       assert(args.size() == 2 && "expected 2 block arguments");
       ValueHandle a(args[0]), b(args[1]);
       linalg_yield(binaryOp(a, b));
     };
     return makeGenericLinalgOp(iterTypes, {I1, I2}, {O}, fun);
   }

 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
                                           ValueHandle vC) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(vC);
   return makeGenericLinalgOp(
-    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+    {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n})},
     {C({m, n})},
     macRegionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
                                           RankedTensorType tC) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(tC);
   return makeGenericLinalgOp(
-    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+    {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n})},
     {C({m, n})},
     mulRegionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_matmul(ValueHandle vA, ValueHandle vB,
                                           ValueHandle vC, RankedTensorType tD) {
   // clang-format off
   AffineExpr m, n, k;
   bindDims(ScopedContext::getContext(), m, n, k);
   StructuredIndexed A(vA), B(vB), C(vC), D(tD);
   return makeGenericLinalgOp(
-    {IterType::Parallel, IterType::Parallel, IterType::Reduction},
+    {IteratorType::Parallel, IteratorType::Parallel, IteratorType::Reduction},
     {A({m, k}), B({k, n}), C({m, n})},
     {D({m, n})},
     macRegionBuilder);
   // clang-format on
 }
 
 Operation *mlir::edsc::ops::linalg_conv_nhwc(ValueHandle vI, ValueHandle vW,
                                              ValueHandle vO,
                                              ArrayRef<int> strides,
                                              ArrayRef<int> dilations) {
   MLIRContext *ctx = ScopedContext::getContext();
   // TODO(ntv) some template magic to make everything rank-polymorphic.
   assert((dilations.empty() || dilations.size() == 2) && "only 2-D conv atm");
   assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");
 
   // Some short names.
-  auto par = IterType::Parallel;
-  auto red = IterType::Reduction;
+  auto par = IteratorType::Parallel;
+  auto red = IteratorType::Reduction;
   auto s = strides;
   auto d = dilations;
 
   AffineExpr b, f, h, w, kh, kw, c;
   bindDims(ctx, b, f, h, w, kh, kw, c);
   unsigned numDims = c.cast<AffineDimExpr>().getPosition() + 1;
   StructuredIndexed I(vI), W(vW), O(vO);
   // clang-format off
Operation *mlir::edsc::ops::linalg_dilated_conv_nhwc(
     ValueHandle vI, ValueHandle vW, ValueHandle vO, int depth_multiplier,
     ArrayRef<int> strides, ArrayRef<int> dilations) {
   MLIRContext *ctx = ScopedContext::getContext();
   // TODO(ntv) some template magic to make everything rank-polymorphic.
   assert((dilations.empty() || dilations.size() == 2) && "only 2-D conv atm");
   assert((strides.empty() || strides.size() == 2) && "only 2-D conv atm");
 
   // Some short names.
-  auto par = IterType::Parallel;
-  auto red = IterType::Reduction;
+  auto par = IteratorType::Parallel;
+  auto red = IteratorType::Reduction;
   auto s = strides;
   auto d = dilations;
 
   // clang-format off
   AffineExpr b, dm, c, h, w, kh, kw;
   bindDims(ctx, b, dm, c, h, w, kh, kw);
   unsigned numDims = kw.cast<AffineDimExpr>().getPosition() + 1;
   StructuredIndexed I(vI), W(vW), O(vO);

mlir/lib/Dialect/VectorOps/CMakeLists.txt

 add_llvm_library(MLIRVectorOps
   DialectRegistration.cpp
   VectorOps.cpp
   VectorTransforms.cpp
   VectorUtils.cpp
+  EDSC/Builders.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/VectorOps
   )
 
 add_dependencies(MLIRVectorOps MLIRVectorOpsIncGen)
 add_dependencies(MLIRVectorOps MLIRVectorTransformPatternsIncGen)
 
 target_link_libraries(MLIRVectorOps MLIRIR MLIRStandardOps MLIRAffineOps MLIRLoopOps MLIRLoopAnalysis)

mlir/lib/Dialect/VectorOps/EDSC/Builders.cpp

+//===- Builders.cpp - MLIR Declarative Linalg Builders --------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/VectorOps/EDSC/Builders.h"
+#include "mlir/Dialect/VectorOps/EDSC/Intrinsics.h"
+#include "mlir/Dialect/VectorOps/VectorOps.h"
+#include "mlir/EDSC/Builders.h"
+#include "mlir/EDSC/Intrinsics.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Support/Functional.h"
+
+using namespace mlir;
+using namespace mlir::edsc;
+using namespace mlir::edsc::intrinsics;
+using namespace mlir::edsc::ops;
+
+Value mlir::edsc::ops::vector_contraction(
+    StructuredIndexed A, StructuredIndexed B, StructuredIndexed C,
+    ArrayRef<IteratorType> iteratorTypes) {
+  using IndexingExprs = ArrayRef<ArrayRef<AffineExpr>>;
+  return vector_contract(
+      A.getValue(), B.getValue(), C.getValue(),
+      IndexingExprs{A.getExprs(), B.getExprs(), C.getExprs()},
+      ArrayRef<StringRef>{functional::map(toString, iteratorTypes)});
+}
+
+Value mlir::edsc::ops::vector_matmul(Value A, Value B, Value C) {
+  AffineExpr m, n, k;
+  bindDims(ScopedContext::getContext(), m, n, k);
+  return vector_contraction(StructuredIndexed(A, {m, k}),
+                            StructuredIndexed(B, {k, n}),
+                            StructuredIndexed(C, {m, n}),
+                            {IteratorType::Parallel, IteratorType::Parallel,
+                             IteratorType::Reduction});
+}

mlir/lib/Dialect/VectorOps/VectorOps.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // ContractionOp
 //===----------------------------------------------------------------------===//
 
 void vector::ContractionOp::build(Builder *builder, OperationState &result,
                                   Value lhs, Value rhs, Value acc,
+                                  ArrayRef<ArrayRef<AffineExpr>> indexingExprs,
+                                  ArrayRef<StringRef> iteratorTypes) {
+  result.addOperands({lhs, rhs, acc});
+  result.addTypes(acc.getType());
+  result.addAttribute(getIndexingMapsAttrName(),
+                      builder->getAffineMapArrayAttr(
+                          AffineMap::inferFromExprList(indexingExprs)));
+  result.addAttribute(getIteratorTypesAttrName(),
+                      builder->getStrArrayAttr(iteratorTypes));
+}
+
+void vector::ContractionOp::build(Builder *builder, OperationState &result,
+                                  Value lhs, Value rhs, Value acc,
                                   ArrayAttr indexingMaps,
                                   ArrayAttr iteratorTypes) {
   result.addOperands({lhs, rhs, acc});
   result.addTypes(acc.getType());
   result.addAttribute(getIndexingMapsAttrName(), indexingMaps);
   result.addAttribute(getIteratorTypesAttrName(), iteratorTypes);
 }
 

mlir/lib/IR/AffineMap.cpp

AffineMap AffineMap::getPermutationMap(ArrayRef<unsigned> permutation,
   for (auto index : permutation)
     affExprs.push_back(getAffineDimExpr(index, context));
   auto m = std::max_element(permutation.begin(), permutation.end());
   auto permutationMap = AffineMap::get(*m + 1, 0, affExprs);
   assert(permutationMap.isPermutation() && "Invalid permutation vector");
   return permutationMap;
 }
 
+static void getMaxDimAndSymbol(ArrayRef<ArrayRef<AffineExpr>> exprsList,
+                               int64_t &maxDim, int64_t &maxSym) {
+  for (auto exprs : exprsList) {
+    for (auto expr : exprs) {
+      expr.walk([&maxDim, &maxSym](AffineExpr e) {
+        if (auto d = e.dyn_cast<AffineDimExpr>())
+          maxDim = std::max(maxDim, static_cast<int64_t>(d.getPosition()));
+        if (auto s = e.dyn_cast<AffineSymbolExpr>())
+          maxSym = std::max(maxSym, static_cast<int64_t>(s.getPosition()));
+      });
+    }
+  }
+}
+
+SmallVector<AffineMap, 4>
+AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList) {
+  int64_t maxDim = -1, maxSym = -1;
+  getMaxDimAndSymbol(exprsList, maxDim, maxSym);
+  SmallVector<AffineMap, 4> maps;
+  maps.reserve(exprsList.size());
+  for (auto exprs : exprsList)
+    maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1,
+                                  /*symbolCount=*/maxSym + 1, exprs));
+  return maps;
+}
+
+SmallVector<AffineMap, 4>
+AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList) {
+  SmallVector<ArrayRef<AffineExpr>, 4> exprsVector;
+  exprsVector.reserve(exprsList.size());
+  for (const auto &el : exprsList)
+    exprsVector.emplace_back(el.begin(), el.end());
+  return inferFromExprList(exprsVector);
+}
+
 AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
                                             MLIRContext *context) {
   SmallVector<AffineExpr, 4> dimExprs;
   dimExprs.reserve(numDims);
   for (unsigned i = 0; i < numDims; ++i)
     dimExprs.push_back(mlir::getAffineDimExpr(i, context));
   return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs);
 }

mlir/test/EDSC/CMakeLists.txt

 add_llvm_executable(mlir-edsc-builder-api-test
   builder-api-test.cpp
 )
 
 llvm_update_compile_flags(mlir-edsc-builder-api-test)
 
 target_link_libraries(mlir-edsc-builder-api-test
   PRIVATE
   MLIRAffineOps
   MLIREDSC
   MLIRIR
   MLIRLinalgEDSC
   MLIRLinalgOps
   MLIRLoopOps
   MLIRStandardOps
   MLIRTransforms
+  MLIRVectorOps
   LLVMCore
   LLVMSupport
 )
 
 target_include_directories(mlir-edsc-builder-api-test PRIVATE ..)
 
 whole_archive_link(mlir-edsc-builder-api-test
   MLIRAffineOps
   MLIRLinalgOps
   MLIRLoopOps
   MLIRStandardOps
+  MLIRVectorOps
   MLIRTransforms
 )

mlir/test/EDSC/builder-api-test.cpp

 //===- builder-api-test.cpp - Tests for Declarative Builder APIs ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 // RUN: mlir-edsc-builder-api-test | FileCheck %s
 
 #include "mlir/Dialect/AffineOps/AffineOps.h"
 #include "mlir/Dialect/Linalg/EDSC/Builders.h"
 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
-#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
+#include "mlir/Dialect/VectorOps/EDSC/Builders.h"
+#include "mlir/Dialect/VectorOps/EDSC/Intrinsics.h"
 #include "mlir/EDSC/Builders.h"
 #include "mlir/EDSC/Helpers.h"
 #include "mlir/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"

 
 TEST_FUNC(zero_and_sign_extendi_op_i1_to_i8) {
   using namespace edsc;
   using namespace edsc::intrinsics;
   using namespace edsc::op;
   auto i1Type = IntegerType::get(1, &globalContext());
   auto i8Type = IntegerType::get(8, &globalContext());
   auto memrefType = MemRefType::get({}, i1Type, {}, 0);
-  auto f = makeFunction("zero_and_sign_extendi_op", {}, {memrefType, memrefType});
+  auto f =
+      makeFunction("zero_and_sign_extendi_op", {}, {memrefType, memrefType});
 
   OpBuilder builder(f.getBody());
   ScopedContext scope(builder, f.getLoc());
   IndexedValue A(f.getArgument(0));
   IndexedValue B(f.getArgument(1));
   // clang-format off
   edsc::intrinsics::zero_extendi(*A, i8Type);
   edsc::intrinsics::sign_extendi(*B, i8Type);
TEST_FUNC(linalg_tensors_test) {
   linalg_pointwise_tanh(SA({i, j}), SC({i, j}));
   Value o1 = linalg_matmul(A, B, tensorType)->getResult(0);
   linalg_matmul(A, B, ValueHandle(o1), tensorType);
 
   f.print(llvm::outs());
   f.erase();
 }
 
+// CHECK-LABEL: func @vector_matmul_test(
+//  CHECK-SAME:   %[[A:.*]]: vector<4x16xf32>,
+//  CHECK-SAME:   %[[B:.*]]: vector<16x8xf32>,
+//  CHECK-SAME:   %[[C:.*]]: vector<4x8xf32>)
+//  CHECK:   vector.contract {{.*}}[affine_map<(d0, d1, d2) -> (d0, d2)>,
+//  CHECK-SAME:                     affine_map<(d0, d1, d2) -> (d2, d1)>,
+//  CHECK-SAME:                     affine_map<(d0, d1, d2) -> (d0, d1)>],
+//  CHECK-SAME:              {{.*}}["parallel", "parallel", "reduction"]
+//  CHECK-SAME: %[[A]], %[[B]], %[[C]]
+//  CHECK-SAME:   vector<4x16xf32>, vector<16x8xf32> into vector<4x8xf32>
+TEST_FUNC(vector_matmul_test) {
+  using namespace edsc;
+  using namespace edsc::ops;
+
+  int64_t M = 4, N = 8, K = 16;
+  auto f32Type = FloatType::getF32(&globalContext());
+  auto mkVectorType = VectorType::get({M, K}, f32Type);
+  auto knVectorType = VectorType::get({K, N}, f32Type);
+  auto mnVectorType = VectorType::get({M, N}, f32Type);
+  auto f = makeFunction("vector_matmul_test", {},
+                        {mkVectorType, knVectorType, mnVectorType});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  ValueHandle A(f.getArgument(0)), B(f.getArgument(1)), C(f.getArgument(2));
+  vector_matmul(A, B, C);
+  f.print(llvm::outs());
+  f.erase();
+}
+
 int main() {
   RUN_TESTS();
   return 0;
 }