QR matrix decomposition and solver for Linalg - #402
Draft
vaijira wants to merge 108 commits into
Draft
Conversation
Initial implementation for QR decomposition by givens rotation and by blocked accelerated householder reflector methods.
* Interpolate migration from burn Co-authored-by: Copilot <copilot@github.com> * Added interpolate benchmarks Co-authored-by: Copilot <copilot@github.com> * Fixed interpolate problems * Fixed interpolate lib lint * fix correctness * export interpolate category * adjust benchmark problems --------- Co-authored-by: Copilot <copilot@github.com> Co-authored-by: louisfd <louisfd94@gmail.com>
* add TiledLayout in cubek-test-utils * cargo fmt --all * remove dead comments * change test name and add simple implementation for layout * rename test name * refactor tests * use RowMajorLayout instead of manually calculating indexes * cargo fmt --all
* Update cubecl rev * Update rev again * Update version to 0.2.0 * Use cubecl commit until published * Use published cubecl
* Interpolate migration from burn Co-authored-by: Copilot <copilot@github.com> * Added interpolate benchmarks Co-authored-by: Copilot <copilot@github.com> * Fixed interpolate problems * Fixed interpolate lib lint * fix correctness * export interpolate category * adjust benchmark problems * add interpolation modes Co-authored-by: Copilot <copilot@github.com> * fix nearest backward algo * Added benchmarks * refactor interpolate benchmarks * fix dtype clone clippy * fix pr comments * fix clippy format * added main cublecl hash * update cubecl hash --------- Co-authored-by: Copilot <copilot@github.com> Co-authored-by: louisfd <louisfd94@gmail.com>
* Bump version to 0.3.0-pre.1 * Update cubecl rev
* migrate pool * add forward tests * added backward pool * add lib wrappers for pool * add backward references and tests * edit cubek toml * edit struct naming * fix pr comments * fix clippy lint
* Add support for f64 tracel-ai#198 * Remove if/else statements with function call. * Add From methods for HostDataType from StorageType
* tile size refactor and interpolate bicubic fix * fix cpu shared memory * resample readme * add comments * fix clippy
…speedup) nsys profiling at 2048x2048 showed 77% of GPU time in the panel factorization: householder_kernel ran on a single thread (serial O(rows) loop, 58us/launch) and apply_householder_kernel on at most 32 threads (one per panel column, 82us/launch). Changes to components/baht_tsqr.rs: - householder_kernel: now uses a full cube with a shared-memory tree reduction for sigma and a parallel normalization. It also writes the reflector directly into v_buf column j (explicit 1 at the head), which removes copy_v_to_buf_kernel and the v_tmp buffer, saving one dispatch per column. - apply_householder_kernel: one cube per remaining panel column, with a cube-wide tree-reduced dot product and parallel rank-1 update (82us -> ~3us per launch). - build_t_tsqr_kernel: rows within each sequential column are now computed in parallel across the cube. Matmul strategy notes (routines/baht_tsqr.rs): - Strategy::Auto was ~30% faster but broke f32 correctness (0.2-16% errors): Auto resolves to SimpleCyclicCmma (tensor cores) and cubek-matmul's adjust_dtypes silently downgrades f32 stage/register types to tf32. There is no full-precision f32 tensor-core path. f64 is unaffected because no f64 CMMA exists and Auto falls back to SimpleUnit. - f32 keeps full-precision unit strategies; strategy_tall upgraded from SimpleUnit to DoubleUnit (121ms vs 131ms at 2048x2048, tied at smaller sizes). Correctness: 14/14 CUDA tests pass (f32 and f64, tolerance 2e-3), plus QR reconstruction checks at 128, 512, 1024, 2048 square and 2048x512 in f32. Benchmarks (CUDA f32, RTX 4070 Laptop, mean of 10 samples): 128x128: 3.10ms -> 1.55ms (2.0x) 512x512: 26.84ms -> 8.87ms (3.0x) 1024x1024: 94.55ms -> 28.8ms (3.3x) 2048x2048: 407.5ms -> ~121ms (3.4x) 2048x512: 152.1ms -> 26.9ms (5.7x) Gap to PyTorch/cuSOLVER f32 shrank from 5-44x to 2.5-9x. Remaining bottleneck is the tall GEMMs (S = W^T*Q^T, Z = V*S) running on full-precision unit matmuls at ~1-2 TFLOPS. Also adds examples/profile_qr.rs, a standalone driver for profiling: cargo build --release -p cubek-linalg --example profile_qr --features cubecl/cuda nsys profile --stats=true target/release/examples/profile_qr [m] [n] [iters]
Review fixes ahead of the linalg PR: - Normalize any input layout on-device: initialize() now runs cubecl's into_contiguous on a transposed view of A (tight row-major [n,m] of A^T == tight col-major [m,n] bytes), replacing the host round-trip (read_one + create_from_slice) that also silently produced wrong Q/R for row-major or pitch-padded inputs. - Validate before launching: reject element-type mismatches between the generic entry points and the tensor dtype (new QRSetupError::TypeMismatch), and reject zero-column matrices (previously a divide-by-zero panic). - Propagate errors instead of panicking: baht_tsqr::launch returns Result and internal matmul setup failures surface as QRSetupError::Matmul. - Overflow-safe indexing: kernels compute flat offsets in usize (r_offset/v_offset params dropped; derived in-kernel from rows/col/j). - Scratch buffers use client.empty instead of TensorHandle::zeros — every region is fully written before read, saving 8 fill dispatches per call; shared-memory tree reduction deduplicated into a tree_reduce_sum cube helper; dead is_col_major flag removed from back-substitution. - Benchmark: col-major input (was factorizing byte-scrambled row-major data), reuse the stored client; profile_qr example gains a dtype arg (f32|f64) and a mean/median summary. - Tests: reuse cubek-test-utils (TestInput builder, HostData comparator), panic on launch errors instead of swallowing them, add row-major input tests and 4 launch-guard tests (20 total, pass on CPU and CUDA); tests/suite/qr/tsqr.rs renamed to baht_tsqr.rs to match the routine. - Deps: drop unused serde, half, pretty_assertions, benchmarks' rand; move num-traits to dev-dependencies. CUDA f32 timings (RTX 4070 Laptop, mean of 10): 128²: 1.49ms (was 1.54), 512²: 8.16ms (was 8.73), 1024²: 22.3ms (was 29.9), 2048x512: 20.5ms (was 27.8) — the removed host round-trip and fill dispatches account for the 25%+ gains at larger sizes.
vaijira
changed the title
Profiling baht_tsqr with nsys showed the Q^T update kernel at 2.7% of GPU kernel time on 1024x1024 and 6.3% on 2048x512. Its two operands (Q^T and Z) are both tight col-major [rows, rows] buffers indexed with the identical flat offset, so it was a plain contiguous elementwise add running scalar under a 2D launch whose second dimension bought nothing. Replace it with add_assign_kernel: Vector<F, N> operands and a 1D launch, with the line size chosen via io_optimized_vector_sizes plus a divisibility check (lands on 4). The kernel itself goes 17,517 -> 7,907 ns at 1024x1024 and 82,284 -> 48,689 ns at 2048x512. Total GPU kernel time drops 1.5% at 1024x1024 and 2.5% at 2048x512. update_trailing_r is deliberately left scalar and 2D. Its operands are transposed relative to each other (R col-major, Z row-major), and declaring Z col-major to fix that moves 124 matmul instances from acc_size_4 to acc_size_1 -- exactly the z_r matmul count -- because it de-vectorizes the matmul's output write. Measured net GPU time was worse than leaving it alone at both shapes, so the reasoning is recorded in a comment instead. Verified on CUDA and wgpu (RTX 4070): 20/20 tests pass on both, f32 and f64.
build_t_tsqr_kernel costs a fixed ~22.9 us per tile regardless of matrix shape (stddev 20 ns across 128 instances at 1024x1024, identical timing at 2048x512) for what is only ~5.5k FMAs. It builds T column by column, and each of the `tile` serialized steps reads back columns written by earlier steps, so building T directly in global memory put a global round trip inside every step of the dependency chain. Stage T and the Gram matrix in shared memory instead and write T out once at the end. The working set is tile * tile elements (4 KiB at f32/tile=32), so it fits easily. `tile` becomes a comptime parameter to size the shared slices. Kernel goes 22,896 -> 20,595 ns, from 3.5% to 3.2% of GPU kernel time at 1024x1024 -- roughly 0.35% end to end, which is below wall-clock noise, so the nsys kernel-time figure is the measurement of record.
Profiling showed the matmuls at ~70% of GPU time, and within them only the trailing-update GEMMs matter. Forcing strategy_tall to Strategy::Auto takes 1024x1024 from 23.5ms to 12.0ms, while forcing the Gram or W matmul to Auto changes nothing end to end. Auto resolves to a tensor-core path that rounds f32 inputs to tf32, which is why it was ruled out wholesale before. Expose it as a knob instead of leaving the speed on the table: - BahtTsqrStrategy gains `allow_tf32` (off by default), mapped into the blueprint and applied to strategy_tall only. Gram and W stay full-precision unconditionally since they cost nothing either way. - New `qr_with_strategy()` entry point; `qr()` delegates with defaults, so the existing API and its accuracy are unchanged. - profile_qr takes an optional 5th arg `tf32` so both paths are profilable. The flag is gated on `supports_type(tf32)`. Without a tensor-core path Strategy::Auto resolves to something *slower* than the full-precision unit routine -- measured on wgpu, 24.1ms -> 30.5ms at 1024x1024 -- so honouring the opt-in there would cost accuracy and speed. Post-gate wgpu measures 21.30/21.56/21.35ms with the flag against 21.30/21.24ms without: a genuine no-op. Also a no-op for f64 (112.7 vs 112.4ms), which has no tf32 path. CUDA f32 medians (RTX 4070 Laptop), full -> tf32: 128x128: 1.34ms -> 1.26ms (1.06x) 512x512: 7.63ms -> 5.31ms (1.44x) 1024x1024: 22.84ms -> 11.90ms (1.92x) 2048x512: 22.01ms -> 11.40ms (1.93x) 2048x2048: 80.22ms -> 48.68ms (1.65x) The gain scales with size -- small problems are dominated by panel dispatch overhead, not matmul -- so the flag is worth enabling for large factorizations and roughly pointless below ~256. Gap to PyTorch/cuSOLVER f32 on the same GPU narrows from 2.2-6.8x to 2.1-4.1x. (cuSOLVER uses tensor cores for f32 too, so the tf32 column is the closer like-for-like comparison.) Tests: adds test_tf32_opt_in (f32 and f64) driving qr_with_strategy at 157x157, asserted at 5e-2 -- the tolerance matching the documented ~1e-2 accuracy rather than pretending it is f32-accurate. 22/22 pass on CUDA and wgpu. Also fixes two unnecessary-cast clippy warnings introduced by the earlier vec_split helper.
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
13 participants
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.
Implementation of QR algorithm based on blocked accelerated householder taking advantage of matmul. There is a solver too.
When calling matmul we are not using most optimized strategies because when using tf32 precision decreases. But with Strategy::auto in matmuls the gain is around 40%.
Validate your PR with burn.
It is important that you make sure that you don't introduce any bugs in burn.
Instructions