Skip to content

Commit 452ec52

Browse files
Add minimal ONNX dialect build option
- Add ONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT CMake flag for minimal builds - Decouple src/Dialect/Mlir from Compiler and Krnl dependencies - Add test-onnx-to-mlir utility for ONNX model import testing - Include documentation and CI workflow for minimal build - Exclude non-essential interfaces and third-party dependencies This enables building only the ONNX dialect and essential utilities without the full onnx-mlir compiler stack, useful for projects that only need ONNX dialect IR generation. Signed-off-by: Anurag <[email protected]>
1 parent 79231a0 commit 452ec52

File tree

11 files changed

+258
-35
lines changed

11 files changed

+258
-35
lines changed
Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
name: Minimal ONNX dialect build
2+
3+
on:
4+
pull_request:
5+
push:
6+
7+
jobs:
8+
build-minimal:
9+
runs-on: ubuntu-22.04
10+
steps:
11+
- uses: actions/checkout@v3
12+
with:
13+
submodules: recursive
14+
- uses: actions/setup-python@v4
15+
with:
16+
python-version: '3.10'
17+
- name: Install prerequisites
18+
run: |
19+
sudo apt-get update
20+
sudo apt-get install -y ninja-build cmake gcc g++ python3-pip
21+
- name: Cache MLIR build
22+
id: cache-mlir
23+
uses: actions/cache@v3
24+
with:
25+
path: llvm-project
26+
key: ${{ runner.os }}-mlir-${{ hashFiles('utils/clone-mlir.sh', 'utils/build-mlir.sh') }}
27+
- name: Clone + build MLIR
28+
if: steps.cache-mlir.outputs.cache-hit != 'true'
29+
run: |
30+
bash utils/clone-mlir.sh
31+
bash utils/build-mlir.sh
32+
- name: Configure minimal build
33+
run: |
34+
cmake -S . -B build -G Ninja \
35+
-DMLIR_DIR=$GITHUB_WORKSPACE/llvm-project/build/lib/cmake/mlir \
36+
-DONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT=ON \
37+
-DCMAKE_BUILD_TYPE=Release
38+
- name: Build test-onnx-to-mlir
39+
run: |
40+
cmake --build build --target test-onnx-to-mlir -j2
41+
- name: Smoke test
42+
run: |
43+
echo 'digraph G { }' > /tmp/empty.onnx || true
44+
# Use a known small ONNX from torch-mlir tests if present in tree; otherwise skip.
45+
if [ -f third_party/onnx/onnx/backend/test/data/node/test_relu/model.onnx ]; then
46+
build/Release/bin/test-onnx-to-mlir third_party/onnx/onnx/backend/test/data/node/test_relu/model.onnx | head -n 20
47+
else
48+
echo 'Minimal build completed.'
49+
fi
50+

CMakeLists.txt

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ option(ONNX_MLIR_ENABLE_WERROR "Enable warnings as errors." OFF)
1212
option(ONNX_MLIR_SUPPRESS_THIRD_PARTY_WARNINGS "Suppress warning in third_party code." ON)
1313
option(ONNX_MLIR_ENABLE_JAVA "Set to ON for building the Java runtime, tools, and tests" ON)
1414
option(ONNX_MLIR_ENABLE_PYRUNTIME_LIGHT "Set to ON for building Python driver of running the compiled model without llvm-project." OFF)
15+
option(ONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT "Build only ONNX dialect without lowering passes and runtime" OFF)
1516

1617
set(CMAKE_CXX_STANDARD 17)
1718

@@ -162,7 +163,10 @@ endif()
162163
set(CMAKE_MESSAGE_LOG_LEVEL NOTICE)
163164

164165
# Add third party subdirectories and define options appropriate to run their cmakes.
165-
if (ONNX_MLIR_ENABLE_PYRUNTIME_LIGHT)
166+
if (ONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT)
167+
# Minimal build: only include ONNX protobuf library.
168+
add_subdirectory(third_party/onnx)
169+
elseif (ONNX_MLIR_ENABLE_PYRUNTIME_LIGHT)
166170
add_subdirectory(third_party/onnx)
167171
add_subdirectory(third_party/pybind11)
168172
else()
@@ -195,13 +199,15 @@ endif()
195199
# compile flags updated via llvm_update_compile_flags, so we need to do that to
196200
# benchmark and rapidcheck as well, so that we can successfully link against them.
197201
# Otherwise, some of the flags for exceptions (among others) are not set correctly.
198-
llvm_update_compile_flags(benchmark)
199-
llvm_update_compile_flags(benchmark_main)
200-
llvm_update_compile_flags(rapidcheck)
202+
if (NOT ONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT)
203+
llvm_update_compile_flags(benchmark)
204+
llvm_update_compile_flags(benchmark_main)
205+
llvm_update_compile_flags(rapidcheck)
201206

202-
if (ONNX_MLIR_ENABLE_STABLEHLO)
203-
llvm_update_compile_flags(stablehlo-opt)
204-
llvm_update_compile_flags(stablehlo-translate)
207+
if (ONNX_MLIR_ENABLE_STABLEHLO)
208+
llvm_update_compile_flags(stablehlo-opt)
209+
llvm_update_compile_flags(stablehlo-translate)
210+
endif()
205211
endif()
206212

207213
# Increase the log level for onnx-mlir
@@ -216,16 +222,24 @@ if (ONNX_MLIR_SUPPRESS_THIRD_PARTY_WARNINGS)
216222
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DSUPPRESS_THIRD_PARTY_WARNINGS")
217223
endif()
218224

219-
if (ONNX_MLIR_ENABLE_STABLEHLO)
225+
if (ONNX_MLIR_ENABLE_STABLEHLO AND NOT ONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT)
220226
add_compile_definitions(ONNX_MLIR_ENABLE_STABLEHLO)
221227
endif()
222228

223-
if (ONNX_MLIR_ENABLE_PYRUNTIME_LIGHT)
229+
if (ONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT)
230+
# Minimal build: only include headers and ONNX dialect source.
231+
add_subdirectory(include)
232+
add_subdirectory(src/Interface)
233+
add_subdirectory(src/Support)
234+
add_subdirectory(src/Builder)
235+
add_subdirectory(src/Dialect/Mlir)
236+
add_subdirectory(src/Dialect/ONNX)
237+
elseif (ONNX_MLIR_ENABLE_PYRUNTIME_LIGHT)
224238
add_subdirectory(src)
225239
else()
226240
add_subdirectory(utils)
227241
add_subdirectory(include)
228242
add_subdirectory(src)
229243
add_subdirectory(docs)
230244
add_subdirectory(test)
231-
endif()
245+
endif()

README.md

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,38 @@ After installation, an `onnx-mlir` executable should appear in the `build/Debug/
7878
If you have difficulties building, rebuilding, or testing `onnx-mlir`, check this [page](docs/TestingHighLevel.md) for helpful hints.
7979

8080

81+
### Minimal build for ONNX dialect only
82+
83+
For consumers that only need the ONNX dialect (without lowering passes, runtime, or non-essential dialects), configure with:
84+
85+
```
86+
cmake -S . -B build -G Ninja \
87+
-DMLIR_DIR=/path/to/llvm-project/build/lib/cmake/mlir \
88+
-DONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT=ON \
89+
-DCMAKE_BUILD_TYPE=Release
90+
91+
cmake --build build --target test-onnx-to-mlir -j$(nproc)
92+
```
93+
94+
This builds only the following components:
95+
- `include/`
96+
- `src/Interface/` (only interfaces required by ONNX)
97+
- `src/Support/`
98+
- `src/Builder/`
99+
- `src/Dialect/Mlir/`
100+
- `src/Dialect/ONNX/`
101+
102+
And provides a small test executable `test-onnx-to-mlir` that loads an ONNX model and prints ONNX dialect IR:
103+
104+
```
105+
build/Release/bin/test-onnx-to-mlir model.onnx > model.mlir
106+
```
107+
108+
Notes:
109+
- Release builds are recommended to reduce memory usage during compilation.
110+
- SpecializedKernel and Krnl dependencies are excluded in this mode.
111+
112+
81113
## Using ONNX-MLIR
82114

83115
The usage of `onnx-mlir` is as such:

docs/MinimalONNXDialectBuild.md

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
<!--- SPDX-License-Identifier: Apache-2.0 -->
2+
3+
## Minimal ONNX dialect build
4+
5+
This document describes how to build a minimal subset of ONNX-MLIR that exposes only the ONNX dialect and the minimal helpers required to parse ONNX models and produce ONNX dialect MLIR.
6+
7+
### Rationale
8+
9+
Some projects want to import the ONNX dialect IR without bringing the full onnx-mlir stack (lowering passes, Krnl dialect, runtime). The CMake option `ONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT` enables such a build.
10+
11+
### What gets built
12+
13+
When `-DONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT=ON` is set, the following subdirectories are included:
14+
15+
- `include/`
16+
- `src/Interface/` (ShapeInference, ResultTypeInference, HasOnnxSubgraph, ShapeHelper)
17+
- `src/Support/`
18+
- `src/Builder/` (for `ImportFrontendModelFile` and builders)
19+
- `src/Dialect/Mlir/` (utility builders only; no Krnl or Compiler dependencies)
20+
- `src/Dialect/ONNX/`
21+
22+
In addition, a small utility `test-onnx-to-mlir` is built to load an ONNX model and print the ONNX dialect IR.
23+
24+
### Configure and build
25+
26+
```
27+
cmake -S . -B build -G Ninja \
28+
-DMLIR_DIR=/path/to/llvm-project/build/lib/cmake/mlir \
29+
-DONNX_MLIR_ENABLE_ONLY_ONNX_DIALECT=ON \
30+
-DCMAKE_BUILD_TYPE=Release
31+
32+
cmake --build build --target test-onnx-to-mlir -j$(nproc)
33+
```
34+
35+
Release builds are recommended to lower memory usage during compilation.
36+
37+
### Using the test helper
38+
39+
```
40+
build/Release/bin/test-onnx-to-mlir /path/to/model.onnx > model.mlir
41+
```
42+
43+
### Notes
44+
45+
- Krnl and Compiler modules are excluded.
46+
- SpecializedKernel interface is excluded.
47+
- The minimal build still requires the upstream LLVM/MLIR toolchain and ONNX protobuf.
48+

src/Builder/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,4 +21,4 @@ add_onnx_mlir_library(OMBuilder
2121
configure_file(OpBuildTable.inc.dc.in
2222
${CMAKE_CURRENT_BINARY_DIR}/OpBuildTable.inc.dc
2323
@ONLY
24-
)
24+
)

src/Dialect/Mlir/CMakeLists.txt

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,7 @@ add_onnx_mlir_library(OMMlirDialects
77
DialectBuilder.cpp
88
VectorMachineSupport.cpp
99

10-
DEPENDS
11-
OMKrnlIncGen
12-
OMSpecializedKernelOpInterface
13-
1410
LINK_LIBS PUBLIC
15-
OMCompilerOptions
1611
MLIRMathDialect
1712
MLIRAffineDialect
1813
MLIRSCFDialect

src/Dialect/Mlir/DialectBuilder.cpp

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
#include "llvm/Support/Debug.h"
2626

2727
// Please do not add dependences on ONNX or KRNL dialects.
28-
#include "src/Compiler/CompilerOptions.hpp"
2928
#include "src/Dialect/Mlir/DialectBuilder.hpp"
3029
#include "src/Dialect/Mlir/VectorMachineSupport.hpp"
3130

@@ -1757,16 +1756,12 @@ Value MemRefBuilder::dim(Value val, Value index) const {
17571756

17581757
void MemRefBuilder::prefetch(Value memref, ValueRange indices, bool isWrite,
17591758
unsigned locality, bool isData) {
1760-
if (disableMemRefPrefetch)
1761-
return;
17621759
b().create<memref::PrefetchOp>(
17631760
loc(), memref, indices, isWrite, locality, isData);
17641761
}
17651762

17661763
void MemRefBuilder::prefetchIE(Value memref, ArrayRef<IndexExpr> indices,
17671764
bool isWrite, unsigned locality, bool isData) {
1768-
if (disableMemRefPrefetch)
1769-
return;
17701765
SmallVector<Value, 4> indexVals;
17711766
IndexExpr::getValues(indices, indexVals);
17721767
prefetch(memref, indexVals, isWrite, locality, isData);

src/Dialect/Mlir/VectorMachineSupport.cpp

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
// =============================================================================
1010

1111
#include "src/Dialect/Mlir/VectorMachineSupport.hpp"
12-
#include "src/Compiler/CompilerOptions.hpp"
1312

1413
#include "mlir/IR/BuiltinTypes.h"
1514
#include "llvm/Support/Debug.h"
@@ -22,6 +21,29 @@ using namespace mlir;
2221

2322
namespace onnx_mlir {
2423

24+
// Minimal version of getZArchNum to avoid dependency on Compiler module
25+
namespace {
26+
int64_t decodeZArchNum(std::string str) {
27+
if (str == "arch12" || str == "z14") // Z14 and equivalents.
28+
return 12;
29+
if (str == "arch13" || str == "z15") // Z15 and equivalents.
30+
return 13;
31+
if (str == "arch14" || str == "z16") // Z16 and equivalents.
32+
return 14;
33+
if (str == "arch15" || str == "z17") // Z17 and equivalents.
34+
return 15;
35+
return -1;
36+
}
37+
38+
int64_t getZArchNumLocal(const std::string &arch, const std::string cpu) {
39+
// Give priority to march, use (deprecated) mcpu if march is not defined.
40+
int64_t num = decodeZArchNum(arch);
41+
if (num == -1)
42+
num = decodeZArchNum(cpu);
43+
return num;
44+
}
45+
} // anonymous namespace
46+
2547
// =============================================================================
2648
// Handling of global vector machine support pointer
2749

@@ -31,7 +53,7 @@ namespace onnx_mlir {
3153
/*static*/ void VectorMachineSupport::setGlobalVectorMachineSupport(
3254
const std::string &arch, const std::string &cpu, const std::string &attr) {
3355
// IBM Z servers use march (deprecated mcpu), process here.
34-
int64_t zArchNum = getZArchNum(arch, cpu);
56+
int64_t zArchNum = getZArchNumLocal(arch, cpu);
3557
if (zArchNum == 12) {
3658
globalVectorMachineSupport = new ZArch12VectorMachineSupport();
3759
} else if (zArchNum == 13) {

src/Dialect/ONNX/CMakeLists.txt

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,22 @@ add_onnx_mlir_library(OMONNXOps
124124
MLIRMemRefTransforms
125125
)
126126

127+
# Add the generated directory to the include path so interface headers can find the generated files
128+
target_include_directories(OMONNXOps PRIVATE ${GEN_DIR})
129+
127130
configure_file(ONNXOps.td.inc.dc.in
128131
${CMAKE_CURRENT_BINARY_DIR}/ONNXOps.td.inc.dc
129132
@ONLY
130133
)
134+
135+
# Add test executable
136+
add_onnx_mlir_executable(test-onnx-to-mlir
137+
test-onnx-to-mlir.cpp
138+
139+
LINK_LIBS PRIVATE
140+
OMONNXOps
141+
OMBuilder
142+
MLIRFuncDialect
143+
MLIRSupport
144+
LLVMSupport
145+
)
Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
// test-onnx-to-mlir.cpp
2+
3+
#include "mlir/Dialect/Func/IR/FuncOps.h"
4+
#include "mlir/IR/BuiltinOps.h"
5+
#include "mlir/IR/DialectRegistry.h"
6+
#include "mlir/IR/MLIRContext.h"
7+
8+
#include "src/Builder/FrontendDialectTransformer.hpp"
9+
#include "src/Dialect/ONNX/ONNXDialect.hpp"
10+
#include "llvm/Support/InitLLVM.h"
11+
#include "llvm/Support/raw_ostream.h"
12+
13+
int main(int argc, char **argv) {
14+
llvm::InitLLVM y(argc, argv);
15+
16+
if (argc != 2) {
17+
llvm::errs() << "Usage: test-onnx-to-mlir <input.onnx>\n";
18+
return 1;
19+
}
20+
21+
mlir::MLIRContext context;
22+
23+
// Register the ONNX dialect.
24+
mlir::DialectRegistry registry;
25+
registry.insert<mlir::func::FuncDialect>();
26+
registry.insert<mlir::ONNXDialect>();
27+
context.appendDialectRegistry(registry);
28+
context.loadAllAvailableDialects();
29+
30+
// Import the ONNX model.
31+
mlir::OwningOpRef<mlir::ModuleOp> module;
32+
std::string errorMessage;
33+
onnx_mlir::ImportOptions options;
34+
35+
int result = onnx_mlir::ImportFrontendModelFile(
36+
argv[1], context, module, &errorMessage, options);
37+
38+
if (result != 0) {
39+
llvm::errs() << "Failed to import model: " << errorMessage << "\n";
40+
return 1;
41+
}
42+
43+
llvm::outs() << "Successfully imported ONNX model!\n";
44+
module->print(llvm::outs());
45+
46+
return 0;
47+
}

0 commit comments

Comments
 (0)