Add runtime check for Gather Op (#3069)

chentong319 · web-flow · commit 584ee43704d4 · 2025-02-07T21:41:28.000-05:00
* implemented

Signed-off-by: Chen Tong &lt;chentong@us.ibm.com&gt;

* add node name

Signed-off-by: Chen Tong &lt;chentong@us.ibm.com&gt;

* format

Signed-off-by: Chen Tong &lt;chentong@us.ibm.com&gt;

---------

Signed-off-by: Chen Tong &lt;chentong@us.ibm.com&gt;
diff --git a/src/Compiler/CompilerOptions.cpp b/src/Compiler/CompilerOptions.cpp
@@ -45,6 +45,7 @@ std::string opsForCall;                                // common for both
 bool disableKrnlOpFusion;                              // common for both
 bool disableQuantZeroPoint;                            // common for both
 bool enableKrnlBufferReuse;                            // common for both
+bool enableSafeCodeGen;                                // common for both
 bool disableMemRefPrefetch;                            // common for both
 uint64_t compilationNumThreads;                        // common for both
 EmissionTargetType emissionTarget;                     // onnx-mlir only
@@ -245,6 +246,16 @@ static llvm::cl::opt<bool, true> enableKrnlBufferReuseOpt(
     llvm::cl::location(enableKrnlBufferReuse), llvm::cl::init(false),
     llvm::cl::cat(OnnxMlirCommonOptions));
 
+static llvm::cl::opt<bool, true> enableSafeCodeGenOpt("enable-safe-code-gen",
+    llvm::cl::desc("enable extra runtime check to be created in code gen. "
+                   "Such check will have cost at runtime, and is not needed if"
+                   "the model and the data are correct."
+                   "Failure of check will trigger assertion error."
+                   "(default=false).\n"
+                   "Set to 'true' if you want to enable the check."),
+    llvm::cl::location(enableSafeCodeGen), 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"
diff --git a/src/Compiler/CompilerOptions.hpp b/src/Compiler/CompilerOptions.hpp
@@ -91,6 +91,7 @@ extern std::string opsForCall;                                // common for both
 extern bool disableKrnlOpFusion;                              // common for both
 extern bool disableQuantZeroPoint;                            // common for both
 extern bool enableKrnlBufferReuse;                            // common for both
+extern bool enableSafeCodeGen;                                // common for both
 extern bool disableMemRefPrefetch;                            // common for both
 extern uint64_t compilationNumThreads;                        // common for both
 extern EmissionTargetType emissionTarget;                     // onnx-mlir only
diff --git a/src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp b/src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp
@@ -12,6 +12,7 @@
 // Krnl IR and standard operations.
 //
 //===----------------------------------------------------------------------===//
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
@@ -367,7 +368,7 @@ void FrontendToKrnlLoweringPass::runOnOperation() {
   target.addLegalDialect<KrnlDialect, affine::AffineDialect,
       arith::ArithDialect, func::FuncDialect, linalg::LinalgDialect,
       math::MathDialect, vector::VectorDialect, memref::MemRefDialect,
-      shape::ShapeDialect, scf::SCFDialect>();
+      shape::ShapeDialect, scf::SCFDialect, cf::ControlFlowDialect>();
   // 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>();
diff --git a/src/Conversion/ONNXToKrnl/Tensor/Gather.cpp b/src/Conversion/ONNXToKrnl/Tensor/Gather.cpp
@@ -12,6 +12,10 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+
+#include "src/Compiler/CompilerOptions.hpp"
 #include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
 
@@ -37,7 +41,8 @@ struct ONNXGatherOpLowering : public OpConversionPattern<ONNXGatherOp> {
     Location loc = ONNXLoc<ONNXGatherOp>(op);
     ValueRange operands = adaptor.getOperands();
 
-    MultiDialectBuilder<KrnlBuilder, IndexExprBuilderForKrnl, MemRefBuilder>
+    MultiDialectBuilder<KrnlBuilder, IndexExprBuilderForKrnl, MemRefBuilder,
+        MathBuilder>
         create(rewriter, loc);
 
     // Get shape.
@@ -122,6 +127,41 @@ struct ONNXGatherOpLowering : public OpConversionPattern<ONNXGatherOp> {
           if (indicesMayBeNegative)
             index = index.selectOrSelf(index < zeroIE, index + axisDim);
 
+          // The Gather op is data dependent: the value of index should be
+          // within the input data size.
+          // Add runtime check if enableSafeCodeGen is set true
+          // Implementation comments vs. createGenerateRuntimeVerificationPass
+          // This check is according to onnx op semantics, not general bound
+          // check for memref. Implementation of RuntimeVerification could be
+          // borrowed. Slightly difference is that onnx semenatics check is for
+          // each dimension independently, not the final address is within
+          // the memref bound.
+          if (enableSafeCodeGen) {
+            // From onnx document:
+            // All index values are expected to be within bounds [-s, s-1]
+            // along axis of size s. It is an error if any of the index values
+            // are out of bounds.
+            // After the negative correction, the range should be [0, s-1]
+            Value upperBound = create.mem.dim(data, axisLit);
+            Value compareUpperBound =
+                create.math.slt(index.getValue(), upperBound);
+            // Report onnx_node_name if the op has the attribute
+            std::string nodeNameStr = op->getName().getStringRef().str() + " ";
+            StringAttr nodeName =
+                op->getAttrOfType<mlir::StringAttr>("onnx_node_name");
+            if (nodeName && !nodeName.getValue().empty()) {
+              nodeNameStr = nodeNameStr + nodeName.getValue().str();
+            }
+            rewriter.create<cf::AssertOp>(loc, compareUpperBound,
+                nodeNameStr +
+                    " indices of GatherOp is larger than the upper bound");
+            Value compareLowerBound =
+                create.math.sge(index.getValue(), zeroIE.getValue());
+            rewriter.create<cf::AssertOp>(loc, compareLowerBound,
+                nodeNameStr +
+                    " indices of GatherOp is less than the lower bound");
+          }
+
           // Compute access function of data: data[ii + (indices[jj],) + kk]
           SmallVector<IndexExpr, 4> dataAccessFct;
           // First add indices iis
diff --git a/src/Conversion/ONNXToKrnl/Tensor/GatherElements.cpp b/src/Conversion/ONNXToKrnl/Tensor/GatherElements.cpp
@@ -12,6 +12,10 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+
+#include "src/Compiler/CompilerOptions.hpp"
 #include "src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
 
@@ -31,7 +35,8 @@ struct ONNXGatherElementsOpLowering
     Location loc = ONNXLoc<ONNXGatherElementsOp>(op);
     ValueRange operands = adaptor.getOperands();
 
-    MultiDialectBuilder<KrnlBuilder, IndexExprBuilderForKrnl, MemRefBuilder>
+    MultiDialectBuilder<KrnlBuilder, IndexExprBuilderForKrnl, MemRefBuilder,
+        MathBuilder>
         create(rewriter, loc);
 
     // Get shape.
@@ -93,6 +98,32 @@ struct ONNXGatherElementsOpLowering
             index = index.selectOrSelf(index < zero, index + axisDim);
           }
 
+          // Check the dynamic requirement of GatherElement Op
+          // Refer to the comments in Gather.cpp
+          if (enableSafeCodeGen) {
+            // From onnx document:
+            // All index values are expected to be within bounds [-s, s-1]
+            // along axis of size s. It is an error if any of the index values
+            // are out of bounds.
+            // After the negative correction, the range should be [0, s-1]
+            Value upperBound = create.mem.dim(data, axis);
+            Value compareUpperBound =
+                create.math.slt(index.getValue(), upperBound);
+            std::string nodeNameStr = op->getName().getStringRef().str() + " ";
+            StringAttr nodeName =
+                op->getAttrOfType<mlir::StringAttr>("onnx_node_name");
+            if (nodeName && !nodeName.getValue().empty()) {
+              nodeNameStr = nodeNameStr + nodeName.getValue().str();
+            }
+            rewriter.create<cf::AssertOp>(loc, compareUpperBound,
+                "indices of GatherOp is larger than the upper bound");
+            LiteralIndexExpr zero(0);
+            Value compareLowerBound =
+                create.math.sge(index.getValue(), zero.getValue());
+            rewriter.create<cf::AssertOp>(loc, compareLowerBound,
+                "indices of GatherOp is less than the lower bound");
+          }
+
           // Access function for the 'data' tensor.
           DimsExpr dataAccessFct;
           for (int64_t i = 0; i < dataRank; ++i)
