VNNI dialect operands extended with memref type

include/TPP/Dialect/VNNI/BufferizableOpInterfaceImpl.h

@@ -0,0 +1,22 @@
+//===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
+//
+// 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_VNNI_BUFFERIZABLEOPINTERFACEIMPL_H
+#define MLIR_DIALECT_VNNI_BUFFERIZABLEOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+} // namespace mlir
+
+namespace mlir {
+namespace vnni {
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace vnni
+} // namespace mlir
+
+#endif // MLIR_DIALECT_VNNI_BUFFERIZABLEOPINTERFACEIMPL_H

include/TPP/Dialect/VNNI/VNNIOps.td

@@ -12,6 +12,8 @@
 
 include "mlir/IR/OpBase.td"
 
+def VNNIOperand : AnyTypeOf<[BF16Tensor, BF16MemRef]>;
+
 def VNNI_Dialect : Dialect {
   let name = "vnni";
   let cppNamespace = "::mlir::vnni";
@@ -37,8 +39,8 @@ def VNNI_MatmulOp : Op<VNNI_Dialect, "matmul",
     ```
   }];
 
-  let arguments = (ins BF16Tensor:$matrixA, BF16Tensor:$matrixB, BF16Tensor:$matrixC);
-  let results = (outs BF16Tensor:$dest);
+  let arguments = (ins VNNIOperand:$matrixA, VNNIOperand:$matrixB, VNNIOperand:$matrixC);
+  let results = (outs VNNIOperand:$dest);
   let assemblyFormat = "`ins` `(` $matrixA `:` type($matrixA) `,` $matrixB `:` type($matrixB) `)` `out` `(` $matrixC `:` type($dest) `)` attr-dict";
 }
 

lib/TPP/Dialect/VNNI/BufferizableOpInterfaceImpl.cpp

@@ -0,0 +1,92 @@
+//===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
+//
+// 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 "TPP/Dialect/VNNI/BufferizableOpInterfaceImpl.h"
+#include "TPP/Dialect/VNNI/VNNIDialect.h"
+#include "TPP/Dialect/VNNI/VNNIOps.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/IR/Operation.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+using namespace mlir::vnni;
+
+namespace mlir {
+namespace vnni {
+namespace {
+
+struct MatmulLayoutInterface
+    : public BufferizableOpInterface::ExternalModel<MatmulLayoutInterface,
+                                                    vnni::MatmulOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    return opOperand.getOperandNumber() == 0 ||
+           opOperand.getOperandNumber() == 1;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    return opOperand.getOperandNumber() == 2;
+  }
+
+  bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
+                            const AnalysisState &state) const {
+    return false;
+  }
+
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    if (opOperand.getOperandNumber() < 1)
+      return {};
+    return {op->getResult(0)};
+  }
+
+  BufferRelation bufferRelation(Operation *op, OpResult opResult,
+                                const AnalysisState &state) const {
+    return BufferRelation::Equivalent;
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options) const {
+    vnni::MatmulOp matmulOp = cast<vnni::MatmulOp>(op);
+
+    FailureOr<Value> maybeDestBuffer =
+        getBuffer(rewriter, matmulOp.getMatrixC(), options);
+    if (failed(maybeDestBuffer))
+      return failure();
+    Value destBuffer = *maybeDestBuffer;
+
+    FailureOr<Value> maybeSrcBufferA =
+        getBuffer(rewriter, matmulOp.getMatrixA(), options);
+    if (failed(maybeSrcBufferA))
+      return failure();
+    Value srcBufferA = *maybeSrcBufferA;
+
+    FailureOr<Value> maybeSrcBufferB =
+        getBuffer(rewriter, matmulOp.getMatrixB(), options);
+    if (failed(maybeSrcBufferB))
+      return failure();
+    Value srcBufferB = *maybeSrcBufferB;
+
+    rewriter.create<vnni::MatmulOp>(op->getLoc(), destBuffer.getType(),
+                                    srcBufferA, srcBufferB, destBuffer);
+    replaceOpWithBufferizedValues(rewriter, op, destBuffer);
+    return success();
+  }
+};
+
+} // namespace
+} // namespace vnni
+} // namespace mlir
+
+void mlir::vnni::registerBufferizableOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, vnni::VNNIDialect *dialect) {
+    MatmulOp::attachInterface<vnni::MatmulLayoutInterface>(*ctx);
+  });
+}

lib/TPP/Dialect/VNNI/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_mlir_dialect_library(TPPVNNIDialect
     # Ops and dialect
+    BufferizableOpInterfaceImpl.cpp
     VNNIDialect.cpp
     VNNIOps.cpp
 

test/BF16/vnni-bufferization.mlir

@@ -0,0 +1,15 @@
+// RUN: tpp-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=identity-layout-map" | FileCheck %s
+
+// CHECK-LABEL: @myfunc(
+// CHECK: %[[ARG0:.+]]: memref<2x2x2xbf16>,
+// CHECK: %[[ARG1:.+]]: memref<2x2xbf16>,
+// CHECK: %[[ARG2:.+]]: memref<4x2xbf16>) -> memref<4x2xbf16> {
+func.func @myfunc(%arg0: tensor<2x2x2xbf16>,
+                  %arg1: tensor<2x2xbf16>,
+		  %arg2: tensor<4x2xbf16>) -> tensor<4x2xbf16> {
+  // CHECK: %[[ALLOC:.+]] = memref.alloc() {alignment = 128 : i64} : memref<4x2xbf16>
+  // CHECK: %[[RET:.+]] = vnni.matmul ins(%[[ARG0]] : memref<2x2x2xbf16>, %[[ARG1]] : memref<2x2xbf16>) out(%[[ALLOC]] : memref<4x2xbf16>)
+  %vnni_result = vnni.matmul ins(%arg0: tensor<2x2x2xbf16>, %arg1: tensor<2x2xbf16>) out(%arg2: tensor<4x2xbf16>)
+  // CHECK: return %[[ALLOC]]
+  return %vnni_result:tensor<4x2xbf16>
+}

test/BF16/vnni-ops.mlir

@@ -8,3 +8,13 @@ func.func @myfunc(%arg0: tensor<2x2x2xbf16>,
   %vnni_result = vnni.matmul ins(%arg0: tensor<2x2x2xbf16>, %arg1: tensor<2x2xbf16>) out(%arg2: tensor<4x2xbf16>)
   return %vnni_result:tensor<4x2xbf16>
 }
+
+
+// CHECK-LABEL: @myfunc2
+func.func @myfunc2(%arg0: memref<2x2x2xbf16>,
+		   %arg1: memref<2x2xbf16>,
+                   %arg2: memref<4x2xbf16>){
+  // CHECK: vnni.matmul
+  vnni.matmul ins(%arg0: memref<2x2x2xbf16>, %arg1: memref<2x2xbf16>) out(%arg2: memref<4x2xbf16>)
+  return
+}

tpp-opt/tpp-opt.cpp

@@ -25,6 +25,7 @@
 #include "TPP/Dialect/LinalgX/LinalgXDialect.h"
 #include "TPP/Dialect/LinalgX/TransformOps/LinalgXTransformOps.h"
 #include "TPP/Dialect/Tpp/TppDialect.h"
+#include "TPP/Dialect/VNNI/BufferizableOpInterfaceImpl.h"
 #include "TPP/Dialect/VNNI/VNNIDialect.h"
 #include "TPP/Dialect/Xsmm/XsmmDialect.h"
 #include "TPP/Passes.h"
@@ -42,6 +43,7 @@ int main(int argc, char **argv) {
   mlir::linalgx::registerTransformDialectExtension(registry);
   mlir::linalgx::registerBufferizableOpInterfaceExternalModels(registry);
   mlir::check::registerBufferizableOpInterfaceExternalModels(registry);
+  mlir::vnni::registerBufferizableOpInterfaceExternalModels(registry);
   // Add the following to include *all* MLIR Core dialects, or selectively
   // include what you need like above. You only need to register dialects that
   // will be *parsed* by the tool, not the one generated