Experimental: convert onnx to linalg

src/Compiler/CompilerOptions.cpp

@@ -210,6 +210,10 @@ llvm::cl::opt<bool> allowSorting("allowSorting",
     llvm::cl::desc("Perform topological sort on onnx graph"),
     llvm::cl::init(true), llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool> enableLinalg("enableLinalg",
+    llvm::cl::desc("Enable ONNX to Linalg conversion and related passes"),
+    llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions));
+
 // Configuration states associated with certain options.
 // For example, when maccel is specified, NNPA can register
 // dependent libdnn.

src/Compiler/CompilerOptions.hpp

@@ -71,6 +71,7 @@ extern llvm::cl::opt<int> onnxOpTransformThreshold;
 extern llvm::cl::opt<bool> onnxOpTransformReport;
 extern llvm::cl::opt<bool> enableParallel;
 extern llvm::cl::opt<bool> enableSimdDataLayout;
+extern llvm::cl::opt<bool> enableLinalg;
 
 // The customEnvFlags must be scanned before the normal options.
 bool parseCustomEnvFlagsCommandLineOption(int argc, const char *const *argv,

src/Compiler/CompilerPasses.cpp

@@ -23,6 +23,7 @@
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Transforms/Passes.h"
@@ -121,17 +122,33 @@ void addONNXToKrnlPasses(mlir::PassManager &pm, int optLevel, bool enableCSE,
   if (enableInstrumentONNXSignature)
     pm.addNestedPass<func::FuncOp>(
         onnx_mlir::createInstrumentONNXSignaturePass());
+  if (enableLinalg) {
+    pm.addPass(onnx_mlir::createLowerONNXToLinalgPass());
+    // Linalg bufferization can be before or after LowerToKrnlPass
+    // However, sice bufferization::AllocTensorOp is in LowerToKrnl temporarily,
+    // these passes have to be called before LowerToKrnlPass
+    pm.addNestedPass<func::FuncOp>(createLinalgBufferizePass());
+    // Convert tensor.EmptyOp to bufferization.alloc_tensor
+    pm.addNestedPass<func::FuncOp>(
+        bufferization::createEmptyTensorToAllocTensorPass());
+  }
   pm.addPass(onnx_mlir::createLowerToKrnlPass(optLevel, enableParallel));
   // An additional pass of canonicalization is helpful because lowering
   // from ONNX dialect to Standard dialect exposes additional canonicalization
   // opportunities.
+
+  // For Linalg and Krnl mixed IR:
+  // Canonicalization pass will clean up bufferization::to_tensor and to_memref
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(
       onnx_mlir::createDisconnectKrnlDimFromAllocPass());
   pm.addPass(mlir::createCanonicalizerPass());
 } // namespace onnx_mlir
 
 void addKrnlToAffinePasses(mlir::PassManager &pm) {
+  if (enableLinalg) {
+    pm.addNestedPass<func::FuncOp>(createConvertLinalgToAffineLoopsPass());
+  }
   pm.addNestedPass<func::FuncOp>(
       onnx_mlir::krnl::createConvertKrnlToAffinePass());
 }

src/Conversion/CMakeLists.txt

@@ -5,6 +5,7 @@ add_subdirectory(KrnlToLLVM)
 add_subdirectory(KrnlToAffine)
 add_subdirectory(KrnlSeqToMemref)
 add_subdirectory(ONNXToTOSA)
+add_subdirectory(ONNXToLinalg)
 
 if (ONNX_MLIR_ENABLE_MHLO)
   add_subdirectory(ONNXToMhlo)

src/Conversion/ONNXToKrnl/CMakeLists.txt

@@ -73,6 +73,7 @@ add_onnx_mlir_library(OMONNXToKrnl
 
   LINK_LIBS PUBLIC
   OMAccelerator
+  OMCompilerOptions
   OMConstPropHelper
   OMONNXOps
   OMSupport

src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp

@@ -13,6 +13,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "src/Compiler/CompilerOptions.hpp"
@@ -24,6 +25,31 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
+struct BufferAllocTensorOpLowering : public ConversionPattern {
+  BufferAllocTensorOpLowering(TypeConverter &typeConverter, MLIRContext *ctx)
+      : ConversionPattern(typeConverter,
+            bufferization::AllocTensorOp::getOperationName(), 1, ctx) {}
+
+  LogicalResult matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const final {
+
+    Location loc = op->getLoc();
+    MultiDialectBuilder<MemRefBuilder> create(rewriter, loc);
+
+    auto alloc =
+        create.mem.alloc(typeConverter->convertType(op->getResult(0).getType())
+                             .cast<MemRefType>(),
+            operands);
+    rewriter.replaceOp(op, alloc->getResults());
+    return success();
+  }
+};
+
+void populateLoweringBufferAllocTensorOpPattern(RewritePatternSet &patterns,
+    TypeConverter &typeConverter, MLIRContext *ctx) {
+  patterns.insert<BufferAllocTensorOpLowering>(typeConverter, ctx);
+}
+
 //===----------------------------------------------------------------------===//
 // EntryPoint Op lowering to Krnl Entry Point.
 //===----------------------------------------------------------------------===//
@@ -256,6 +282,8 @@ void populateONNXToKrnlConversionPattern(RewritePatternSet &patterns,
   populateLoweringONNXSequenceLengthOpPattern(patterns, typeConverter, ctx);
   // Entry point
   patterns.insert<ONNXEntryPointLowering>(ctx);
+  // Experiment: lowering bufferization::AllocTensorOp to memref alloc
+  populateLoweringBufferAllocTensorOpPattern(patterns, typeConverter, ctx);
 }
 
 //===----------------------------------------------------------------------===//
@@ -350,6 +378,7 @@ void FrontendToKrnlLoweringPass::runOnOperation() {
   target.addIllegalOp<mlir::AffineLoadOp>();
   target.addIllegalOp<mlir::memref::StoreOp>();
   target.addIllegalOp<mlir::AffineStoreOp>();
+  target.addIllegalOp<bufferization::AllocTensorOp>();
 
   // If `emitDealloc` is turned off, make sure we don't have buffer deallocation
   // at this level. Will use MLIR buffer-deallocation for this purpose instead.

src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp

@@ -13,8 +13,11 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+
 #include "src/Accelerators/Accelerator.hpp"
+#include "src/Compiler/CompilerOptions.hpp"
+#include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
 #include "src/Dialect/Krnl/DialectBuilder.hpp"
 #include "src/Dialect/Mlir/DialectBuilder.hpp"
 #include "src/Dialect/ONNX/ONNXOps/OpHelper.hpp"
@@ -724,8 +727,16 @@ KrnlTypeConverter::KrnlTypeConverter() {
     if (inputs.size() != 1)
       return llvm::None;
 
-    return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
-        .getResult(0);
+    // Use ToTensorOp instead of UnrealizedConversionCastOp
+    // because Linalg use ToTensor, though they are the same in semantic
+    // Since UnrealizedConversionCastOp is used in other places and will not
+    // be replaced in this PR
+    if (enableLinalg)
+      return builder.create<bufferization::ToTensorOp>(loc, resultType, inputs)
+          .getResult();
+    else
+      return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
+          .getResult(0);
   });
 
   addTargetMaterialization([&](OpBuilder &builder, Type resultType,
@@ -734,8 +745,13 @@ KrnlTypeConverter::KrnlTypeConverter() {
     if (inputs.size() != 1)
       return llvm::None;
 
-    return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
-        .getResult(0);
+    // Replace UnrealizedConversionCastOp
+    if (enableLinalg)
+      return builder.create<bufferization::ToMemrefOp>(loc, resultType, inputs)
+          .getResult();
+    else
+      return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
+          .getResult(0);
   });
 }
 

src/Conversion/ONNXToLinalg/CMakeLists.txt

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# Please keep in alphabetical order.
+add_onnx_mlir_library(OMONNXToLinalg
+  ConvertONNXToLinalg.cpp
+  ONNXToLinalgCommon.cpp
+  Math/MatMul.cpp
+
+  LINK_LIBS PUBLIC
+  OMAccelerator
+  OMConstPropHelper
+  OMONNXOps
+  OMSupport
+  MLIRFuncDialect
+  MLIRFuncTransforms
+  )

src/Conversion/ONNXToLinalg/ConvertONNXToLinalg.cpp

@@ -0,0 +1,132 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//====------ ConvertONNXToLinalg.cpp - ONNX dialects to Krnl lowering -----===//
+//
+// Copyright 2019-2022 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file implements the lowering of frontend operations to a combination of
+// Krnl IR and standard operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "src/Compiler/CompilerOptions.hpp"
+
+#include "src/Accelerators/Accelerator.hpp"
+#include "src/Conversion/ONNXToLinalg/ONNXToLinalgCommon.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+void populateONNXToLinalgConversionPattern(RewritePatternSet &patterns,
+    TypeConverter &typeConverter, MLIRContext *ctx) {
+
+  // Math
+  populateLoweringONNXMatMulOpLinalgPattern(patterns, typeConverter, ctx);
+}
+
+//===----------------------------------------------------------------------===//
+// ONNX to Krnl Dialect lowering pass
+//===----------------------------------------------------------------------===//
+
+/// This is a partial lowering to Krnl loops of the ONNX operations.
+struct ONNXToLinalgLoweringPass
+    : public PassWrapper<ONNXToLinalgLoweringPass, OperationPass<ModuleOp>> {
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ONNXToLinalgLoweringPass)
+
+  StringRef getArgument() const override { return "convert-onnx-to-linalg"; }
+
+  StringRef getDescription() const override {
+    return "Lower ONNX ops to Linalg dialect.";
+  }
+
+  // Make sure that we have a valid default constructor and copy
+  // constructor to make sure that the options are initialized properly.
+  ONNXToLinalgLoweringPass() = default;
+  ONNXToLinalgLoweringPass(const ONNXToLinalgLoweringPass &pass)
+      : PassWrapper<ONNXToLinalgLoweringPass, OperationPass<ModuleOp>>() {}
+
+  void runOnOperation() final;
+};
+
+void ONNXToLinalgLoweringPass::runOnOperation() {
+  ModuleOp module = getOperation();
+
+  // The first thing to define is the conversion target. This will define the
+  // final target for this lowering.
+  ConversionTarget target(getContext());
+
+  // We define the specific operations, or dialects, that are legal targets for
+  // this lowering.
+  target.addLegalDialect<KrnlDialect, AffineDialect, arith::ArithDialect,
+      func::FuncDialect, linalg::LinalgDialect, math::MathDialect,
+      memref::MemRefDialect, shape::ShapeDialect, scf::SCFDialect,
+      tensor::TensorDialect>();
+  // Needed to support unsigned int computations. To be removed if we use a
+  // scheme that does not rely on the UnrealizedConversionCastOp.
+  target.addLegalOp<::mlir::UnrealizedConversionCastOp>();
+  // Make ONNXNoneOp legal so that other ONNX ops can use it during the
+  // lowering. ONNXNoneOp will be dangling and removed by calling
+  // canonicalization after the lowering.
+  target.addLegalOp<::mlir::ONNXNoneOp>();
+  target.addLegalOp<linalg::MatmulOp>();
+  target.addLegalOp<tensor::EmptyOp>();
+
+  // The following requirements are from Krnl and they are kept if ONNXToKrnl
+  // is after this pass.
+  // If the Linalg is on tensor instead of memref, this lowering will not
+  // generate memref or Affine load/store. However, these requiremnts will may
+  // be an issue if Ops are lowered other than Krnl Use krnl.load/store instead
+  // of std.load/store and affine.load/store. krnl.load/store will be lowered to
+  // std.load/store and affine.load/store by `convert-krnl-to-affine` pass.
+  target.addIllegalOp<mlir::memref::LoadOp>();
+  target.addIllegalOp<mlir::AffineLoadOp>();
+  target.addIllegalOp<mlir::memref::StoreOp>();
+  target.addIllegalOp<mlir::AffineStoreOp>();
+
+  target.addIllegalOp<ONNXMatMulOp>();
+
+  // TODO: add any other ops which are considered legal.
+  // Some operations can be marked as being still legal.
+  // Example: target.addLegalOp<mlir::OpName>();
+
+  // For future: Handle the accelerator target.
+  // for (auto *accel : onnx_mlir::accel::Accelerator::getAccelerators())
+  // accel->conversionTargetONNXToLinalg(target);
+
+  // Now that the conversion target has been defined, we just need to provide
+  // the set of patterns that will lower the frontend operations.
+  RewritePatternSet patterns(&getContext());
+
+  // Convert types to legal types for the Krnl dialect.
+  LinalgTypeConverter linalgTypeConverter;
+
+  // Define patterns.
+  populateONNXToLinalgConversionPattern(
+      patterns, linalgTypeConverter, &getContext());
+
+  // For future: Rewrite patterns for accelerators.
+  // for (auto *accel : onnx_mlir::accel::Accelerator::getAccelerators())
+  //  accel->rewritePatternONNXToLinalg(patterns, krnlTypeConverter,
+  //  &getContext());
+
+  // With the target and rewrite patterns defined, we can now attempt the
+  // conversion. The conversion will signal failure if any of our `illegal`
+  // operations were not converted successfully.
+  if (failed(applyPartialConversion(module, target, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+std::unique_ptr<Pass> createLowerONNXToLinalgPass() {
+  return std::make_unique<ONNXToLinalgLoweringPass>();
+}
+
+} // namespace onnx_mlir

src/Conversion/ONNXToLinalg/Math/MatMul.cpp

@@ -0,0 +1,55 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//===----------------- Matmul.cpp - Lowering Matmul Op --------------------===//
+//
+// Copyright 2019-2022 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file lowers the ONNX Matmul Operator to Linalg dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/Debug.h"
+
+#include "src/Conversion/ONNXToLinalg/ONNXToLinalgCommon.hpp"
+#include "src/Dialect/Mlir/DialectBuilder.hpp"
+
+#define DEBUG_TYPE "matmul"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+struct ONNXMatMulOpLinalgLowering : public ConversionPattern {
+  ONNXMatMulOpLinalgLowering(TypeConverter &typeConverter, MLIRContext *ctx)
+      : ConversionPattern(
+            typeConverter, mlir::ONNXMatMulOp::getOperationName(), 1, ctx) {}
+  LogicalResult matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+      ConversionPatternRewriter &rewriter) const final {
+    Location loc = op->getLoc();
+
+    auto outputType = op->getResult(0).getType().cast<ShapedType>();
+
+    // ToFix: dimension size is assumed to be static
+    SmallVector<Value> newDynamicSizes;
+    auto outV = rewriter.create<tensor::EmptyOp>(loc, outputType.getShape(),
+        outputType.getElementType(), newDynamicSizes);
+
+    SmallVector<Value, 1> outputs;
+    outputs.emplace_back(outV);
+    auto newOp =
+        rewriter.create<linalg::MatmulOp>(loc, outputType, operands, outputs);
+    rewriter.replaceOp(op, newOp.getResults());
+    return success();
+  }
+}; // namespace onnx_mlir
+
+void populateLoweringONNXMatMulOpLinalgPattern(RewritePatternSet &patterns,
+    TypeConverter &typeConverter, MLIRContext *ctx) {
+  patterns.insert<ONNXMatMulOpLinalgLowering>(typeConverter, ctx);
+}
+
+} // namespace onnx_mlir