Added address sanitization feature for translating ML models to LLVM modules

iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp

@@ -230,6 +230,15 @@ class LLVMAOTTargetBackend final : public TargetBackend {
                                      "dialect to the native llvm::Module";
     }
 
+    switch (options_.sanitizerKind) {
+      case SanitizerKind::kNone:
+        break;
+      case SanitizerKind::kAddress: {
+        for (auto &function : llvmModule->getFunctionList())
+          function.addFnAttr(llvm::Attribute::SanitizeAddress);
+      } break;
+    }
+
     // Try to grab a linker tool based on the options (and target environment).
     auto linkerTool = LinkerTool::getForTarget(targetTriple, options_);
     if (!linkerTool) {
@@ -340,6 +349,8 @@ class LLVMAOTTargetBackend final : public TargetBackend {
         builder, debugDatabaseFilenameRef);
     iree_DyLibExecutableDef_debug_database_embedded_add(builder,
                                                         debugDatabaseRef);
+    iree_DyLibExecutableDef_sanitized_kind_add(
+        builder, static_cast<unsigned char>(options_.sanitizerKind));
     iree_DyLibExecutableDef_end_as_root(builder);
 
     // Add the binary data to the target executable.

iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.cpp

@@ -27,6 +27,7 @@
 #include "llvm/Support/Host.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
 
 namespace mlir {
 namespace iree_compiler {
@@ -87,6 +88,31 @@ LogicalResult runLLVMIRPasses(const LLVMTargetOptions &options,
   passBuilder.registerLoopAnalyses(loopAnalysisManager);
   passBuilder.crossRegisterProxies(loopAnalysisManager, functionAnalysisManager,
                                    cGSCCAnalysisManager, moduleAnalysisManager);
+
+  switch (options.sanitizerKind) {
+    case SanitizerKind::kNone:
+      break;
+    case SanitizerKind::kAddress: {
+      passBuilder.registerOptimizerLastEPCallback(
+          [](llvm::ModulePassManager &modulePassManager,
+             llvm::PassBuilder::OptimizationLevel Level) {
+            bool compileKernel = false;
+            bool recover = false;
+            bool useAfterScope = true;
+            bool moduleUseAfterScope = false;
+            bool useOdrIndicator = false;
+            modulePassManager.addPass(
+                llvm::RequireAnalysisPass<llvm::ASanGlobalsMetadataAnalysis,
+                                          llvm::Module>());
+            modulePassManager.addPass(llvm::ModuleAddressSanitizerPass(
+                compileKernel, recover, moduleUseAfterScope, useOdrIndicator));
+            modulePassManager.addPass(
+                createModuleToFunctionPassAdaptor(llvm::AddressSanitizerPass(
+                    compileKernel, recover, useAfterScope)));
+          });
+    } break;
+  }
+
   if (options.optLevel != llvm::PassBuilder::OptimizationLevel::O0) {
     llvm::ModulePassManager modulePassManager;
     modulePassManager =

iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.cpp

@@ -82,6 +82,13 @@ LLVMTargetOptions getLLVMTargetOptionsFromFlags() {
     llvmTargetOptions.targetCPUFeatures = clTargetCPUFeatures;
   }
 
+  static llvm::cl::opt<SanitizerKind> clSanitizerKind(
+      "iree-llvm-sanitize", llvm::cl::desc("Apply LLVM sanitize feature"),
+      llvm::cl::init(SanitizerKind::kNone),
+      llvm::cl::values(clEnumValN(SanitizerKind::kAddress, "address",
+                                  "Address sanitizer support")));
+  llvmTargetOptions.sanitizerKind = clSanitizerKind;
+
   static llvm::cl::opt<std::string> clTargetABI(
       "iree-llvm-target-abi",
       llvm::cl::desc("LLVM target machine ABI; specify for -mabi"),

iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.h

@@ -23,6 +23,13 @@ namespace iree_compiler {
 namespace IREE {
 namespace HAL {
 
+// Defines kinds of Sanitizer
+// The order in enum class should be same as one in flat buffer schema
+enum class SanitizerKind {
+  kNone = 0,
+  kAddress,
+};
+
 struct LLVMTargetOptions {
   // Target machine configuration.
   std::string targetTriple;
@@ -39,6 +46,9 @@ struct LLVMTargetOptions {
   // and benchmarking
   bool debugSymbols = true;
 
+  // Sanitizer Kind for CPU Kernels
+  SanitizerKind sanitizerKind = SanitizerKind::kNone;
+
   // Link any required runtime libraries into the produced binaries statically.
   // This increases resulting binary size but enables the binaries to be used on
   // any machine without requiring matching system libraries to be installed.

iree/hal/local/loaders/legacy_library_loader.cc

@@ -17,6 +17,7 @@
 #include "iree/base/dynamic_library.h"
 #include "iree/base/internal/file_io.h"
 #include "iree/base/internal/file_path.h"
+#include "iree/base/target_platform.h"
 #include "iree/base/tracing.h"
 #include "iree/hal/local/local_executable.h"
 
@@ -81,6 +82,20 @@ static iree_status_t iree_hal_dylib_executable_flatbuffer_verify(
                             "executable library_embedded is missing/empty");
   }
 
+  switch (iree_DyLibExecutableDef_sanitized_kind_get(executable_def)) {
+    case iree_SanitizerKind_None:
+      break;
+    case iree_SanitizerKind_Address: {
+      if (!IREE_SANITIZER_ADDRESS) {
+        return iree_make_status(
+            IREE_STATUS_UNAVAILABLE,
+            "Dynamic library executable is compiled with ASAN support, but "
+            "this host application failed to enable ASAN to load this "
+            "executable");
+      }
+    } break;
+  }
+
   return iree_ok_status();
 }
 

iree/schemas/dylib_executable_def.fbs

@@ -14,14 +14,22 @@
 
 namespace iree;
 
+// Define kinds of sanitizers (the order in enum should be same as enum in LLVM
+// TargetOptions)
+enum SanitizerKind : ubyte {
+  None = 0,
+  Address
+}
+
 // 'Dynamic Library (dylib) Executable'.
 
 file_identifier "DLIB";
 file_extension "dlib";
 
 // Dynamic library (.so/.dll/.dylib) executable module.
 table DyLibExecutableDef {
-  // A map of entry points to string names with the same order as in the executable op.
+  // A map of entry points to string names with the same order as in the
+  // executable op.
   entry_points:[string];
   // An embedded (as opposed to external) dynamic library file.
   // TODO(scotttodd): List of embedded files?
@@ -30,6 +38,7 @@ table DyLibExecutableDef {
 
   debug_database_filename:string;
   debug_database_embedded:[ubyte];
+  sanitized_kind:SanitizerKind = None;
 
   // TODO(scotttodd): Relative file path from this flatbuffer file
 }