feat(ai): add ROCm and MIGraphX execution providers for AMD GPUs#335
Draft
Benehiko wants to merge 1 commit into
Draft
feat(ai): add ROCm and MIGraphX execution providers for AMD GPUs#335Benehiko wants to merge 1 commit into
Benehiko wants to merge 1 commit into
Conversation
Collaborator
|
This looks interesting. CC: @Ayobami-00 |
750bffe to
27ce036
Compare
517c297 to
5cb85ec
Compare
Wires ort's ROCmExecutionProvider and MIGraphXExecutionProvider through the same path as the existing CUDA / TensorRT / DirectML / CoreML providers so ahnlich-ai can target AMD GPUs on Linux when the host has a matching ROCm or MIGraphX runtime and an ORT build that supports either provider. Two variants are included because upstream onnxruntime removed the ROCm execution provider in release 1.23 and recommends MIGraphX as its replacement. ahnlich currently pins ort to 2.0.0-rc.5 (against ORT 1.19, which still ships ROCm), so both variants stay useful until the ORT pin moves past 1.23. - protos/ai/execution_provider.proto: add ROCM = 4 and MIGRAPHX = 5 - ahnlich/types/src/ai/execution_provider.rs: regenerated Rust enum - ahnlich/ai/src/engine/ai/providers/ort/mod.rs: register ROCmExecutionProvider and MIGraphXExecutionProvider via InnerAIExecutionProvider::ROCm and ::MIGraphX - ahnlich/dsl/src/ai.rs: accept "rocm" and "migraphx" in parse_to_execution_provider - ahnlich/dsl/src/syntax/syntax.pest: extend the execution_provider rule to tokenise "rocm" and "migraphx" (without this the DSL rejects the new keywords at the parser layer before parse_to_execution_provider runs) - ahnlich/dsl/src/tests/ai.rs: add round-trip tests for "rocm" and "migraphx" through parse_ai_query (mirrors the existing TensorRT and CUDA tests) - sdk/ahnlich-client-go/grpc/ai/execution_provider/execution_provider.pb.go: regenerate via `buf generate` - sdk/ahnlich-client-node/grpc/ai/execution_provider_pb.ts: regenerate via `buf generate` - sdk/ahnlich-client-py/ahnlich_client_py/grpc/ai/execution_provider/__init__.py: add ROCM and MIGRAPHX variants (matches betterproto's emit; full `make grpc-update-python` regen left to the maintainer's environment) - README.md / protos/README.md: document ROCm and MIGraphX prerequisites and the ORT 1.23 ROCm removal Generated as a reference patch with Claude — not validated against AMD hardware. Verified locally with: - `cargo test -p dsl` (30 / 30 passing, +2 new tests) - `cargo check -p ai` against the actual libonnxruntime.so 1.19.0 bundled in the official ahnlich-ai image - end-to-end DSL smoke test via ahnlich-cli: the patched binary boots, accepts `executionprovider rocm` and `executionprovider migraphx` through the gRPC layer, and rejects unknown tokens at the parser - proto regen confirmed reproducible (build.rs round-trips execution_provider.rs cleanly, `buf generate` produces the same Go/Node stubs committed here)
5cb85ec to
e63d12e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds two opt-in execution providers so ahnlich-ai can target AMD GPUs:
ROCM— ONNX Runtime's ROCm execution provider (works ononnxruntime < 1.23)MIGRAPHX— AMD's recommended replacement afteronnxruntimeremoved the ROCm provider in 1.23 (ROCm EP removal note)Both wire through the same path the existing
CUDA,TENSOR_RT,DIRECT_ML, andCORE_MLproviders use — no new dependencies, no behaviour change for existing providers.The
ortcrate (2.0.0-rc.5, already pinned inahnlich/ai/Cargo.toml) exposes both asROCmExecutionProviderandMIGraphXExecutionProvider.Why two providers
ahnlich currently pins
ortto2.0.0-rc.5(against ONNX Runtime 1.19), where the ROCm execution provider still ships. AMD removed the ROCm provider from onnxruntime in 1.23 and recommends MIGraphX for new builds. To stay useful both today and after an ORT bump, this PR adds both variants. Maintainers can keep both or dropROCMonce ahnlich bumps the ORT pin past 1.23.Changes
Protocol + Rust core
protos/ai/execution_provider.proto— addROCM = 4andMIGRAPHX = 5with doc comments that flow through to every language binding via the generatorsprotos/README.md— mentionROCMandMIGRAPHXalongsideCUDA/TENSOR_RTahnlich/types/src/ai/execution_provider.rs— regenerated Rust enum (mirrors whatbuild.rsproduces from the updated.proto)ahnlich/ai/src/engine/ai/providers/ort/mod.rsROCmExecutionProviderandMIGraphXExecutionProviderfromortInnerAIExecutionProvider::ROCmandInnerAIExecutionProvider::MIGraphXFrom<AIExecutionProvider>impl for both variantsregister_providerto callROCmExecutionProvider::default().register(...)andMIGraphXExecutionProvider::default().register(...)DSL grammar + parser + tests
ahnlich/dsl/src/syntax/syntax.pest— extend theexecution_providerrule so the pest tokeniser actually emitsrocm/migraphx. Without this the DSL rejects the new keywords at the parser layer beforeparse_to_execution_providerever sees them.ahnlich/dsl/src/ai.rs— accept"rocm"and"migraphx"inparse_to_execution_providerahnlich/dsl/src/tests/ai.rs— addtest_get_sim_n_parse_rocm_execution_providerandtest_get_sim_n_parse_migraphx_execution_provider, mirroring the existing TensorRT / CUDA round-tripsSDK regeneration
sdk/ahnlich-client-go/grpc/ai/execution_provider/execution_provider.pb.go— regenerated viabuf generatesdk/ahnlich-client-node/grpc/ai/execution_provider_pb.ts— regenerated viabuf generatesdk/ahnlich-client-py/ahnlich_client_py/grpc/ai/execution_provider/__init__.py—ROCM+MIGRAPHXvariants added in the shapebetterprotoemits. A fullmake grpc-update-pythonrun in the maintainer environment is recommended so any incidental codegen drift (formatter / generator version) is captured properly.Docs
README.md— document ROCm and MIGraphX prerequisites under "Execution Providers", including the note that upstream ORT removed the ROCm provider in 1.23Verification done locally
cargo check -p ahnlich_types -p dsl— passescargo test -p dsl— 30 / 30 passing (the two new tests round-trip"rocm"and"migraphx"end-to-end throughparse_ai_query)cargo check -p ai— passes against the actuallibonnxruntime.so 1.19.0shipped in the official ahnlich-ai image (used asORT_LIB_LOCATION), confirming the patchedort/mod.rscompiles cleanlyahnlich-cli:executionprovider rocmandexecutionprovider migraphxparse successfully (server responds with the expected store-lookup error against an empty store)executionprovider bogusis rejected at the parser layer (the negative test that surfaced the missingsyntax.pestrule in the first place)build.rsregeneratesahnlich/types/src/ai/execution_provider.rsfrom the updated.prototo the same Rust shape this PR ships; confirmedbuf generatereproduces the committed Go and Node stubsWhat is intentionally not included
Dockerfilebuilds againstonnxruntime-linux-x64-gpu-1.19.0.tgz, which is the CUDA flavour. To actually exercise ROCm or MIGraphX at runtime, ahnlich-ai needs an image built against a ROCm-enabled or MIGraphX-enabled ORT (either built from source with--use_rocm/--use_migraphx, or pulled from AMD'srocm-onnxruntimepackage). That is a packaging concern that deserves its own PR + maintainer input — probablyDockerfile.rocmand/orDockerfile.migraphxplus a CI matrix entry.register_providerarms can only be exercised through the DSL → gRPC layer (which is done above), not through actual model inference.make grpc-update-pythonrun. The Python__init__.pychange is a minimal manual mirror of whatbetterprotowould emit. A maintainer-sidepoetry run generate_from_protosis recommended to catch any formatter drift.Manual checklist for whoever picks this up
make grpc-update-pythonfrom a clean environment to confirm the Python stub matches the patchROCMandMIGRAPHXor dropROCMafter the next ORT bump past 1.23Dockerfile.rocm(orDockerfile.migraphx) bundling a ROCm-enabled / MIGraphX-enabled ORT, plus a CI matrix entryInnerAIExecutionProvider::CPUpath is unchanged, so this should be free, but worth a sanity run)References
ortcrateROCmExecutionProvider: https://docs.rs/ort/2.0.0-rc.5/ort/struct.ROCmExecutionProvider.htmlortcrateMIGraphXExecutionProvider: https://docs.rs/ort/2.0.0-rc.5/ort/struct.MIGraphXExecutionProvider.html