Added address sanitization feature for translating ML models to LLVM modules (#4676)
1. Added an option, -iree-llvm-address-sanitizer. Users can turn it on and it injects ASAN to binary of the module.
2. If ASAN is used on compilation, it changes an ubyte to a enum value in flat buffer to mark it is compiled with ASAN enabled. IREE runtime can read the ubyte before running the module. And, it will make an error message if runtime does not have ASAN feature when running address sanitized module.
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp
index 7bb954a..d0dbde6 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp
@@ -230,6 +230,15 @@
"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 @@
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.
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.cpp
index 467ebfd..87c2f79 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.cpp
+++ b/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 @@
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 =
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.cpp
index 33f1416..b6c18e0 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.cpp
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.cpp
@@ -82,6 +82,13 @@
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"),
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.h b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.h
index 1ac53bc..3e85ed3 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.h
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.h
@@ -23,6 +23,13 @@
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 @@
// 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.
diff --git a/iree/hal/local/loaders/legacy_library_loader.cc b/iree/hal/local/loaders/legacy_library_loader.cc
index ef2a7e9..7e59cfb 100644
--- a/iree/hal/local/loaders/legacy_library_loader.cc
+++ b/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 @@
"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();
}
diff --git a/iree/schemas/dylib_executable_def.fbs b/iree/schemas/dylib_executable_def.fbs
index 8c30493..38ee55b 100644
--- a/iree/schemas/dylib_executable_def.fbs
+++ b/iree/schemas/dylib_executable_def.fbs
@@ -14,6 +14,13 @@
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";
@@ -21,7 +28,8 @@
// 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 @@
debug_database_filename:string;
debug_database_embedded:[ubyte];
+ sanitized_kind:SanitizerKind = None;
// TODO(scotttodd): Relative file path from this flatbuffer file
}
