Skip to content

Commit

Permalink
Support large models with lrodata linkage (onnx#2289)
Browse files Browse the repository at this point in the history
* - Support large models with lrodata linkage (thanks Andreas Krebbel)
- Fix bug in setXoptOption that was setting Xllc instead

Signed-off-by: Gong Su <[email protected]>
  • Loading branch information
gongsu832 authored Jun 1, 2023
1 parent 667011c commit 431e6dc
Show file tree
Hide file tree
Showing 9 changed files with 290 additions and 2 deletions.
8 changes: 8 additions & 0 deletions src/Compiler/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -1,5 +1,13 @@
# SPDX-License-Identifier: Apache-2.0

# Include "Tools" before "Compiler" so that LRODATA_SCRIPT
# is set before ExternalUtil.hpp.in is processed.
if (LRODATA_SCRIPT)
file(READ ${LRODATA_SCRIPT} ONNX_MLIR_LRODATA_SCRIPT)
else()
set(ONNX_MLIR_LRODATA_SCRIPT "")
endif()

configure_file(
${CMAKE_CURRENT_SOURCE_DIR}/ExternalUtil.hpp.in
${CMAKE_CURRENT_BINARY_DIR}/ExternalUtil.hpp.cfg
Expand Down
5 changes: 4 additions & 1 deletion src/Compiler/CompilerOptions.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,9 @@ llvm::cl::opt<std::string> reportHeapAfter("report-heap-after",
// replaced by a class of its own.
std::map<std::string, std::vector<std::string>> CompilerConfigMap;

// Must match ModelSize enum
const std::string modelSizeStr[] = {"small", "medium", "large", "huge"};

// =============================================================================
// Methods for setting and getting compiler variables.

Expand Down Expand Up @@ -430,7 +433,7 @@ std::string getOptimizationLevelOption() {
// Support for Xopt.
void setXoptOption(const std::vector<std::string> &flags) {
for (const std::string &flag : flags)
Xllc.addValue(flag);
Xopt.addValue(flag);
}

void clearXoptOption() { Xopt.clear(); }
Expand Down
2 changes: 2 additions & 0 deletions src/Compiler/CompilerOptions.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,8 @@ typedef enum {
// clang-format on
} ModelSize;

extern const std::string modelSizeStr[];

// Options for onnx-mlir only.
extern llvm::cl::OptionCategory OnnxMlirOptions;
// Common options shared between onnx-mlir and onnx-mlir-opt.
Expand Down
20 changes: 20 additions & 0 deletions src/Compiler/CompilerUtils.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
#include "src/Dialect/ONNX/ONNXDialect.hpp"
#include "src/Version/Version.hpp"

#include <fstream>
#include <regex>

#define DEBUG_TYPE "compiler_utils"
Expand Down Expand Up @@ -428,6 +429,7 @@ static int genLLVMBitcode(const mlir::OwningOpRef<ModuleOp> &module,
// Use the LLVM's 'opt' command to optimize the bitcode.
std::string optPath = getToolPath("opt", kOptPath);
Command optBitcode(/*exePath=*/optPath);
setXoptOption({"--code-model", modelSizeStr[modelSize]});
int rc = optBitcode.appendStr(getOptimizationLevelOption())
.appendStr(getTargetTripleOption())
.appendStr(getTargetArchOption())
Expand All @@ -447,6 +449,7 @@ static int genModelObject(

std::string llcPath = getToolPath("llc", kLlcPath);
Command llvmToObj(/*exePath=*/llcPath);
setXllcOption({"--code-model", modelSizeStr[modelSize]});
int rc = llvmToObj.appendStr(getOptimizationLevelOption())
.appendStr(getTargetTripleOption())
.appendStr(getTargetArchOption())
Expand Down Expand Up @@ -497,6 +500,23 @@ static int genSharedLib(std::string sharedLibNameWithExt,
std::vector<std::string> sharedLibOpts = {"-shared", "-fPIC"};
llvm::for_each(libs, [](std::string &lib) { lib = "-l" + lib; });
llvm::for_each(libDirs, [](std::string &libDir) { libDir = "-L" + libDir; });
#ifdef __s390x__
llvm::SmallString<64> lds;
if (modelSize == ModelSize::large) {
if (auto ec =
llvm::sys::fs::createTemporaryFile("s390x-lrodata", "ld", lds)) {
llvm::errs() << ec.message() << "\n";
return CompilerFailureInObjToLib;
}

std::string ldScript = std::string(lds);
std::ofstream ofs(ldScript);
ofs << kLrodataScript;
ofs.close();
sharedLibOpts.push_back("-Wl,-T," + ldScript);
}
llvm::FileRemover ldsRemover(lds);
#endif
#endif

Command link(kCxxPath);
Expand Down
1 change: 1 addition & 0 deletions src/Compiler/ExternalUtil.hpp.in
Original file line number Diff line number Diff line change
Expand Up @@ -18,4 +18,5 @@ const std::string kObjCopyPath = "@CMAKE_OBJCOPY@";
const std::string kArPath = "@CMAKE_AR@";
const std::string kJarPath = "@Java_JAR_EXECUTABLE@";
const std::string kDefaultTriple = "@ONNX_MLIR_DEFAULT_TRIPLE@";
const std::string kLrodataScript = R"(@ONNX_MLIR_LRODATA_SCRIPT@)";
} // namespace onnx_mlir
1 change: 1 addition & 0 deletions src/Tools/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,4 @@

add_subdirectory(onnx-mlir-opt)
add_subdirectory(onnx-mlir-reduce)
add_subdirectory(lrodata-script)
10 changes: 10 additions & 0 deletions src/Tools/lrodata-script/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
# SPDX-License-Identifier: Apache-2.0

# file(READ ...) cannot set global variable so we only set the
# script filename here and let the "Compiler" subdirectory do
# the file(READ ...)
unset(LRODATA_SCRIPT CACHE)
if (${CMAKE_SYSTEM_PROCESSOR} STREQUAL "s390x")
set(LRODATA_SCRIPT "${CMAKE_CURRENT_SOURCE_DIR}/s390x-lrodata.ld"
CACHE INTERNAL "")
endif()
243 changes: 243 additions & 0 deletions src/Tools/lrodata-script/s390x-lrodata.ld
Original file line number Diff line number Diff line change
@@ -0,0 +1,243 @@
/* Script for -z combreloc -z separate-code -z relro */
/* Copyright (C) 2014-2022 Free Software Foundation, Inc.
Copying and distribution of this script, with or without modification,
are permitted in any medium without royalty provided the copyright
notice and this notice are preserved. */
OUTPUT_FORMAT("elf64-s390", "elf64-s390",
"elf64-s390")
OUTPUT_ARCH(s390:64-bit)
ENTRY(_start)
SEARCH_DIR("=/usr/s390x-redhat-linux/lib64"); SEARCH_DIR("=/usr/lib64"); SEARCH_DIR("=/usr/local/lib64"); SEARCH_DIR("=/lib64"); SEARCH_DIR("=/usr/s390x-redhat-linux/lib"); SEARCH_DIR("=/usr/local/lib"); SEARCH_DIR("=/lib"); SEARCH_DIR("=/usr/lib");
SECTIONS
{
PROVIDE (__executable_start = SEGMENT_START("text-segment", 0x1000000)); . = SEGMENT_START("text-segment", 0x1000000) + SIZEOF_HEADERS;
.interp : { *(.interp) }
.note.gnu.build-id : { *(.note.gnu.build-id) }
.hash : { *(.hash) }
.gnu.hash : { *(.gnu.hash) }
.dynsym : { *(.dynsym) }
.dynstr : { *(.dynstr) }
.gnu.version : { *(.gnu.version) }
.gnu.version_d : { *(.gnu.version_d) }
.gnu.version_r : { *(.gnu.version_r) }
.rela.dyn :
{
*(.rela.init)
*(.rela.text .rela.text.* .rela.gnu.linkonce.t.*)
*(.rela.fini)
*(.rela.rodata .rela.rodata.* .rela.gnu.linkonce.r.*)
*(.rela.data .rela.data.* .rela.gnu.linkonce.d.*)
*(.rela.tdata .rela.tdata.* .rela.gnu.linkonce.td.*)
*(.rela.tbss .rela.tbss.* .rela.gnu.linkonce.tb.*)
*(.rela.ctors)
*(.rela.dtors)
*(.rela.got)
*(.rela.bss .rela.bss.* .rela.gnu.linkonce.b.*)
*(.rela.ldata .rela.ldata.* .rela.gnu.linkonce.l.*)
*(.rela.lbss .rela.lbss.* .rela.gnu.linkonce.lb.*)
*(.rela.lrodata .rela.lrodata.* .rela.gnu.linkonce.lr.*)
*(.rela.ifunc)
}
.rela.plt :
{
*(.rela.plt)
PROVIDE_HIDDEN (__rela_iplt_start = .);
*(.rela.iplt)
PROVIDE_HIDDEN (__rela_iplt_end = .);
}
. = ALIGN(CONSTANT (MAXPAGESIZE));
.init :
{
KEEP (*(SORT_NONE(.init)))
} =0x07070707
.plt : { *(.plt) *(.iplt) }
.text :
{
*(.text.unlikely .text.*_unlikely .text.unlikely.*)
*(.text.exit .text.exit.*)
*(.text.startup .text.startup.*)
*(.text.hot .text.hot.*)
*(SORT(.text.sorted.*))
*(.text .stub .text.* .gnu.linkonce.t.*)
/* .gnu.warning sections are handled specially by elf.em. */
*(.gnu.warning)
} =0x07070707
.fini :
{
KEEP (*(SORT_NONE(.fini)))
} =0x07070707
PROVIDE (__etext = .);
PROVIDE (_etext = .);
PROVIDE (etext = .);
. = ALIGN(CONSTANT (MAXPAGESIZE));
/* Adjust the address for the rodata segment. We want to adjust up to
the same address within the page on the next page up. */
. = SEGMENT_START("rodata-segment", ALIGN(CONSTANT (MAXPAGESIZE)) + (. & (CONSTANT (MAXPAGESIZE) - 1)));
.rodata : { *(.rodata .rodata.* .gnu.linkonce.r.*) }
.rodata1 : { *(.rodata1) }
.eh_frame_hdr : { *(.eh_frame_hdr) *(.eh_frame_entry .eh_frame_entry.*) }
.eh_frame : ONLY_IF_RO { KEEP (*(.eh_frame)) *(.eh_frame.*) }
.gcc_except_table : ONLY_IF_RO { *(.gcc_except_table .gcc_except_table.*) }
.gnu_extab : ONLY_IF_RO { *(.gnu_extab*) }
/* These sections are generated by the Sun/Oracle C++ compiler. */
.exception_ranges : ONLY_IF_RO { *(.exception_ranges*) }
/* Adjust the address for the data segment. We want to adjust up to
the same address within the page on the next page up. */
. = DATA_SEGMENT_ALIGN (CONSTANT (MAXPAGESIZE), CONSTANT (COMMONPAGESIZE));
/* Exception handling */
.eh_frame : ONLY_IF_RW { KEEP (*(.eh_frame)) *(.eh_frame.*) }
.gnu_extab : ONLY_IF_RW { *(.gnu_extab) }
.gcc_except_table : ONLY_IF_RW { *(.gcc_except_table .gcc_except_table.*) }
.exception_ranges : ONLY_IF_RW { *(.exception_ranges*) }
/* Thread Local Storage sections */
.tdata :
{
PROVIDE_HIDDEN (__tdata_start = .);
*(.tdata .tdata.* .gnu.linkonce.td.*)
}
.tbss : { *(.tbss .tbss.* .gnu.linkonce.tb.*) *(.tcommon) }
.preinit_array :
{
PROVIDE_HIDDEN (__preinit_array_start = .);
KEEP (*(.preinit_array))
PROVIDE_HIDDEN (__preinit_array_end = .);
}
.init_array :
{
PROVIDE_HIDDEN (__init_array_start = .);
KEEP (*(SORT_BY_INIT_PRIORITY(.init_array.*) SORT_BY_INIT_PRIORITY(.ctors.*)))
KEEP (*(.init_array EXCLUDE_FILE (*crtbegin.o *crtbegin?.o *crtend.o *crtend?.o ) .ctors))
PROVIDE_HIDDEN (__init_array_end = .);
}
.fini_array :
{
PROVIDE_HIDDEN (__fini_array_start = .);
KEEP (*(SORT_BY_INIT_PRIORITY(.fini_array.*) SORT_BY_INIT_PRIORITY(.dtors.*)))
KEEP (*(.fini_array EXCLUDE_FILE (*crtbegin.o *crtbegin?.o *crtend.o *crtend?.o ) .dtors))
PROVIDE_HIDDEN (__fini_array_end = .);
}
.ctors :
{
/* gcc uses crtbegin.o to find the start of
the constructors, so we make sure it is
first. Because this is a wildcard, it
doesn't matter if the user does not
actually link against crtbegin.o; the
linker won't look for a file to match a
wildcard. The wildcard also means that it
doesn't matter which directory crtbegin.o
is in. */
KEEP (*crtbegin.o(.ctors))
KEEP (*crtbegin?.o(.ctors))
/* We don't want to include the .ctor section from
the crtend.o file until after the sorted ctors.
The .ctor section from the crtend file contains the
end of ctors marker and it must be last */
KEEP (*(EXCLUDE_FILE (*crtend.o *crtend?.o ) .ctors))
KEEP (*(SORT(.ctors.*)))
KEEP (*(.ctors))
}
.dtors :
{
KEEP (*crtbegin.o(.dtors))
KEEP (*crtbegin?.o(.dtors))
KEEP (*(EXCLUDE_FILE (*crtend.o *crtend?.o ) .dtors))
KEEP (*(SORT(.dtors.*)))
KEEP (*(.dtors))
}
.jcr : { KEEP (*(.jcr)) }
.data.rel.ro : { *(.data.rel.ro.local* .gnu.linkonce.d.rel.ro.local.*) *(.data.rel.ro .data.rel.ro.* .gnu.linkonce.d.rel.ro.*) }
.dynamic : { *(.dynamic) }
.got : { *(.got) *(.igot) }
. = DATA_SEGMENT_RELRO_END (0, .);
.got.plt : { *(.got.plt) *(.igot.plt) }
.data :
{
*(.data .data.* .gnu.linkonce.d.*)
SORT(CONSTRUCTORS)
}
.data1 : { *(.data1) }
_edata = .; PROVIDE (edata = .);
. = .;
__bss_start = .;
.bss :
{
*(.dynbss)
*(.bss .bss.* .gnu.linkonce.b.*)
*(COMMON)
/* Align here to ensure that the .bss section occupies space up to
_end. Align after .bss to ensure correct alignment even if the
.bss section disappears because there are no input sections.
FIXME: Why do we need it? When there is no .bss section, we do not
pad the .data section. */
. = ALIGN(. != 0 ? 64 / 8 : 1);
}
.lbss :
{
*(.dynlbss)
*(.lbss .lbss.* .gnu.linkonce.lb.*)
*(LARGE_COMMON)
}
. = ALIGN(64 / 8);
. = SEGMENT_START("ldata-segment", .);
.lrodata ALIGN(CONSTANT (MAXPAGESIZE)) + (. & (CONSTANT (MAXPAGESIZE) - 1)) :
{
*(.lrodata .lrodata.* .gnu.linkonce.lr.*)
}
.ldata ALIGN(CONSTANT (MAXPAGESIZE)) + (. & (CONSTANT (MAXPAGESIZE) - 1)) :
{
*(.ldata .ldata.* .gnu.linkonce.l.*)
. = ALIGN(. != 0 ? 64 / 8 : 1);
}
. = ALIGN(64 / 8);
_end = .; PROVIDE (end = .);
. = DATA_SEGMENT_END (.);
/* Stabs debugging sections. */
.stab 0 : { *(.stab) }
.stabstr 0 : { *(.stabstr) }
.stab.excl 0 : { *(.stab.excl) }
.stab.exclstr 0 : { *(.stab.exclstr) }
.stab.index 0 : { *(.stab.index) }
.stab.indexstr 0 : { *(.stab.indexstr) }
.comment 0 : { *(.comment) }
.gnu.build.attributes : { *(.gnu.build.attributes .gnu.build.attributes.*) }
/* DWARF debug sections.
Symbols in the DWARF debugging sections are relative to the beginning
of the section so we begin them at 0. */
/* DWARF 1. */
.debug 0 : { *(.debug) }
.line 0 : { *(.line) }
/* GNU DWARF 1 extensions. */
.debug_srcinfo 0 : { *(.debug_srcinfo) }
.debug_sfnames 0 : { *(.debug_sfnames) }
/* DWARF 1.1 and DWARF 2. */
.debug_aranges 0 : { *(.debug_aranges) }
.debug_pubnames 0 : { *(.debug_pubnames) }
/* DWARF 2. */
.debug_info 0 : { *(.debug_info .gnu.linkonce.wi.*) }
.debug_abbrev 0 : { *(.debug_abbrev) }
.debug_line 0 : { *(.debug_line .debug_line.* .debug_line_end) }
.debug_frame 0 : { *(.debug_frame) }
.debug_str 0 : { *(.debug_str) }
.debug_loc 0 : { *(.debug_loc) }
.debug_macinfo 0 : { *(.debug_macinfo) }
/* SGI/MIPS DWARF 2 extensions. */
.debug_weaknames 0 : { *(.debug_weaknames) }
.debug_funcnames 0 : { *(.debug_funcnames) }
.debug_typenames 0 : { *(.debug_typenames) }
.debug_varnames 0 : { *(.debug_varnames) }
/* DWARF 3. */
.debug_pubtypes 0 : { *(.debug_pubtypes) }
.debug_ranges 0 : { *(.debug_ranges) }
/* DWARF 5. */
.debug_addr 0 : { *(.debug_addr) }
.debug_line_str 0 : { *(.debug_line_str) }
.debug_loclists 0 : { *(.debug_loclists) }
.debug_macro 0 : { *(.debug_macro) }
.debug_names 0 : { *(.debug_names) }
.debug_rnglists 0 : { *(.debug_rnglists) }
.debug_str_offsets 0 : { *(.debug_str_offsets) }
.debug_sup 0 : { *(.debug_sup) }
.gnu.attributes 0 : { KEEP (*(.gnu.attributes)) }
/DISCARD/ : { *(.note.GNU-stack) *(.gnu_debuglink) *(.gnu.lto_*) }
}
2 changes: 1 addition & 1 deletion test/mlir/driver/check_passthrough_options.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
// RUN: onnx-mlir -Xllc --data-sections -v %s -o %t 2>&1 | FileCheck --check-prefix=LLC %s
// RUN: onnx-mlir -mllvm --data-sections -v %s -o %t 2>&1 | FileCheck --check-prefix=LLVM %s

// OPT: opt{{.*}} --data-sections -o {{.*}}check_passthrough_options{{.*}}.bc
// OPT: opt{{.*}} --data-sections {{.*}} -o {{.*}}check_passthrough_options{{.*}}.bc
// OPT-NOT: llc{{.*}} --data-sections {{.*}}
// LLC-NOT: opt{{.*}} --data-sections -o {{.*}}check_passthrough_options{{.*}}.bc
// LLC: llc{{.*}} --data-sections {{.*}}
Expand Down

0 comments on commit 431e6dc

Please sign in to comment.