diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp index 7bb954ad7589..d0dbde6b585f 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp +++ b/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(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 467ebfdc195b..87c2f797d9ce 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 @@ 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()); + 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 33f1416789c7..b6c18e0904b9 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.cpp +++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMTargetOptions.cpp @@ -82,6 +82,13 @@ LLVMTargetOptions getLLVMTargetOptionsFromFlags() { llvmTargetOptions.targetCPUFeatures = clTargetCPUFeatures; } + static llvm::cl::opt 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 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 1ac53bc9ad06..3e85ed34b9ad 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_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. diff --git a/iree/hal/local/loaders/legacy_library_loader.cc b/iree/hal/local/loaders/legacy_library_loader.cc index 2a03f86ae9c0..0b78f7b230a7 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 @@ 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(); } diff --git a/iree/schemas/dylib_executable_def.fbs b/iree/schemas/dylib_executable_def.fbs index 8c3049330cee..38ee55b76145 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 @@ 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 }