feat: fp and cuda-graph support#346
Open
devillove084 wants to merge 9 commits into
Open
Conversation
These storage-only types map to i8/i16 in LLVM IR since pliron lacks dedicated types for them. FP8 and BF16 are used by Blackwell tensor cores (tcgen05) and device-extern FFI declarations.
Introduces CudaGraph (RAII wrapper for CUgraph), CudaGraphExec (instantiated executable graph), CaptureMode, CaptureModeGuard, and CudaStreamCaptureExt for stream-based graph capture. Graph tests require serial execution due to shared primary context; run with --test-threads=1.
- Tcgen05ElementType gains E4M3 and E5M2 variants for FP8 tensor cores - New tcgen05_mma_ws_e4m3 / tcgen05_mma_ws_e5m2 device intrinsics - CudaGraph: add_kernel_node, add_empty_node, add_dependencies - CudaGraphNode: RAII wrapper for CUgraphNode - KernelNodeParams: safe kernel param builder with push_param - CudaGraphExec: update() and set_kernel_node_params() - 19 graph integration tests (run with --test-threads=1)
- dialect-nvvm: Tcgen05MmaWsE4M3Op, Tcgen05MmaWsE5M2Op - mir-importer: emit_tcgen05_mma_ws_e4m3, emit_tcgen05_mma_ws_e5m2 - mir-lower: PTX lowering via convert_mma_ws with kind::f16 (FP8 element type is encoded in the idesc, not PTX mnemonic)
Inspired by llama.cpp's CUDA graph pattern: - GraphStrategy: Direct, AutoCapture, Prebuilt - AutoCapture warmup: first run direct, second triggers capture - Property-change detection via FNV hash of data ptrs + shapes - CachedGraphExec::launch_or_capture() single entry point
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 FP8 (E4M3/E5M2), FP6 (E3M2/E2M3), and FP4 (E2M1) type support through the full compiler pipeline (types → dialect → MIR translation → PTX lowering), plus a complete CUDA Graph API with stream capture, manual node construction, executable graph update, and an auto-warmup strategy inspired by llama.cpp. This is foundational work for Blackwell (sm_100a+) tensor core operations and low-overhead kernel launch.
Changes
Sub-byte float type system (4 crates)
DeviceExternTypegainsBFloat16,Float8E4M3,Float8E5M2,Float4E2M1,Float6E3M2,Float6E2M3— all sub-byte float types map toi8/i16in LLVM IR since no native types exist in pliron/LLVMmodule.rs,ops.rs,lower.rs,export_test.rsFP8/FP6/FP4 tcgen05 device intrinsics + compiler pipeline (5 crates)
Tcgen05ElementTypeextended withE4M3,E5M2,E2M3,E3M2,E2M1+idesc_value()method for kind-dependent idesc encoding (enum discriminants are identifiers;build()maps to correct PTX ISA bit patterns per kind family)tcgen05_mma_ws_e4m3/e5m2/e2m3/e3m2/e2m1(), plusnew_e4m3()/new_e5m2()/new_e2m3()/new_e3m2()/new_e2m1()descriptor buildersTcgen05MmaWsE2M3Op,Tcgen05MmaWsE3M2Op,Tcgen05MmaWsE2M1Opregistered indialect-nvvmemit_tcgen05_mma_ws_e2m3/e3m2/e2m1()inmir-importer+ dispatch entries for all f8f6f4 typeskind::f8f6f4; F16/BF16 usekind::f16; TF32 useskind::tf32— the element type is encoded in the instruction descriptor (idesc), not the PTX mnemonicCUDA Graph API (1 crate, 1 new module)
CudaStreamCaptureExttrait —begin_capture,end_capture,is_capturingCudaGraph(RAII forCUgraph),CudaGraphExec(instantiate/launch/upload/destroy)CudaGraphNode,add_kernel_node,add_empty_node,add_dependenciesKernelNodeParamswith alignedpush_param<T>()builderCudaGraphExec::update()returningGraphUpdateResult,set_kernel_node_params()for hot param updatesCachedGraphExecwithGraphStrategy(Direct / AutoCapture / Prebuilt), FNV-hash-based property change detection, two-phase warmup protocolCaptureModeGuardRAII forcuThreadExchangeStreamCaptureModeTesting
22 new tests added across 3 test suites. The
test_scalar_vs_tensor_core_fp4_instruction_counttest demonstrates that a singletcgen05.mma.ws.kind::f8f6f4instruction replaces ~8.4M scalar instructions for a 128×256×32 FP4 tile (~1,000,000× speedup).cargo test --workspacepassescargo oxide run <example>— N/A (library-level additions, no new examples)