[mlir] Refactor code out of BufferPlacement.cpp

[silvas]

Summary

Now BufferPlacement.cpp doesn't depend on Bufferize.h.

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

[rriddle]

mlir/lib/Transforms/Bufferize.cpp

23–31

39

Are the conversions cheep to copy? If not please use & here. May be more clear if you use the proper type instead of auto, same comment applied to several places below.

60

nit: Can you merge this variable definition into the if condition? (I'm assuming the result here is Optional<T>?

98

If there is a comment in the body, please use braces.

https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

152

Please prefer array_pod_sort over std::sort when possible.

157

Why this instead of a normal range based for? This seems more verbose.

192

Is the this-> necessary here?

[silvas]

Thanks River for the comments. To keep this patch a pure move, I've put the stylistic fixes in https://reviews.llvm.org/D89272

mlir/lib/Transforms/Bufferize.cpp

152

llvm::sort has a specialization for POD that does the right thing.

mlir/include/mlir/Transforms/Bufferize.h

 // pattern needs to be written. The infrastructure in this file assists in
 // defining these conversion patterns in a composable way.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_BUFFERIZE_H
 #define MLIR_TRANSFORMS_BUFFERIZE_H
 
-#include "mlir/Analysis/Liveness.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Builders.h"
-#include "mlir/IR/Dominance.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 
 /// A helper type converter class for using inside Buffer Assignment operation
 /// conversion patterns. The default constructor keeps all the types intact

mlir/lib/Transforms/BufferPlacement.cpp

 // The current implementation does not support explicit-control-flow loops and
 // the resulting code will be invalid with respect to program semantics.
 // However, structured control-flow loops are fully supported. Furthermore, it
 // doesn't accept functions which return buffers already.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
+#include "mlir/Analysis/Liveness.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/Bufferize.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SetOperations.h"
 
 using namespace mlir;
 
 /// Walks over all immediate return-like terminators in the given region.
 template <typename FuncT>
 static void walkReturnOperations(Region *region, const FuncT &func) {
void runOnFunction() override {
     BufferPlacement placement(getFunction());
     placement.place();
   }
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
-// BufferAssignmentTypeConverter
-//===----------------------------------------------------------------------===//
-
-/// Registers conversions into BufferAssignmentTypeConverter
-BufferAssignmentTypeConverter::BufferAssignmentTypeConverter() {
-  // Keep all types unchanged.
-  addConversion([](Type type) { return type; });
-  // Convert RankedTensorType to MemRefType.
-  addConversion([](RankedTensorType type) {
-    return (Type)MemRefType::get(type.getShape(), type.getElementType());
-  });
-  // Convert UnrankedTensorType to UnrankedMemRefType.
-  addConversion([](UnrankedTensorType type) {
-    return (Type)UnrankedMemRefType::get(type.getElementType(), 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 BufferAssignmentTypeConverter::tryDecomposeValue(
-    OpBuilder &builder, Location loc, Type type, Value value,
-    SmallVectorImpl<Value> &results) {
-  for (auto conversion : decomposeValueConversions)
-    if (conversion(builder, loc, type, value, results) != llvm::None)
-      return;
-  results.push_back(value);
-}
-
-/// This method tries to decompose a type using provided decompose callback
-/// functions. If it is unable to do so, the original type is returned.
-void BufferAssignmentTypeConverter::tryDecomposeType(
-    Type type, SmallVectorImpl<Type> &types) {
-  for (auto conversion : decomposeTypeConversions)
-    if (conversion(type, types) != llvm::None)
-      return;
-  types.push_back(type);
-}
-
-/// This method returns ResultConversionKind for the input type.
-BufferAssignmentTypeConverter::ResultConversionKind
-BufferAssignmentTypeConverter::getResultConversionKind(Type origin,
-                                                       Type converted) {
-  for (auto conversion : resultTypeConversions) {
-    auto res = conversion(origin, converted);
-    if (res != llvm::None)
-      return res.getValue();
-  }
-  return KeepAsFunctionResult;
-}
-
-//===----------------------------------------------------------------------===//
-// BufferAssignmentFuncOpConverter
-//===----------------------------------------------------------------------===//
-
-/// Performs the actual function signature rewriting step.
-LogicalResult BufferAssignmentFuncOpConverter::matchAndRewrite(
-    mlir::FuncOp funcOp, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
-  auto funcType = funcOp.getType();
-
-  // Convert function arguments using the provided TypeConverter.
-  TypeConverter::SignatureConversion conversion(funcType.getNumInputs());
-  for (auto argType : llvm::enumerate(funcType.getInputs())) {
-    SmallVector<Type, 2> decomposedTypes, convertedTypes;
-    converter.tryDecomposeType(argType.value(), decomposedTypes);
-    converter.convertTypes(decomposedTypes, convertedTypes);
-    conversion.addInputs(argType.index(), convertedTypes);
-  }
-
-  // Convert the result types of the function.
-  SmallVector<Type, 2> newResultTypes;
-  newResultTypes.reserve(funcOp.getNumResults());
-  for (Type resultType : funcType.getResults()) {
-    SmallVector<Type, 2> originTypes;
-    converter.tryDecomposeType(resultType, originTypes);
-    for (auto origin : originTypes) {
-      Type converted = converter.convertType(origin);
-      auto kind = converter.getResultConversionKind(origin, converted);
-      if (kind == BufferAssignmentTypeConverter::AppendToArgumentsList)
-        conversion.addInputs(converted);
-      else
-        // kind = BufferAssignmentTypeConverter::KeepAsFunctionResult
-        newResultTypes.push_back(converted);
-    }
-  }
-
-  if (failed(rewriter.convertRegionTypes(&funcOp.getBody(), converter,
-                                         &conversion)))
-    return failure();
-
-  // Update the signature of the function.
-  rewriter.updateRootInPlace(funcOp, [&] {
-    funcOp.setType(rewriter.getFunctionType(conversion.getConvertedTypes(),
-                                            newResultTypes));
-  });
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
-// BufferAssignmentCallOpConverter
-//===----------------------------------------------------------------------===//
-
-namespace {
-// This class represents a mapping from a result to a list of values and some
-// results that have not yet constructed. Instead, the indices of these
-// results in the operation that will be constructed are known. They will be
-// replaced with the actual values when they are available. The order of
-// adding to this mapping is important.
-class CallOpResultMapping {
-public:
-  CallOpResultMapping() { order = 0; };
-
-  /// Add an available value to the mapping.
-  void addMapping(Value value) { toValuesMapping.push_back({order++, value}); }
-
-  /// Add the index of unavailble result value to the mapping.
-  void addMapping(unsigned index) {
-    toIndicesMapping.push_back({order++, index});
-  }
-
-  /// This method returns the mapping values list. The unknown result values
-  /// that only their indicies are available are replaced with their values.
-  void getMappingValues(ValueRange valuesToReplaceIndices,
-                        SmallVectorImpl<Value> &values) {
-    // Append available values to the list.
-    SmallVector<std::pair<unsigned, Value>, 2> res(toValuesMapping.begin(),
-                                                   toValuesMapping.end());
-    // Replace the indices with the actual values.
-    llvm::for_each(
-        toIndicesMapping, [&](const std::pair<unsigned, unsigned> &entry) {
-          assert(entry.second < valuesToReplaceIndices.size() &&
-                 "The value index is out of range.");
-          res.push_back({entry.first, valuesToReplaceIndices[entry.second]});
-        });
-    // Sort the values based on their adding orders.
-    llvm::sort(res, [](const std::pair<unsigned, Value> &v1,
-                       const std::pair<unsigned, Value> &v2) {
-      return v1.first < v2.first;
-    });
-    // Fill the values.
-    llvm::for_each(res, [&](const std::pair<unsigned, Value> &entry) {
-      values.push_back(entry.second);
-    });
-  }
-
-private:
-  /// Keeping the inserting order of mapping values.
-  int order;
-
-  /// Containing the mapping values with their inserting orders.
-  SmallVector<std::pair<unsigned, Value>, 2> toValuesMapping;
-
-  /// Containing the indices of result values with their inserting orders.
-  SmallVector<std::pair<unsigned, unsigned>, 2> toIndicesMapping;
-};
-} // namespace
-
-/// Performs the actual rewriting step.
-LogicalResult BufferAssignmentCallOpConverter::matchAndRewrite(
-    CallOp callOp, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
-
-  Location loc = callOp.getLoc();
-  OpBuilder builder(callOp);
-  SmallVector<Value, 2> newOperands;
-
-  // TODO: if the CallOp references a FuncOp that only has a declaration (e.g.
-  // to an externally defined symbol like an external library calls), only
-  // convert if some special attribute is set.
-  // This will allow more control of interop across ABI boundaries.
-
-  // Create the operands list of the new `CallOp`. It unpacks the decomposable
-  // values if a decompose callback function has been provided by the user.
-  for (auto operand : operands) {
-    SmallVector<Value, 2> values;
-    this->converter.tryDecomposeValue(builder, loc, operand.getType(), operand,
-                                      values);
-    newOperands.append(values.begin(), values.end());
-  }
-
-  // Create the new result types for the new `CallOp` and a mapping from the old
-  // result to new value(s).
-  SmallVector<Type, 2> newResultTypes;
-  SmallVector<CallOpResultMapping, 4> mappings;
-  mappings.resize(callOp.getNumResults());
-  for (auto result : llvm::enumerate(callOp.getResults())) {
-    SmallVector<Type, 2> originTypes;
-    converter.tryDecomposeType(result.value().getType(), originTypes);
-    auto &resultMapping = mappings[result.index()];
-    for (Type origin : originTypes) {
-      Type converted = converter.convertType(origin);
-      auto kind = converter.getResultConversionKind(origin, converted);
-      if (kind == BufferAssignmentTypeConverter::KeepAsFunctionResult) {
-        newResultTypes.push_back(converted);
-        // The result value is not yet available. Its index is kept and it is
-        // replaced with the actual value of the new `CallOp` later.
-        resultMapping.addMapping(newResultTypes.size() - 1);
-      } else {
-        // kind = BufferAssignmentTypeConverter::AppendToArgumentsList
-        MemRefType memref = converted.dyn_cast<MemRefType>();
-        if (!memref)
-          return callOp.emitError("Cannot allocate for a non-Memref type");
-        Value alloc = rewriter.create<AllocOp>(loc, memref);
-        newOperands.push_back(alloc);
-        resultMapping.addMapping(alloc);
-      }
-    }
-  }
-
-  CallOp newCallOp = rewriter.create<CallOp>(loc, callOp.getCallee(),
-                                             newResultTypes, newOperands);
-
-  // Build a replacing value for each result to replace its uses. If a result
-  // has multiple mapping values, it needs to be packed to a single value.
-  OpBuilder nextBuilder(callOp.getOperation()->getNextNode());
-  SmallVector<Value, 2> replacedValues;
-  replacedValues.reserve(callOp.getNumResults());
-  for (unsigned i = 0, e = callOp.getNumResults(); i < e; ++i) {
-    SmallVector<Value, 2> valuesToPack;
-    mappings[i].getMappingValues(newCallOp.getResults(), valuesToPack);
-    if (valuesToPack.empty()) {
-      // No replacement is required.
-      replacedValues.push_back(nullptr);
-    } else if (valuesToPack.size() == 1) {
-      replacedValues.push_back(valuesToPack.front());
-    } else {
-      // Values need to be packed using callback function. The same callback
-      // that is used for materializeArgumentConversion is used for packing.
-      Value packed = converter.materializeArgumentConversion(
-          nextBuilder, loc, callOp.getType(i), valuesToPack);
-      replacedValues.push_back(packed);
-    }
-  }
-  rewriter.replaceOp(callOp, replacedValues);
-  return success();
-}
-
-//===----------------------------------------------------------------------===//
 // BufferPlacementPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::createBufferPlacementPass() {
   return std::make_unique<BufferPlacementPass>();
 }

mlir/lib/Transforms/Bufferize.cpp

+//===- Bufferize.cpp - Bufferization utilities ----------------------------===//
+//
+// 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/Transforms/Bufferize.h"
+#include "mlir/IR/Operation.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// BufferAssignmentTypeConverter
+//===----------------------------------------------------------------------===//
+
+/// Registers conversions into BufferAssignmentTypeConverter
+BufferAssignmentTypeConverter::BufferAssignmentTypeConverter() {
+  // Keep all types unchanged.
+  addConversion([](Type type) { return type; });
+  // Convert RankedTensorType to MemRefType.
+  addConversion([](RankedTensorType type) {
+    return (Type)MemRefType::get(type.getShape(), type.getElementType());
+  });
+  // Convert UnrankedTensorType to UnrankedMemRefType.
+  addConversion([](UnrankedTensorType type) {
+    return (Type)UnrankedMemRefType::get(type.getElementType(), 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 BufferAssignmentTypeConverter::tryDecomposeValue(
+    OpBuilder &builder, Location loc, Type type, Value value,
+    SmallVectorImpl<Value> &results) {
+  for (auto conversion : decomposeValueConversions)
+    if (conversion(builder, loc, type, value, results) != llvm::None)
+      return;
+  results.push_back(value);
+}
+
+/// This method tries to decompose a type using provided decompose callback
+/// functions. If it is unable to do so, the original type is returned.
+void BufferAssignmentTypeConverter::tryDecomposeType(
+    Type type, SmallVectorImpl<Type> &types) {
+  for (auto conversion : decomposeTypeConversions)
+    if (conversion(type, types) != llvm::None)
+      return;
+  types.push_back(type);
+}
+
+/// This method returns ResultConversionKind for the input type.
+BufferAssignmentTypeConverter::ResultConversionKind
+BufferAssignmentTypeConverter::getResultConversionKind(Type origin,
+                                                       Type converted) {
+  for (auto conversion : resultTypeConversions) {
+    auto res = conversion(origin, converted);
+    if (res != llvm::None)
+      return res.getValue();
+  }
+  return KeepAsFunctionResult;
+}
+
+//===----------------------------------------------------------------------===//
+// BufferAssignmentFuncOpConverter
+//===----------------------------------------------------------------------===//
+
+/// Performs the actual function signature rewriting step.
+LogicalResult BufferAssignmentFuncOpConverter::matchAndRewrite(
+    mlir::FuncOp funcOp, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+  auto funcType = funcOp.getType();
+
+  // Convert function arguments using the provided TypeConverter.
+  TypeConverter::SignatureConversion conversion(funcType.getNumInputs());
+  for (auto argType : llvm::enumerate(funcType.getInputs())) {
+    SmallVector<Type, 2> decomposedTypes, convertedTypes;
+    converter.tryDecomposeType(argType.value(), decomposedTypes);
+    converter.convertTypes(decomposedTypes, convertedTypes);
+    conversion.addInputs(argType.index(), convertedTypes);
+  }
+
+  // Convert the result types of the function.
+  SmallVector<Type, 2> newResultTypes;
+  newResultTypes.reserve(funcOp.getNumResults());
+  for (Type resultType : funcType.getResults()) {
+    SmallVector<Type, 2> originTypes;
+    converter.tryDecomposeType(resultType, originTypes);
+    for (auto origin : originTypes) {
+      Type converted = converter.convertType(origin);
+      auto kind = converter.getResultConversionKind(origin, converted);
+      if (kind == BufferAssignmentTypeConverter::AppendToArgumentsList)
+        conversion.addInputs(converted);
+      else
+        // kind = BufferAssignmentTypeConverter::KeepAsFunctionResult
+        newResultTypes.push_back(converted);
+    }
+  }
+
+  if (failed(rewriter.convertRegionTypes(&funcOp.getBody(), converter,
+                                         &conversion)))
+    return failure();
+
+  // Update the signature of the function.
+  rewriter.updateRootInPlace(funcOp, [&] {
+    funcOp.setType(rewriter.getFunctionType(conversion.getConvertedTypes(),
+                                            newResultTypes));
+  });
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// BufferAssignmentCallOpConverter
+//===----------------------------------------------------------------------===//
+
+namespace {
+// This class represents a mapping from a result to a list of values and some
+// results that have not yet constructed. Instead, the indices of these
+// results in the operation that will be constructed are known. They will be
+// replaced with the actual values when they are available. The order of
+// adding to this mapping is important.
+class CallOpResultMapping {
+public:
+  CallOpResultMapping() { order = 0; };
+
+  /// Add an available value to the mapping.
+  void addMapping(Value value) { toValuesMapping.push_back({order++, value}); }
+
+  /// Add the index of unavailble result value to the mapping.
+  void addMapping(unsigned index) {
+    toIndicesMapping.push_back({order++, index});
+  }
+
+  /// This method returns the mapping values list. The unknown result values
+  /// that only their indicies are available are replaced with their values.
+  void getMappingValues(ValueRange valuesToReplaceIndices,
+                        SmallVectorImpl<Value> &values) {
+    // Append available values to the list.
+    SmallVector<std::pair<unsigned, Value>, 2> res(toValuesMapping.begin(),
+                                                   toValuesMapping.end());
+    // Replace the indices with the actual values.
+    llvm::for_each(
+        toIndicesMapping, [&](const std::pair<unsigned, unsigned> &entry) {
+          assert(entry.second < valuesToReplaceIndices.size() &&
+                 "The value index is out of range.");
+          res.push_back({entry.first, valuesToReplaceIndices[entry.second]});
+        });
+    // Sort the values based on their adding orders.
+    llvm::sort(res, [](const std::pair<unsigned, Value> &v1,
+                       const std::pair<unsigned, Value> &v2) {
+      return v1.first < v2.first;
+    });
+    // Fill the values.
+    llvm::for_each(res, [&](const std::pair<unsigned, Value> &entry) {
+      values.push_back(entry.second);
+    });
+  }
+
+private:
+  /// Keeping the inserting order of mapping values.
+  int order;
+
+  /// Containing the mapping values with their inserting orders.
+  SmallVector<std::pair<unsigned, Value>, 2> toValuesMapping;
+
+  /// Containing the indices of result values with their inserting orders.
+  SmallVector<std::pair<unsigned, unsigned>, 2> toIndicesMapping;
+};
+} // namespace
+
+/// Performs the actual rewriting step.
+LogicalResult BufferAssignmentCallOpConverter::matchAndRewrite(
+    CallOp callOp, ArrayRef<Value> operands,
+    ConversionPatternRewriter &rewriter) const {
+
+  Location loc = callOp.getLoc();
+  OpBuilder builder(callOp);
+  SmallVector<Value, 2> newOperands;
+
+  // TODO: if the CallOp references a FuncOp that only has a declaration (e.g.
+  // to an externally defined symbol like an external library calls), only
+  // convert if some special attribute is set.
+  // This will allow more control of interop across ABI boundaries.
+
+  // Create the operands list of the new `CallOp`. It unpacks the decomposable
+  // values if a decompose callback function has been provided by the user.
+  for (auto operand : operands) {
+    SmallVector<Value, 2> values;
+    this->converter.tryDecomposeValue(builder, loc, operand.getType(), operand,
+                                      values);
+    newOperands.append(values.begin(), values.end());
+  }
+
+  // Create the new result types for the new `CallOp` and a mapping from the old
+  // result to new value(s).
+  SmallVector<Type, 2> newResultTypes;
+  SmallVector<CallOpResultMapping, 4> mappings;
+  mappings.resize(callOp.getNumResults());
+  for (auto result : llvm::enumerate(callOp.getResults())) {
+    SmallVector<Type, 2> originTypes;
+    converter.tryDecomposeType(result.value().getType(), originTypes);
+    auto &resultMapping = mappings[result.index()];
+    for (Type origin : originTypes) {
+      Type converted = converter.convertType(origin);
+      auto kind = converter.getResultConversionKind(origin, converted);
+      if (kind == BufferAssignmentTypeConverter::KeepAsFunctionResult) {
+        newResultTypes.push_back(converted);
+        // The result value is not yet available. Its index is kept and it is
+        // replaced with the actual value of the new `CallOp` later.
+        resultMapping.addMapping(newResultTypes.size() - 1);
+      } else {
+        // kind = BufferAssignmentTypeConverter::AppendToArgumentsList
+        MemRefType memref = converted.dyn_cast<MemRefType>();
+        if (!memref)
+          return callOp.emitError("Cannot allocate for a non-Memref type");
+        Value alloc = rewriter.create<AllocOp>(loc, memref);
+        newOperands.push_back(alloc);
+        resultMapping.addMapping(alloc);
+      }
+    }
+  }
+
+  CallOp newCallOp = rewriter.create<CallOp>(loc, callOp.getCallee(),
+                                             newResultTypes, newOperands);
+
+  // Build a replacing value for each result to replace its uses. If a result
+  // has multiple mapping values, it needs to be packed to a single value.
+  OpBuilder nextBuilder(callOp.getOperation()->getNextNode());
+  SmallVector<Value, 2> replacedValues;
+  replacedValues.reserve(callOp.getNumResults());
+  for (unsigned i = 0, e = callOp.getNumResults(); i < e; ++i) {
+    SmallVector<Value, 2> valuesToPack;
+    mappings[i].getMappingValues(newCallOp.getResults(), valuesToPack);
+    if (valuesToPack.empty()) {
+      // No replacement is required.
+      replacedValues.push_back(nullptr);
+    } else if (valuesToPack.size() == 1) {
+      replacedValues.push_back(valuesToPack.front());
+    } else {
+      // Values need to be packed using callback function. The same callback
+      // that is used for materializeArgumentConversion is used for packing.
+      Value packed = converter.materializeArgumentConversion(
+          nextBuilder, loc, callOp.getType(i), valuesToPack);
+      replacedValues.push_back(packed);
+    }
+  }
+  rewriter.replaceOp(callOp, replacedValues);
+  return success();
+}

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   BufferPlacement.cpp
+  Bufferize.cpp
   Canonicalizer.cpp
   CopyRemoval.cpp
   CSE.cpp
   DialectConversion.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp