Add a flag to turn on/off the lowering of scalar broadcasting binary ops to NNPA (#2778) (#2782)

* Add a flag to turn on/off scalar broadcasting binary op in NNPA

Signed-off-by: Tung D. Le <[email protected]>

---------

Signed-off-by: Tung D. Le <[email protected]>
Co-authored-by: Alexandre Eichenberger <[email protected]>
(cherry picked from commit 08d4fed)

Co-authored-by: Tung D. Le <[email protected]>

Author: cjvolzka

src/Accelerators/NNPA/Compiler/CMakeLists.txt

@@ -1,5 +1,3 @@
-get_property(OMLibs GLOBAL PROPERTY ONNX_MLIR_LIBS)
-
 add_onnx_mlir_library(OMNNPACompilerOptions
   NNPACompilerOptions.cpp
 
@@ -12,7 +10,6 @@ add_onnx_mlir_library(OMNNPACompilerOptions
   ${NNPA_ONNX_MLIR_BIN_ROOT}
 
   LINK_LIBS PUBLIC
-  ${OMLibs}
   OMCompilerOptions
 
   ACCEL_INCLUDE_DIRS PRIVATE
@@ -32,7 +29,6 @@ add_onnx_mlir_library(OMNNPACompilerUtils
   ${NNPA_ONNX_MLIR_BIN_ROOT}
 
   LINK_LIBS PUBLIC
-  ${OMLibs}
   OMNNPACompilerOptions
   OMCompilerPasses
 

src/Accelerators/NNPA/Compiler/NNPACompilerOptions.cpp

@@ -55,6 +55,13 @@ llvm::cl::opt<bool> nnpaEnableCompilerStickUnstick(
                    "stick/unstick code. Default is false."),
     llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions));
 
+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::init(false), llvm::cl::cat(OnnxMlirCommonOptions));
+
 llvm::cl::opt<std::string> nnpaLoadDevicePlacementFile{
     "nnpa-load-device-placement-file",
     llvm::cl::desc(

src/Accelerators/NNPA/Compiler/NNPACompilerOptions.hpp

@@ -49,11 +49,13 @@ typedef enum {
 } NNPAPlacementHeuristic;
 
 extern llvm::cl::OptionCategory OnnxMlirOptions;
+extern llvm::cl::OptionCategory OnnxMlirCommonOptions;
 extern llvm::cl::opt<onnx_mlir::NNPAEmissionTargetType> nnpaEmissionTarget;
 extern llvm::cl::opt<bool> nnpaClipToDLFloatRange;
 extern llvm::cl::opt<bool> nnpaEnableZHighToOnnx;
 extern llvm::cl::opt<bool> nnpaEnableZHighDecomposeStickUnstick;
 extern llvm::cl::opt<bool> nnpaEnableCompilerStickUnstick;
+extern llvm::cl::opt<bool> nnpaEnableScalarBcastBinary;
 extern llvm::cl::opt<NNPAPlacementHeuristic> nnpaPlacementHeuristic;
 extern llvm::cl::opt<bool> profileZHighIR;
 extern llvm::cl::opt<std::string> nnpaLoadDevicePlacementFile;

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

@@ -11,7 +11,7 @@ add_onnx_mlir_library(OMONNXToZHigh
   libzdnn
 
   LINK_LIBS PUBLIC
-  OMCompilerOptions
+  OMNNPACompilerOptions
   OMONNXOps
   OMONNXToKrnl
   OMZHighOps
@@ -32,7 +32,7 @@ add_onnx_mlir_library(OMRewriteONNXForZHigh
   libzdnn
 
   LINK_LIBS PUBLIC
-  OMCompilerOptions
+  OMNNPACompilerOptions
   OMONNXOps
   OMONNXToKrnl
   OMZHighOps

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

@@ -324,12 +324,19 @@ bool isSuitableForZDNN<ONNXDivOp>(
   // Check NNPA level.
   if (!isCompatibleWithNNPALevel(NNPA_Z16))
     return false;
-  if (!isF32ScalarConstantTensor(A) && !isValidElementTypeAndRank(A))
+  // Broadcast with a scalar operand.
+  if (isEnableScalarBcastBinary()) {
+    if (isF32ScalarConstantTensor(A) && isValidElementTypeAndRank(B))
+      return true;
+    if (isF32ScalarConstantTensor(B) && isValidElementTypeAndRank(A))
+      return true;
+  }
+  // Non-broadcast cases.
+  if (!isValidElementTypeAndRank(A))
     return false;
-  if (!isF32ScalarConstantTensor(B) && !isValidElementTypeAndRank(B))
+  if (!isValidElementTypeAndRank(B))
     return false;
-  return isF32ScalarConstantTensor(A) || isF32ScalarConstantTensor(B) ||
-         dimAnalysis->sameShape(A, B);
+  return dimAnalysis->sameShape(A, B);
 }
 
 /// Check legality for ONNXSum.

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

@@ -29,6 +29,8 @@ include "src/Accelerators/NNPA/Conversion/ONNXToZHigh/ONNXToZHighCommon.td"
 ///    dag benefitsAdded = (addBenefit 0)
 /// >;
 
+def IsEnableScalarBcastBinary: Constraint<CPred<"isEnableScalarBcastBinary()">>;
+
 def IsNoneType : Constraint<CPred<"(($_self).getType().isa<NoneType>())">>;
 
 def IsNotNoneType : Constraint<CPred<"(!($_self).getType().isa<NoneType>())">>;
@@ -227,7 +229,7 @@ def replaceONNXDivBroadcastPattern1 : Pat<
         (GetScalarF32AttrFromConstant $y),
         (NoneLayoutAttr)),
       (returnType $s_x))),
-  [(IsF32ScalarConstantTensor $y)], [],
+  [(IsEnableScalarBcastBinary), (IsF32ScalarConstantTensor $y)], [],
   (addBenefit 1)
 >;
 
@@ -241,7 +243,7 @@ def replaceONNXDivBroadcastPattern2 : Pat<
         (NoneLayoutAttr)),
       (ZHighStickOp:$s_y $y, (NoneLayoutAttr)),
       (returnType $s_y))),
-  [(IsF32ScalarConstantTensor $x)], [],
+  [(IsEnableScalarBcastBinary), (IsF32ScalarConstantTensor $x)], [],
   (addBenefit 1)
 >;
 

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

@@ -14,11 +14,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/Accelerators/NNPA/Conversion/ONNXToZHigh/ONNXToZHighCommon.hpp"
+#include "src/Accelerators/NNPA/Compiler/NNPACompilerOptions.hpp"
 #include "src/Dialect/ONNX/DialectBuilder.hpp"
 
 using namespace mlir;
 namespace onnx_mlir {
 
+bool isEnableScalarBcastBinary() { return nnpaEnableScalarBcastBinary; }
+
 /// Get transposed tensor by using a permutation array.
 Value emitONNXTranspose(
     Location loc, PatternRewriter &rewriter, Value x, ArrayRef<int64_t> perms) {

src/Accelerators/NNPA/Conversion/ONNXToZHigh/ONNXToZHighCommon.hpp

@@ -27,6 +27,8 @@ const std::string DEVICE_ATTRIBUTE = "device";
 const std::string CPU_DEVICE = "cpu";
 const std::string NNPA_DEVICE = "nnpa";
 
+bool isEnableScalarBcastBinary();
+
 template <typename OP_TYPE>
 void addDynamicallyLegalOpFor(mlir::ConversionTarget *target,
     const onnx_mlir::DimAnalysis *dimAnalysis,

src/Accelerators/NNPA/Pass/NNPAPasses.hpp

@@ -29,11 +29,9 @@ std::unique_ptr<mlir::Pass> createDevicePlacementPass(
 
 /// Add pass for lowering ONNX ops to ZHigh ops.
 std::unique_ptr<mlir::Pass> createONNXToZHighPass();
-std::unique_ptr<mlir::Pass> createONNXToZHighPass();
 
 /// Add pass for rewriting ONNX ops for ZHigh.
 std::unique_ptr<mlir::Pass> createRewriteONNXForZHighPass();
-std::unique_ptr<mlir::Pass> createRewriteONNXForZHighPass();
 
 /// Add pass for re-construct ONNX ops from ZHigh ops.
 std::unique_ptr<mlir::Pass> createZHighToONNXPass();

src/Conversion/KrnlToAffine/KrnlGetLinearOffsetIndex.cpp

@@ -53,7 +53,7 @@ class KrnlGetLinearOffsetIndexLowering : public ConversionPattern {
 
     auto memrefTy = llvm::dyn_cast<MemRefType>(memref.getType());
     int64_t rank = memrefTy.getRank();
-    assert(mapResults.value().size() == rank && "Invalid indices");
+    assert((int64_t)mapResults.value().size() == rank && "Invalid indices");
 
     // Only lower this op after the memref is normalized.
     if (!memrefTy.getLayout().isIdentity())