From 7b2dbaae04b2dccfc84f89f461f23517d78a741d Mon Sep 17 00:00:00 2001
From: kurisu6912 <227995639+kurisu6912@users.noreply.github.com>
Date: Thu, 5 Feb 2026 15:36:00 +0800
Subject: [PATCH] [Feature] Add option to disable out-of-bound access warnings
 in safe memory access legalization

---
 src/op/builtin.cc                            |  1 +
 src/op/builtin.h                             |  3 +++
 src/transform/legalize_safe_memory_access.cc | 13 ++++++++++---
 tilelang/transform/pass_config.py            |  3 +++
 4 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/op/builtin.cc b/src/op/builtin.cc
index 7343d9558..d8122aa77 100644
--- a/src/op/builtin.cc
+++ b/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); }
 
diff --git a/src/op/builtin.h b/src/op/builtin.h
index 429a8db34..5da3b521a 100644
--- a/src/op/builtin.h
+++ b/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
  *
diff --git a/src/transform/legalize_safe_memory_access.cc b/src/transform/legalize_safe_memory_access.cc
index 4e49c7aab..d483aa294 100644
--- a/src/transform/legalize_safe_memory_access.cc
+++ b/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);
     }
diff --git a/tilelang/transform/pass_config.py b/tilelang/transform/pass_config.py
index 5a1fd61de..5a98c06d0 100644
--- a/tilelang/transform/pass_config.py
+++ b/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"""