[mlir] Add std.tensor_to_memref op and teach the infra about it

[silvas]

Summary

The opposite of tensor_to_memref is tensor_load.

• Add some basic tensor_load/tensor_to_memref folding.
• Add source/target materializations to BufferizeTypeConverter.
• Add an example std bufferization pattern/pass that shows how the materialiations work together (more std bufferization patterns to come in subsequent commits).
  • In coming commits, I'll document how to write composable bufferization passes/patterns and update the other in-tree bufferization passes to match this convention. The populate* functions will of course continue to be exposed for power users.

The naming on tensor_load/tensor_to_memref and their pretty forms are
not very intuitive. I'm open to any suggestions here. One key
observation is that the memref type must always be the one specified in
the pretty form, since the tensor type can be inferred from the memref
type but not vice-versa.

With this, I've been able to replace all my custom bufferization type
converters in npcomp with BufferizeTypeConverter!

Part of the plan discussed in:
https://llvm.discourse.group/t/what-is-the-strategy-for-tensor-memref-conversion-bufferization/1938/17

[silvas]

update description

[silvas]

update

[herhut]

mlir/include/mlir/Transforms/Bufferize.h

16

Should we make this distinction here? Why not just have all derive from BufferizeConversionPattern? I am not strongly opposed, just wondering what the motivation is.

18

nit: is sufficient

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3610

More of a question than a comment: this seems like a canonicalization rather than constant folding. Should it be expressed that way (as a canonicalization pattern?)

3621

Same comment as above: should this be a canonicalization?

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3610

Folding isn't about "constant folding", but when you can infer a value from the existing IR, i.e. the main difference between canonicalization and folding is that folding cannot create new IR.

[jurahul]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3610

Make sense. Thanks!

[silvas]

mlir/include/mlir/Transforms/Bufferize.h

16

Good point! I will be writing more docs elaborating on this soon (I felt the need to put this advice in this patch since I started applying it).

One fundamental reason to avoid BufferizationTypeConverter is that many type conversions that are useful for bufferization also apply to any type. E.g. the patterns for select (with scalar predicate) or scf.if. These are often "structural" ops whose operands/results behave much as a block argument, where we structurally know that it is a pure passthrough from some other SSA value, so any type conversion is ok (e.g. index->i64). A recent example of this is https://reviews.llvm.org/D88083 in which we inadvertently specialized the shape.assuming type conversion to bufferization, when it need not be limited to that (and I will be refactoring that soon).

There are a couple more "stylistic" reasons as well:

• a principle of least requirements -- if you are writing an STL algorithm that only use ++ on an iterator, it's overkill to require a RandomAccessIterator. There is both an cover-more-use-cases advantage and a clarity-for-reader advantage to know that the less-powerful ForwardIterator abstraction is sufficient for the algorithm. This is similar to how one would use a RewritePattern instead of a ConversionRewritePattern if the pattern does not need type conversions.
• a (very?) small minority of bufferization conversions actually require BufferizeTypeConverter's extra methods. So using a regular ConversionPattern helps to avoid introducing new concepts for readers.

Actually, I have been considering making BufferizeTypeConverter not inherit from TypeConverter at all. I think it could be split into a free function populateBufferizeTypeConversions(TypeConverter &) + a new BufferizeConversionPolicy with all the extra methods. BufferizeConversionPattern would then change to hold a BufferizeConversionPolicy, and any type conversions would happen through the TypeConverter obtained from getTypeConverter().

[herhut]

mlir/include/mlir/Transforms/Bufferize.h

16

Thanks for clarifying. Splitting into general _type conversion_ and _bufferize_ patterns seems useful. I also like the idea of splitting the bufferization policy out form type conversion, as thinking about it they are separate things.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

TypesMatchWith<"result type matches tensor equivalent of 'memref'",
                     "memref", "result",
                     "getTensorTypeFromMemRefType($_self)">]> {
   let summary = "tensor load operation";
   let description = [{
     Create a tensor from a memref, making an independent copy of the element
     data. The result value is a tensor whose shape and element type match the
     memref operand.
 
+    The opposite of this op is tensor_to_memref. Together, these two ops are
+    useful for source/target materializations when doing type conversions
+    involving tensors and memrefs.
+
     Example:
 
     ```mlir
     // Produces a value of tensor<4x?xf32> type.
     %12 = tensor_load %10 : memref<4x?xf32, #layout, memspace0>
     ```
   }];
 
TensorType getType() {
       Type resultType = getResult().getType();
       if (resultType.isa<TensorType>())
         return resultType.cast<TensorType>();
       return {};
     }
   }];
 
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
+
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TensorStoreOp
 //===----------------------------------------------------------------------===//
 
 def TensorStoreOp : Std_Op<"tensor_store",
     [SameOperandsShape, SameOperandsElementType,
let arguments = (ins AnyTensor:$tensor, Arg<AnyRankedOrUnrankedMemRef,
                        "the reference to store to", [MemWrite]>:$memref);
   // TensorStoreOp is fully verified by traits.
   let verifier = ?;
 
   let assemblyFormat = "$tensor `,` $memref attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
+// TensorToMemrefOp
+//===----------------------------------------------------------------------===//
+
+def TensorToMemrefOp : Std_Op<"tensor_to_memref",
+    [SameOperandsAndResultShape, SameOperandsAndResultElementType,
+     TypesMatchWith<"type of 'tensor' is the tensor equivalent of 'memref'",
+                    "memref", "tensor",
+                    "getTensorTypeFromMemRefType($_self)">]> {
+  let summary = "tensor to memref operation";
+  let description = [{
+    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.
+
+    The opposite of this op is tensor_load. Together, these two ops are useful
+    for source/target materializations when doing type conversions involving
+    tensors and memrefs.
+
+    Note: This op takes the memref type in its pretty form because the tensor
+    type can always be inferred from the memref type, but the reverse is not
+    true. For example, the memref might have a layout map or memory space which
+    cannot be inferred from the tensor type.
+
+    ```mlir
+    // Result type is tensor<4x?xf32>
+    %12 = tensor_to_memref %10 :  memref<4x?xf32, #map0, 42>
+    ```
+  }];
+
+  let arguments = (ins AnyTensor:$tensor);
+  let results = (outs Res<AnyRankedOrUnrankedMemRef,
+                      "the memref to create", [MemAlloc]>:$memref);
+  // This op is fully verified by traits.
+  let verifier = ?;
+
+  let assemblyFormat = "$tensor attr-dict `:` type($memref)";
+
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // TransposeOp
 //===----------------------------------------------------------------------===//
 
 def TransposeOp : Std_Op<"transpose", [NoSideEffect]>,
     Arguments<(ins AnyStridedMemRef:$in, AffineMapAttr:$permutation)>,
     Results<(outs AnyStridedMemRef)> {
   let summary = "`transpose` produces a new strided memref (metadata-only)";
   let description = [{

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

 // transformation library.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Bufferize.h"
 
 namespace mlir {
 
 class OwningRewritePatternList;
 
 /// Creates an instance of the ExpandAtomic pass.
 std::unique_ptr<Pass> createExpandAtomicPass();
 
 void populateExpandTanhPattern(OwningRewritePatternList &patterns,
                                MLIRContext *ctx);
 
+void populateStdBufferizePatterns(MLIRContext *context,
+                                  BufferizeTypeConverter &typeConverter,
+                                  OwningRewritePatternList &patterns);
+
+/// Creates an instance of the StdBufferize pass.
+std::unique_ptr<Pass> createStdBufferizePass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES_H_

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

 
 include "mlir/Pass/PassBase.td"
 
 def ExpandAtomic : FunctionPass<"expand-atomic"> {
   let summary = "Expands AtomicRMWOp into GenericAtomicRMWOp.";
   let constructor = "mlir::createExpandAtomicPass()";
 }
 
+def StdBufferize : FunctionPass<"std-bufferize"> {
+  let summary = "Bufferize the std dialect";
+  let constructor = "mlir::createStdBufferizePass()";
+}
+
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

mlir/include/mlir/Transforms/Bufferize.h

 //===- Bufferize.h - Bufferization utilities --------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // We use the term "bufferize" to mean conversion from tensor types to
 // memref types.
 //
 // Generally speaking, for each op that operates on tensor types, a conversion
 // pattern needs to be written. The infrastructure in this file assists in
 // defining these conversion patterns in a composable way.
 //
+// Bufferization conversion patterns should generally use the ordinary
+// conversion pattern classes (e.g. OpConversionPattern). A TypeConverter
+// (accessible with getTypeConverter()) available on such patterns in sufficient
+// for most cases (if needed at all).
+//
+// But some patterns require access to the extra functions on
+// BufferizeTypeConverter that don't exist on the base TypeConverter class. For
+// those cases, BufferizeConversionPattern and its related classes should be
+// used, which provide access to a BufferizeTypeConverter directly.
+//
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_BUFFERIZE_H
 #define MLIR_TRANSFORMS_BUFFERIZE_H
 
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 } // namespace
 
 void TensorCastOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<ChainedTensorCast>(context);
 }
 
 //===----------------------------------------------------------------------===//
-// Helpers for Tensor[Load|Store]Op
+// Helpers for Tensor[Load|Store]Op and TensorToMemrefOp
 //===----------------------------------------------------------------------===//
 
 static Type getTensorTypeFromMemRefType(Type type) {
   if (auto memref = type.dyn_cast<MemRefType>())
     return RankedTensorType::get(memref.getShape(), memref.getElementType());
   if (auto memref = type.dyn_cast<UnrankedMemRefType>())
     return UnrankedTensorType::get(memref.getElementType());
   return NoneType::get(type.getContext());
 }
 
 //===----------------------------------------------------------------------===//
+// TensorLoadOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult TensorLoadOp::fold(ArrayRef<Attribute>) {
+  if (auto tensorToMemref = memref().getDefiningOp<TensorToMemrefOp>())
+    return tensorToMemref.tensor();
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
+// TensorToMemrefOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult TensorToMemrefOp::fold(ArrayRef<Attribute>) {
+  if (auto tensorLoad = tensor().getDefiningOp<TensorLoadOp>())
+    if (tensorLoad.memref().getType() == getType())
+      return tensorLoad.memref();
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // TransposeOp
 //===----------------------------------------------------------------------===//
 
 /// Build a strided memref type by applying `permutationMap` tp `memRefType`.
 static MemRefType inferTransposeResultType(MemRefType memRefType,
                                            AffineMap permutationMap) {
   auto rank = memRefType.getRank();
   auto originalSizes = memRefType.getShape();

mlir/lib/Dialect/StandardOps/Transforms/Bufferize.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 std ops.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/Bufferize.h"
+#include "PassDetail.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+class BufferizeTensorCastOp : public OpConversionPattern<TensorCastOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(TensorCastOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto resultType = getTypeConverter()->convertType(op.getType());
+    rewriter.replaceOpWithNewOp<MemRefCastOp>(op, resultType, operands[0]);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateStdBufferizePatterns(MLIRContext *context,
+                                        BufferizeTypeConverter &typeConverter,
+                                        OwningRewritePatternList &patterns) {
+  patterns.insert<BufferizeTensorCastOp>(typeConverter, context);
+}
+
+namespace {
+struct StdBufferizePass : public StdBufferizeBase<StdBufferizePass> {
+  void runOnFunction() override {
+    auto *context = &getContext();
+    BufferizeTypeConverter typeConverter;
+    OwningRewritePatternList patterns;
+    ConversionTarget target(*context);
+
+    target.addLegalDialect<StandardOpsDialect>();
+
+    populateStdBufferizePatterns(context, typeConverter, patterns);
+    target.addIllegalOp<TensorCastOp>();
+
+    if (failed(mlir::applyPartialConversion(getFunction(), target, patterns)))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createStdBufferizePass() {
+  return std::make_unique<StdBufferizePass>();
+}

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

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

mlir/lib/Transforms/Bufferize.cpp

BufferizeTypeConverter::BufferizeTypeConverter() {
   // Convert RankedTensorType to MemRefType.
   addConversion([](RankedTensorType type) -> Type {
     return MemRefType::get(type.getShape(), type.getElementType());
   });
   // Convert UnrankedTensorType to UnrankedMemRefType.
   addConversion([](UnrankedTensorType type) -> Type {
     return UnrankedMemRefType::get(type.getElementType(), 0);
   });
+  addSourceMaterialization([](OpBuilder &builder, RankedTensorType type,
+                              ValueRange inputs, Location loc) -> Value {
+    assert(inputs.size() == 1);
+    assert(inputs[0].getType().isa<BaseMemRefType>());
+    return builder.create<TensorLoadOp>(loc, type, inputs[0]);
+  });
+  addTargetMaterialization([](OpBuilder &builder, MemRefType type,
+                              ValueRange inputs, Location loc) -> Value {
+    assert(inputs.size() == 1);
+    assert(inputs[0].getType().isa<TensorType>());
+    return builder.create<TensorToMemrefOp>(loc, type, inputs[0]);
+  });
 }
 
 /// This method tries to decompose a value of a certain type using provided
 /// decompose callback functions. If it is unable to do so, the original value
 /// is returned.
 void BufferizeTypeConverter::tryDecomposeValue(
     OpBuilder &builder, Location loc, Type type, Value value,
     SmallVectorImpl<Value> &results) {

mlir/test/Dialect/Standard/bufferize.mlir

+// RUN: mlir-opt %s -std-bufferize | FileCheck %s
+
+// CHECK-LABEL:   func @tensor_cast(
+// CHECK-SAME:                      %[[TENSOR:.*]]: tensor<?xindex>) -> tensor<2xindex> {
+// CHECK:           %[[MEMREF:.*]] = tensor_to_memref %[[TENSOR]]
+// CHECK:           %[[CASTED:.*]] = memref_cast %[[MEMREF]] : memref<?xindex> to memref<2xindex>
+// CHECK:           %[[RET:.*]] = tensor_load %[[CASTED]]
+// CHECK:           return %[[RET]] : tensor<2xindex>
+func @tensor_cast(%arg0: tensor<?xindex>) -> tensor<2xindex> {
+  %0 = tensor_cast %arg0 : tensor<?xindex> to tensor<2xindex>
+  return %0 : tensor<2xindex>
+}

mlir/test/Dialect/Standard/canonicalize.mlir

+// RUN: mlir-opt %s -canonicalize | FileCheck %s
+
+// Test case: Basic folding of tensor_load(tensor_to_memref(t)) -> t
+// CHECK-LABEL:   func @tensor_load_of_tensor_to_memref(
+// CHECK-SAME:                                          %[[TENSOR:.*]]: tensor<?xf32>) -> tensor<?xf32> {
+// CHECK:           return %[[TENSOR]]
+func @tensor_load_of_tensor_to_memref(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+    %0 = tensor_to_memref %arg0 : memref<?xf32>
+    %1 = tensor_load %0 : memref<?xf32>
+    return %1 : tensor<?xf32>
+}
+
+// Test case: Basic folding of tensor_to_memref(tensor_load(m)) -> m
+// CHECK-LABEL:   func @tensor_to_memref_of_tensor_load(
+// CHECK-SAME:                                          %[[MEMREF:.*]]: memref<?xf32>) -> memref<?xf32> {
+// CHECK:           return %[[MEMREF]]
+func @tensor_to_memref_of_tensor_load(%arg0: memref<?xf32>) -> memref<?xf32> {
+    %0 = tensor_load %arg0 : memref<?xf32>
+    %1 = tensor_to_memref %0 : memref<?xf32>
+    return %1 : memref<?xf32>
+}
+
+// Test case: If the memrefs are not the same type, don't fold them.
+// CHECK-LABEL:   func @no_fold_tensor_to_memref_of_tensor_load(
+// CHECK-SAME:                                                  %[[MEMREF_ADDRSPACE2:.*]]: memref<?xf32, 2>) -> memref<?xf32, 7> {
+// CHECK:           %[[TENSOR:.*]] = tensor_load %[[MEMREF_ADDRSPACE2]] : memref<?xf32, 2>
+// CHECK:           %[[MEMREF_ADDRSPACE7:.*]] = tensor_to_memref %[[TENSOR]] : memref<?xf32, 7>
+// CHECK:           return %[[MEMREF_ADDRSPACE7]]
+func @no_fold_tensor_to_memref_of_tensor_load(%arg0: memref<?xf32, 2>) -> memref<?xf32, 7> {
+    %0 = tensor_load %arg0 : memref<?xf32, 2>
+    %1 = tensor_to_memref %0 : memref<?xf32, 7>
+    return %1 : memref<?xf32, 7>
+}

mlir/test/Dialect/Standard/ops.mlir

 }
 
 // CHECK-LABEL: test_index_cast_tensor_reverse
 func @test_index_cast_tensor_reverse(%arg0 : tensor<i64>) -> tensor<index> {
   %0 = index_cast %arg0 : tensor<i64> to tensor<index>
   return %0 : tensor<index>
 }
 
+// CHECK-LABEL: test_tensor_to_memref
+func @test_tensor_to_memref(%arg0: tensor<?xi64>, %arg1: tensor<*xi64>) -> (memref<?xi64, affine_map<(d0) -> (d0 + 7)>>, memref<*xi64, 1>) {
+  %0 = tensor_to_memref %arg0 : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>
+  %1 = tensor_to_memref %arg1 : memref<*xi64, 1>
+  return %0, %1 : memref<?xi64, affine_map<(d0) -> (d0 + 7)>>, memref<*xi64, 1>
+}
+
 // CHECK-LABEL: @assert
 func @assert(%arg : i1) {
   assert %arg, "Some message in case this assertion fails."
   return
 }
 
 // CHECK-LABEL: @dynamic_tensor_from_elements
 func @dynamic_tensor_from_elements(%m : index, %n : index)