try to use new buffer deallocation (#2919)

* implementation

Signed-off-by: Chen Tong <[email protected]>

* comments

Signed-off-by: Chen Tong <[email protected]>

* format

Signed-off-by: Chen Tong <[email protected]>

---------

Signed-off-by: Chen Tong <[email protected]>
Co-authored-by: Tung D. Le <[email protected]>
Co-authored-by: Alexandre Eichenberger <[email protected]>

Author: 3 people

src/Compiler/CompilerDialects.cpp

@@ -46,8 +46,9 @@ DialectRegistry registerDialects(ArrayRef<accel::Accelerator::Kind> accels) {
   for (auto *accel : accel::Accelerator::getAccelerators())
     accel->registerDialects(registry);
 
-  if (useOldBufferization)
-    memref::registerAllocationOpInterfaceExternalModels(registry);
+  // Register interface needed by both old and new buffer deallocation pass.
+  memref::registerAllocationOpInterfaceExternalModels(registry);
+  arith::registerBufferDeallocationOpInterfaceExternalModels(registry);
 
   return registry;
 }

src/Compiler/CompilerOptions.cpp

@@ -42,6 +42,7 @@ bool enableONNXHybridPass;                             // common for both
 std::vector<std::string> functionsToDecompose;         // common for both
 std::string opsForCall;                                // common for both
 bool disableKrnlOpFusion;                              // common for both
+bool disableMemRefPrefetch;                            // common for both
 EmissionTargetType emissionTarget;                     // onnx-mlir only
 bool invokeOnnxVersionConverter;                       // onnx-mlir only
 bool preserveLocations;                                // onnx-mlir only
@@ -211,6 +212,13 @@ static llvm::cl::opt<bool, true> disableKrnlOpFusionOpt(
     llvm::cl::location(disableKrnlOpFusion), llvm::cl::init(false),
     llvm::cl::cat(OnnxMlirCommonOptions));
 
+static llvm::cl::opt<bool, true> disableMemRefPrefetchOpt(
+    "disable-memref-prefetch",
+    llvm::cl::desc("disable generation of memref.prefetch (default=false)\n"
+                   "Set to 'true' if you want to disable prefetch."),
+    llvm::cl::location(disableMemRefPrefetch), llvm::cl::init(false),
+    llvm::cl::cat(OnnxMlirCommonOptions));
+
 static llvm::cl::opt<bool, true> disableRecomposeOptionOpt("disable-recompose",
     llvm::cl::desc("Disable recomposition of ONNX operations."),
     llvm::cl::location(disableRecomposeOption), llvm::cl::init(false),

src/Compiler/CompilerOptions.hpp

@@ -87,6 +87,7 @@ extern bool enableONNXHybridPass;                             // common for both
 extern std::vector<std::string> functionsToDecompose;         // common for both
 extern std::string opsForCall;                                // common for both
 extern bool disableKrnlOpFusion;                              // common for both
+extern bool disableMemRefPrefetch;                            // common for both
 extern EmissionTargetType emissionTarget;                     // onnx-mlir only
 extern bool invokeOnnxVersionConverter;                       // onnx-mlir only
 extern bool preserveLocations;                                // onnx-mlir only

src/Dialect/Krnl/CMakeLists.txt

@@ -20,6 +20,7 @@ add_onnx_mlir_library(OMKrnlOps
   OMSpecializedKernelOpInterface
 
   LINK_LIBS PUBLIC
+  OMCompilerOptions
   OMONNXOps
   MLIRLLVMCommonConversion
   MLIRAffineDialect

src/Dialect/Krnl/DialectBuilder.cpp

@@ -14,6 +14,7 @@
 
 #include "llvm/ADT/TypeSwitch.h"
 
+#include "src/Compiler/CompilerOptions.hpp"
 #include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
 #include "src/Dialect/Krnl/DialectBuilder.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -94,12 +95,16 @@ Value KrnlBuilder::getLinearOffsetIndexIE(
 
 void KrnlBuilder::prefetch(Value memref, ValueRange indices, bool isWrite,
     unsigned localityHint, bool isDataCache) {
+  if (disableMemRefPrefetch)
+    return;
   b().create<KrnlPrefetchOp>(
       loc(), memref, indices, isWrite, localityHint, isDataCache);
 }
 
 void KrnlBuilder::prefetchIE(Value memref, ArrayRef<IndexExpr> indices,
     bool isWrite, unsigned localityHint, bool isDataCache) {
+  if (disableMemRefPrefetch)
+    return;
   SmallVector<Value, 4> indexValues;
   IndexExpr::getValues(indices, indexValues);
   b().create<KrnlPrefetchOp>(

src/Dialect/Mlir/CMakeLists.txt

@@ -12,6 +12,7 @@ add_onnx_mlir_library(OMMlirDialects
   OMSpecializedKernelOpInterface
 
   LINK_LIBS PUBLIC
+  OMCompilerOptions
   MLIRMathDialect
   MLIRAffineDialect
   MLIRSCFDialect

src/Dialect/Mlir/DialectBuilder.cpp

@@ -25,6 +25,7 @@
 #include "llvm/Support/Debug.h"
 
 // Please do not add dependences on ONNX or KRNL dialects.
+#include "src/Compiler/CompilerOptions.hpp"
 #include "src/Dialect/Mlir/DialectBuilder.hpp"
 #include "src/Dialect/Mlir/VectorMachineSupport.hpp"
 
@@ -1657,12 +1658,16 @@ Value MemRefBuilder::dim(Value val, Value index) const {
 
 void MemRefBuilder::prefetch(Value memref, ValueRange indices, bool isWrite,
     unsigned locality, bool isData) {
+  if (disableMemRefPrefetch)
+    return;
   b().create<memref::PrefetchOp>(
       loc(), memref, indices, isWrite, locality, isData);
 }
 
 void MemRefBuilder::prefetchIE(Value memref, ArrayRef<IndexExpr> indices,
     bool isWrite, unsigned locality, bool isData) {
+  if (disableMemRefPrefetch)
+    return;
   SmallVector<Value, 4> indexVals;
   IndexExpr::getValues(indices, indexVals);
   prefetch(memref, indexVals, isWrite, locality, isData);