Merge branch 'main' into pr_extend_dynamic_backend_test

Author: negiyas

docs/BuildOnLinuxOSX.md

@@ -15,7 +15,7 @@ Firstly, install MLIR (as a part of LLVM-Project):
 ``` bash
 git clone -n https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX-MLIR.
-cd llvm-project && git checkout a4ca07f13b560b4f6fa5459eef7159e4f9ee9a6b && cd ..
+cd llvm-project && git checkout 7ac7d418ac2b16fd44789dcf48e2b5d73de3e715 && cd ..
 ```
 
 [same-as-file]: <> (utils/build-mlir.sh)

docs/BuildOnWindows.md

@@ -52,7 +52,7 @@ Install MLIR (as a part of LLVM-Project):
 ```shell
 git clone -n https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX-MLIR.
-cd llvm-project && git checkout a4ca07f13b560b4f6fa5459eef7159e4f9ee9a6b && cd ..
+cd llvm-project && git checkout 7ac7d418ac2b16fd44789dcf48e2b5d73de3e715 && cd ..
 ```
 
 [same-as-file]: <> (utils/build-mlir.cmd)

src/Accelerators/NNPA/Compiler/NNPACompilerOptions.cpp

@@ -57,9 +57,9 @@ llvm::cl::opt<bool> nnpaEnableCompilerStickUnstick(
 
 llvm::cl::opt<bool> nnpaEnableScalarBcastBinary(
     "nnpa-enable-scalar-bcast-binary",
-    llvm::cl::desc("Enable the lowering to NNPA the broadcasting binary ops "
-                   "whose one of the operands is scalar. Currently support "
-                   "ONNXDiv only. Default is false."),
+    llvm::cl::desc("Enable the lowering to NNPA of binary operations with "
+                   "broadcasting of a scalar operand."
+                   "Currently only enable ONNXDiv. Default is false."),
     llvm::cl::init(false), llvm::cl::cat(OnnxMlirCommonOptions));
 
 llvm::cl::opt<std::string> nnpaLoadDevicePlacementFile{

src/Accelerators/NNPA/Conversion/ONNXToZHigh/CMakeLists.txt

@@ -14,6 +14,7 @@ add_onnx_mlir_library(OMONNXToZHigh
   OMNNPACompilerOptions
   OMONNXOps
   OMONNXToKrnl
+  OMShapeInferencePass
   OMZHighOps
 
   ACCEL_INCLUDE_DIRS PRIVATE

src/Accelerators/NNPA/Conversion/ONNXToZHigh/DevicePlacement.cpp

@@ -196,7 +196,7 @@ void DevicePlacementPass::runOnOperation() {
   // Call ONNXToZHigh pass for lowering multiple ONNX ops at once to ZHigh.
   // E.g. `onnx.ReLu (onnx.Conv)` to zhigh.Conv.
   RewritePatternSet Patterns2(context);
-  getONNXToZHighOneOpPatterns(Patterns2);
+  getONNXToZHighMultipleOpPatterns(Patterns2);
   (void)applyAnalysisConversion(module, target, std::move(Patterns2),
       ConversionConfig{.legalizableOps = &legalizedOps2});
 

src/Accelerators/NNPA/Conversion/ONNXToZHigh/ONNXToZHigh.cpp

@@ -22,6 +22,7 @@
 #include "src/Dialect/ONNX/ONNXDimAnalysis.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
+#include "src/Dialect/ONNX/Transforms/ShapeInference.hpp"
 
 using namespace mlir;
 
@@ -328,16 +329,13 @@ void getONNXToZHighMultipleOpPatterns(RewritePatternSet &patterns) {
   patterns.insert<replaceONNXMatMulAddPattern2>(context);
   patterns.insert<replaceONNXReluConvPattern>(context);
   patterns.insert<replaceONNXLogSoftmaxPattern>(context);
+  // Shape inference for newly-added operations.
+  getShapeInferencePatterns(patterns);
 }
 
 void ONNXToZHighLoweringPass::runOnOperation() {
   ModuleOp module = getOperation();
 
-  // Run the unknown dimension analysis to help check equality of unknown
-  // dimensions at compile time.
-  onnx_mlir::DimAnalysis dimAnalysis(module);
-  dimAnalysis.analyze();
-
   // The first thing to define is the conversion target. This will define the
   // final target for this lowering.
   ConversionTarget target(getContext());
@@ -363,6 +361,11 @@ void ONNXToZHighLoweringPass::runOnOperation() {
   // It's ok to fail.
   (void)applyPatternsAndFoldGreedily(module, std::move(combinedPatterns));
 
+  // Run the unknown dimension analysis to help check equality of unknown
+  // dimensions at compile time.
+  onnx_mlir::DimAnalysis dimAnalysis(module);
+  dimAnalysis.analyze();
+
   // Single ONNX to ZHigh operation lowering.
   RewritePatternSet patterns(&getContext());
   onnx_mlir::getONNXToZHighOneOpPatterns(patterns);

src/Accelerators/NNPA/Conversion/ONNXToZHigh/ZHighToONNX.cpp

@@ -59,6 +59,8 @@ void ZHighToONNXLoweringPass::runOnOperation() {
 
   RewritePatternSet patterns(&getContext());
   populateWithGenerated(patterns);
+  zhigh::ZHighStickOp::getCanonicalizationPatterns(patterns, &getContext());
+  zhigh::ZHighUnstickOp::getCanonicalizationPatterns(patterns, &getContext());
 
   (void)applyPatternsAndFoldGreedily(function, std::move(patterns));
 }

src/Accelerators/NNPA/Conversion/ONNXToZHigh/ZHighToONNX.td

@@ -37,62 +37,107 @@ def CreateONNXMaxOp : NativeCodeCall<"$_builder.create<ONNXMaxOp>($_loc, $0.getT
 // ONNXAddOp %X = ZHighUnstickOp (ZHighAddOp (ZHighStickOp %X),
 // (ZHighStickOp %Y))
 //===----------------------------------------------------------------------===//
-def replaceZHighAddPattern : Pat<
-  (ZHighUnstickOp (ZHighAddOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
-  (ONNXAddOp $x, $y),
-  [(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
+def replaceZHighAddPattern1 : Pat<
+  (ZHighUnstickOp (ZHighAddOp (ZHighStickOp:$s_x $x, $_), $y)),
+  (ONNXAddOp $x, (ZHighUnstickOp $y)),
+  [(NotBlockArgument:$x), (HasOneUse:$s_x)]
 >;
 
+def replaceZHighAddPattern2 : Pat<
+  (ZHighUnstickOp (ZHighAddOp $x, (ZHighStickOp:$s_y $y, $_))),
+  (ONNXAddOp (ZHighUnstickOp $x), $y),
+  [(NotBlockArgument:$y), (HasOneUse:$s_y)]
+>;
 
 //===----------------------------------------------------------------------===//
 // ONNXMulOp %X = ZHighUnstickOp (ZHighMulOp (ZHighStickOp %X),
 // (ZHighStickOp %Y))
 //===----------------------------------------------------------------------===//
-def replaceZHighMulPattern : Pat<
-  (ZHighUnstickOp (ZHighMulOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
-  (ONNXMulOp $x, $y),
-  [(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
+def replaceZHighMulPattern1 : Pat<
+  (ZHighUnstickOp (ZHighMulOp (ZHighStickOp:$s_x $x, $_), $y)),
+  (ONNXMulOp $x, (ZHighUnstickOp $y)),
+  [(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
+  (addBenefit 1)
+>;
+
+def replaceZHighMulPattern2 : Pat<
+  (ZHighUnstickOp (ZHighMulOp $x, (ZHighStickOp:$s_y $y, $_))),
+  (ONNXMulOp (ZHighUnstickOp $x), $y),
+  [(NotBlockArgument:$y), (HasOneUse:$s_y)], [],
+  (addBenefit 0)
 >;
 
 //===----------------------------------------------------------------------===//
 // ONNXSubOp %X = ZHighUnstickOp (ZHighSubOp (ZHighStickOp %X),
 // (ZHighStickOp %Y))
 //===----------------------------------------------------------------------===//
-def replaceZHighSubPattern : Pat<
-  (ZHighUnstickOp (ZHighSubOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
-  (ONNXSubOp $x, $y),
-  [(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
+def replaceZHighSubPattern1 : Pat<
+  (ZHighUnstickOp (ZHighSubOp (ZHighStickOp:$s_x $x, $_), $y)),
+  (ONNXSubOp $x, (ZHighUnstickOp $y)),
+  [(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
+  (addBenefit 1)
+>;
+
+def replaceZHighSubPattern2 : Pat<
+  (ZHighUnstickOp (ZHighSubOp $x, (ZHighStickOp:$s_y $y, $_))),
+  (ONNXSubOp (ZHighUnstickOp $x), $y),
+  [(NotBlockArgument:$y), (HasOneUse:$s_y)], [ ],
+  (addBenefit 0)
 >;
 
 //===----------------------------------------------------------------------===//
 // ONNXDivOp %X = ZHighUnstickOp (ZHighDivOp (ZHighStickOp
 // %X),(ZHighStickOp %Y))
 // Note: turn off this pattern since NNPA is faster at this moment.
 //===----------------------------------------------------------------------===//
-// def replaceZHighDivPattern : Pat<
-//   (ZHighUnstickOp (ZHighDivOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
-//   (ONNXDivOp $x, $y),
-//   [(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
-// >;
+//def replaceZHighDivPattern1 : Pat<
+//  (ZHighUnstickOp (ZHighDivOp (ZHighStickOp:$s_x $x, $_), $y)),
+//  (ONNXDivOp $x, (ZHighUnstickOp $y)),
+//  [(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
+//  (addBenefit 1)
+//>;
+//
+//def replaceZHighDivPattern2 : Pat<
+//  (ZHighUnstickOp (ZHighDivOp $x, (ZHighStickOp:$s_y $y, $_))),
+//  (ONNXDivOp (ZHighUnstickOp $x), $y),
+//  [(NotBlockArgument:$y), (HasOneUse:$s_y)], [ ],
+//  (addBenefit 0)
+//>;
 
 //===----------------------------------------------------------------------===//
 // ONNXMinOp %X = ZHighUnstickOp (ZHighMinOp (ZHighStickOp %X),
 // (ZHighStickOp %Y))
 //===----------------------------------------------------------------------===//
-def replaceZHighMinPattern : Pat<
-  (ZHighUnstickOp:$u (ZHighMinOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
-  (CreateONNXMinOp $u, $x, $y),
-  [(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
+def replaceZHighMinPattern1 : Pat<
+  (ZHighUnstickOp:$u (ZHighMinOp (ZHighStickOp:$s_x $x, $_), $y)),
+  (CreateONNXMinOp $u, $x, (ZHighUnstickOp $y)),
+  [(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
+  (addBenefit 1)
+>;
+
+def replaceZHighMinPattern2 : Pat<
+  (ZHighUnstickOp:$u (ZHighMinOp $x, (ZHighStickOp:$s_y $y, $_))),
+  (CreateONNXMinOp $u, (ZHighUnstickOp $x), $y),
+  [(NotBlockArgument:$y), (HasOneUse:$s_y)], [ ],
+  (addBenefit 0)
 >;
 
 //===----------------------------------------------------------------------===//
 // ONNXMaxOp %X = ZHighUnstickOp (ZHighMaxOp (ZHighStickOp %X),
 // (ZHighStickOp %Y))
 //===----------------------------------------------------------------------===//
-def replaceZHighMaxPattern : Pat<
-  (ZHighUnstickOp:$u (ZHighMaxOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
-  (CreateONNXMaxOp $u, $x, $y),
-  [(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
+def replaceZHighMaxPattern1 : Pat<
+  (ZHighUnstickOp:$u (ZHighMaxOp (ZHighStickOp:$s_x $x, $_), $y)),
+  (CreateONNXMaxOp $u, $x, (ZHighUnstickOp $y)),
+  [(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
+  (addBenefit 1)
+>;
+
+def replaceZHighMaxPattern2 : Pat<
+  (ZHighUnstickOp:$u (ZHighMaxOp $x, (ZHighStickOp:$s_y $y, $_))),
+  (CreateONNXMaxOp $u, (ZHighUnstickOp $x), $y),
+  [(NotBlockArgument:$y), (HasOneUse:$s_y)], [ ],
+  (addBenefit 0)
 >;
 
 //===----------------------------------------------------------------------===//