[hybrid] Hybrid API: native library tasks with cuDNN provider#888
Conversation
There was a problem hiding this comment.
Pull request overview
Adds a “Hybrid API” execution path so a TaskGraph can mix JIT-compiled kernels with calls into native vendor libraries (cuBLAS/cuFFT/cuDNN) on the same device buffers/stream, plus new CUDA library modules, JNI drivers, tests, docs, and packaging updates to ship/run them.
Changes:
- Introduce library-task descriptors + runtime SPI (
TornadoLibraryProvider) and runtime plumbing to create per-(device, plan) native contexts and dispatch library calls from the interpreter. - Expose CUDA native stream/context handles for ordering native library work with TornadoVM kernels/transfers (incl. CUDA Graph capture constraints).
- Add new modules (
tornado-cublas,tornado-cufft,tornado-cudnn), JNI driver submodules, unit tests/benchmarks, documentation, and distribution wiring.
Reviewed changes
Copilot reviewed 89 out of 91 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| tornado-unittests/src/main/java/uk/ac/manchester/tornado/unittests/cufft/TestCuFft.java | Adds cuFFT unit tests using library tasks with a Java DFT reference. |
| tornado-unittests/src/main/java/module-info.java | Adds module requires/exports for new CUDA library test packages. |
| tornado-unittests/pom.xml | Adds dependencies on tornado-cublas, tornado-cufft, tornado-cudnn. |
| tornado-runtime/src/main/java/uk/ac/manchester/tornado/runtime/tasks/TornadoTaskGraph.java | Extends task-graph runtime to accept/track library tasks and cleanup contexts. |
| tornado-runtime/src/main/java/uk/ac/manchester/tornado/runtime/tasks/LibraryTask.java | Introduces a schedulable task type representing a native library call. |
| tornado-runtime/src/main/java/uk/ac/manchester/tornado/runtime/library/spi/TornadoNativeStreamSupport.java | Defines backend capability to expose native stream/context handles. |
| tornado-runtime/src/main/java/uk/ac/manchester/tornado/runtime/library/spi/TornadoLibraryProvider.java | Adds provider SPI for native library bindings. |
| tornado-runtime/src/main/java/uk/ac/manchester/tornado/runtime/library/spi/LibraryInvocation.java | Defines the runtime-resolved argument/pointer bundle passed to providers. |
| tornado-runtime/src/main/java/uk/ac/manchester/tornado/runtime/library/spi/LibraryContext.java | Defines per-(library, device, plan) opaque provider context. |
| tornado-runtime/src/main/java/uk/ac/manchester/tornado/runtime/library/LibraryRegistry.java | Adds provider discovery and context caching/teardown registry. |
| tornado-runtime/src/main/java/uk/ac/manchester/tornado/runtime/interpreter/TornadoVMInterpreter.java | Adds interpreter support to prepare and dispatch library tasks at launch time. |
| tornado-runtime/src/main/java/module-info.java | Exports/uses library SPI packages from tornado.runtime. |
| tornado-drivers/pom.xml | Adds new JNI driver submodules under the CUDA profile. |
| tornado-drivers/cufft-jni/src/main/cpp/source/tornado-cufft.cpp | Implements JNI wrappers around cuFFT plan/exec/stream binding. |
| tornado-drivers/cufft-jni/src/main/cpp/CMakeLists.txt | Builds tornado-cufft JNI shared library and locates cuFFT. |
| tornado-drivers/cufft-jni/pom.xml | Maven build for cuFFT JNI driver (cmake + packaging). |
| tornado-drivers/cudnn-jni/src/main/cpp/source/tornado-cudnn-sdpa.cpp | Adds cuDNN graph-API SDPA (flash attention) JNI implementation. |
| tornado-drivers/cudnn-jni/src/main/cpp/CMakeLists.txt | Builds tornado-cudnn JNI shared library and fetches cudnn-frontend headers. |
| tornado-drivers/cudnn-jni/pom.xml | Maven build for cuDNN JNI driver (cmake + packaging). |
| tornado-drivers/cuda/src/main/java/uk/ac/manchester/tornado/drivers/cuda/runtime/CUDATornadoDevice.java | Implements TornadoNativeStreamSupport for CUDA devices. |
| tornado-drivers/cuda/src/main/java/uk/ac/manchester/tornado/drivers/cuda/CUDADeviceContextInterface.java | Adds default native stream/context interop hooks to device contexts. |
| tornado-drivers/cuda/src/main/java/uk/ac/manchester/tornado/drivers/cuda/CUDADeviceContext.java | Exposes CUstream/CUcontext handles from CUDA command queues. |
| tornado-drivers/cuda/src/main/java/uk/ac/manchester/tornado/drivers/cuda/CUDACommandQueue.java | Adds Java natives to fetch raw stream/context pointers. |
| tornado-drivers/cuda-jni/src/main/cpp/source/CUDAGraph.cpp | Implements JNI natives to return queue CUstream/CUcontext pointers. |
| tornado-drivers/cuda-jni/src/main/cpp/CMakeLists.txt | Adjusts CUDA include/lib hint ordering for consistent toolkit selection. |
| tornado-drivers/cublas-jni/src/main/cpp/CMakeLists.txt | Adds build for tornado-cublas JNI shared library. |
| tornado-drivers/cublas-jni/pom.xml | Maven build for cuBLAS JNI driver (cmake + packaging). |
| tornado-cufft/src/main/resources/META-INF/services/uk.ac.manchester.tornado.runtime.library.spi.TornadoLibraryProvider | Registers cuFFT provider via ServiceLoader. |
| tornado-cufft/src/main/java/uk/ac/manchester/tornado/cufft/tests/BenchmarkFft.java | Adds FFT benchmark comparing Java/JIT DFT vs cuFFT library task. |
| tornado-cufft/src/main/java/uk/ac/manchester/tornado/cufft/provider/CuFftNativeLib.java | Adds JNI loader + cuFFT native method declarations and error decoding. |
| tornado-cufft/src/main/java/uk/ac/manchester/tornado/cufft/provider/CuFftLibraryProvider.java | Implements cuFFT provider and plan caching per shape/context. |
| tornado-cufft/src/main/java/uk/ac/manchester/tornado/cufft/CuFft.java | Adds public factory methods for cuFFT library task descriptors. |
| tornado-cufft/src/main/java/module-info.java | Declares tornado.cufft module and provides provider implementation. |
| tornado-cufft/README.md | Documents cuFFT library task usage and build/run instructions. |
| tornado-cufft/pom.xml | Maven module for cuFFT Java bindings. |
| tornado-cudnn/src/main/resources/META-INF/services/uk.ac.manchester.tornado.runtime.library.spi.TornadoLibraryProvider | Registers cuDNN provider via ServiceLoader. |
| tornado-cudnn/src/main/java/uk/ac/manchester/tornado/cudnn/tests/BenchmarkSdpa.java | Adds SDPA benchmark (JIT vs fused cuDNN graph API). |
| tornado-cudnn/src/main/java/uk/ac/manchester/tornado/cudnn/tests/BenchmarkConv2d.java | Adds conv2d benchmark (JIT vs cuDNN). |
| tornado-cudnn/src/main/java/uk/ac/manchester/tornado/cudnn/provider/CuDnnNativeLib.java | Adds JNI loader + native method declarations for cuDNN ops and plans. |
| tornado-cudnn/src/main/java/uk/ac/manchester/tornado/cudnn/CuDnn.java | Adds public factory methods for cuDNN library task descriptors. |
| tornado-cudnn/src/main/java/module-info.java | Declares tornado.cudnn module and provides provider implementation. |
| tornado-cudnn/README.md | Documents cuDNN library task usage, requirements, and benchmarks. |
| tornado-cudnn/pom.xml | Maven module for cuDNN Java bindings. |
| tornado-cublas/src/main/resources/META-INF/services/uk.ac.manchester.tornado.runtime.library.spi.TornadoLibraryProvider | Registers cuBLAS and cuBLASLt providers via ServiceLoader. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/TestCuBlasWorkspace.java | Adds cuBLAS workspace allocation/reuse test. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/TestCuBlasSgemvWithTornadoVMTasksPOST.java | Adds mixed JIT + cuBLAS + JIT pipeline test. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/TestCuBlasSgemvWithTasksCudaGraph.java | Adds CUDA Graph capture/replay integration test for mixed graphs. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/TestCuBlasSgemvBeta.java | Adds SGEMV beta!=0 test validating READ_WRITE semantics. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/TestCuBlasSgemv.java | Adds basic SGEMV correctness test. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/TestCuBlasSgemmTF32.java | Adds TF32-vs-FP32 SGEMM validation + timing. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/TestCuBlasSgemm.java | Adds SGEMM correctness test using row-major trick. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/TestCuBlasGemmExFP16.java | Adds mixed-precision GEMMEx FP16 tests. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/tests/BenchmarkLtFusedMlp.java | Adds cuBLASLt fused MLP benchmark (GELU_BIAS epilogue). |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/provider/CuBlasNativeLib.java | Adds JNI loader + native method declarations for cuBLAS calls. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/provider/CuBlasLtNativeLib.java | Adds JNI declarations for cuBLASLt handle/plan/execute. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/provider/CuBlasLtLibraryProvider.java | Implements cuBLASLt provider with plan cache and workspace management. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/enums/CudaDataType.java | Adds CUDA data type enum mapping for Ex APIs. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/enums/CuBlasOperation.java | Adds cuBLAS transpose-op enum mapping. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/enums/CuBlasMathMode.java | Adds cuBLAS math-mode enum mapping (incl. TF32). |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/enums/CuBlasLtEpilogue.java | Adds cuBLASLt epilogue enum mapping (bias/GELU). |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/enums/CublasComputeType.java | Adds cuBLAS compute-type enum mapping. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/CuBlasOptions.java | Adds per-call tuning options container for cuBLAS tasks. |
| tornado-cublas/src/main/java/uk/ac/manchester/tornado/cublas/CuBlasLt.java | Adds public factory methods for cuBLASLt library task descriptors. |
| tornado-cublas/src/main/java/module-info.java | Declares tornado.cublas module and provides provider implementations. |
| tornado-cublas/pom.xml | Maven module for cuBLAS Java bindings. |
| tornado-assembly/src/bin/tornado.py | Adds CUDA library modules to --add-modules when CUDA backend is enabled. |
| tornado-assembly/src/bin/tornado-test | Adds cuBLAS/cuFFT/cuDNN tests to the “test the world” list. |
| tornado-assembly/pom.xml | Adds assembly dependencies for CUDA library modules + JNI libs. |
| tornado-assembly/assembly.xml | Includes new CUDA library modules and JNI artifacts in the distribution. |
| tornado-api/src/main/java/uk/ac/manchester/tornado/api/TornadoTaskGraphInterface.java | Adds API hook for adding a library task to a task graph. |
| tornado-api/src/main/java/uk/ac/manchester/tornado/api/common/TornadoFunctions.java | Adds functional interfaces for libraryTask(...) factories. |
| tornado-api/src/main/java/uk/ac/manchester/tornado/api/common/LibraryTaskDescriptor.java | Introduces the public descriptor type for a library task call. |
| pom.xml | Adds new top-level Maven modules for cuBLAS/cuFFT/cuDNN. |
| docs/source/index.rst | Adds the hybrid API doc page to the docs index. |
| docs/source/hybrid-api.rst | Documents hybrid API concepts, usage, and provider SPI. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| public static void destroyContexts(long executionPlanId) { | ||
| Iterator<Map.Entry<ContextKey, LibraryContext>> iterator = CONTEXTS.entrySet().iterator(); | ||
| while (iterator.hasNext()) { | ||
| Map.Entry<ContextKey, LibraryContext> entry = iterator.next(); | ||
| if (entry.getKey().executionPlanId() == executionPlanId) { | ||
| for (TornadoLibraryProvider provider : getProviders()) { | ||
| if (provider.libraryName().equals(entry.getKey().libraryName())) { | ||
| provider.destroyContext(entry.getValue()); | ||
| break; | ||
| } | ||
| } | ||
| iterator.remove(); | ||
| } | ||
| } | ||
| } |
| @Override | ||
| public boolean equals(Object obj) { | ||
| if (obj instanceof LibraryTask other) { | ||
| return getId().equals(other.getId()); | ||
| } | ||
| return false; | ||
| } | ||
|
|
||
| @Override | ||
| public int hashCode() { | ||
| int hash = 7; | ||
| hash = 71 * hash + Objects.hashCode(getId()); | ||
| hash = 71 * hash + Objects.hashCode(descriptor.getLibraryName()); | ||
| hash = 71 * hash + Objects.hashCode(descriptor.getFunctionName()); | ||
| return hash; | ||
| } |
| final int n = (int) invocation.getArg(2); | ||
| final int batch = (int) invocation.getArg(3); |
| final int transa = (int) invocation.getArg(0); | ||
| final int transb = (int) invocation.getArg(1); | ||
| final int m = (int) invocation.getArg(2); | ||
| final int n = (int) invocation.getArg(3); | ||
| final int k = (int) invocation.getArg(4); |
| final float alpha = (float) invocation.getArg(5); | ||
| final int lda = (int) invocation.getArg(7); | ||
| final int ldb = (int) invocation.getArg(9); | ||
| final float beta = (float) invocation.getArg(10); | ||
| final int ldc = (int) invocation.getArg(12); |
| target_link_libraries(tornado-cudnn ${JNI_LIB_DIRS} "-lcudnn -lcudart") | ||
| endif() | ||
|
|
||
| message(STATUS "TornadoVM cuFFT: include=${CUDA_INCLUDE_DIR}; cudnn=${CUDNN_LIB}") |
| </parent> | ||
| <artifactId>tornado-cudnn</artifactId> | ||
| <name>tornado-cudnn</name> | ||
| <description>TornadoVM cuDNN bindings: NVIDIA cuFFT library tasks for TornadoVM task graphs</description> |
|
Cherry-picking commit |
|
Cherry-picking commit |
Replace non-ASCII characters (em dashes, middle dots) in javadoc, add trailing periods to arg-list javadoc, wrap an over-long line, and remove ParenPad whitespace in Map.of( // dispatch registries. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Cherry-picking commit |
|
Cherry-picking commit |
| /** (input, output, n, batch). */ | ||
| private static void execC2C(CuFftContext context, LibraryInvocation invocation, int direction) { | ||
| final int n = (int) invocation.getArg(2); | ||
| final int batch = (int) invocation.getArg(3); | ||
|
|
|
|
||
| private static long getOrCreatePlan(CuFftContext context, PlanKind kind, int dim1, int dim2) { | ||
| String planKey = kind + ":" + dim1 + ":" + dim2; | ||
| /** (input, output, n, batch). */ |
Description
Adds the cuDNN provider to the Hybrid API — CNN primitives plus fused scaled-dot-product attention (flash attention, FP16) via the cuDNN graph API.
What's included
Legacy-API primitives (FP32, NCHW) —
tornado-cudnn+tornado-drivers/cudnn-jni(requires cuDNN 9):cudnnSoftmax(in, out, rows, cols)cudnnSoftmaxForward(ACCURATE/INSTANCE)cudnnRelu/cudnnSigmoid/cudnnTanhcudnnActivationForwardcudnnMaxPool2dcudnnPoolingForward(MAX)cudnnConv2dcudnnConvolutionForwardThese legacy functions are deprecated-since-cuDNN-9.0 (fully functional, validated on 9.23);
tornado-cudnn/ROADMAP.mdtracks their migration to graph-API patterns behind the same factories.Fused SDPA / flash attention (the flagship):
CuDnn.sdpaForward(q, k, v, o, b, h, sQ, sKv, d, scale, causal)— FP16HalfFloatArrays in packed BHSD layout, FP32 accumulation, inference, optional causal mask (d multiple of 8, ≤ 256, Ampere+). Implemented with the header-only cudnn-frontend v1.25.0 (NVIDIA-recommended over raw backend descriptors), fetched at build time; graph → heuristic → execution plan bundled per shape and created in the provider'sprepare()hook, so SDPA is safe with CUDA graph capture.Core fix:
TornadoTaskGraph.isArgumentIgnorablenow coversBoolean,HalfFloatArray, andByteArraytask arguments — aBooleantask argument previously crashedlockInPendingFieldsObjectswithIndexOutOfBoundsException(surfaced by the SDPAcausalflag; latent for the array types).Tests
uk.ac.manchester.tornado.unittests.cudnn.TestCuDnn— 9/9 viatornado-test(auto-skip[UNSUPPORTED]without CUDA backend /libtornado-cudnn): softmax/relu/tanh/maxpool/conv2d vs Java references; a CNN mixed pipeline (conv → JIT bias task → ReLU → max pool on shared device buffers); conv underwithCUDAGraph; SDPA plain and causal vs a Java attention reference. Cross-suite regression after the core fix: cuBLAS 11/11, cuBLASLt 5/5, cuFFT 8/8, TestArrays 23/23.Benchmarks (RTX 4090, cuDNN 9.23, cross-validated)
Convolution (
BenchmarkConv2d, FP32):Fused causal flash attention (
BenchmarkSdpa, FP16 vs the FP32 JIT attention kernel on identical FP16-rounded inputs):Docs
tornado-cudnn/README.md(supported ops, deprecation note, benchmark tables) andtornado-cudnn/ROADMAP.md(coverage tracker derived from the cuDNN docs: legacy completion, FP16 Tensor Core conv + autotuning, RMSNorm/LayerNorm, GQA/paged-KV attention, graph-API migration; FP8 items marked hardware-gated on Hopper).Backend/s tested
Verified on the CUDA backend (
make BACKEND=cuda), RTX 4090; suites skip cleanly on other backends.