[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

49

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

49

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();
+    if (failed(options.createMemCpy(rewriter, op->getLoc(), *srcBuffer,
+                                    tensorStoreOp.getMemref())))
+      return failure();
+    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/transform-op-bufferize-to-allocation.mlir

 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = transform.get_result %0[0] : (!pdl.operation) -> !transform.any_value
   %2 = transform.structured.bufferize_to_allocation %1
 }
 
 // -----
 
+// CHECK-LABEL: func @tensor_pad_constant(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
+//       CHECK:   %[[src:.*]] = bufferization.to_memref %[[t]]
+//       CHECK:   %[[alloc:.*]] = memref.alloc
+//       CHECK:   %[[subview:.*]] = memref.subview %[[alloc]]
+//       CHECK:   memref.copy %[[src]], %[[subview]]
+//       CHECK:   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>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.get_result %0[0] : (!pdl.operation) -> !transform.any_value
+  %2 = transform.structured.bufferize_to_allocation %1
+  // Make sure that One-Shot Bufferize can bufferize the rest.
+  transform.bufferization.one_shot_bufferize %arg1
+}
+
+// -----
+
 // CHECK-LABEL: func @materialization_of_bbarg(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[c0:.*]] = arith.constant 0 : index
 //       CHECK:   %[[dim:.*]] = tensor.dim %[[t]], %[[c0]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%[[dim]]) : memref<?x10xindex, 4>
 //       CHECK:   memref.tensor_store %[[t]], %[[alloc]]
 //       CHECK:   %[[alloc_t:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
 //       CHECK:   %[[r:.*]] = tensor.extract %[[alloc_t]]
 //       CHECK:   memref.dealloc %[[alloc]]
 //       CHECK:   return %[[r]]
 func.func @materialization_of_bbarg(%t: tensor<?x10xindex>, %idx: index) -> index {
   %r = tensor.extract %t[%idx, %idx] : tensor<?x10xindex>
   return %r : index
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = test_produce_value_handle_to_argument_of_parent_block %0, 0 : (!pdl.operation) -> !transform.any_value
   %2 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
 }
+
+// -----
+
+// CHECK-LABEL: func @materialization_of_bbarg(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
+//       CHECK:   %[[m:.*]] = bufferization.to_memref %[[t]]
+//       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4>
+//       CHECK:   memref.copy %[[m]], %[[alloc]]
+//       CHECK:   %[[r:.*]] = memref.load %[[alloc]]
+//       CHECK:   return %[[r]]
+func.func @materialization_of_bbarg(%t: tensor<?x10xindex>, %idx: index) -> index {
+  %r = tensor.extract %t[%idx, %idx] : tensor<?x10xindex>
+  return %r : index
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = test_produce_value_handle_to_argument_of_parent_block %0, 0 : (!pdl.operation) -> !transform.any_value
+  %2 = transform.structured.bufferize_to_allocation %1 {memory_space = 4}
+  // Make sure that One-Shot Bufferize can bufferize the rest.
+  transform.bufferization.one_shot_bufferize %arg1
+}

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
+}

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",