Enable LLVM AddressSanitizer for triton-shared (#294)

Implemented support for debugging Triton programs with sanitizers:
- New MLIR pass added to triton-shared-opt: --add-debug-info, appending
missing #di attributes for useful error logs
- New LLVM pass ran using opt: --sanitizer-attributes, appending missing
sanitizer attributes to functions in LLVM IR
- Adding -fsanitize flags during compilation + linking. Doing so
required changing llc and g++ to clang++
- Introducing environment variable SANITIZER_TYPE to enable the
sanitizers during JIT compilation
- 3 scripts to facilitate simple build, runtime setup, and running a
Triton program with sanitizers enabled (see usage)

Usage for anyone who wants to debug a Triton program they have written:
1. /triton_shared/scripts/build_triton_shared_for_sanitizers.sh: one
time build script. Sets up a virtual environment, installs a custom LLVM
with compiler-rt, openmp, clang, and mlir. Then builds triton_shared
using this custom LLVM
2. source /triton_shared/scripts/setup_runtime_for_sanitizers.sh: after
installation, prepare the runtime once to activate the virtual
environment and set environment variables. Should be called with source
when a new terminal is used
3. /triton_shared/scripts/run_triton.sh [asan|tsan] python
triton_program.py: executes the Triton program with ASan or TSan
enabled, with the required environment variables

---------

Co-authored-by: Tim Lu <t-timlu@microsoft.com>

Author: timthlu

backend/compiler.py

@@ -10,6 +10,7 @@
 import shutil
 import subprocess
 import functools
+import triton
 from pathlib import Path
 
 def _get_triton_shared_opt_path() -> str:
@@ -33,6 +34,16 @@ def _dump_ir_if_needed(files):
     for f in files:
         shutil.copy(f, os.path.join(path, os.path.basename(f)))
 
+def _get_sanitizer_type():
+    # returns "" if not set
+    # throws error if set to something other than "asan"
+    sanitizer_type = os.getenv("TRITON_SHARED_SANITIZER_TYPE", "")
+
+    if sanitizer_type != "" and sanitizer_type != "asan":
+        # throw error
+        raise Exception(f"TRITON_SHARED_SANITIZER_TYPE {sanitizer_type} is invalid.")
+    
+    return sanitizer_type
 
 def _ttir_to_ttsharedir(mod):
     # Get Triton-MLIR as string
@@ -43,7 +54,16 @@ def _ttir_to_ttsharedir(mod):
         Path(src_path).write_text(ttir_code)
         _dump_ir_if_needed([src_path])
         triton_shared_opt_path = _get_triton_shared_opt_path()
-        subprocess.check_call([triton_shared_opt_path, src_path, "--triton-to-linalg-experimental", "--mlir-print-debuginfo", "-o", dst_path])
+
+        subprocess_args = [triton_shared_opt_path, src_path, "--triton-to-linalg-experimental", "--mlir-print-debuginfo", "-o", dst_path]
+
+        if _get_sanitizer_type() != "":
+            print("Building with sanitizer support...")
+
+            # has to run before the other passes as operates on the tt dialect
+            subprocess_args.insert(2, "--add-llvm-debug-info")
+
+        subprocess.check_call(subprocess_args)
         return Path(dst_path).read_text()
 
 
@@ -118,8 +138,38 @@ def _llir_to_bin(llir: str, metadata):
         src_path = os.path.join(tmpdir, "kernel.ll")
         dst_path = os.path.join(tmpdir, "kernel.o")
         Path(src_path).write_text(llir)
-        llc_path = _get_llvm_bin_path("llc")
-        subprocess.check_call([llc_path, src_path, "-filetype=obj", "-o", dst_path])
+
+        sanitizer_type = _get_sanitizer_type()
+
+        if sanitizer_type != "":
+            # using a sanitizer
+            # invoke pass to append sanitizer attributes
+            instrumented_src_path = os.path.join(tmpdir, "kernel-instrumented.ll")
+        
+            opt_path = _get_llvm_bin_path("opt")
+            top_level_triton_path = os.path.dirname(triton.__file__)
+            sanitizer_attributes_pass_path = str(next(Path(top_level_triton_path).rglob("libSanitizerAttributes.so"), None))
+
+            if not sanitizer_attributes_pass_path:
+                raise Exception(f"libSanitizerAttributes.so does not exist.")
+
+            subprocess.check_call([opt_path, "-load-pass-plugin", sanitizer_attributes_pass_path, 
+                "-passes=sanitizer-attributes", f"-sanitizer-type={sanitizer_type}", "-S", src_path, 
+                "-o", instrumented_src_path])
+
+            # compile to object file
+            clang_path = _get_llvm_bin_path("clang++")
+
+            subprocess_args = [clang_path, "-c", instrumented_src_path, "-o", dst_path]
+
+            if sanitizer_type == "asan":
+                subprocess_args.extend(["-g", "-fsanitize=address", "-mllvm", "-asan-stack=0"])
+                
+            subprocess.check_call(subprocess_args)
+        else:
+            llc_path = _get_llvm_bin_path("llc")
+            subprocess.check_call([llc_path, src_path, "-filetype=obj", "-o", dst_path])
+        
         return Path(dst_path).read_bytes()
 
 

backend/driver.py

@@ -12,6 +12,31 @@
 from triton.backends.driver import DriverBase
 from triton.backends.compiler import GPUTarget
 
+def _get_llvm_bin_path(bin_name: str) -> str:
+    path = os.getenv("LLVM_BINARY_DIR", "")
+    if path == "":
+        raise Exception("LLVM_BINARY_DIR is not set.")
+    return os.path.join(path, bin_name)
+
+def _get_sanitizer_type():
+    # returns "" if not set
+    # throws error if set to something other than "asan"
+    sanitizer_type = os.getenv("TRITON_SHARED_SANITIZER_TYPE", "")
+
+    if sanitizer_type != "" and sanitizer_type != "asan":
+        # throw error
+        raise Exception(f"TRITON_SHARED_SANITIZER_TYPE {sanitizer_type} is invalid.")
+    
+    return sanitizer_type
+
+def _sanitizer_available(sanitizer_type):
+    if "LD_PRELOAD" not in os.environ:
+        return False
+    if f"libclang_rt.{sanitizer_type}.so" not in os.environ["LD_PRELOAD"]:
+        return False
+    
+    return True
+
 # -------------------- Launcher ----------------------------
 def _ty_to_cpp(ty):
     if ty[0] == '*':
@@ -253,7 +278,12 @@ def launch(
 
         if cache_path is None:
           with tempfile.TemporaryDirectory() as tmpdir:
+              sanitizer_type = _get_sanitizer_type()
+
               if platform.system() == "Windows":
+                  if sanitizer_type != "":
+                      raise Exception("Sanitizers are not supported on Windows with triton-shared.")
+
                   obj_path = os.path.join(tmpdir, "kernel.obj")
                   launcher_src_path = os.path.join(tmpdir, "main.cxx")
                   so_path = os.path.join(tmpdir, "kernel.pyd")
@@ -271,12 +301,30 @@ def launch(
                   so_path = os.path.join(tmpdir, "kernel.so")
                   Path(obj_path).write_bytes(kernel_obj)
                   Path(launcher_src_path).write_text(src)
+
                   # Compile it together.
-                  subprocess.check_call([
-                    "g++", "-std=c++17", launcher_src_path, obj_path,
-                    f"-I{py_include_dir}", f"-I{include_dir}", f"-L{py_lib_dir}",
-                    "-shared", f"-l{py_lib}", "-fPIC", "-o", so_path
-                  ])
+                  if sanitizer_type != "":
+                      clang_path = _get_llvm_bin_path("clang++")
+
+                      subprocess_args = [
+                          clang_path, "-std=c++17", launcher_src_path, obj_path,
+                          f"-I{py_include_dir}", f"-I{include_dir}", f"-L{py_lib_dir}",
+                          "-shared", f"-l{py_lib}", "-fPIC", "-o", so_path
+                      ]
+
+                      if not _sanitizer_available(sanitizer_type):
+                          raise Exception(f"Use LD_PRELOAD=\"path/to/libclang_rt.{sanitizer_type}.so\" TRITON_SHARED_SANITIZER_TYPE={sanitizer_type} python ...")
+
+                      if sanitizer_type == "asan":
+                          subprocess_args.extend(["-g", "-fsanitize=address", "-mllvm", "-asan-stack=0"])
+                      
+                      subprocess.check_call(subprocess_args)
+                  else:
+                      subprocess.check_call([
+                        "g++", "-std=c++17", launcher_src_path, obj_path,
+                        f"-I{py_include_dir}", f"-I{include_dir}", f"-L{py_lib_dir}",
+                        "-shared", f"-l{py_lib}", "-fPIC", "-o", so_path
+                      ])
 
               with open(so_path, "rb") as f:
                 cache_path = cache.put(f.read(), filename, binary=True)

include/triton-shared/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
+add_subdirectory(Transform)

include/triton-shared/Transform/AddLLVMDebugInfo/AddLLVMDebugInfo.h

@@ -0,0 +1,27 @@
+//===----------------------------------------------------------------------===//
+//
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TRITON_TRANSFORM_ADDLLVMDEBUGINFO_ADDLLVMDEBUGINFO_H
+#define TRITON_TRANSFORM_ADDLLVMDEBUGINFO_ADDLLVMDEBUGINFO_H
+
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+namespace mlir {
+namespace triton {
+
+std::unique_ptr<OperationPass<ModuleOp>> createAddLLVMDebugInfoPass();
+
+
+} // namespace triton
+} // namespace mlir
+
+#endif // TRITON_TRANSFORM_ADDLLVMDEBUGINFO_ADDLLVMDEBUGINFO_H

include/triton-shared/Transform/AddLLVMDebugInfo/CMakeLists.txt

@@ -0,0 +1,9 @@
+#===------------------------------------------------------------------------===#
+#
+# Copyright (c) Triton Project Contributors.
+#
+#===------------------------------------------------------------------------===#
+
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls --name AddLLVMDebugInfo)
+add_public_tablegen_target(AddLLVMDebugInfoTransformPassIncGen)

include/triton-shared/Transform/AddLLVMDebugInfo/Passes.h

@@ -0,0 +1,22 @@
+//===----------------------------------------------------------------------===//
+//
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ADD_LLVM_DEBUG_INFO_TRANSFORM_PASSES_H
+#define ADD_LLVM_DEBUG_INFO_TRANSFORM_PASSES_H
+
+#include "triton-shared/Transform/AddLLVMDebugInfo/AddLLVMDebugInfo.h"
+
+namespace mlir {
+namespace triton {
+
+#define GEN_PASS_REGISTRATION
+#include "triton-shared/Transform/AddLLVMDebugInfo/Passes.h.inc"
+
+} // namespace triton
+} // namespace mlir
+
+#endif

include/triton-shared/Transform/AddLLVMDebugInfo/Passes.td

@@ -0,0 +1,18 @@
+//===----------------------------------------------------------------------===//
+//
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ADD_LLVM_DEBUG_INFO_TRANSFORM_PASSES
+#define ADD_LLVM_DEBUG_INFO_TRANSFORM_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def AddLLVMDebugInfo : Pass<"add-llvm-debug-info", "mlir::ModuleOp"> {
+  let summary = "Add LLVM debug info";
+  let constructor = "triton::createAddLLVMDebugInfoPass()";
+}
+
+#endif

include/triton-shared/Transform/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(AddLLVMDebugInfo)

lib/CMakeLists.txt

@@ -2,4 +2,6 @@ add_subdirectory(Analysis)
 add_subdirectory(AnalysisStructured)
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
+add_subdirectory(Sanitizer)
+add_subdirectory(Transform)
 add_subdirectory(Utils)

lib/Sanitizer/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(SanitizerAttributes)