[mlir] Bufferize tensor constant ops

[silvas]

Summary

We lower them to a std.global_memref (uniqued by constant value) + a
std.get_global_memref to produce the corresponding memref value.
This allows removing Linalg's somewhat hacky lowering of tensor
constants, now that std properly supports this.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp

312

I think this proposal has far reaching consequences on transformations that are not yet well understood.
In the short-term I'd favor a first class "bufferizable_in_place" attribute on subtensor/subtensor_insert to carry the information.

That attribute may be either:

1. obtained as the result of some analysis / canonicalization
2. as a transformation invariant (e.g. the transformation that introduces it decides whether it has enough information/guarantees to add the attribute)
3. directly written in the IR for testing purposes.

Irrespective of this discussion, I think the root of the problem still lies on the tensor_load + tensor_to_memref canonicalization (see below).

328

I reverted this line locally after patching this revision and here is what I see:

./bin/mlir-opt -linalg-bufferize -std-bufferize -tensor-constant-bufferize -func-bufferize -convert-linalg-to-loops -print-ir-after-all ../mlir/integration_test/Dialect/Linalg/CPU/test-subtensor-insert.mlir

// *** IR Dump After TensorConstantBufferize ***
module {
  global_memref "private" constant @__constant_1xf32 : memref<1xf32> = dense<2.000000e+01>
  global_memref "private" constant @__constant_2xf32 : memref<2xf32> = dense<1.000000e+01>
  func @main() {
    %0 = get_global_memref @__constant_2xf32 : memref<2xf32>
    %1 = tensor_load %0 : memref<2xf32>
    %2 = get_global_memref @__constant_1xf32 : memref<1xf32>
    %3 = tensor_load %2 : memref<1xf32>
    %4 = tensor_to_memref %3 : memref<1xf32>
    %5 = tensor_to_memref %1 : memref<2xf32>
    %6 = subview %5[0] [1] [1]  : memref<2xf32> to memref<1xf32>
    linalg.copy(%4, %6) : memref<1xf32>, memref<1xf32>
    %7 = tensor_load %5 : memref<2xf32>
    %8 = tensor_to_memref %7 : memref<2xf32>
    %9 = memref_cast %8 : memref<2xf32> to memref<*xf32>
    %10 = tensor_load %9 : memref<*xf32>
    call @print_memref_f32(%10) : (tensor<*xf32>) -> ()
    return
  }
  func @print_memref_f32(tensor<*xf32>)
}

// *** IR Dump After FuncBufferize ***
module {
  global_memref "private" constant @__constant_1xf32 : memref<1xf32> = dense<2.000000e+01>
  global_memref "private" constant @__constant_2xf32 : memref<2xf32> = dense<1.000000e+01>
  func @main() {
    %0 = get_global_memref @__constant_2xf32 : memref<2xf32>
    %1 = get_global_memref @__constant_1xf32 : memref<1xf32>
    %2 = subview %0[0] [1] [1]  : memref<2xf32> to memref<1xf32>
    linalg.copy(%1, %2) : memref<1xf32>, memref<1xf32>
    %3 = memref_cast %0 : memref<2xf32> to memref<*xf32>
    call @print_memref_f32(%3) : (memref<*xf32>) -> ()
    return
  }
  func @print_memref_f32(memref<*xf32>)
}

The problem seems to be again coming from the tensor_load + tensor_to_memref canonicalization.

The tensor_to_memref semantics is:

Create a memref from a tensor. This is equivalent to allocating a new
memref of the appropriate (possibly dynamic) shape, and then copying the
elements (as if by a tensor_store op) into the newly allocated memref.

Under these semantics, I would argue the IR is correct after TensorConstantBufferize and incorrect after FuncBufferize.

Relaxing the semantics of tensor_to_memref to not guarantee an alloc + copy would make it reasonable to require the bufferization to insert alloc + copy when it does not know better.

However, I am not sure how to relax tensor_to_memref or even phrase it:

Create a memref from a tensor %t. The resulting memref aliases the memref operand of all tensor_load operations that flows into %t .. ?

[silvas]

mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp

312

First of all, this is not a proposal. This is a fact, and this fixes a miscompile that was already present in the code, even before this patch. For clarity, I'll split that out.

I think that adding a BufferizesToSomethingThatWritesIntoAnOperand trait would help. Then we could have a copy insertion pass that inserts copies at the tensor level (if needed), and the bufferization patterns will assume that those copies have been inserted.

"bufferizable in place" is a global property, so transformations cannot possibly know when creating IR whether it is safe to do in place. See e.g. test-tensor-matmul -- that one fails because we attempt to do a subtensor_insert into a std.constant, which happens to bufferize to something read only. A similar problem trivially occurs if the std.constant is used as the init tensor of two different linalg ops (possibly obscured by arbitrary control flow).

328

Agreed that we aren't obeying the "copy" aspect of tensor_to_memref. But that's not the problem here. The current way that we are lowering tensor_to_memref is simply emulating what would happen if you were running dialect conversion in one big pass (in that situation, tensors would be replaced directly by memrefs everywhere, and no tensor_to_memref/tensor_load pairs would be inserted to keep the IR legal). So the problem is independent of tensor_to_memref.

It's a fundamental aspect of simply converting a tensor Value to a memref Value without doing anything else -- the correct semantics is that every *use* of a tensor Value that has been converted to memref needs to see a separate copy, *if* that use might mutate its operand in place.

So we either need to update the dialect conversion framework to respect that, or we need to do a separate copy insertion phase that establishes the invariants.

[silvas]

Split out fix to subtensor bufferization into https://reviews.llvm.org/D91371

[silvas]

mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp

312

I split out https://reviews.llvm.org/D91371 which fixes just subtensor_insert in isolation.

[nicolasvasilache]

LGTM once rebased on the addressed https://reviews.llvm.org/D91371.

Looking forward to the next steps :) !

[silvas]

update

mlir/include/mlir/Dialect/Linalg/Passes.td

def LinalgLowerToLoops : FunctionPass<"convert-linalg-to-loops"> {
   let summary = "Lower the operations from the linalg dialect into loops";
   let constructor = "mlir::createConvertLinalgToLoopsPass()";
   let dependentDialects = ["linalg::LinalgDialect", "scf::SCFDialect", "AffineDialect"];
 }
 
 def LinalgBufferize : Pass<"linalg-bufferize", "ModuleOp"> {
   let summary = "Bufferize the linalg dialect";
   let constructor = "mlir::createLinalgBufferizePass()";
-  let dependentDialects = ["linalg::LinalgDialect", "vector::VectorDialect"];
+  let dependentDialects = ["linalg::LinalgDialect"];
 }
 
 def LinalgLowerToParallelLoops
     : FunctionPass<"convert-linalg-to-parallel-loops"> {
   let summary = "Lower the operations from the linalg dialect into parallel "
                 "loops";
   let constructor = "mlir::createConvertLinalgToParallelLoopsPass()";
   let dependentDialects = ["AffineDialect", "linalg::LinalgDialect", "scf::SCFDialect"];

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.h

void populateStdBufferizePatterns(MLIRContext *context,
                                   OwningRewritePatternList &patterns);
 
 /// Creates an instance of std bufferization pass.
 std::unique_ptr<Pass> createStdBufferizePass();
 
 /// Creates an instance of func bufferization pass.
 std::unique_ptr<Pass> createFuncBufferizePass();
 
+/// Creates an instance of tensor constant bufferization pass.
+std::unique_ptr<Pass> createTensorConstantBufferizePass();
+
 /// Creates an instance of the StdExpand pass that legalizes Std
 /// dialect ops to be convertible to LLVM. For example,
 /// `std.ceildivi_signed` gets transformed to a number of std operations,
 /// which can be lowered to LLVM; `memref_reshape` gets converted to
 /// `memref_reinterpret_cast`.
 std::unique_ptr<Pass> createStdExpandOpsPass();
 
 /// Collects a set of patterns to rewrite ops within the Std dialect.

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

let description = [{
     - block arguments
     - the result of tensor_load
     Other values of tensor type should be eliminated by earlier
     bufferization passes.
   }];
   let constructor = "mlir::createFuncBufferizePass()";
 }
 
+def TensorConstantBufferize : Pass<"tensor-constant-bufferize", "ModuleOp"> {
+  let summary = "Bufferize tensor constants.";
+  let description = [{
+    This pass bufferizes tensor constants.
+
+    This pass needs to be a module pass because it inserts std.global_memref
+    ops into the module, which cannot be done safely from a function pass due to
+    multi-threading. Most other bufferization passes can run in parallel at
+    function granularity.
+  }];
+  let constructor = "mlir::createTensorConstantBufferizePass()";
+}
+
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

mlir/integration_test/Dialect/Linalg/CPU/test-elementwise.mlir

-// RUN: mlir-opt %s -convert-elementwise-to-linalg -std-bufferize -linalg-bufferize -func-bufferize -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
+// RUN: mlir-opt %s -convert-elementwise-to-linalg -std-bufferize -tensor-constant-bufferize -linalg-bufferize -func-bufferize -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @main() {
   %a = constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
   %b = constant dense<[10.0, 20.0, 30.0]> : tensor<3xf32>
 

mlir/integration_test/Dialect/Linalg/CPU/test-subtensor-insert.mlir

+// RUN: mlir-opt %s -linalg-bufferize -std-bufferize -tensor-constant-bufferize -func-bufferize \
+// RUN: -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
+// RUN: mlir-cpu-runner -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
+// RUN: | FileCheck %s
+
+func @main() {
+  %const = constant dense<10.0> : tensor<2xf32>
+  %insert_val = constant dense<20.0> : tensor<1xf32>
+  %inserted = subtensor_insert %insert_val into %const[0][1][1] : tensor<1xf32> into tensor<2xf32>
+
+  %unranked = tensor_cast %inserted : tensor<2xf32> to tensor<*xf32>
+  call @print_memref_f32(%unranked) : (tensor<*xf32>) -> ()
+
+  //      CHECK: Unranked Memref base@ = {{0x[-9a-f]*}}
+  // CHECK-SAME: rank = 1 offset = 0 sizes = [2] strides = [1] data =
+  // CHECK-NEXT: [20, 10]
+
+  return
+}
+
+func @print_memref_f32(%ptr : tensor<*xf32>)

mlir/integration_test/Dialect/Linalg/CPU/test-tensor-e2e.mlir

-// RUN: mlir-opt %s -std-bufferize -linalg-bufferize -func-bufferize -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
+// RUN: mlir-opt %s -tensor-constant-bufferize -std-bufferize -linalg-bufferize -func-bufferize -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @foo() -> tensor<4xf32> {
   %0 = constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
   return %0 : tensor<4xf32>
 }

mlir/integration_test/Dialect/Linalg/CPU/test-tensor-matmul.mlir

-// RUN: mlir-opt %s -linalg-bufferize -std-bufferize -func-bufferize \
+// RUN: mlir-opt %s -linalg-bufferize -std-bufferize -tensor-constant-bufferize -func-bufferize \
 // RUN: -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 // RUN: mlir-opt %s  -linalg-tile="linalg-tile-sizes=1,2,3" -linalg-bufferize \
-// RUN: -scf-bufferize -std-bufferize -func-bufferize -convert-linalg-to-loops \
+// RUN: -scf-bufferize -std-bufferize -tensor-constant-bufferize -func-bufferize -convert-linalg-to-loops \
 // RUN:  -convert-scf-to-std -convert-linalg-to-llvm | \
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func @main() {
   %A = constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf32>
   %B = constant dense<[[1.0, 2.0, 3.0, 4.0],

mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp

 }
 
 static Value maybeConvertToIndex(Location loc, Value val, OpBuilder &b) {
   if (val.getType().isIndex())
     return val;
   return b.create<IndexCastOp>(loc, val, b.getIndexType());
 }
 
+static Value cloneMemref(Location loc, Value memref, OpBuilder &b) {
+  auto memrefType = memref.getType().cast<MemRefType>();
+  SmallVector<Value, 4> dynOperands;
+  for (auto dim : llvm::enumerate(memrefType.getShape())) {
+    if (dim.value() == TensorType::kDynamicSize) {
+      dynOperands.push_back(b.create<DimOp>(loc, memref, dim.index()));
+    }
+  }
+  auto alloc = b.create<AllocOp>(loc, memrefType, dynOperands);
+  b.create<linalg::CopyOp>(loc, memref, alloc);
+  return alloc;
+}
+
 static LogicalResult
 allocateBuffersForResults(Location loc, LinalgOp linalgOp,
                           linalg::GenericOpAdaptor &adaptor,
                           SmallVectorImpl<Value> &resultBuffers, OpBuilder &b) {
   // Lazily compute loopRanges.
   SmallVector<Range, 4> loopRanges;
 
   // Allocate a buffer for every tensor result.
for (auto en : llvm::enumerate(linalgOp.getOperation()->getResultTypes())) {
     }
     auto tensorShape = tensorType.getShape();
     auto memrefType = MemRefType::get(tensorShape, tensorType.getElementType());
 
     // Allocate buffers for init tensors that are assumed to fold onto the first
     // results.
     // TODO: update this assumption because the reality is more complex
     // under linalg on tensor based transformations.
-    bool foldedInitTensor = resultIndex < linalgOp.getNumInitTensors();
-    if (foldedInitTensor) {
-      Value initTensor = linalgOp.getInitTensor(resultIndex);
-      Value initBuffer = adaptor.init_tensors()[resultIndex];
-      SmallVector<Value, 4> dynOperands;
-      for (auto dim : llvm::enumerate(tensorShape)) {
-        if (dim.value() == TensorType::kDynamicSize) {
-          dynOperands.push_back(b.create<DimOp>(loc, initTensor, dim.index()));
-        }
-      }
-      auto alloc = b.create<AllocOp>(loc, memrefType, dynOperands);
-      b.create<linalg::CopyOp>(loc, initBuffer, alloc);
-      resultBuffers.push_back(alloc);
+    bool hasInitTensor = resultIndex < linalgOp.getNumInitTensors();
+    if (hasInitTensor) {
+      resultBuffers.push_back(
+          cloneMemref(loc, adaptor.init_tensors()[resultIndex], b));
       continue;
     }
 
     // Allocate buffers for statically-shaped results.
     if (memrefType.hasStaticShape()) {
       resultBuffers.push_back(b.create<AllocOp>(loc, memrefType));
       continue;
     }
public:
   LogicalResult
   matchAndRewrite(SubTensorInsertOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     SubTensorInsertOpAdaptor adaptor(operands,
                                      op.getOperation()->getAttrDictionary());
     Value sourceMemRef = adaptor.source();
     assert(sourceMemRef.getType().isa<MemRefType>());
 
-    Value destMemRef = adaptor.dest();
+    // TODO: We need a much more coherent story for tensor-level code to become
+    // "in-place" sensibly.
+    //
+    // In general, the converted memref here could be aliased or could point
+    // into constant memory, which causes miscompilations.
+    // However, it kind of defeats the purpose of a subtensor_insert if it
+    // doesn't lower into someting that "inserts" but rather "copy + inserts".
+    // There are multiple avenues of fixing this, such as capturing the desire
+    // for an in-place update more with a more structured op like:
+    // subtensor_update %whole_tensor[0, %i0][2, %i0][1, 2] {
+    // ^bb0(%subtensor: tensor<2x?xf32>):
+    //   %updated = linalg.matmul ins(%A, %B) init(%C)
+    //   yield %updated
+    // } : tensor<?x?xf32> updating tensor<2x?xf32>
+    // This is much easier to reason about than loose subtensor/subtensor_insert
+    // ops. We can bufferize the subtensor_update op to a single copy of
+    // %whole_tensor with in-place updates to a subview, which then with some
+    // analysis we can tractably hoist out of loop nests to get the final
+    // desired in-place code.
+    Value destMemRef = cloneMemref(op.getLoc(), adaptor.dest(), rewriter);
     assert(destMemRef.getType().isa<MemRefType>());
 
     // Take a subview to copy the small memref.
     Value subview = rewriter.create<SubViewOp>(
         op.getLoc(), destMemRef, extractFromI64ArrayAttr(op.static_offsets()),
         extractFromI64ArrayAttr(op.static_sizes()),
         extractFromI64ArrayAttr(op.static_strides()), adaptor.offsets(),
         adaptor.sizes(), adaptor.strides());
     // Copy the small memref.
     rewriter.create<linalg::CopyOp>(op.getLoc(), sourceMemRef, subview);
     rewriter.replaceOp(op, destMemRef);
     return success();
   }
 };
-
-/// TensorConstantOp conversion inserts a linearized 1-D vector constant that is
-/// stored in memory. A linalg.reshape is introduced to convert to the desired
-/// n-D buffer form.
-class TensorConstantOpConverter : public OpConversionPattern<ConstantOp> {
-public:
-  using OpConversionPattern::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(ConstantOp op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const final {
-
-    RankedTensorType rankedTensorType =
-        op.getType().dyn_cast<RankedTensorType>();
-    if (!rankedTensorType)
-      return failure();
-    if (llvm::any_of(rankedTensorType.getShape(), [](int64_t s) {
-          return s == 0 || ShapedType::isDynamic(s);
-        }))
-      return failure();
-
-    int64_t nElements = 1;
-    for (int64_t s : rankedTensorType.getShape())
-      nElements *= s;
-    Type elementType = rankedTensorType.getElementType();
-    MemRefType memrefType =
-        getTypeConverter()->convertType(op.getType()).cast<MemRefType>();
-    VectorType flatVectorType = VectorType::get({nElements}, elementType);
-    MemRefType memrefOfFlatVectorType = MemRefType::get({}, flatVectorType);
-    MemRefType flatMemrefType = MemRefType::get({nElements}, elementType);
-
-    Location loc = op.getLoc();
-    auto attr = op.getValue().cast<DenseElementsAttr>();
-    Value alloc =
-        rewriter.create<AllocOp>(loc, memrefOfFlatVectorType, ValueRange{});
-    Value cstVec = rewriter.create<ConstantOp>(loc, flatVectorType,
-                                               attr.reshape(flatVectorType));
-    rewriter.create<StoreOp>(loc, cstVec, alloc);
-
-    Value memref =
-        rewriter.create<vector::TypeCastOp>(loc, flatMemrefType, alloc);
-    if (rankedTensorType.getRank() > 1) {
-      // Introduce a linalg.reshape to flatten the memref.
-      AffineMap collapseAllDims = AffineMap::getMultiDimIdentityMap(
-          /*numDims=*/rankedTensorType.getRank(), op.getContext());
-      memref = rewriter.create<linalg::ReshapeOp>(
-          loc, memrefType, memref,
-          rewriter.getAffineMapArrayAttr(collapseAllDims));
-    }
-    rewriter.replaceOp(op, memref);
-
-    return success();
-  }
-};
 } // namespace
 
 namespace {
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     BufferizeTypeConverter typeConverter;
 
     // Mark all Standard operations legal.
-    target.addLegalDialect<StandardOpsDialect, vector::VectorDialect>();
+    target.addLegalDialect<StandardOpsDialect>();
     target.addIllegalOp<SubTensorOp, SubTensorInsertOp>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return typeConverter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);
     target.addDynamicallyLegalOp<ConstantOp>(isLegalOperation);

 void mlir::linalg::populateLinalgBufferizePatterns(
     MLIRContext *context, BufferizeTypeConverter &typeConverter,
     OwningRewritePatternList &patterns) {
   patterns.insert<BufferizeAnyLinalgOp>(typeConverter);
   // TODO: Drop this once tensor constants work in standard.
   patterns.insert<
       // clang-format off
       SubTensorOpConverter,
-      SubTensorInsertOpConverter,
-      TensorConstantOpConverter
+      SubTensorInsertOpConverter
       // clang-format on
       >(typeConverter, context);
 }

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

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   Bufferize.cpp
   DecomposeCallGraphTypes.cpp
   ExpandOps.cpp
   ExpandTanh.cpp
   FuncBufferize.cpp
   FuncConversions.cpp
+  TensorConstantBufferize.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms
 
   DEPENDS
   MLIRStandardTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
   MLIRSCF
   MLIRStandard
   MLIRTransforms
   )

mlir/lib/Dialect/StandardOps/Transforms/TensorConstantBufferize.cpp

+//===- Bufferize.cpp - Bufferization for std ops --------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements bufferization of tensor-valued std.constant ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/Transforms/Bufferize.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+// This class creates global ops for all tensor-valued constants in the program.
+// It creates them with pretty names and makes sure that duplicate globals
+// aren't created.
+class GlobalCreator {
+public:
+  explicit GlobalCreator(ModuleOp module);
+  GlobalMemrefOp getGlobalFor(Attribute attr) {
+    assert(globals.find(attr) != globals.end() && "unknown constant attr");
+    return globals[attr];
+  }
+
+private:
+  DenseMap<Attribute, GlobalMemrefOp> globals;
+};
+
+GlobalCreator::GlobalCreator(ModuleOp module) {
+  BufferizeTypeConverter typeConverter;
+  // Create a builder without an insertion point. We will insert using the
+  // symbol table to guarantee unique names.
+  OpBuilder globalBuilder(module.getContext());
+  SymbolTable symbolTable(module);
+  module.walk([&](ConstantOp op) {
+    // We only want tensor constants for now.
+    auto type = op.getType().dyn_cast<RankedTensorType>();
+    if (!type)
+      return;
+    // If we already have a global for this constant value, no need to do
+    // anything else.
+    auto it = globals.find(op.getValue());
+    if (it != globals.end())
+      return;
+
+    // Create a pretty name.
+    SmallString<64> buf;
+    llvm::raw_svector_ostream os(buf);
+    interleave(type.getShape(), os, "x");
+    os << "x" << type.getElementType();
+
+    auto global = globalBuilder.create<GlobalMemrefOp>(
+        op.getLoc(), (Twine("__constant_") + os.str()).str(),
+        /*sym_visibility=*/globalBuilder.getStringAttr("private"),
+        /*type=*/
+        TypeAttr::get(typeConverter.convertType(type)), /*initial_value=*/
+        op.getValue().cast<ElementsAttr>(), /*constant=*/true);
+    symbolTable.insert(global);
+    // The symbol table inserts at the end of the module, but globals are a bit
+    // nicer if they are at the beginning.
+    global.getOperation()->moveBefore(&module.front());
+    globals[op.getValue()] = global;
+  });
+}
+} // namespace
+
+namespace {
+class BufferizeTensorConstantOp : public OpConversionPattern<ConstantOp> {
+public:
+  BufferizeTensorConstantOp(GlobalCreator &globals,
+                            TypeConverter &typeConverter, MLIRContext *context)
+      : OpConversionPattern<ConstantOp>(typeConverter, context, /*benefit=*/1),
+        globals(globals) {}
+
+  LogicalResult
+  matchAndRewrite(ConstantOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto type = op.getType().dyn_cast<RankedTensorType>();
+    if (!type)
+      return failure();
+
+    auto globalMemref = globals.getGlobalFor(op.value());
+    rewriter.replaceOpWithNewOp<GetGlobalMemrefOp>(op, globalMemref.type(),
+                                                   globalMemref.getName());
+    return success();
+  }
+  GlobalCreator &globals;
+};
+} // namespace
+
+namespace {
+struct TensorConstantBufferizePass
+    : public TensorConstantBufferizeBase<TensorConstantBufferizePass> {
+  void runOnOperation() override {
+    auto module = getOperation();
+    GlobalCreator globals(module);
+
+    auto *context = &getContext();
+    BufferizeTypeConverter typeConverter;
+    OwningRewritePatternList patterns;
+    ConversionTarget target(*context);
+
+    target.addLegalDialect<StandardOpsDialect>();
+    patterns.insert<BufferizeTensorConstantOp>(globals, typeConverter, context);
+    target.addDynamicallyLegalOp<ConstantOp>(
+        [&](ConstantOp op) { return typeConverter.isLegal(op.getType()); });
+    if (failed(applyPartialConversion(module, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createTensorConstantBufferizePass() {
+  return std::make_unique<TensorConstantBufferizePass>();
+}

mlir/test/Dialect/Linalg/bufferize.mlir

%0, %1 = linalg.generic {
         %tmp1 = exp %gen_arg1 : f32
         linalg.yield %tmp1, %tmp1 : f32, f32
     } -> tensor<?x?xf32>, tensor<?x?xf32>
     return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32>
 }
 
 // -----
 
-// Check lowering of tensor-valued std.constant's
-// TODO: Move this to std-bufferize.
-
-// CHECK-LABEL:   func @constant() -> tensor<2x3xf32> {
-// CHECK:           %[[VECTOR_MEMREF:.*]] = alloc() : memref<vector<6xf32>>
-// CHECK:           %[[VECTOR_CONST:.*]] = constant dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00]> : vector<6xf32>
-// CHECK:           store %[[VECTOR_CONST]], %[[VECTOR_MEMREF]][] : memref<vector<6xf32>>
-// CHECK:           %[[MEMREF:.*]] = vector.type_cast %[[VECTOR_MEMREF]] : memref<vector<6xf32>> to memref<6xf32>
-// CHECK:           %[[FINAL_SHAPE:.*]] = linalg.reshape %[[MEMREF]] [#map] : memref<6xf32> into memref<2x3xf32>
-// CHECK:           %[[RESULT:.*]] = tensor_load %[[FINAL_SHAPE]] : memref<2x3xf32>
-// CHECK:           return %[[RESULT]] : tensor<2x3xf32>
-func @constant() -> tensor<2x3xf32> {
-  %0 = constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf32>
-  return %0: tensor<2x3xf32>
-}
-
-// -----
-
 #accesses = [
   affine_map<(i, j, k) -> (j, i, k)>,
   affine_map<(i, j, k) -> (i, j)>
 ]
 
 #trait = {
   indexing_maps = #accesses,
   iterator_types = ["parallel", "parallel", "reduction"]

 //  CHECK-SAME:   %[[ST1:[0-9a-z]*]]: tensor<2x?xf32>
 func @bufferize_subtensor_insert(%t : tensor<?x?xf32>, %st0 : tensor<2x3xf32>, %st1 : tensor<2x?xf32>) ->
     (tensor<?x?xf32>, tensor<?x?xf32>) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   //      CHECK: %[[IDX:.*]] = call @make_index() : () -> index
   %i0 = call @make_index() : () -> index
 
+
   //  CHECK-DAG: %[[M0:.*]] = tensor_to_memref %[[T]] : memref<?x?xf32>
   //  CHECK-DAG: %[[SM0:.*]] = tensor_to_memref %[[ST0]] : memref<2x3xf32>
-  // CHECK-NEXT: %[[SUBVIEW0:.*]] = subview %[[M0]][0, 0] [2, 3] [1, 1]
+  // CHECK-NEXT: %[[C0:.*]] = constant 0 : index
+  // CHECK-NEXT: %[[DIM0:.*]] = dim %[[M0]], %[[C0]] : memref<?x?xf32>
+  // CHECK-NEXT: %[[C1:.*]] = constant 1 : index
+  // CHECK-NEXT: %[[DIM1:.*]] = dim %[[M0]], %[[C1]] : memref<?x?xf32>
+  // CHECK-NEXT: %[[M0_COPY:.*]] = alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
+  // CHECK-NEXT: linalg.copy(%[[M0]], %[[M0_COPY]]) : memref<?x?xf32>, memref<?x?xf32>
+  // CHECK-NEXT: %[[SUBVIEW0:.*]] = subview %[[M0_COPY]][0, 0] [2, 3] [1, 1]
   // CHECK-SAME:   memref<?x?xf32> to memref<2x3xf32, #[[$MAP0]]>
   // CHECK-NEXT: linalg.copy(%[[SM0]], %[[SUBVIEW0]]) : memref<2x3xf32>, memref<2x3xf32, #[[$MAP0]]>
-  // CHECK-NEXT: %[[RT0:.*]] = tensor_load %[[M0]] : memref<?x?xf32>
+  // CHECK-NEXT: %[[RT0:.*]] = tensor_load %[[M0_COPY]] : memref<?x?xf32>
   %t0 = subtensor_insert %st0 into %t[0, 0][2, 3][1, 1] : tensor<2x3xf32> into tensor<?x?xf32>
 
   //  CHECK-DAG: %[[M1:.*]] = tensor_to_memref %[[T]] : memref<?x?xf32>
   //  CHECK-DAG: %[[SM1:.*]] = tensor_to_memref %[[ST1]] : memref<2x?xf32>
-  // CHECK-NEXT: %[[SUBVIEW1:.*]] = subview %[[M1]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2]
+  // CHECK-NEXT: %[[C0:.*]] = constant 0 : index
+  // CHECK-NEXT: %[[DIM0:.*]] = dim %[[M1]], %[[C0]] : memref<?x?xf32>
+  // CHECK-NEXT: %[[C1:.*]] = constant 1 : index
+  // CHECK-NEXT: %[[DIM1:.*]] = dim %[[M1]], %[[C1]] : memref<?x?xf32>
+  // CHECK-NEXT: %[[M1_COPY:.*]] = alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
+  // CHECK-NEXT: linalg.copy(%[[M1]], %[[M1_COPY]]) : memref<?x?xf32>, memref<?x?xf32>
+  // CHECK-NEXT: %[[SUBVIEW1:.*]] = subview %[[M1_COPY]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2]
   // CHECK-SAME:   memref<?x?xf32> to memref<2x?xf32, #[[$MAP1]]>
   // CHECK-NEXT: linalg.copy(%[[SM1]], %[[SUBVIEW1]]) : memref<2x?xf32>, memref<2x?xf32, #[[$MAP1]]>
-  // CHECK-NEXT: %[[RT1:.*]] = tensor_load %[[M1]] : memref<?x?xf32>
+  // CHECK-NEXT: %[[RT1:.*]] = tensor_load %[[M1_COPY]] : memref<?x?xf32>
   %t1 = subtensor_insert %st1 into %t[0, %i0][2, %i0][1, 2] : tensor<2x?xf32> into tensor<?x?xf32>
 
   //     CHECK: return %[[RT0]], %[[RT1]]
   return %t0, %t1: tensor<?x?xf32>, tensor<?x?xf32>
 }

mlir/test/Dialect/Standard/tensor-constant-bufferize.mlir

+// RUN: mlir-opt %s -tensor-constant-bufferize -split-input-file
+
+// CHECK-LABEL: module {
+// We check the debug name too since we put some effort into making that readable.
+// The name isn't load-bearing though.
+// CHECK: global_memref "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00>
+// CHECK: @basic
+func @basic() -> tensor<3x4xf32> {
+  // CHECK: %[[MEMREF:.*]] = get_global_memref @__constant_3x4xf32 : memref<3x4xf32>
+  // CHECK: %[[TENSOR:.*]] = tensor_load %[[MEMREF]]
+  %0 = constant dense<7.0> : tensor<3x4xf32>
+  // CHECK: return %[[TENSOR]]
+  return %0 : tensor<3x4xf32>
+}
+
+// CHECK: }
+
+// -----
+
+// CHECK-LABEL: module {
+
+// Only one global is created.
+// CHECK: global_memref
+// CHECK-NOT: global_memref
+func @duplicate_constants() -> (tensor<3x4xf32>, tensor<3x4xf32>) {
+  %0 = constant dense<7.0> : tensor<3x4xf32>
+  %1 = constant dense<7.0> : tensor<3x4xf32>
+  return %0, %1 : tensor<3x4xf32>, tensor<3x4xf32>
+}
+
+// CHECK: }
+
+// -----
+
+// CHECK-LABEL: module {
+
+// Two globals are created.
+// CHECK: global_memref
+// CHECK: global_memref
+// CHECK-NOT: global_memref
+func @multiple_constants() -> (tensor<3x4xf32>, tensor<3x4xf32>) {
+  %0 = constant dense<7.0> : tensor<3x4xf32>
+  %1 = constant dense<8.0> : tensor<3x4xf32>
+  return %0, %1 : tensor<3x4xf32>, tensor<3x4xf32>
+}
+
+// CHECK: }
+
+// -----
+
+// CHECK-LABEL: module {
+// We don't convert non-tensor globals.
+// CHECK-NOT: global_memref
+func @non_tensor() {
+    %0 = constant 7 : i32
+    return
+}
+
+// CHECK: }