Added address sanitization feature for translating ML models to LLVM modules

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

@@ -229,6 +229,10 @@ class LLVMAOTTargetBackend final : public TargetBackend {
       return targetOp.emitError() << "failed to translate the MLIR LLVM "
                                      "dialect to the native llvm::Module";
     }
+    if (options_.sanitizerKind == LLVMTargetOptions::SanitizerKind::Address) {
+      for (auto &function : llvmModule->getFunctionList())
+        function.addFnAttr(llvm::Attribute::SanitizeAddress);
+    }
 
     // Try to grab a linker tool based on the options (and target environment).
     auto linkerTool = LinkerTool::getForTarget(targetTriple, options_);
@@ -340,6 +344,8 @@ class LLVMAOTTargetBackend final : public TargetBackend {
         builder, debugDatabaseFilenameRef);
     iree_DyLibExecutableDef_debug_database_embedded_add(builder,
                                                         debugDatabaseRef);
+    iree_DyLibExecutableDef_sanitized_kind_add(
+        builder, (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,28 @@ LogicalResult runLLVMIRPasses(const LLVMTargetOptions &options,
   passBuilder.registerLoopAnalyses(loopAnalysisManager);
   passBuilder.crossRegisterProxies(loopAnalysisManager, functionAnalysisManager,
                                    cGSCCAnalysisManager, moduleAnalysisManager);
+
+  if (options.sanitizerKind == LLVMTargetOptions::SanitizerKind::Address) {
+    bool compileKernel = false;
+    bool recover = false;
+    bool useAfterScope = true;
+    bool moduleUseAfterScope = false;
+    bool useOdrIndicator = false;
+    passBuilder.registerOptimizerLastEPCallback(
+        [compileKernel, recover, useAfterScope, moduleUseAfterScope,
+         useOdrIndicator](llvm::ModulePassManager &modulePassManager,
+                          llvm::PassBuilder::OptimizationLevel Level) {
+          modulePassManager.addPass(
+              llvm::RequireAnalysisPass<llvm::ASanGlobalsMetadataAnalysis,
+                                        llvm::Module>());
+          modulePassManager.addPass(llvm::ModuleAddressSanitizerPass(
+              compileKernel, recover, moduleUseAfterScope, useOdrIndicator));
+          modulePassManager.addPass(
+              createModuleToFunctionPassAdaptor(llvm::AddressSanitizerPass(
+                  compileKernel, recover, useAfterScope)));
+        });
+  }
+
   if (options.optLevel != llvm::PassBuilder::OptimizationLevel::O0) {
     llvm::ModulePassManager modulePassManager;
     modulePassManager =

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

@@ -82,6 +82,18 @@ LLVMTargetOptions getLLVMTargetOptionsFromFlags() {
     llvmTargetOptions.targetCPUFeatures = clTargetCPUFeatures;
   }
 
+  static llvm::cl::opt<std::string> clSanitizerKind(
+      "iree-llvm-sanitize",
+      llvm::cl::desc("Specify LLVM sanitize feature such as 'address' ETC"),
+      llvm::cl::init(""));
+  if (clSanitizerKind == "address")
+    llvmTargetOptions.sanitizerKind = LLVMTargetOptions::SanitizerKind::Address;
+  else if (clSanitizerKind != "") {
+    llvm::errs() << "For now, you can use 'address' for llvm-sanitize, other "
+                    "values are ignored\n";
+    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

@@ -39,6 +39,11 @@ struct LLVMTargetOptions {
   // and benchmarking
   bool debugSymbols = true;
 
+  // Define and declare kinds of Sanitizer as enum
+  // the order in enum should be same as one in flat buffer schema
+  enum SanitizerKind { None = 0, Address };
+  SanitizerKind sanitizerKind = SanitizerKind::None;
+
   // 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,15 @@ static iree_status_t iree_hal_dylib_executable_flatbuffer_verify(
                             "executable library_embedded is missing/empty");
   }
 
+  if (iree_DyLibExecutableDef_sanitized_kind_get(executable_def) ==
+          iree_Sanitizer_Address &&
+      !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");
+  }
+
   return iree_ok_status();
 }
 

iree/schemas/dylib_executable_def.fbs

@@ -13,6 +13,12 @@
 // limitations under the License.
 
 namespace iree;
+// Define kinds of sanitizers (the order in enum should be same as enum in LLVM
+// TargetOptions)
+enum Sanitizer : ubyte {
+  None = 0,
+  Address
+}
 
 // 'Dynamic Library (dylib) Executable'.
 
@@ -21,15 +27,22 @@ 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.
-  entry_points:[string];
-  // An embedded (as opposed to external) dynamic library file.
-  // TODO(scotttodd): List of embedded files?
-  // TODO(scotttodd): Format of files, platform information (x86/arm/etc.)
-  library_embedded:[ubyte];
+// 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?
+// TODO(scotttodd): Format of files, platform information (x86/arm/etc.)
+library_embedded:
+  [ubyte];
 
-  debug_database_filename:string;
-  debug_database_embedded:[ubyte];
+debug_database_filename:
+  string;
+debug_database_embedded:
+  [ubyte];
+sanitized_kind:
+  Sanitizer = None;
 
   // TODO(scotttodd): Relative file path from this flatbuffer file
 }