[TEST] Add ASan/LeakSanitizer harness and fix memory bugs it found

scripts/asan/lsan.supp

@@ -0,0 +1,26 @@
+# LeakSanitizer suppressions for host-side Triton tools (triton-opt, C++ unittests).
+#
+# These are known one-time / process-lifetime allocations in LLVM and MLIR that
+# are never freed by design (managed statics, command-line option registries,
+# one-time target/pass/dialect registration). They are not real leaks; keeping
+# them suppressed lets genuine leaks in the Intel backend stand out.
+#
+# This list is seeded conservatively and refined empirically from real
+# LeakSanitizer output -- add an entry only for allocations confirmed to be
+# benign LLVM/MLIR globals, never for anything under third_party/intel.
+#
+# See https://github.com/intel/intel-xpu-backend-for-triton/issues/5029
+#
+# Entries are matched as substrings against the (demangled) frames of each
+# leak's allocation stack. Format: `leak:<substring>`.
+
+# LLVM command-line option machinery: cl::opt / cl::list objects are registered
+# as globals and intentionally live for the whole process.
+leak:llvm::cl::
+
+# LLVM ManagedStatic-backed singletons (registries, thread pools, etc.).
+leak:llvm::ManagedStaticBase::RegisterManagedStatic
+
+# One-time target / pass registration kept alive for the process lifetime.
+leak:llvm::PassRegistry
+leak:llvm::TargetRegistry

scripts/test-triton.sh

@@ -67,6 +67,8 @@ OPTION:
     --warning-reports
     --ignore-errors
     --run-all
+    --asan            run host-only unit + LIT tests under AddressSanitizer/LeakSanitizer;
+                      requires a build produced with TRITON_BUILD_WITH_ASAN=1
     --skip-list SKIPLIST
     --extra-skip-list-suffixes SEMICOLON-SEPARATED LIST OF SUFFIXES
     --select-from-file SELECTFILE
@@ -139,6 +141,7 @@ TRITON_TEST_RUN_ALL=false
 SKIP_PIP=false
 SKIP_PYTORCH=false
 TEST_UNSKIP=false
+TEST_ASAN=false
 
 while (( $# != 0 )); do
   case "$1" in
@@ -438,6 +441,10 @@ while (( $# != 0 )); do
       TRITON_TEST_RUN_ALL=true
       shift
       ;;
+    --asan)
+      TEST_ASAN=true
+      shift
+      ;;
     --skip-list)
       # Must be absolute
       TRITON_TEST_SKIPLIST_DIR="$(mkdir -p "$2" && cd "$2" && pwd)"
@@ -484,6 +491,53 @@ source "$SCRIPTS_DIR/pytest-utils.sh"
 # Provides the `pip` wrapper (pip or `uv pip`).
 source "$SCRIPTS_DIR/pip-utils.sh"
 
+# AddressSanitizer / LeakSanitizer mode. ASan does not work with binaries that
+# run code on the GPU, so this is restricted to the host-side test surface:
+# the C++ unittests and the LIT tests, both of which exercise triton-opt without
+# touching the device. Requires a build produced with TRITON_BUILD_WITH_ASAN=1.
+# See https://github.com/intel/intel-xpu-backend-for-triton/issues/5029
+if [ "$TEST_ASAN" = true ]; then
+  # Restrict to the host-only unit-test suite (CXX unittests + LIT).
+  TEST_DEFAULT=false
+  TEST_UNIT=true
+  TEST_CORE=false
+  TEST_TUTORIAL=false
+  TEST_MICRO_BENCHMARKS=false
+  TEST_TRITON_KERNELS=false
+  # No Python/GPU tests run under ASan, so skip the pip/pytorch install steps.
+  SKIP_PIP=true
+  SKIP_PYTORCH=true
+
+  # Fail early with a clear message if the build is not ASan-instrumented.
+  # Note: the script runs under `set -o pipefail`, so avoid `... | grep -q`
+  # patterns -- grep closes the pipe on first match and the upstream command
+  # dies with SIGPIPE (141), which pipefail would report as failure. Capture
+  # the output first, then match.
+  ASAN_TRITON_OPT=$(ls -1 "$TRITON_PROJ"/build/cmake*/bin/triton-opt 2>/dev/null || true)
+  ASAN_TRITON_OPT=${ASAN_TRITON_OPT%%$'\n'*}
+  if [ -z "$ASAN_TRITON_OPT" ]; then
+    err "****** ERROR: triton-opt not found. Build Triton first (with TRITON_BUILD_WITH_ASAN=1). ******"
+  fi
+  ASAN_SYMS=$(nm "$ASAN_TRITON_OPT" 2>/dev/null | grep -c '__asan_init' || true)
+  if [ "$ASAN_SYMS" -eq 0 ]; then
+    err "****** ERROR: $ASAN_TRITON_OPT is not ASan-instrumented. Rebuild with TRITON_BUILD_WITH_ASAN=1. ******"
+  fi
+
+  # LeakSanitizer runs at exit; suppress known-benign LLVM/MLIR global leaks so
+  # real leaks stand out. detect_leaks defaults to on for Linux ASan; set it
+  # explicitly for clarity.
+  #
+  # allow_user_poisoning=0 is required: the prebuilt LLVM we link against is not
+  # ASan-instrumented, but its allocators emit __asan_poison_memory_region calls
+  # that resolve against our runtime, poisoning buffers that instrumented Triton
+  # code (e.g. MLIR SmallVector move-assignment in Dialect::addType) then writes
+  # to -- a false use-after-poison that aborts every MLIR tool at static init.
+  # Disabling user poisoning makes those manual poison calls no-ops; it does NOT
+  # weaken leak detection or ASan redzone checks for heap overflow/use-after-free.
+  export ASAN_OPTIONS="detect_leaks=1:allow_user_poisoning=0:${ASAN_OPTIONS:-}"
+  export LSAN_OPTIONS="suppressions=$SCRIPTS_DIR/asan/lsan.supp:print_suppressions=0:${LSAN_OPTIONS:-}"
+fi
+
 if [ "$TRITON_TEST_REPORTS" == true ]; then
     capture_runtime_env
 fi

third_party/intel/lib/Dialect/Triton/Transforms/StrideVersioning.cpp

@@ -199,37 +199,43 @@ struct TritonIntelStrideVersioning
   void runOnOperation() final {
     ModuleOp moduleOp = getOperation();
 
+    // Collect candidate loops during the walk, but do NOT version them here:
+    // LoopVersioner::version() erases the loop, and mutating the IR while the
+    // walker still holds iterators into it is a use-after-free (the walker
+    // dereferences the erased op's regions and the parent block iterator steps
+    // past the freed node). Only non-nested loops are candidates, so they are
+    // top-level siblings and versioning one never invalidates another's handle.
+    SmallVector<scf::ForOp> candidateLoops;
     moduleOp->walk<WalkOrder::PreOrder>([&](Operation *op) {
-      if (auto forOp = dyn_cast<scf::ForOp>(op)) {
-        if (!isCandidateLoop(forOp))
-          return WalkResult::advance();
-
-        // Collect candidate operations. These are descriptor load/store
-        // operations with no stride equal to one (at compile time).
-        OpSelector selector;
-        OpsCollector collector(forOp, selector);
-        if (collector.collectOps() == 0)
-          return WalkResult::advance();
-
-        OpBuilder builder(forOp);
-        SmallVector<Operation *> selectedOps;
-        std::unordered_map<Operation *, Value> selectedOpToStride;
-        for (Operation *op : collector.getOps()) {
-          TypeSwitch<Operation *>(op)
-              .Case<tt::DescriptorLoadOp>([&](auto loadOp) {
-                processDescLoad(loadOp, selectedOps, selectedOpToStride);
-              })
-              .Default([](auto) { return false; });
-        }
+      if (auto forOp = dyn_cast<scf::ForOp>(op))
+        if (isCandidateLoop(forOp))
+          candidateLoops.push_back(forOp);
+      return WalkResult::advance();
+    });
 
-        if (!selectedOpToStride.empty()) {
-          LoopVersioner loopVersioner;
-          loopVersioner.version(forOp, selectedOps, selectedOpToStride);
-        }
+    for (scf::ForOp forOp : candidateLoops) {
+      // Collect candidate operations. These are descriptor load/store
+      // operations with no stride equal to one (at compile time).
+      OpSelector selector;
+      OpsCollector collector(forOp, selector);
+      if (collector.collectOps() == 0)
+        continue;
+
+      SmallVector<Operation *> selectedOps;
+      std::unordered_map<Operation *, Value> selectedOpToStride;
+      for (Operation *op : collector.getOps()) {
+        TypeSwitch<Operation *>(op)
+            .Case<tt::DescriptorLoadOp>([&](auto loadOp) {
+              processDescLoad(loadOp, selectedOps, selectedOpToStride);
+            })
+            .Default([](auto) { return false; });
       }
 
-      return WalkResult::advance();
-    });
+      if (!selectedOpToStride.empty()) {
+        LoopVersioner loopVersioner;
+        loopVersioner.version(forOp, selectedOps, selectedOpToStride);
+      }
+    }
 
     LLVM_DEBUG(llvm::dbgs() << "After versioning:\n" << moduleOp << "\n");
     assert(succeeded(verify(moduleOp)) && "Module verification failed");