Skip to content

[hybrid] Hybrid API: native library tasks with cuDNN provider#888

Merged
mikepapadim merged 2 commits into
beehive-lab:developfrom
mikepapadim:hybrid-cudnn
Jul 3, 2026
Merged

[hybrid] Hybrid API: native library tasks with cuDNN provider#888
mikepapadim merged 2 commits into
beehive-lab:developfrom
mikepapadim:hybrid-cudnn

Conversation

@mikepapadim

@mikepapadim mikepapadim commented Jul 2, 2026

Copy link
Copy Markdown
Member

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.

Stacked on #887 (this branch includes its commits: the library-task SPI, cuBLAS/cuBLASLt, and cuFFT providers). Suggest reviewing #887 first and re-targeting this PR's base to develop for consistency.

What's included

Legacy-API primitives (FP32, NCHW)tornado-cudnn + tornado-drivers/cudnn-jni (requires cuDNN 9):

Factory cuDNN call Semantics
cudnnSoftmax(in, out, rows, cols) cudnnSoftmaxForward (ACCURATE/INSTANCE) numerically stable per-row softmax
cudnnRelu / cudnnSigmoid / cudnnTanh cudnnActivationForward element-wise activations
cudnnMaxPool2d cudnnPoolingForward (MAX) square window/stride
cudnnConv2d cudnnConvolutionForward 2D cross-correlation, NCHW/KCRS; per-shape plan (descriptors + algorithm + workspace size) with a grow-only device workspace

These legacy functions are deprecated-since-cuDNN-9.0 (fully functional, validated on 9.23); tornado-cudnn/ROADMAP.md tracks 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) — FP16 HalfFloatArrays 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's prepare() hook, so SDPA is safe with CUDA graph capture.

Core fix: TornadoTaskGraph.isArgumentIgnorable now covers Boolean, HalfFloatArray, and ByteArray task arguments — a Boolean task argument previously crashed lockInPendingFieldsObjects with IndexOutOfBoundsException (surfaced by the SDPA causal flag; latent for the array types).

Tests

uk.ac.manchester.tornado.unittests.cudnn.TestCuDnn9/9 via tornado-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 under withCUDAGraph; 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):

Layer (NCHW, 3×3, pad 1) TornadoVM JIT direct conv cuDNN Speedup
8×64×56×56, k=64 3.8 TFLOP/s 11.6 TFLOP/s 3.1×
16×128×28×28, k=128 4.0 TFLOP/s 21.1 TFLOP/s 5.3×

Fused causal flash attention (BenchmarkSdpa, FP16 vs the FP32 JIT attention kernel on identical FP16-rounded inputs):

Shape (b, h, s, d) JIT attention kernel cuDNN fused SDPA Speedup
4, 16, 1024, 128 90.7 ms 0.268 ms (64.2 TFLOP/s) 339×
1, 32, 2048, 128 203.3 ms 0.407 ms (84.3 TFLOP/s) 499×

Docs

tornado-cudnn/README.md (supported ops, deprecation note, benchmark tables) and tornado-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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +95 to +109
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();
}
}
}
Comment on lines +109 to +124
@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;
}
Comment on lines +81 to +82
final int n = (int) invocation.getArg(2);
final int batch = (int) invocation.getArg(3);
Comment on lines +111 to +115
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);
Comment on lines +116 to +120
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}")
Comment thread tornado-cudnn/pom.xml
</parent>
<artifactId>tornado-cudnn</artifactId>
<name>tornado-cudnn</name>
<description>TornadoVM cuDNN bindings: NVIDIA cuFFT library tasks for TornadoVM task graphs</description>
@mikepapadim mikepapadim changed the title [hybrid] Hybrid API: native library tasks with cuddn provider [hybrid] Hybrid API: native library tasks with cuDNN provider Jul 3, 2026
@mikepapadim
mikepapadim marked this pull request as ready for review July 3, 2026 09:30
@stratika
stratika changed the base branch from master to develop July 3, 2026 10:26
@stratika
stratika marked this pull request as draft July 3, 2026 15:25
@stratika
stratika marked this pull request as ready for review July 3, 2026 15:25
@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

⚠️ Could not automatically backport this PR to jdk25.

Cherry-picking commit 4c67d0b64b865870b82045f1cc92ae64a09b66e8 onto jdk25 produced conflicts, so no jdk25 PR was created. Please backport these changes manually.

@stratika stratika closed this Jul 3, 2026
@stratika stratika reopened this Jul 3, 2026
@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

⚠️ Could not automatically backport this PR to jdk25.

Cherry-picking commit 4c67d0b64b865870b82045f1cc92ae64a09b66e8 onto jdk25 produced conflicts, so no jdk25 PR was created. Please backport these changes manually.

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>
@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

⚠️ Could not automatically backport this PR to jdk25.

Cherry-picking commit 4c67d0b64b865870b82045f1cc92ae64a09b66e8 onto jdk25 produced conflicts, so no jdk25 PR was created. Please backport these changes manually.

Copilot AI review requested due to automatic review settings July 3, 2026 16:10
@github-actions

github-actions Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

⚠️ Could not automatically backport this PR to jdk25.

Cherry-picking commit e4220c2fb05a1b8cb4e1cdfd407359c9e06ba92a onto jdk25 produced conflicts, so no jdk25 PR was created. Please backport these changes manually.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated 2 comments.

Comment on lines +128 to +132
/** (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). */
@mikepapadim
mikepapadim merged commit 3cb4c0a into beehive-lab:develop Jul 3, 2026
3 checks passed
@yrq0208 yrq0208 moved this from Needs backport to Backport PR Open in JDK PR Backport Tracking Project Jul 6, 2026
@stratika stratika mentioned this pull request Jul 6, 2026
10 tasks
@stratika stratika moved this from Backport PR Open to Backporting in JDK PR Backport Tracking Project Jul 6, 2026
@stratika stratika moved this from Backporting to Merged to JDK25 in JDK PR Backport Tracking Project Jul 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CUDA enhancement New feature or request

Projects

Status: Merged to JDK25

Development

Successfully merging this pull request may close these issues.

4 participants