[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/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/Dialect/VectorOps/VectorOps.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/EDSC/Builders.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

struct StructuredIndexed : public ValueHandle {
 
   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");
+            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;

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 a vector of AffineMaps; each 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

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<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;

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;
 }
 
+template <typename AffineExprContainer>
+static void getMaxDimAndSymbol(ArrayRef<AffineExprContainer> exprsList,
+                               int64_t &maxDim, int64_t &maxSym) {
+  for (const 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()));
+      });
+    }
+  }
+}
+
+template <typename AffineExprContainer>
+SmallVector<AffineMap, 4>
+inferFromExprList(ArrayRef<AffineExprContainer> exprsList) {
+  int64_t maxDim = -1, maxSym = -1;
+  getMaxDimAndSymbol(exprsList, maxDim, maxSym);
+  SmallVector<AffineMap, 4> maps;
+  maps.reserve(exprsList.size());
+  for (const auto &exprs : exprsList)
+    maps.push_back(AffineMap::get(/*dimCount=*/maxDim + 1,
+                                  /*symbolCount=*/maxSym + 1, exprs));
+  return maps;
+}
+
+SmallVector<AffineMap, 4>
+AffineMap::inferFromExprList(ArrayRef<ArrayRef<AffineExpr>> exprsList) {
+  return ::inferFromExprList(exprsList);
+}
+
+SmallVector<AffineMap, 4>
+AffineMap::inferFromExprList(ArrayRef<SmallVector<AffineExpr, 4>> exprsList) {
+  return ::inferFromExprList(exprsList);
+}
+
 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/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
 #include "mlir/Dialect/LoopOps/EDSC/Builders.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
+#include "mlir/Dialect/VectorOps/EDSC/Intrinsics.h"
 #include "mlir/EDSC/Builders.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"
 #include "mlir/IR/StandardTypes.h"
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;
 }