[mlir][memref] Bufferize memref.tensor_store op

[springerm]

Summary

This change adds the BufferizableOpInterface implementation for memref.tensor_store.

Depends On: D144022

[nicolasvasilache]

mlir/lib/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.cpp

48

Can we use a linalg.copy here?
Experience shows that going to memref.copy prematurely can yield very bad perf.
There are simple woarkarounds like rewriting a memref.copy into a linalg.copy but it is better to avoid lowering prematurely.

mlir/test/Dialect/MemRef/bufferize.mlir

12

nit:nl

[springerm]

mlir/lib/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.cpp

48

I changed it options.createMemCpy then it's using whatever "copy" callback we configure the bufferization with.

mlir/include/mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h

+//===- 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_MEMREF_BUFFERIZABLEOPINTERFACEIMPL_H
+#define MLIR_DIALECT_MEMREF_BUFFERIZABLEOPINTERFACEIMPL_H
+
+namespace mlir {
+
+class DialectRegistry;
+
+namespace memref {
+void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace memref
+} // namespace mlir
+
+#endif // MLIR_DIALECT_MEMREF_BUFFERIZABLEOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"
 #include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
+#include "mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
inline void registerAllDialects(DialectRegistry &registry) {
   vector::registerTransformDialectExtension(registry);
 
   // Register all external models.
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
+  memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);
   shape::registerBufferizableOpInterfaceExternalModels(registry);
   sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerBufferizableOpInterfaceExternalModels(registry);
   tensor::registerInferTypeOpInterfaceExternalModels(registry);
   tensor::registerTilingInterfaceExternalModels(registry);
   vector::registerBufferizableOpInterfaceExternalModels(registry);

mlir/lib/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.cpp

+//===- 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 "mlir/Dialect/MemRef/Transforms/BufferizableOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Operation.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+
+namespace {
+/// Bufferization of memref.tensor_store. Replace with memref.copy.
+struct TensorStoreOpInterface
+    : public BufferizableOpInterface::ExternalModel<TensorStoreOpInterface,
+                                                    memref::TensorStoreOp> {
+  AliasingOpResultList getAliasingOpResults(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
+    return {};
+  }
+
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    assert(opOperand.getOperandNumber() == 0 && "expected src operand");
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    // The memref operand is written but not the tensor operand.
+    assert(opOperand.getOperandNumber() == 0 && "expected src operand");
+    return false;
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options) const {
+    auto tensorStoreOp = cast<memref::TensorStoreOp>(op);
+    auto srcBuffer = getBuffer(rewriter, tensorStoreOp.getTensor(), options);
+    if (failed(srcBuffer))
+      return failure();
+    rewriter.create<memref::CopyOp>(op->getLoc(), *srcBuffer,
+                                    tensorStoreOp.getMemref());
+    rewriter.eraseOp(tensorStoreOp);
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::memref::registerBufferizableOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, MemRefDialect *dialect) {
+    TensorStoreOp::attachInterface<TensorStoreOpInterface>(*ctx);
+  });
+}

mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefTransforms
+  BufferizableOpInterfaceImpl.cpp
   ComposeSubView.cpp
   ExpandOps.cpp
   ExpandStridedMetadata.cpp
   EmulateWideInt.cpp
   FoldMemRefAliasOps.cpp
   MultiBuffer.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
   RuntimeOpVerification.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineUtils
   MLIRArithDialect
   MLIRArithTransforms
+  MLIRBufferizationDialect
   MLIRFuncDialect
   MLIRInferTypeOpInterface
   MLIRLoopLikeInterface
   MLIRMemRefDialect
   MLIRPass
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
 )
 

mlir/test/Dialect/Linalg/bufferize-non-dps-ops.mlir

 // RUN: mlir-opt -split-input-file \
 // RUN:   -test-linalg-transform-patterns=test-bufferize-non-dps-ops-patterns \
 // RUN:   -canonicalize %s | FileCheck %s
 
+// RUN: mlir-opt -split-input-file \
+// RUN:   -test-linalg-transform-patterns=test-bufferize-non-dps-ops-patterns \
+// RUN:   -canonicalize -one-shot-bufferize %s | \
+// RUN: FileCheck %s --check-prefix=CHECK-BUFFERIZE-ALL
+
 // CHECK:       #[[$map:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 5)>
 // CHECK:       #[[$map1:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 10)>
 // CHECK-LABEL: func @tensor_pad_constant(
 //  CHECK-SAME:   %[[t:.*]]: tensor<?x10xindex>, %[[l2:.*]]: index, %[[h1:.*]]: index, %[[h2:.*]]: index
 //   CHECK-DAG:   %[[c0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[c50:.*]] = arith.constant 50 : index
 //   CHECK-DAG:   %[[dim0:.*]] = tensor.dim %[[t]], %[[c0]]
 //   CHECK-DAG:   %[[size0:.*]] = affine.apply #[[$map]]()[%[[h1]], %[[dim0]]]
 //   CHECK-DAG:   %[[size1:.*]] = affine.apply #[[$map1]]()[%[[l2]], %[[h2]]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%[[size0]], %[[size1]]) : memref<?x?xindex>
 //       CHECK:   linalg.fill ins(%[[c50]] : index) outs(%[[alloc]] : memref<?x?xindex>)
 //       CHECK:   %[[dim0:.*]] = tensor.dim %[[t]], %[[c0]]
 //       CHECK:   %[[subview:.*]] = memref.subview %[[alloc]][5, %[[l2]]] [%[[dim0]], 10] [1, 1]
 //       CHECK:   memref.tensor_store %[[t]], %[[subview]]
 //       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable : memref<?x?xindex>
 //       CHECK:   return %[[r]]
+
+// CHECK-BUFFERIZE-ALL-LABEL: func @tensor_pad_constant(
+//  CHECK-BUFFERIZE-ALL-SAME:     %[[t:.*]]: tensor<?x10xindex>
+//       CHECK-BUFFERIZE-ALL:   %[[src:.*]] = bufferization.to_memref %[[t]]
+//       CHECK-BUFFERIZE-ALL:   %[[alloc:.*]] = memref.alloc
+//       CHECK-BUFFERIZE-ALL:   %[[subview:.*]] = memref.subview %[[alloc]]
+//       CHECK-BUFFERIZE-ALL:   memref.copy %[[src]], %[[subview]]
+//       CHECK-BUFFERIZE-ALL:   bufferization.to_tensor %[[alloc]] restrict writable
 func.func @tensor_pad_constant(%t: tensor<?x10xindex>, %l2: index, %h1: index,
                                %h2: index) -> tensor<?x?xindex> {
   %0 = tensor.pad %t low[5, %l2] high[%h1, %h2] {
   ^bb0(%arg0: index, %arg1: index):
     %c = arith.constant 50 : index
     tensor.yield %c : index
   } : tensor<?x10xindex> to tensor<?x?xindex>
   return %0 : tensor<?x?xindex>
 }

mlir/test/Dialect/MemRef/bufferize.mlir

+// RUN: mlir-opt -one-shot-bufferize %s | FileCheck %s
+
+// CHECK-LABEL: func @tensor_store(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>, %[[m:.*]]: memref<?xf32>
+//       CHECK:   %[[src:.*]] = bufferization.to_memref %[[t]]
+//       CHECK:   memref.copy %[[src]], %[[m]]
+//       CHECK:   return
+func.func @tensor_store(%t: tensor<?xf32>, %m: memref<?xf32>) {
+  memref.tensor_store %t, %m : memref<?xf32>
+  return
+}
 No newline at end of file

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     hdrs = glob(["include/mlir/Dialect/MemRef/Transforms/*.h"]),
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":AffineUtils",
         ":ArithDialect",
         ":ArithTransforms",
         ":ArithUtils",
+        ":BufferizationDialect",
         ":ControlFlowDialect",
         ":DialectUtils",
         ":FuncDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefPassIncGen",