[Feature] Add option to disable out-of-bound access warnings in safe memory access legalization

src/op/builtin.cc

@@ -41,6 +41,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION(kDisableDataRaceCheck, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kEnableLowerLDGSTG, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kEnableLowerLDGSTGPredicated, Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION(kDisableLoopUnswitching, Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION(kDisableOutOfBoundWarning, Bool);
 
 DataType cuTensorMapType() { return DataType::UInt(8, 128); }
 

src/op/builtin.h

@@ -114,6 +114,9 @@ static constexpr const char *kDisableThreadStorageSync =
  */
 static constexpr const char *kForceLetInline = "tl.force_let_inline";
 
+static constexpr const char *kDisableOutOfBoundWarning =
+    "tl.disable_out_of_bound_warning";
+
 /*!
  * \brief Get the type of the CUDA tensor map
  *

src/transform/legalize_safe_memory_access.cc

@@ -31,9 +31,16 @@ using arith::IRMutatorWithAnalyzer;
 //    log a warning (local/shared) or handle accordingly (global).
 struct SafeMemChecker : public StmtExprVisitor {
 
+  bool disableOOBWarning = false;
+
   SafeMemChecker(arith::Analyzer *analyzer, bool recursively_collect_conds)
       : analyzer_(analyzer),
-        recursively_collect_conds_(recursively_collect_conds) {}
+        recursively_collect_conds_(recursively_collect_conds) {
+    disableOOBWarning =
+        tvm::transform::PassContext::Current()
+            ->GetConfig(kDisableOutOfBoundWarning, Optional<Bool>())
+            .value_or(false);
+  }
   void VisitExpr_(const BufferLoadNode *op) final {
     // If the buffer is in global scope, we will check its indices and add
     // corresponding bound checks.
@@ -42,7 +49,7 @@ struct SafeMemChecker : public StmtExprVisitor {
     // we are writing TilePrograms. Therefore we only log warnings if there
     // are possible out-of-bounds.
     CheckBufferIndices(op->buffer, op->indices, /*is_load=*/true,
-                       !IsGlobalBuffer(op->buffer));
+                       !disableOOBWarning && !IsGlobalBuffer(op->buffer));
     if (recursively_collect_conds_) {
       StmtExprVisitor::VisitExpr_(op);
     }
@@ -51,7 +58,7 @@ struct SafeMemChecker : public StmtExprVisitor {
   void VisitStmt_(const BufferStoreNode *op) final {
     // Check if the buffer is in global scope
     CheckBufferIndices(op->buffer, op->indices, /*is_load=*/false,
-                       !IsGlobalBuffer(op->buffer));
+                       !disableOOBWarning && !IsGlobalBuffer(op->buffer));
     if (recursively_collect_conds_) {
       StmtExprVisitor::VisitStmt_(op);
     }

tilelang/transform/pass_config.py

@@ -216,3 +216,6 @@ class PassConfigKey(str, Enum):
 
     CUDA_KERNELS_OUTPUT_DIR = "cuda.kernels_output_dir"
     """Output directory for generated CUDA kernels. Default: empty string"""
+
+    TL_DISABLE_OUT_OF_BOUND_WARNING = "tl.disable_out_of_bound_warning"
+    """Disable out-of-bound access warnings in safe memory access legalization. Default: False"""