Generate messages for legality check for MatMul op.

Signed-off-by: Haruki Imai <[email protected]>

Author: imaihal

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

@@ -20,6 +20,9 @@
 #include "src/Conversion/ONNXToKrnl/RNN/RNNBase.hpp"
 #include "src/Dialect/ONNX/ONNXDimAnalysis.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "legality-check"
 
 using namespace mlir;
 using namespace onnx_mlir;
@@ -266,6 +269,11 @@ bool meetPoolParamRestrictions(int64_t inputShape, int64_t kernelShape,
   return true;
 }
 
+void emitLegalityCheckMessage(Operation *op, StringRef message) {
+  LLVM_DEBUG(llvm::outs() << "[NNPA Legality Check] Warning: " << op->getLoc()
+                          << " runs on CPU. Reason: " << message << "\n");
+}
+
 /// Default legality check.
 template <typename OP_TYPE>
 bool isSuitableForZDNN(OP_TYPE op, const DimAnalysis *dimAnalysis) {
@@ -489,23 +497,37 @@ bool isSuitableForZDNN<ONNXExpOp>(
 template <>
 bool isSuitableForZDNN<ONNXMatMulOp>(
     ONNXMatMulOp op, const DimAnalysis *dimAnalysis) {
+
   // Check NNPA level.
-  if (!isCompatibleWithNNPALevel(NNPA_Z16))
+  if (!isCompatibleWithNNPALevel(NNPA_Z16)) {
+    emitLegalityCheckMessage(
+        op.getOperation(), "Not compatible with NNPA level.");
     return false;
+  }
   int64_t opnum = op.getNumOperands();
   if (opnum != 2) {
+    emitLegalityCheckMessage(
+        op.getOperation(), "The number of operands not 2.");
     return false;
   }
-  if (!isValidElementTypeAndRank(op.getOperand(0)))
+  if (!isValidElementTypeAndRank(op.getOperand(0))) {
+    emitLegalityCheckMessage(
+        op.getOperation(), "Operand 0 not valid element type and rank.");
     return false;
-  if (!isValidElementTypeAndRank(op.getOperand(1)))
+  }
+  if (!isValidElementTypeAndRank(op.getOperand(1))) {
+    emitLegalityCheckMessage(
+        op.getOperation(), "Operand 1 not valid element type and rank.");
     return false;
+  }
   ShapedType aType = op.getOperand(0).getType().cast<ShapedType>();
   ShapedType bType = op.getOperand(1).getType().cast<ShapedType>();
 
   // Illegal if A or B is unranked.
-  if (!aType.hasRank() || !bType.hasRank())
+  if (!aType.hasRank() || !bType.hasRank()) {
+    emitLegalityCheckMessage(op.getOperation(), "A or B is unranked.");
     return false;
+  }
 
   auto shapeA = aType.getShape();
   auto shapeB = bType.getShape();
@@ -518,21 +540,34 @@ bool isSuitableForZDNN<ONNXMatMulOp>(
   // by using broadcasting etc.
   if ((shapeA.size() == 2) && (shapeB.size() == 2)) {
     // unstacked case
-    if (aType.hasStaticShape() && bType.hasStaticShape())
-      return (shapeA[1] == shapeB[0]);
-    else
+    if (aType.hasStaticShape() && bType.hasStaticShape()) {
+      bool returnVal = (shapeA[1] == shapeB[0]);
+      if (!returnVal)
+        emitLegalityCheckMessage(
+            op.getOperation(), "Unstacked case, dim A 1 not equal dim B 0.");
+      return returnVal;
+    } else
       return true;
   } else if ((shapeA.size() == 3) && (shapeB.size() == 3)) {
     // stacked w/o bcast case
-    if (aType.hasStaticShape() && bType.hasStaticShape())
-      return ((shapeA[0] == shapeB[0]) && (shapeA[2] == shapeB[1]));
-    else
+    if (aType.hasStaticShape() && bType.hasStaticShape()) {
+      bool returnVal = ((shapeA[0] == shapeB[0]) && (shapeA[2] == shapeB[1]));
+      if (!returnVal)
+        emitLegalityCheckMessage(
+            op.getOperation(), "Stacked w/o bcast, dim A 0 not equal dim B 0 "
+                               "or dim A 2 not equal dim B 1.");
+      return returnVal;
+    } else
       return true;
   } else if ((shapeA.size() == 3) && (shapeB.size() == 2)) {
     // stacked w/ bcast
-    if (aType.hasStaticShape() && bType.hasStaticShape())
-      return (shapeA[2] == shapeB[0]);
-    else
+    if (aType.hasStaticShape() && bType.hasStaticShape()) {
+      bool returnVal = (shapeA[2] == shapeB[0]);
+      if (!returnVal)
+        emitLegalityCheckMessage(
+            op.getOperation(), "Stacked w/ bcast, dim A 2 not equal dim B 0.");
+      return returnVal;
+    } else
       return true;
   }
   return false; // unsupported case

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

@@ -45,3 +45,5 @@ mlir::StringRef getStrPaddingType(OP op);
 /// See "MaxPool2D/AvgPool2D Parameter Restrictions" in "zDNN API Reference"
 bool meetPoolParamRestrictions(int64_t inputShape, int64_t kernelShape,
     int64_t strides, int64_t outputShape, mlir::StringRef paddingType);
+
+void emitLegalityCheckMessage(mlir::Operation *op, mlir::StringRef message);

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

@@ -589,14 +589,24 @@ void getRewriteONNXForZHighDynamicallyLegal(
         // No input is N-D (N > 3) but dimension N or M (NxK * KxM) is dynamic
         // or exceeds NNPA limitation.
         bool nExceeded, mExceeded;
-        if (SplitLargeMatMulPattern::canBeRewritten(op, nExceeded, mExceeded))
+        if (SplitLargeMatMulPattern::canBeRewritten(op, nExceeded, mExceeded)) {
+          emitLegalityCheckMessage(op.getOperation(),
+              "No input is N-D (N > 3) but dimension N or M (NxK * KxM) is "
+              "dynamic or exceeds NNPA limitation.");
           return false;
+        }
 
         // - one input is N-D (N > 3) and the other is 2-D.
-        if (aRank == 2 && bRank > 3)
+        if (aRank == 2 && bRank > 3) {
+          emitLegalityCheckMessage(op.getOperation(),
+              "one input is N-D (N > 3) and the other is 2-D.");
           return false;
-        if (bRank == 2 && aRank > 3)
+        }
+        if (bRank == 2 && aRank > 3) {
+          emitLegalityCheckMessage(op.getOperation(),
+              "one input is N-D (N > 3) and the other is 2-D.");
           return false;
+        }
 
         // - both inputs are *the same* N-D, N > 3 and there is no broadcasting
         if (aRank > 3 && (aRank == bRank)) {
@@ -607,6 +617,10 @@ void getRewriteONNXForZHighDynamicallyLegal(
               sameBatchDims =
                   dimAnalysis->sameDynDim(op.getA(), i, op.getB(), i);
           }
+          if (!sameBatchDims)
+            emitLegalityCheckMessage(
+                op.getOperation(), "both inputs are *the same* N-D, N > 3 and "
+                                   "there is no broadcasting.");
           return !sameBatchDims;
         }