diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index ab8db6433..10846d548 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -160,7 +160,7 @@ jobs: ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }} FUSOR_CONFORMANCE_REQUIRE_GPU: "1" - run: cargo test --lib --bins --tests --examples --workspace --features kalosm/language,kalosm/sound,kalosm/vision,kalosm/remote,kalosm/scrape --exclude fusor-core --exclude rbert --exclude rwhisper + run: cargo test --lib --bins --tests --examples --workspace --features kalosm/language,kalosm/sound,kalosm/vision,kalosm/remote,kalosm/scrape --exclude rbert --exclude rwhisper doc: if: github.event.pull_request.draft == false diff --git a/.gitignore b/.gitignore index 58bdb0970..2da34edc3 100644 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,4 @@ rust-analyzer /.claude **/.claude/ profiles +/fusor-ml/fusor/examples/data diff --git a/Cargo.lock b/Cargo.lock index 79ac05f6b..2f1e0e16c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4207,6 +4207,26 @@ dependencies = [ "rustfft", ] +[[package]] +name = "egg" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd40cfd4196d7a8f882ace95d623d4d6734588e502b4e188675a7c2f55eb4fb4" +dependencies = [ + "env_logger", + "hashbrown 0.15.5", + "indexmap 2.14.0", + "log", + "num-bigint", + "num-traits", + "quanta", + "rustc-hash 2.1.2", + "smallvec 1.15.1", + "symbol_table", + "symbolic_expressions", + "thiserror 1.0.69", +] + [[package]] name = "ego-tree" version = "0.6.3" @@ -4371,6 +4391,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "env_logger" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a12e6657c4c97ebab115a42dcee77225f7f482cdd841cf7088c657a42e9e00e7" +dependencies = [ + "log", +] + [[package]] name = "equator" version = "0.4.2" @@ -4404,7 +4433,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -4674,11 +4703,13 @@ dependencies = [ "candle-nn", "candle-transformers", "criterion 0.8.2", + "flate2", "fusor-core", "fusor-cpu", "fusor-gguf", "fusor-types", "futures", + "futures-channel", "half", "kalosm-common", "kalosm-model-types", @@ -4688,6 +4719,9 @@ dependencies = [ "rand 0.9.4", "tokio", "tracing", + "tracing-subscriber 0.3.23", + "ureq 3.3.0", + "wgpu", ] [[package]] @@ -4723,6 +4757,7 @@ name = "fusor-core" version = "0.1.0" dependencies = [ "bytemuck", + "egg", "fusor-gguf", "fusor-tile-ir", "fusor-tile-ir-kernels", @@ -4803,10 +4838,14 @@ dependencies = [ name = "fusor-tile-ir-runtime" version = "0.1.0" dependencies = [ + "bincode 1.3.3", "fusor-tile-ir", + "libc", "lru 0.14.0", + "naga", "parking_lot", "rustc-hash 2.1.2", + "serde", "tracing", "wgpu", ] @@ -5351,7 +5390,7 @@ dependencies = [ "log", "presser", "thiserror 2.0.18", - "windows 0.62.2", + "windows 0.54.0", ] [[package]] @@ -6203,7 +6242,7 @@ checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" dependencies = [ "hermit-abi 0.5.2", "libc", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -7339,7 +7378,6 @@ dependencies = [ [[package]] name = "naga" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "arrayvec", "bit-set 0.9.1", @@ -7356,6 +7394,7 @@ dependencies = [ "num-traits", "once_cell", "rustc-hash 1.1.0", + "serde", "spirv", "thiserror 2.0.18", "unicode-ident", @@ -8671,6 +8710,21 @@ dependencies = [ "bytemuck", ] +[[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi", + "web-sys", + "winapi", +] + [[package]] name = "quick-error" version = "2.0.1" @@ -9622,7 +9676,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -10687,6 +10741,23 @@ dependencies = [ "vart 0.9.3", ] +[[package]] +name = "symbol_table" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f19bffd69fb182e684d14e3c71d04c0ef33d1641ac0b9e81c712c734e83703bc" +dependencies = [ + "crossbeam-utils", + "foldhash 0.1.5", + "hashbrown 0.15.5", +] + +[[package]] +name = "symbolic_expressions" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c68d531d83ec6c531150584c42a4290911964d5f0d79132b193b67252a23b71" + [[package]] name = "symphonia" version = "0.5.5" @@ -10932,7 +11003,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -12297,7 +12368,6 @@ checksum = "a28ac98ddc8b9274cb41bb4d9d4d5c425b6020c50c46f25559911905610b4a88" [[package]] name = "wgpu" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "arrayvec", "bitflags 2.13.0", @@ -12326,7 +12396,6 @@ dependencies = [ [[package]] name = "wgpu-core" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "arrayvec", "bit-set 0.9.1", @@ -12358,7 +12427,6 @@ dependencies = [ [[package]] name = "wgpu-core-deps-apple" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "wgpu-hal", ] @@ -12366,7 +12434,6 @@ dependencies = [ [[package]] name = "wgpu-core-deps-emscripten" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "wgpu-hal", ] @@ -12374,7 +12441,6 @@ dependencies = [ [[package]] name = "wgpu-core-deps-windows-linux-android" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "wgpu-hal", ] @@ -12382,7 +12448,6 @@ dependencies = [ [[package]] name = "wgpu-hal" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "android_system_properties", "arrayvec", @@ -12435,7 +12500,6 @@ dependencies = [ [[package]] name = "wgpu-naga-bridge" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "naga", "wgpu-types", @@ -12444,7 +12508,6 @@ dependencies = [ [[package]] name = "wgpu-types" version = "29.0.3" -source = "git+https://github.com/ealmloff/wgpu?branch=yield-now#477d335af5f2bd3f5c25017db67ac7062e53f69a" dependencies = [ "bitflags 2.13.0", "bytemuck", diff --git a/Cargo.toml b/Cargo.toml index 3c868ee93..e35dc013f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,7 +44,9 @@ members = [ ] [patch.crates-io] -wgpu = { git = "https://github.com/ealmloff/wgpu", branch = "yield-now" } +# Local checkout of the wgpu fork while the mixed-precision cooperative-MMA +# extension bakes (f16 operands with f32 accumulators in the MSL writer). +wgpu = { path = "../wgpu-fusor/wgpu" } [workspace.dependencies] kalosm = { path = "./interfaces/kalosm", version = "0.4.0", default-features = false } @@ -390,3 +392,7 @@ codegen-units = 1 [profile.dev.package."*"] opt-level = 3 + +[patch."https://github.com/ealmloff/wgpu"] +naga = { path = "../wgpu-fusor/naga" } +wgpu = { path = "../wgpu-fusor/wgpu" } diff --git a/fusor-ml/conformance/goldens/attention_gqa_causal.txt b/fusor-ml/conformance/goldens/attention_gqa_causal.txt new file mode 100644 index 000000000..f26b3d953 --- /dev/null +++ b/fusor-ml/conformance/goldens/attention_gqa_causal.txt @@ -0,0 +1,4 @@ +out 0xa45dec1e98442bd9 +dq 0x9d3eefd367974641 +dk 0x712914ed2040211b +dv 0xeedffc029ea157a5 diff --git a/fusor-ml/conformance/goldens/bilstm_trace.txt b/fusor-ml/conformance/goldens/bilstm_trace.txt new file mode 100644 index 000000000..8a80e7e08 --- /dev/null +++ b/fusor-ml/conformance/goldens/bilstm_trace.txt @@ -0,0 +1 @@ +out 0xab42f2b1f7f2e4b7 diff --git a/fusor-ml/conformance/goldens/decode_trace_dispatch.txt b/fusor-ml/conformance/goldens/decode_trace_dispatch.txt new file mode 100644 index 000000000..0d339623f --- /dev/null +++ b/fusor-ml/conformance/goldens/decode_trace_dispatch.txt @@ -0,0 +1,5 @@ +399 [("flash_attention", 6), ("matmul_f32", 19), ("merged_matmul", 6), ("nary_direct", 33), ("row_program", 13)] +1 [("flash_attention", 6), ("matmul_f32", 19), ("merged_matmul", 6), ("nary_direct", 32), ("row_program", 13)] +1 [("flash_attention", 6), ("matmul_f32", 19), ("merged_matmul", 6), ("merged_region", 1), ("merged_row", 1), ("nary_direct", 35), ("row_program", 15)] +1 [("flash_attention", 24), ("matmul_f32", 59), ("merged_matmul", 18), ("merged_region", 40), ("merged_row", 32), ("nary_direct", 114), ("row_program", 27)] +1 [("flash_attention", 24), ("matmul_f32", 59), ("merged_matmul", 18), ("merged_region", 39), ("merged_row", 26), ("nary_direct", 114), ("row_program", 27)] diff --git a/fusor-ml/conformance/goldens/kernel_bench_dispatch.txt b/fusor-ml/conformance/goldens/kernel_bench_dispatch.txt new file mode 100644 index 000000000..72a97417e --- /dev/null +++ b/fusor-ml/conformance/goldens/kernel_bench_dispatch.txt @@ -0,0 +1,9 @@ +wgrad merged_matmul 2 +wgrad256 merged_matmul 2 +wgrad256m merged_matmul 2 +fwd merged_matmul 2 +fwd256 merged_matmul 2 +fwdup merged_matmul 2 +attn merged_matmul 2 +softmax merged_row 1 +softmax row_program 1 diff --git a/fusor-ml/conformance/goldens/qgemv_decode_ggml.txt b/fusor-ml/conformance/goldens/qgemv_decode_ggml.txt new file mode 100644 index 000000000..b16cf13ea --- /dev/null +++ b/fusor-ml/conformance/goldens/qgemv_decode_ggml.txt @@ -0,0 +1,4 @@ +q4k_4096x8192 0x7f4870e80e75707a +q4k_4096x8193 0x92adc0e985bb5224 +q4k_4096x5120 0x253507257e08480d +q6k_4096x8192 0x5394b93db2f8ac10 diff --git a/fusor-ml/conformance/src/bench/burn.rs b/fusor-ml/conformance/src/bench/burn.rs index f2bd3e057..e3f81d5d6 100644 --- a/fusor-ml/conformance/src/bench/burn.rs +++ b/fusor-ml/conformance/src/bench/burn.rs @@ -1,8 +1,12 @@ //! Burn WGPU benchmark cases that mirror Fusor WebGPU cases where possible. +//! +//! Like `webgpu.rs`, each case body is parameterized by the sizes that vary +//! between the fixed registry entry and the per-size sweep, and the +//! `fixed_cases!` invocation at the bottom pins the registry sizes. use burn::{ backend::{Wgpu, wgpu::WgpuDevice}, - nn::{RmsNormConfig, RotaryEncodingConfig}, + nn::{LayerNormConfig, RmsNormConfig, RotaryEncodingConfig}, tensor::{ Tensor, TensorData, activation, module, ops::{AttentionModuleOptions, ConvOptions}, @@ -87,806 +91,565 @@ async fn materialize(tensor: BurnTensor) -> BenchmarkResult<( Ok(()) } -async fn materialize_inputs(inputs: &[BurnTensor]) -> BenchmarkResult<()> { - for input in inputs { - materialize(input.clone()).await?; - } - Ok(()) +async fn values_input( + device: &WgpuDevice, + shape: [usize; R], + values: Vec, +) -> BenchmarkResult> { + let tensor = burn_tensor(values, shape, device); + materialize(tensor.clone()).await?; + Ok(tensor) +} + +async fn input_tensor( + device: &WgpuDevice, + shape: [usize; R], + seed: usize, + scale: f32, +) -> BenchmarkResult> { + values_input( + device, + shape, + deterministic_values(elements(&shape), seed, scale), + ) + .await } fn bench_case( name: &'static str, - run: impl for<'a> FnOnce(&'a Device, BenchmarkConfig) -> super::CaseFuture<'a> + 'static, + run: impl FnOnce(BenchmarkConfig, String) -> super::CaseFuture<'static> + 'static, ) -> BenchmarkCase { - BenchmarkCase::new(name, run) -} - -pub fn elementwise_add_square() -> BenchmarkCase { - bench_case("burn::elementwise_add_square", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [512usize, 512usize]; - let lhs = burn_tensor( - deterministic_values(elements(&shape), 1, 0.01), - shape, - &device, - ); - let rhs = burn_tensor( - deterministic_values(elements(&shape), 2, 0.008), - shape, - &device, - ); - materialize_inputs(&[lhs.clone(), rhs.clone()]).await?; - - let samples = time_samples(config, || { - let output = lhs.clone() + rhs.clone(); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::elementwise_add_square", - config, - samples, - format!("{} f32 add", shape_label(&shape)), - )) - }) + BenchmarkCase::new(name, move |_fusor_device: &Device, config| { + run(config, name.to_string()) }) } -pub fn elementwise_mul_rank4() -> BenchmarkCase { - bench_case("burn::elementwise_mul_rank4", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [9usize, 11usize, 32usize, 16usize]; - let lhs = burn_tensor( - deterministic_values(elements(&shape), 3, 0.012), - shape, - &device, - ); - let rhs = burn_tensor( - deterministic_values(elements(&shape), 4, 0.009), - shape, - &device, - ); - materialize_inputs(&[lhs.clone(), rhs.clone()]).await?; - - let samples = time_samples(config, || { - let output = lhs.clone() * rhs.clone(); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::elementwise_mul_rank4", - config, - samples, - format!("{} f32 mul", shape_label(&shape)), - )) - }) +pub(super) async fn elementwise_add_square_case( + config: BenchmarkConfig, + name: String, + size: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let shape = [size, size]; + let lhs = input_tensor(&device, shape, 1, 0.01).await?; + let rhs = input_tensor(&device, shape, 2, 0.008).await?; + let samples = time_samples(config, || { + let output = lhs.clone() + rhs.clone(); + async move { materialize(output).await } }) -} - -pub fn unary_trig_chain() -> BenchmarkCase { - bench_case("burn::unary_trig_chain", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [384usize, 384usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 10, 0.01), - shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = input.clone().sin() + input.clone().cos(); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::unary_trig_chain", - config, - samples, - format!("{} sin+cos", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} f32 add", shape_label(&shape)), + )) +} + +pub(super) async fn elementwise_mul_rank4_case( + config: BenchmarkConfig, + name: String, + shape: [usize; 4], +) -> BenchmarkResult { + let device = initialized_device().await; + let lhs = input_tensor(&device, shape, 3, 0.012).await?; + let rhs = input_tensor(&device, shape, 4, 0.009).await?; + let samples = time_samples(config, || { + let output = lhs.clone() * rhs.clone(); + async move { materialize(output).await } }) -} - -pub fn activation_gelu() -> BenchmarkCase { - bench_case("burn::activation_gelu", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [512usize, 256usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 11, 0.015), - shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = activation::gelu(input.clone()); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::activation_gelu", - config, - samples, - format!("{} gelu", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} f32 mul", shape_label(&shape)), + )) +} + +pub(super) async fn unary_trig_chain_case( + config: BenchmarkConfig, + name: String, + size: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let shape = [size, size]; + let input = input_tensor(&device, shape, 10, 0.01).await?; + let samples = time_samples(config, || { + let output = input.clone().sin() + input.clone().cos(); + async move { materialize(output).await } }) -} - -pub fn broadcast_add() -> BenchmarkCase { - bench_case("burn::broadcast_add", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let matrix_shape = [256usize, 512usize]; - let vector_shape = [512usize]; - let matrix = burn_tensor( - deterministic_values(elements(&matrix_shape), 12, 0.006), - matrix_shape, - &device, - ); - let vector = burn_tensor( - deterministic_values(elements(&vector_shape), 13, 0.01), - vector_shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&matrix)).await?; - materialize_inputs(std::slice::from_ref(&vector)).await?; - - let samples = time_samples(config, || { - let output = matrix.clone() + vector.clone().reshape([1, vector_shape[0]]); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::broadcast_add", - config, - samples, - "256x512 + broadcast 512", - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} sin+cos", shape_label(&shape)), + )) +} + +pub(super) async fn activation_gelu_case( + config: BenchmarkConfig, + name: String, + shape: [usize; 2], +) -> BenchmarkResult { + let device = initialized_device().await; + let input = input_tensor(&device, shape, 11, 0.015).await?; + let samples = time_samples(config, || { + let output = activation::gelu(input.clone()); + async move { materialize(output).await } }) -} - -pub fn transpose_then_elementwise() -> BenchmarkCase { - bench_case( - "burn::transpose_then_elementwise", - |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [256usize, 384usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 14, 0.01), - shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let transposed = input.clone().transpose(); - let output = transposed.clone() * transposed; - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::transpose_then_elementwise", - config, - samples, - "256x384 transpose, square", - )) - }) - }, - ) -} - -pub fn reduction_sum_last_dim() -> BenchmarkCase { - bench_case("burn::reduction_sum_last_dim", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [256usize, 512usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 15, 0.004), - shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = input.clone().sum_dim(1); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::reduction_sum_last_dim", - config, - samples, - format!("{} sum axis 1", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} gelu", shape_label(&shape)), + )) +} + +pub(super) async fn broadcast_add_case( + config: BenchmarkConfig, + name: String, + rows: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let matrix_shape = [rows, 512usize]; + let vector_shape = [512usize]; + let matrix = input_tensor(&device, matrix_shape, 12, 0.006).await?; + let vector = input_tensor(&device, vector_shape, 13, 0.01).await?; + let samples = time_samples(config, || { + let output = matrix.clone() + vector.clone().reshape([1, vector_shape[0]]); + async move { materialize(output).await } }) -} - -pub fn reduction_max_middle_axis() -> BenchmarkCase { - bench_case( - "burn::reduction_max_middle_axis", - |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [64usize, 128usize, 64usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 16, 0.004), - shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = input.clone().max_dim(1); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::reduction_max_middle_axis", - config, - samples, - format!("{} max axis 1", shape_label(&shape)), - )) - }) - }, - ) -} - -pub fn softmax_last_dim() -> BenchmarkCase { - bench_case("burn::softmax_last_dim", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [512usize, 256usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 5, 0.006), - shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = activation::softmax(input.clone(), 1); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::softmax_last_dim", - config, - samples, - format!("{} last-axis softmax", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} + broadcast 512", shape_label(&matrix_shape)), + )) +} + +pub(super) async fn transpose_then_elementwise_case( + config: BenchmarkConfig, + name: String, + shape: [usize; 2], +) -> BenchmarkResult { + let device = initialized_device().await; + let input = input_tensor(&device, shape, 14, 0.01).await?; + let samples = time_samples(config, || { + let transposed = input.clone().transpose(); + let output = transposed.clone() * transposed; + async move { materialize(output).await } }) -} - -pub fn softmax_middle_axis() -> BenchmarkCase { - bench_case("burn::softmax_middle_axis", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [32usize, 128usize, 64usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 17, 0.004), - shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = activation::softmax(input.clone(), 1); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::softmax_middle_axis", - config, - samples, - format!("{} softmax axis 1", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} transpose, square", shape_label(&shape)), + )) +} + +pub(super) async fn reduction_sum_last_dim_case( + config: BenchmarkConfig, + name: String, + rows: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let shape = [rows, 512usize]; + let input = input_tensor(&device, shape, 15, 0.004).await?; + let samples = time_samples(config, || { + let output = input.clone().sum_dim(1); + async move { materialize(output).await } }) -} - -pub fn layer_norm_last_dim() -> BenchmarkCase { - bench_case("burn::layer_norm_last_dim", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [8usize, 128usize, 512usize]; - let last_dim = shape[2]; - let input = burn_tensor( - deterministic_values(elements(&shape), 18, 0.01), - shape, - &device, - ); - let layer = burn::nn::LayerNormConfig::new(last_dim) - .with_epsilon(1.0e-5) - .init::(&device); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = layer.clone().forward(input.clone()); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::layer_norm_last_dim", - config, - samples, - format!("{} layer norm", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} sum axis 1", shape_label(&shape)), + )) +} + +pub(super) async fn reduction_max_middle_axis_case( + config: BenchmarkConfig, + name: String, + shape: [usize; 3], +) -> BenchmarkResult { + let device = initialized_device().await; + let input = input_tensor(&device, shape, 16, 0.004).await?; + let samples = time_samples(config, || { + let output = input.clone().max_dim(1); + async move { materialize(output).await } }) -} - -pub fn rms_norm_fused() -> BenchmarkCase { - bench_case("burn::rms_norm_fused", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [8usize, 128usize, 512usize]; - let last_dim = shape[2]; - let input = burn_tensor( - deterministic_values(elements(&shape), 21, 0.01), - shape, - &device, - ); - let rms = RmsNormConfig::new(last_dim) - .with_epsilon(1.0e-5) - .init::(&device); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = rms.clone().forward(input.clone()); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::rms_norm_fused", - config, - samples, - format!("{} rms norm", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} max axis 1", shape_label(&shape)), + )) +} + +pub(super) async fn softmax_last_dim_case( + config: BenchmarkConfig, + name: String, + rows: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let shape = [rows, 256usize]; + let input = input_tensor(&device, shape, 5, 0.006).await?; + let samples = time_samples(config, || { + let output = activation::softmax(input.clone(), 1); + async move { materialize(output).await } }) -} - -pub fn dense_matmul_square() -> BenchmarkCase { - bench_case("burn::dense_matmul_square", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let lhs_shape = [256usize, 256usize]; - let rhs_shape = [256usize, 256usize]; - let lhs = burn_tensor( - deterministic_values(elements(&lhs_shape), 6, 0.004), - lhs_shape, - &device, - ); - let rhs = burn_tensor( - deterministic_values(elements(&rhs_shape), 7, 0.004), - rhs_shape, - &device, - ); - materialize_inputs(&[lhs.clone(), rhs.clone()]).await?; - - let samples = time_samples(config, || { - let output = lhs.clone().matmul(rhs.clone()); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::dense_matmul_square", - config, - samples, - "256x256 @ 256x256 f32", - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} last-axis softmax", shape_label(&shape)), + )) +} + +pub(super) async fn softmax_middle_axis_case( + config: BenchmarkConfig, + name: String, + shape: [usize; 3], +) -> BenchmarkResult { + let device = initialized_device().await; + let input = input_tensor(&device, shape, 17, 0.004).await?; + let samples = time_samples(config, || { + let output = activation::softmax(input.clone(), 1); + async move { materialize(output).await } }) -} - -pub fn dense_batched_matmul() -> BenchmarkCase { - bench_case("burn::dense_batched_matmul", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let lhs_shape = [8usize, 64usize, 96usize]; - let rhs_shape = [8usize, 96usize, 64usize]; - let lhs = burn_tensor( - deterministic_values(elements(&lhs_shape), 23, 0.004), - lhs_shape, - &device, - ); - let rhs = burn_tensor( - deterministic_values(elements(&rhs_shape), 24, 0.004), - rhs_shape, - &device, - ); - materialize_inputs(&[lhs.clone(), rhs.clone()]).await?; - - let samples = time_samples(config, || { - let output = lhs.clone().matmul(rhs.clone()); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::dense_batched_matmul", - config, - samples, - "8x64x96 @ 8x96x64 f32", - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} softmax axis 1", shape_label(&shape)), + )) +} + +pub(super) async fn layer_norm_last_dim_case( + config: BenchmarkConfig, + name: String, + shape: [usize; 3], +) -> BenchmarkResult { + let device = initialized_device().await; + let last_dim = shape[2]; + let input = input_tensor(&device, shape, 18, 0.01).await?; + let layer = LayerNormConfig::new(last_dim) + .with_epsilon(1.0e-5) + .init::(&device); + let samples = time_samples(config, || { + let output = layer.clone().forward(input.clone()); + async move { materialize(output).await } }) -} - -pub fn conv1d_small() -> BenchmarkCase { - bench_case("burn::conv1d_small", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let input_shape = [4usize, 8usize, 256usize]; - let weight_shape = [16usize, 8usize, 5usize]; - let bias_shape = [16usize]; - let input = burn_tensor( - deterministic_values(elements(&input_shape), 25, 0.01), - input_shape, - &device, - ); - let weight = burn_tensor( - deterministic_values(elements(&weight_shape), 26, 0.01), - weight_shape, - &device, - ); - let bias = burn_tensor( - deterministic_values(elements(&bias_shape), 27, 0.001), - bias_shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - materialize_inputs(std::slice::from_ref(&weight)).await?; - materialize_inputs(std::slice::from_ref(&bias)).await?; - - let samples = time_samples(config, || { - let output = module::conv1d( - input.clone(), - weight.clone(), - Some(bias.clone()), - ConvOptions::new([2], [1], [1], 1), - ); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::conv1d_small", - config, - samples, - "4x8x256 conv 16x8x5", - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} layer norm", shape_label(&shape)), + )) +} + +pub(super) async fn rms_norm_fused_case( + config: BenchmarkConfig, + name: String, + shape: [usize; 3], +) -> BenchmarkResult { + let device = initialized_device().await; + let last_dim = shape[2]; + let input = input_tensor(&device, shape, 21, 0.01).await?; + let rms = RmsNormConfig::new(last_dim) + .with_epsilon(1.0e-5) + .init::(&device); + let samples = time_samples(config, || { + let output = rms.clone().forward(input.clone()); + async move { materialize(output).await } }) -} - -pub fn top_k_large() -> BenchmarkCase { - bench_case("burn::top_k_large", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let input_len = 65_537usize; - let k = 64usize; - let values = (0..input_len) - .map(|index| { - let base = ((index * 67 + 29) % 10_007) as f32 * 0.001; - let bump = if index % 4099 == 0 { 20.0 } else { 0.0 }; - base + bump - (index % 13) as f32 * 0.0001 - }) - .collect::>(); - let input = burn_tensor(values, [input_len], &device); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = input.clone().topk(k, 0); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::top_k_large", - config, - samples, - format!("{input_len} logits, k={k}"), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} rms norm", shape_label(&shape)), + )) +} + +pub(super) async fn dense_matmul_square_case( + config: BenchmarkConfig, + name: String, + size: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let lhs_shape = [size, size]; + let rhs_shape = [size, size]; + let lhs = input_tensor(&device, lhs_shape, 6, 0.004).await?; + let rhs = input_tensor(&device, rhs_shape, 7, 0.004).await?; + let samples = time_samples(config, || { + let output = lhs.clone().matmul(rhs.clone()); + async move { materialize(output).await } }) -} - -pub fn top_k_qwen_vocab() -> BenchmarkCase { - bench_case("burn::top_k_qwen_vocab", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let input_len = 151_936usize; - let k = 40usize; - let input = burn_tensor( - deterministic_values(input_len, 28, 0.01), - [input_len], - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = input.clone().topk(k, 0); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::top_k_qwen_vocab", - config, - samples, - format!("{input_len} logits, k={k}"), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "{} @ {} f32", + shape_label(&lhs_shape), + shape_label(&rhs_shape) + ), + )) +} + +pub(super) async fn dense_batched_matmul_case( + config: BenchmarkConfig, + name: String, + batch: usize, + m: usize, + k: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let lhs_shape = [batch, m, k]; + let rhs_shape = [batch, k, m]; + let lhs = input_tensor(&device, lhs_shape, 23, 0.004).await?; + let rhs = input_tensor(&device, rhs_shape, 24, 0.004).await?; + let samples = time_samples(config, || { + let output = lhs.clone().matmul(rhs.clone()); + async move { materialize(output).await } }) -} - -pub fn q8_0_qgemv() -> BenchmarkCase { - bench_case("burn::q8_0_qgemv", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let weight_shape = [4096usize, 896usize]; - let input_shape = [1usize, weight_shape[1]]; - let dense_weight_shape = [weight_shape[1], weight_shape[0]]; - let input = burn_tensor( - deterministic_values(elements(&input_shape), 8, 0.003), - input_shape, - &device, - ); - let weights = burn_tensor( - deterministic_values(elements(&dense_weight_shape), 80, 0.003), - dense_weight_shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - materialize_inputs(std::slice::from_ref(&weights)).await?; - - let samples = time_samples(config, || { - let output = input.clone().matmul(weights.clone()); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::q8_0_qgemv", - config, - samples, - "1x896 @ dense f32 896x4096 baseline", - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "{} @ {} f32", + shape_label(&lhs_shape), + shape_label(&rhs_shape) + ), + )) +} + +pub(super) async fn conv1d_small_case( + config: BenchmarkConfig, + name: String, + len: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let input_shape = [4usize, 8usize, len]; + let weight_shape = [16usize, 8usize, 5usize]; + let bias_shape = [16usize]; + let input = input_tensor(&device, input_shape, 25, 0.01).await?; + let weight = input_tensor(&device, weight_shape, 26, 0.01).await?; + let bias = input_tensor(&device, bias_shape, 27, 0.001).await?; + let samples = time_samples(config, || { + let output = module::conv1d( + input.clone(), + weight.clone(), + Some(bias.clone()), + ConvOptions::new([2], [1], [1], 1), + ); + async move { materialize(output).await } }) -} - -pub fn q4k_qgemv() -> BenchmarkCase { - bench_case("burn::q4k_qgemv", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let weight_shape = [2048usize, 1024usize]; - let input_shape = [1usize, weight_shape[1]]; - let dense_weight_shape = [weight_shape[1], weight_shape[0]]; - let input = burn_tensor( - deterministic_values(elements(&input_shape), 29, 0.003), - input_shape, - &device, - ); - let weights = burn_tensor( - deterministic_values(elements(&dense_weight_shape), 81, 0.003), - dense_weight_shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - materialize_inputs(std::slice::from_ref(&weights)).await?; - - let samples = time_samples(config, || { - let output = input.clone().matmul(weights.clone()); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::q4k_qgemv", - config, - samples, - "1x1024 @ dense f32 1024x2048 baseline", - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "{} conv {}", + shape_label(&input_shape), + shape_label(&weight_shape) + ), + )) +} + +pub(super) async fn top_k_case( + config: BenchmarkConfig, + name: String, + input_len: usize, + k: usize, + values: Vec, +) -> BenchmarkResult { + let device = initialized_device().await; + let input = values_input(&device, [input_len], values).await?; + let samples = time_samples(config, || { + let output = input.clone().topk(k, 0); + async move { materialize(output).await } }) -} - -pub fn q4k_paired_silu() -> BenchmarkCase { - bench_case("burn::q4k_paired_silu", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let weight_shape = [2048usize, 1024usize]; - let input_shape = [1usize, weight_shape[1]]; - let dense_weight_shape = [weight_shape[1], weight_shape[0]]; - let pair_len = weight_shape[0] / 2; - let input = burn_tensor( - deterministic_values(elements(&input_shape), 30, 0.003), - input_shape, - &device, - ); - let weights = burn_tensor( - deterministic_values(elements(&dense_weight_shape), 82, 0.003), - dense_weight_shape, - &device, - ); - materialize_inputs(std::slice::from_ref(&input)).await?; - materialize_inputs(std::slice::from_ref(&weights)).await?; - - let samples = time_samples(config, || { - let projected = input.clone().matmul(weights.clone()); - let gate = projected.clone().narrow(1, 0, pair_len); - let up = projected.narrow(1, pair_len, pair_len); - let output = activation::silu(gate) * up; - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::q4k_paired_silu", - config, - samples, - "1x1024 @ dense f32 1024x2048 + paired SiLU baseline", - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{input_len} logits, k={k}"), + )) +} + +pub(super) async fn qgemv_dense_case( + config: BenchmarkConfig, + name: String, + weight_shape: [usize; 2], + input_seed: usize, + weight_seed: usize, + detail_suffix: &'static str, +) -> BenchmarkResult { + let device = initialized_device().await; + let input_shape = [1usize, weight_shape[1]]; + let dense_weight_shape = [weight_shape[1], weight_shape[0]]; + let input = input_tensor(&device, input_shape, input_seed, 0.003).await?; + let weights = input_tensor(&device, dense_weight_shape, weight_seed, 0.003).await?; + let samples = time_samples(config, || { + let output = input.clone().matmul(weights.clone()); + async move { materialize(output).await } }) -} - -pub fn flash_attention_small() -> BenchmarkCase { - bench_case("burn::flash_attention_small", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [1usize, 4usize, 128usize, 64usize]; - let q = burn_tensor( - deterministic_values(elements(&shape), 31, 0.003), - shape, - &device, - ); - let k = burn_tensor( - deterministic_values(elements(&shape), 32, 0.003), - shape, - &device, - ); - let v = burn_tensor( - deterministic_values(elements(&shape), 33, 0.003), - shape, - &device, - ); - materialize_inputs(&[q.clone(), k.clone(), v.clone()]).await?; - - let samples = time_samples(config, || { - let output = module::attention( - q.clone(), - k.clone(), - v.clone(), - None, - None, - AttentionModuleOptions { - scale: Some(1.0 / (64.0f64).sqrt()), - softcap: None, - is_causal: false, - }, - ); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::flash_attention_small", - config, - samples, - format!("{} scaled dot-product attention", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "1x{} @ dense f32 {}{detail_suffix}", + weight_shape[1], + shape_label(&dense_weight_shape) + ), + )) +} + +pub(super) async fn q4k_paired_silu_case( + config: BenchmarkConfig, + name: String, + weight_shape: [usize; 2], + detail_suffix: &'static str, +) -> BenchmarkResult { + let device = initialized_device().await; + let input_shape = [1usize, weight_shape[1]]; + let dense_weight_shape = [weight_shape[1], weight_shape[0]]; + let pair_len = weight_shape[0] / 2; + let input = input_tensor(&device, input_shape, 30, 0.003).await?; + let weights = input_tensor(&device, dense_weight_shape, 82, 0.003).await?; + let samples = time_samples(config, || { + let projected = input.clone().matmul(weights.clone()); + let gate = projected.clone().narrow(1, 0, pair_len); + let up = projected.narrow(1, pair_len, pair_len); + let output = activation::silu(gate) * up; + async move { materialize(output).await } }) -} - -pub fn flash_attention_causal_small() -> BenchmarkCase { - bench_case( - "burn::flash_attention_causal_small", - |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [1usize, 4usize, 128usize, 64usize]; - let q = burn_tensor( - deterministic_values(elements(&shape), 34, 0.003), - shape, - &device, - ); - let k = burn_tensor( - deterministic_values(elements(&shape), 35, 0.003), - shape, - &device, - ); - let v = burn_tensor( - deterministic_values(elements(&shape), 36, 0.003), - shape, - &device, - ); - materialize_inputs(&[q.clone(), k.clone(), v.clone()]).await?; - - let samples = time_samples(config, || { - let output = module::attention( - q.clone(), - k.clone(), - v.clone(), - None, - None, - AttentionModuleOptions { - scale: Some(1.0 / (64.0f64).sqrt()), - softcap: None, - is_causal: true, - }, - ); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::flash_attention_causal_small", - config, - samples, - format!( - "{} causal scaled dot-product attention", - shape_label(&shape) - ), - )) - }) - }, - ) -} - -pub fn rope_fused_decode() -> BenchmarkCase { - bench_case("burn::rope_fused_decode", |_fusor_device, config| { - Box::pin(async move { - let device = initialized_device().await; - let shape = [1usize, 8usize, 256usize, 64usize]; - let [batch, heads, seq_len, head_dim] = shape; - let input = burn_tensor( - deterministic_values(batch * heads * seq_len * head_dim, 9, 0.01), - shape, - &device, - ); - let rope = RotaryEncodingConfig::new(seq_len * 2, head_dim).init::(&device); - materialize_inputs(std::slice::from_ref(&input)).await?; - - let samples = time_samples(config, || { - let output = rope.clone().forward(input.clone()); - async move { materialize(output).await } - }) - .await?; - - Ok(BenchmarkReport::new( - "burn::rope_fused_decode", - config, - samples, - format!("{} rotary encoding", shape_label(&shape)), - )) - }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "1x{} @ dense f32 {} + paired SiLU{detail_suffix}", + weight_shape[1], + shape_label(&dense_weight_shape) + ), + )) +} + +pub(super) async fn attention_case( + config: BenchmarkConfig, + name: String, + seq_len: usize, + seeds: [usize; 3], + causal: bool, + detail_op: &'static str, +) -> BenchmarkResult { + let device = initialized_device().await; + let shape = [1usize, 4usize, seq_len, 64usize]; + let q = input_tensor(&device, shape, seeds[0], 0.003).await?; + let k = input_tensor(&device, shape, seeds[1], 0.003).await?; + let v = input_tensor(&device, shape, seeds[2], 0.003).await?; + let samples = time_samples(config, || { + let output = module::attention( + q.clone(), + k.clone(), + v.clone(), + None, + None, + AttentionModuleOptions { + scale: Some(1.0 / (64.0f64).sqrt()), + softcap: None, + is_causal: causal, + }, + ); + async move { materialize(output).await } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} {detail_op}", shape_label(&shape)), + )) +} + +pub(super) async fn rope_fused_decode_case( + config: BenchmarkConfig, + name: String, + seq_len: usize, +) -> BenchmarkResult { + let device = initialized_device().await; + let shape = [1usize, 8usize, seq_len, 64usize]; + let [_, _, _, head_dim] = shape; + let input = input_tensor(&device, shape, 9, 0.01).await?; + let rope = RotaryEncodingConfig::new(seq_len * 2, head_dim).init::(&device); + let samples = time_samples(config, || { + let output = rope.clone().forward(input.clone()); + async move { materialize(output).await } + }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} rotary encoding", shape_label(&shape)), + )) +} + +macro_rules! fixed_cases { + ($($case:ident => $body:ident($($arg:expr),* $(,)?);)*) => { + $( + pub fn $case() -> BenchmarkCase { + bench_case(concat!("burn::", stringify!($case)), |config, name| { + Box::pin($body(config, name, $($arg),*)) + }) + } + )* + }; +} + +fixed_cases! { + elementwise_add_square => elementwise_add_square_case(512); + elementwise_mul_rank4 => elementwise_mul_rank4_case([9, 11, 32, 16]); + unary_trig_chain => unary_trig_chain_case(384); + activation_gelu => activation_gelu_case([512, 256]); + broadcast_add => broadcast_add_case(256); + transpose_then_elementwise => transpose_then_elementwise_case([256, 384]); + reduction_sum_last_dim => reduction_sum_last_dim_case(256); + reduction_max_middle_axis => reduction_max_middle_axis_case([64, 128, 64]); + softmax_last_dim => softmax_last_dim_case(512); + softmax_middle_axis => softmax_middle_axis_case([32, 128, 64]); + layer_norm_last_dim => layer_norm_last_dim_case([8, 128, 512]); + rms_norm_fused => rms_norm_fused_case([8, 128, 512]); + dense_matmul_square => dense_matmul_square_case(256); + dense_batched_matmul => dense_batched_matmul_case(8, 64, 96); + conv1d_small => conv1d_small_case(256); + top_k_large => top_k_case(65_537, 64, crate::bench::webgpu::topk_values(65_537)); + top_k_qwen_vocab => top_k_case(151_936, 40, deterministic_values(151_936, 28, 0.01)); + q8_0_qgemv => qgemv_dense_case([4096, 896], 8, 80, " baseline"); + q4k_qgemv => qgemv_dense_case([2048, 1024], 29, 81, " baseline"); + q4k_paired_silu => q4k_paired_silu_case([2048, 1024], " baseline"); + attention_small => attention_case(128, [31, 32, 33], false, "scaled dot-product attention"); + attention_causal_small => attention_case(128, [34, 35, 36], true, "causal scaled dot-product attention"); + rope_fused_decode => rope_fused_decode_case(256); } diff --git a/fusor-ml/conformance/src/bench/registry.rs b/fusor-ml/conformance/src/bench/registry.rs index 2582d556d..547b9306c 100644 --- a/fusor-ml/conformance/src/bench/registry.rs +++ b/fusor-ml/conformance/src/bench/registry.rs @@ -64,6 +64,23 @@ macro_rules! registry { } } + #[cfg(feature = "burn-bench")] + fn burn_cases_for_suite(name: &str) -> Option> { + match name { + $( + concat!("burn::", stringify!($case)) => Some(vec![ + crate::bench::burn::$case(), + ]), + )* + _ => None, + } + } + + #[cfg(not(feature = "burn-bench"))] + fn burn_cases_for_suite(_name: &str) -> Option> { + None + } + #[cfg(test)] mod generated_tests { use super::*; @@ -112,66 +129,21 @@ macro_rules! registry { } )* } - }; -} - -#[cfg(feature = "burn-bench")] -fn burn_cases_for_suite(name: &str) -> Option> { - match name { - "burn::elementwise_add_square" => Some(vec![crate::bench::burn::elementwise_add_square()]), - "burn::elementwise_mul_rank4" => Some(vec![crate::bench::burn::elementwise_mul_rank4()]), - "burn::unary_trig_chain" => Some(vec![crate::bench::burn::unary_trig_chain()]), - "burn::activation_gelu" => Some(vec![crate::bench::burn::activation_gelu()]), - "burn::broadcast_add" => Some(vec![crate::bench::burn::broadcast_add()]), - "burn::transpose_then_elementwise" => { - Some(vec![crate::bench::burn::transpose_then_elementwise()]) - } - "burn::reduction_sum_last_dim" => Some(vec![crate::bench::burn::reduction_sum_last_dim()]), - "burn::reduction_max_middle_axis" => { - Some(vec![crate::bench::burn::reduction_max_middle_axis()]) - } - "burn::softmax_last_dim" => Some(vec![crate::bench::burn::softmax_last_dim()]), - "burn::softmax_middle_axis" => Some(vec![crate::bench::burn::softmax_middle_axis()]), - "burn::layer_norm_last_dim" => Some(vec![crate::bench::burn::layer_norm_last_dim()]), - "burn::rms_norm_fused" => Some(vec![crate::bench::burn::rms_norm_fused()]), - "burn::dense_matmul_square" => Some(vec![crate::bench::burn::dense_matmul_square()]), - "burn::dense_batched_matmul" => Some(vec![crate::bench::burn::dense_batched_matmul()]), - "burn::conv1d_small" => Some(vec![crate::bench::burn::conv1d_small()]), - "burn::top_k_large" => Some(vec![crate::bench::burn::top_k_large()]), - "burn::top_k_qwen_vocab" => Some(vec![crate::bench::burn::top_k_qwen_vocab()]), - "burn::q8_0_qgemv" => Some(vec![crate::bench::burn::q8_0_qgemv()]), - "burn::q4k_qgemv" => Some(vec![crate::bench::burn::q4k_qgemv()]), - "burn::q4k_paired_silu" => Some(vec![crate::bench::burn::q4k_paired_silu()]), - "burn::flash_attention_small" => Some(vec![crate::bench::burn::flash_attention_small()]), - "burn::flash_attention_causal_small" => { - Some(vec![crate::bench::burn::flash_attention_causal_small()]) - } - "burn::rope_fused_decode" => Some(vec![crate::bench::burn::rope_fused_decode()]), - _ => None, - } -} -#[cfg(not(feature = "burn-bench"))] -fn burn_cases_for_suite(_name: &str) -> Option> { - None -} - -#[cfg(all(test, feature = "burn-bench"))] -mod burn_generated_tests { - use super::*; + #[cfg(all(test, feature = "burn-bench"))] + mod burn_generated_tests { + use super::*; - async fn gpu_device() -> Option { - match Device::gpu().await { - Ok(device) => Some(device), - Err(err) => { - tracing::warn!("skipping Burn benchmark smoke test: {err}"); - None + async fn gpu_device() -> Option { + match Device::gpu().await { + Ok(device) => Some(device), + Err(err) => { + tracing::warn!("skipping Burn benchmark smoke test: {err}"); + None + } + } } - } - } - macro_rules! burn_tests { - ($($case:ident),* $(,)?) => { $( #[allow(clippy::await_holding_lock)] #[tokio::test] @@ -201,34 +173,8 @@ mod burn_generated_tests { } } )* - }; - } - - burn_tests! { - elementwise_add_square, - elementwise_mul_rank4, - unary_trig_chain, - activation_gelu, - broadcast_add, - transpose_then_elementwise, - reduction_sum_last_dim, - reduction_max_middle_axis, - softmax_last_dim, - softmax_middle_axis, - layer_norm_last_dim, - rms_norm_fused, - dense_matmul_square, - dense_batched_matmul, - conv1d_small, - top_k_large, - top_k_qwen_vocab, - q8_0_qgemv, - q4k_qgemv, - q4k_paired_silu, - flash_attention_small, - flash_attention_causal_small, - rope_fused_decode, - } + } + }; } registry! { @@ -252,7 +198,7 @@ registry! { q8_0_qgemv, q4k_qgemv, q4k_paired_silu, - flash_attention_small, - flash_attention_causal_small, + attention_small, + attention_causal_small, rope_fused_decode, } diff --git a/fusor-ml/conformance/src/bench/sweep.rs b/fusor-ml/conformance/src/bench/sweep.rs index fe881d7cb..a4073d22f 100644 --- a/fusor-ml/conformance/src/bench/sweep.rs +++ b/fusor-ml/conformance/src/bench/sweep.rs @@ -1,10 +1,12 @@ //! Per-benchmark size sweeps used by the web runner detail route. +//! +//! Each sweep point runs the shared case bodies in `webgpu.rs` / `burn.rs` +//! with sweep-derived sizes; only the size tables and the size-to-parameter +//! mapping live here. -use fusor::{Device, GgmlType, MaskKind, QMatrix, Tensor as FusorTensor}; +use fusor::Device; -use crate::common::quantized::{q4k_raw_bytes, q8_0_raw_bytes, qmatrix_from_raw_bytes}; - -use super::{BenchmarkConfig, BenchmarkReport, BenchmarkResult, time_samples}; +use super::{BenchmarkConfig, BenchmarkReport, BenchmarkResult}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct BenchmarkSweepSize { @@ -41,355 +43,241 @@ pub enum BenchmarkSweepEvent { }, } +const fn size(label: &'static str, value: usize) -> BenchmarkSweepSize { + BenchmarkSweepSize { label, value } +} + const SQUARE_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "128x128", - value: 128, - }, - BenchmarkSweepSize { - label: "256x256", - value: 256, - }, - BenchmarkSweepSize { - label: "512x512", - value: 512, - }, - BenchmarkSweepSize { - label: "768x768", - value: 768, - }, + size("128x128", 128), + size("256x256", 256), + size("512x512", 512), + size("768x768", 768), ]; const RANK4_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "3x5x16x8", - value: 16, - }, - BenchmarkSweepSize { - label: "6x8x24x12", - value: 24, - }, - BenchmarkSweepSize { - label: "9x11x32x16", - value: 32, - }, - BenchmarkSweepSize { - label: "12x16x48x24", - value: 48, - }, + size("3x5x16x8", 16), + size("6x8x24x12", 24), + size("9x11x32x16", 32), + size("12x16x48x24", 48), ]; const ROW_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "64 rows", - value: 64, - }, - BenchmarkSweepSize { - label: "128 rows", - value: 128, - }, - BenchmarkSweepSize { - label: "256 rows", - value: 256, - }, - BenchmarkSweepSize { - label: "512 rows", - value: 512, - }, + size("64 rows", 64), + size("128 rows", 128), + size("256 rows", 256), + size("512 rows", 512), ]; const MID_AXIS_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "16x64x32", - value: 64, - }, - BenchmarkSweepSize { - label: "24x96x48", - value: 96, - }, - BenchmarkSweepSize { - label: "32x128x64", - value: 128, - }, - BenchmarkSweepSize { - label: "48x192x96", - value: 192, - }, + size("16x64x32", 64), + size("24x96x48", 96), + size("32x128x64", 128), + size("48x192x96", 192), ]; const SEQ_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "seq 32", - value: 32, - }, - BenchmarkSweepSize { - label: "seq 64", - value: 64, - }, - BenchmarkSweepSize { - label: "seq 128", - value: 128, - }, - BenchmarkSweepSize { - label: "seq 256", - value: 256, - }, + size("seq 32", 32), + size("seq 64", 64), + size("seq 128", 128), + size("seq 256", 256), ]; const MATMUL_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "64", - value: 64, - }, - BenchmarkSweepSize { - label: "128", - value: 128, - }, - BenchmarkSweepSize { - label: "256", - value: 256, - }, - BenchmarkSweepSize { - label: "384", - value: 384, - }, + size("64", 64), + size("128", 128), + size("256", 256), + size("384", 384), ]; const TOPK_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "8k", - value: 8_192, - }, - BenchmarkSweepSize { - label: "32k", - value: 32_768, - }, - BenchmarkSweepSize { - label: "64k", - value: 65_537, - }, - BenchmarkSweepSize { - label: "128k", - value: 131_072, - }, + size("8k", 8_192), + size("32k", 32_768), + size("64k", 65_537), + size("128k", 131_072), ]; const QWEN_TOPK_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "32k", - value: 32_768, - }, - BenchmarkSweepSize { - label: "65k", - value: 65_536, - }, - BenchmarkSweepSize { - label: "100k", - value: 100_000, - }, - BenchmarkSweepSize { - label: "151936", - value: 151_936, - }, + size("32k", 32_768), + size("65k", 65_536), + size("100k", 100_000), + size("151936", 151_936), ]; const Q8_GEMV_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "1024x512", - value: 1024, - }, - BenchmarkSweepSize { - label: "2048x768", - value: 2048, - }, - BenchmarkSweepSize { - label: "4096x896", - value: 4096, - }, - BenchmarkSweepSize { - label: "6144x1024", - value: 6144, - }, + size("1024x512", 1024), + size("2048x768", 2048), + size("4096x896", 4096), + size("6144x1024", 6144), ]; const Q4_GEMV_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "512x512", - value: 512, - }, - BenchmarkSweepSize { - label: "1024x768", - value: 1024, - }, - BenchmarkSweepSize { - label: "2048x1024", - value: 2048, - }, - BenchmarkSweepSize { - label: "4096x1024", - value: 4096, - }, + size("512x512", 512), + size("1024x768", 1024), + size("2048x1024", 2048), + size("4096x1024", 4096), ]; const PAIRED_SILU_SIZES: [BenchmarkSweepSize; 4] = [ - BenchmarkSweepSize { - label: "pair 256", - value: 256, - }, - BenchmarkSweepSize { - label: "pair 512", - value: 512, - }, - BenchmarkSweepSize { - label: "pair 1024", - value: 1024, - }, - BenchmarkSweepSize { - label: "pair 1536", - value: 1536, - }, + size("pair 256", 256), + size("pair 512", 512), + size("pair 1024", 1024), + size("pair 1536", 1536), ]; +const fn desc( + case: &'static str, + title: &'static str, + detail: &'static str, + sizes: &'static [BenchmarkSweepSize], +) -> BenchmarkSweepDescriptor { + BenchmarkSweepDescriptor { + case, + title, + detail, + sizes, + } +} + pub fn descriptor(case: &str) -> Option { let descriptor = match case { - "elementwise_add_square" => BenchmarkSweepDescriptor { - case: "elementwise_add_square", - title: "Elementwise add", - detail: "F32 add over square tensors.", - sizes: &SQUARE_SIZES, - }, - "elementwise_mul_rank4" => BenchmarkSweepDescriptor { - case: "elementwise_mul_rank4", - title: "Elementwise mul rank 4", - detail: "F32 multiply over rank-4 tensors.", - sizes: &RANK4_SIZES, - }, - "unary_trig_chain" => BenchmarkSweepDescriptor { - case: "unary_trig_chain", - title: "Unary trig chain", - detail: "sin(x) + cos(x) over square tensors.", - sizes: &SQUARE_SIZES, - }, - "activation_gelu" => BenchmarkSweepDescriptor { - case: "activation_gelu", - title: "GELU activation", - detail: "F32 GELU over square tensors.", - sizes: &SQUARE_SIZES, - }, - "broadcast_add" => BenchmarkSweepDescriptor { - case: "broadcast_add", - title: "Broadcast add", - detail: "Matrix plus broadcast row vector.", - sizes: &ROW_SIZES, - }, - "transpose_then_elementwise" => BenchmarkSweepDescriptor { - case: "transpose_then_elementwise", - title: "Transpose then elementwise", - detail: "Transpose followed by elementwise square.", - sizes: &ROW_SIZES, - }, - "reduction_sum_last_dim" => BenchmarkSweepDescriptor { - case: "reduction_sum_last_dim", - title: "Reduction sum", - detail: "Sum over the last matrix axis.", - sizes: &ROW_SIZES, - }, - "reduction_max_middle_axis" => BenchmarkSweepDescriptor { - case: "reduction_max_middle_axis", - title: "Reduction max", - detail: "Max over the middle rank-3 axis.", - sizes: &MID_AXIS_SIZES, - }, - "softmax_last_dim" => BenchmarkSweepDescriptor { - case: "softmax_last_dim", - title: "Softmax last axis", - detail: "Softmax over the last matrix axis.", - sizes: &ROW_SIZES, - }, - "softmax_middle_axis" => BenchmarkSweepDescriptor { - case: "softmax_middle_axis", - title: "Softmax middle axis", - detail: "Softmax over the middle rank-3 axis.", - sizes: &MID_AXIS_SIZES, - }, - "layer_norm_last_dim" => BenchmarkSweepDescriptor { - case: "layer_norm_last_dim", - title: "Layer norm", - detail: "Layer normalization over the last dimension.", - sizes: &SEQ_SIZES, - }, - "rms_norm_fused" => BenchmarkSweepDescriptor { - case: "rms_norm_fused", - title: "RMS norm", - detail: "RMS normalization over the last dimension.", - sizes: &SEQ_SIZES, - }, - "dense_matmul_square" => BenchmarkSweepDescriptor { - case: "dense_matmul_square", - title: "Dense matmul", - detail: "Square F32 matrix multiplication.", - sizes: &MATMUL_SIZES, - }, - "dense_batched_matmul" => BenchmarkSweepDescriptor { - case: "dense_batched_matmul", - title: "Batched matmul", - detail: "Batched F32 matrix multiplication.", - sizes: &MATMUL_SIZES, - }, - "conv1d_small" => BenchmarkSweepDescriptor { - case: "conv1d_small", - title: "Conv1D", - detail: "Small 1D convolution with fixed channels.", - sizes: &ROW_SIZES, - }, - "top_k_large" => BenchmarkSweepDescriptor { - case: "top_k_large", - title: "Top K", - detail: "Top-k selection over a logits vector.", - sizes: &TOPK_SIZES, - }, - "top_k_qwen_vocab" => BenchmarkSweepDescriptor { - case: "top_k_qwen_vocab", - title: "Top K Qwen vocab", - detail: "Top-k selection over vocabulary-scale logits.", - sizes: &QWEN_TOPK_SIZES, - }, - "q8_0_qgemv" => BenchmarkSweepDescriptor { - case: "q8_0_qgemv", - title: "Q8_0 GEMV", - detail: "Fusor Q8_0 GEMV against a Burn dense-f32 baseline.", - sizes: &Q8_GEMV_SIZES, - }, - "q4k_qgemv" => BenchmarkSweepDescriptor { - case: "q4k_qgemv", - title: "Q4K GEMV", - detail: "Fusor Q4K GEMV against a Burn dense-f32 baseline.", - sizes: &Q4_GEMV_SIZES, - }, - "q4k_paired_silu" => BenchmarkSweepDescriptor { - case: "q4k_paired_silu", - title: "Q4K paired SiLU", - detail: "Fusor fused paired SiLU GEMV against a Burn dense-f32 baseline.", - sizes: &PAIRED_SILU_SIZES, - }, - "flash_attention_small" => BenchmarkSweepDescriptor { - case: "flash_attention_small", - title: "Attention", - detail: "Scaled dot-product attention across sequence lengths.", - sizes: &SEQ_SIZES, - }, - "flash_attention_causal_small" => BenchmarkSweepDescriptor { - case: "flash_attention_causal_small", - title: "Causal attention", - detail: "Causal scaled dot-product attention across sequence lengths.", - sizes: &SEQ_SIZES, - }, - "rope_fused_decode" => BenchmarkSweepDescriptor { - case: "rope_fused_decode", - title: "RoPE", - detail: "Rotary positional encoding across sequence lengths.", - sizes: &SEQ_SIZES, - }, + "elementwise_add_square" => desc( + "elementwise_add_square", + "Elementwise add", + "F32 add over square tensors.", + &SQUARE_SIZES, + ), + "elementwise_mul_rank4" => desc( + "elementwise_mul_rank4", + "Elementwise mul rank 4", + "F32 multiply over rank-4 tensors.", + &RANK4_SIZES, + ), + "unary_trig_chain" => desc( + "unary_trig_chain", + "Unary trig chain", + "sin(x) + cos(x) over square tensors.", + &SQUARE_SIZES, + ), + "activation_gelu" => desc( + "activation_gelu", + "GELU activation", + "F32 GELU over square tensors.", + &SQUARE_SIZES, + ), + "broadcast_add" => desc( + "broadcast_add", + "Broadcast add", + "Matrix plus broadcast row vector.", + &ROW_SIZES, + ), + "transpose_then_elementwise" => desc( + "transpose_then_elementwise", + "Transpose then elementwise", + "Transpose followed by elementwise square.", + &ROW_SIZES, + ), + "reduction_sum_last_dim" => desc( + "reduction_sum_last_dim", + "Reduction sum", + "Sum over the last matrix axis.", + &ROW_SIZES, + ), + "reduction_max_middle_axis" => desc( + "reduction_max_middle_axis", + "Reduction max", + "Max over the middle rank-3 axis.", + &MID_AXIS_SIZES, + ), + "softmax_last_dim" => desc( + "softmax_last_dim", + "Softmax last axis", + "Softmax over the last matrix axis.", + &ROW_SIZES, + ), + "softmax_middle_axis" => desc( + "softmax_middle_axis", + "Softmax middle axis", + "Softmax over the middle rank-3 axis.", + &MID_AXIS_SIZES, + ), + "layer_norm_last_dim" => desc( + "layer_norm_last_dim", + "Layer norm", + "Layer normalization over the last dimension.", + &SEQ_SIZES, + ), + "rms_norm_fused" => desc( + "rms_norm_fused", + "RMS norm", + "RMS normalization over the last dimension.", + &SEQ_SIZES, + ), + "dense_matmul_square" => desc( + "dense_matmul_square", + "Dense matmul", + "Square F32 matrix multiplication.", + &MATMUL_SIZES, + ), + "dense_batched_matmul" => desc( + "dense_batched_matmul", + "Batched matmul", + "Batched F32 matrix multiplication.", + &MATMUL_SIZES, + ), + "conv1d_small" => desc( + "conv1d_small", + "Conv1D", + "Small 1D convolution with fixed channels.", + &ROW_SIZES, + ), + "top_k_large" => desc( + "top_k_large", + "Top K", + "Top-k selection over a logits vector.", + &TOPK_SIZES, + ), + "top_k_qwen_vocab" => desc( + "top_k_qwen_vocab", + "Top K Qwen vocab", + "Top-k selection over vocabulary-scale logits.", + &QWEN_TOPK_SIZES, + ), + "q8_0_qgemv" => desc( + "q8_0_qgemv", + "Q8_0 GEMV", + "Fusor Q8_0 GEMV against a Burn dense-f32 baseline.", + &Q8_GEMV_SIZES, + ), + "q4k_qgemv" => desc( + "q4k_qgemv", + "Q4K GEMV", + "Fusor Q4K GEMV against a Burn dense-f32 baseline.", + &Q4_GEMV_SIZES, + ), + "q4k_paired_silu" => desc( + "q4k_paired_silu", + "Q4K paired SiLU", + "Fusor fused paired SiLU GEMV against a Burn dense-f32 baseline.", + &PAIRED_SILU_SIZES, + ), + "attention_small" => desc( + "attention_small", + "Attention", + "Scaled dot-product attention across sequence lengths.", + &SEQ_SIZES, + ), + "attention_causal_small" => desc( + "attention_causal_small", + "Causal attention", + "Causal scaled dot-product attention across sequence lengths.", + &SEQ_SIZES, + ), + "rope_fused_decode" => desc( + "rope_fused_decode", + "RoPE", + "Rotary positional encoding across sequence lengths.", + &SEQ_SIZES, + ), _ => return None, }; Some(descriptor) @@ -446,37 +334,6 @@ pub async fn run_sweep( Ok(points) } -fn deterministic_values(len: usize, seed: usize, scale: f32) -> Vec { - (0..len) - .map(|index| { - let bucket = (index - .wrapping_mul(37) - .wrapping_add(seed.wrapping_mul(17)) - .wrapping_add(11)) - % 211; - (bucket as f32 - 105.0) * scale - }) - .collect() -} - -fn shape_label(shape: &[usize]) -> String { - shape - .iter() - .map(usize::to_string) - .collect::>() - .join("x") -} - -fn elements(shape: &[usize]) -> usize { - shape.iter().product() -} - -async fn materialize_inputs(inputs: &[&FusorTensor]) { - for input in inputs { - input.materialize().await; - } -} - fn rank4_shape(size: usize) -> [usize; 4] { [size / 4, size / 3, size, size / 2] } @@ -504,1328 +361,135 @@ fn q4_shape(value: usize) -> [usize; 2] { [value, k] } -fn sweep_report( - suite: &str, - case: &str, - size: BenchmarkSweepSize, - config: BenchmarkConfig, - samples: Vec, - detail: impl Into, -) -> BenchmarkReport { - BenchmarkReport::new( - format!("{suite}::{case}@{}", size.label), - config, - samples, - detail, - ) -} - async fn run_webgpu_case( case: &str, device: &Device, size: BenchmarkSweepSize, config: BenchmarkConfig, ) -> BenchmarkResult { + use crate::bench::webgpu as cases; + let name = format!("webgpu::{case}@{}", size.label); + let value = size.value; match case { "elementwise_add_square" => { - let shape = [size.value, size.value]; - let lhs: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 1, 0.01), - ); - let rhs: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 2, 0.008), - ); - materialize_inputs(&[&lhs, &rhs]).await; - let samples = time_samples(config, || { - let output = (&lhs + &rhs).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} f32 add", shape_label(&shape)), - )) + cases::elementwise_add_square_case(device, config, name, value).await } "elementwise_mul_rank4" => { - let shape = rank4_shape(size.value); - let lhs: FusorTensor<4, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 3, 0.012), - ); - let rhs: FusorTensor<4, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 4, 0.009), - ); - materialize_inputs(&[&lhs, &rhs]).await; - let samples = time_samples(config, || { - let output = (&lhs * &rhs).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} f32 mul", shape_label(&shape)), - )) - } - "unary_trig_chain" => { - let shape = [size.value, size.value]; - let input: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 10, 0.01), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let output = (input.sin() + input.cos()).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} sin+cos", shape_label(&shape)), - )) + cases::elementwise_mul_rank4_case(device, config, name, rank4_shape(value)).await } + "unary_trig_chain" => cases::unary_trig_chain_case(device, config, name, value).await, "activation_gelu" => { - let shape = [size.value, size.value]; - let input: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 11, 0.015), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let output = input.gelu(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} gelu", shape_label(&shape)), - )) - } - "broadcast_add" => { - let matrix_shape = [size.value, 512usize]; - let vector_shape = [512usize]; - let matrix: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - matrix_shape, - &deterministic_values(elements(&matrix_shape), 12, 0.006), - ); - let vector: FusorTensor<1, f32> = FusorTensor::from_slice( - device, - vector_shape, - &deterministic_values(elements(&vector_shape), 13, 0.01), - ); - materialize_inputs(&[&matrix]).await; - materialize_inputs(&[&vector]).await; - let samples = time_samples(config, || { - let vector_row = vector.reshape([1, vector_shape[0]]); - let output = (&matrix + vector_row.broadcast_as(matrix_shape)).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} + broadcast 512", shape_label(&matrix_shape)), - )) + cases::activation_gelu_case(device, config, name, [value, value]).await } + "broadcast_add" => cases::broadcast_add_case(device, config, name, value).await, "transpose_then_elementwise" => { - let shape = [size.value, size.value + size.value / 2]; - let input: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 14, 0.01), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let transposed = input.transpose(0, 1); - let output = (transposed.clone() * transposed).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} transpose, square", shape_label(&shape)), - )) + cases::transpose_then_elementwise_case(device, config, name, [value, value + value / 2]) + .await } "reduction_sum_last_dim" => { - let shape = [size.value, 512usize]; - let input: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 15, 0.004), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let output = input.sum::<1>(1); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} sum axis 1", shape_label(&shape)), - )) + cases::reduction_sum_last_dim_case(device, config, name, value).await } "reduction_max_middle_axis" => { - let shape = middle_shape(size.value); - let input: FusorTensor<3, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 16, 0.004), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let output = input.max::<2>(1); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} max axis 1", shape_label(&shape)), - )) - } - "softmax_last_dim" => { - let shape = [size.value, 256usize]; - let input: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 5, 0.006), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let output = input.softmax_last_dim::<1>(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} last-axis softmax", shape_label(&shape)), - )) + cases::reduction_max_middle_axis_case(device, config, name, middle_shape(value)).await } + "softmax_last_dim" => cases::softmax_last_dim_case(device, config, name, value).await, "softmax_middle_axis" => { - let shape = middle_shape(size.value); - let input: FusorTensor<3, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 17, 0.004), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let output = input.softmax::<2>(1); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} softmax axis 1", shape_label(&shape)), - )) + cases::softmax_middle_axis_case(device, config, name, middle_shape(value)).await } "layer_norm_last_dim" => { - let shape = [4usize, size.value, 512usize]; - let last_dim = shape[2]; - let input: FusorTensor<3, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 18, 0.01), - ); - let weight_values = deterministic_values(last_dim, 19, 0.002) - .into_iter() - .map(|value| value + 1.0) - .collect::>(); - let bias_values = deterministic_values(last_dim, 20, 0.001); - let weight: FusorTensor<1, f32> = - FusorTensor::from_slice(device, [last_dim], &weight_values); - let bias: FusorTensor<1, f32> = - FusorTensor::from_slice(device, [last_dim], &bias_values); - materialize_inputs(&[&input]).await; - materialize_inputs(&[&weight, &bias]).await; - let samples = time_samples(config, || { - let output = - input.layer_norm_last_dim_fused::<2, 1, _, _>(&weight, Some(&bias), 1.0e-5); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} layer norm", shape_label(&shape)), - )) - } - "rms_norm_fused" => { - let shape = [4usize, size.value, 512usize]; - let last_dim = shape[2]; - let input: FusorTensor<3, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 21, 0.01), - ); - let weight_values = deterministic_values(last_dim, 22, 0.002) - .into_iter() - .map(|value| value + 1.0) - .collect::>(); - let weight: FusorTensor<1, f32> = - FusorTensor::from_slice(device, [last_dim], &weight_values); - materialize_inputs(&[&input]).await; - materialize_inputs(&[&weight]).await; - let samples = time_samples(config, || { - let output = input.rms_norm_fused_no_bias::<1, 2>(&weight, 1.0e-5); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} rms norm", shape_label(&shape)), - )) - } - "dense_matmul_square" => { - let lhs_shape = [size.value, size.value]; - let rhs_shape = [size.value, size.value]; - let lhs: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - lhs_shape, - &deterministic_values(elements(&lhs_shape), 6, 0.004), - ); - let rhs: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - rhs_shape, - &deterministic_values(elements(&rhs_shape), 7, 0.004), - ); - materialize_inputs(&[&lhs, &rhs]).await; - let samples = time_samples(config, || { - let output = lhs.matmul(&rhs); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!( - "{} @ {} f32", - shape_label(&lhs_shape), - shape_label(&rhs_shape) - ), - )) + cases::layer_norm_last_dim_case(device, config, name, [4, value, 512]).await } + "rms_norm_fused" => cases::rms_norm_fused_case(device, config, name, [4, value, 512]).await, + "dense_matmul_square" => cases::dense_matmul_square_case(device, config, name, value).await, "dense_batched_matmul" => { - let batch = 4usize; - let k = size.value + 32; - let lhs_shape = [batch, size.value, k]; - let rhs_shape = [batch, k, size.value]; - let lhs: FusorTensor<3, f32> = FusorTensor::from_slice( - device, - lhs_shape, - &deterministic_values(elements(&lhs_shape), 23, 0.004), - ); - let rhs: FusorTensor<3, f32> = FusorTensor::from_slice( - device, - rhs_shape, - &deterministic_values(elements(&rhs_shape), 24, 0.004), - ); - materialize_inputs(&[&lhs, &rhs]).await; - let samples = time_samples(config, || { - let output = lhs.matmul(&rhs); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!( - "{} @ {} f32", - shape_label(&lhs_shape), - shape_label(&rhs_shape) - ), - )) + cases::dense_batched_matmul_case(device, config, name, 4, value, value + 32).await } - "conv1d_small" => { - let input_shape = [4usize, 8usize, size.value]; - let weight_shape = [16usize, 8usize, 5usize]; - let bias_shape = [16usize]; - let input: FusorTensor<3, f32> = FusorTensor::from_slice( - device, - input_shape, - &deterministic_values(elements(&input_shape), 25, 0.01), - ); - let weight: FusorTensor<3, f32> = FusorTensor::from_slice( - device, - weight_shape, - &deterministic_values(elements(&weight_shape), 26, 0.01), - ); - let bias: FusorTensor<1, f32> = FusorTensor::from_slice( - device, - bias_shape, - &deterministic_values(elements(&bias_shape), 27, 0.001), - ); - materialize_inputs(&[&input]).await; - materialize_inputs(&[&weight]).await; - materialize_inputs(&[&bias]).await; - let samples = time_samples(config, || { - let output = input.conv(&weight, Some(&bias), [2], [1]); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!( - "{} conv {}", - shape_label(&input_shape), - shape_label(&weight_shape) - ), - )) + "conv1d_small" => cases::conv1d_small_case(device, config, name, value).await, + "top_k_large" => { + cases::top_k_case(device, config, name, value, 64, cases::topk_values(value)).await } - "top_k_large" | "top_k_qwen_vocab" => { - let input_len = size.value; - let k = if case == "top_k_qwen_vocab" { 40 } else { 64 }; - let input: FusorTensor<1, f32> = - FusorTensor::from_slice(device, [input_len], &topk_values(input_len)); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || async { - let top = input.top_k_pairs(k).await?; - if top.len() != k { - return Err(format!("top_k returned {} pairs, expected {k}", top.len()).into()); - } - Ok(()) - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{input_len} logits, k={k}"), - )) - } - "q8_0_qgemv" => { - let weight_shape = q8_shape(size.value); - let input_shape = [1usize, weight_shape[1]]; - let raw_bytes = q8_0_raw_bytes(weight_shape); - let matrix: QMatrix = - qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q8_0); - let input: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - input_shape, - &deterministic_values(elements(&input_shape), 8, 0.003), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let output = input.q_mat_mul(&matrix); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!( - "1x{} @ Q8_0 {}", - weight_shape[1], - shape_label(&weight_shape) - ), - )) - } - "q4k_qgemv" => { - let weight_shape = q4_shape(size.value); - let input_shape = [1usize, weight_shape[1]]; - let raw_bytes = q4k_raw_bytes(weight_shape); - let matrix: QMatrix = - qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q4K); - let input: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - input_shape, - &deterministic_values(elements(&input_shape), 29, 0.003), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let output = input.q_mat_mul(&matrix); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("1x{} @ Q4K {}", weight_shape[1], shape_label(&weight_shape)), - )) + "top_k_qwen_vocab" => { + cases::top_k_case(device, config, name, value, 40, cases::topk_values(value)).await } + "q8_0_qgemv" => cases::q8_0_qgemv_case(device, config, name, q8_shape(value)).await, + "q4k_qgemv" => cases::q4k_qgemv_case(device, config, name, q4_shape(value)).await, "q4k_paired_silu" => { - let weight_shape = [size.value * 2, 1024usize]; - let input_shape = [1usize, weight_shape[1]]; - let raw_bytes = q4k_raw_bytes(weight_shape); - let matrix: QMatrix = - qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q4K); - let input: FusorTensor<2, f32> = FusorTensor::from_slice( - device, - input_shape, - &deterministic_values(elements(&input_shape), 30, 0.003), - ); - materialize_inputs(&[&input]).await; - let samples = time_samples(config, || { - let pair_len = weight_shape[0] / 2; - let projected = input.q_mat_mul(&matrix); - let gate = projected - .narrow(fusor::D::Minus1, 0, pair_len) - .to_concrete(); - let up = projected - .narrow(fusor::D::Minus1, pair_len, pair_len) - .to_concrete(); - let output = (gate.silu() * up).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("1x1024 @ paired Q4K {}", shape_label(&weight_shape)), - )) + cases::q4k_paired_silu_case(device, config, name, [value * 2, 1024]).await } - "flash_attention_small" | "flash_attention_causal_small" => { - let seq_len = size.value; - let shape = [1usize, 4usize, seq_len, 64usize]; - let q: FusorTensor<4, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 31, 0.003), - ); - let k: FusorTensor<4, f32> = FusorTensor::from_slice( + "attention_small" => { + cases::attention_case( device, - shape, - &deterministic_values(elements(&shape), 32, 0.003), - ); - let v: FusorTensor<4, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(elements(&shape), 33, 0.003), - ); - let mask_shape = [seq_len, seq_len]; - let mask_values = vec![0.0f32; elements(&mask_shape)]; - let mask: FusorTensor<2, f32> = - FusorTensor::from_slice(device, mask_shape, &mask_values); - materialize_inputs(&[&q, &k, &v]).await; - if case == "flash_attention_causal_small" { - materialize_inputs(&[&mask]).await; - } - let samples = time_samples(config, || { - let mask_arg = if case == "flash_attention_causal_small" { - Some((&mask, MaskKind::Causal)) - } else { - None - }; - let output = q.flash_attention(&k, &v, 1.0 / (64.0f32).sqrt(), mask_arg); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, config, - samples, - format!("{} attention", shape_label(&shape)), - )) + name, + value, + [31, 32, 33], + false, + "attention", + ) + .await } - "rope_fused_decode" => { - let seq_len = size.value; - let shape = [1usize, 8usize, seq_len, 64usize]; - let [batch, heads, _, head_dim] = shape; - let pos_shape = [seq_len * 2, head_dim / 2]; - let cos_values = rope_values(pos_shape, head_dim, true); - let sin_values = rope_values(pos_shape, head_dim, false); - let input: FusorTensor<4, f32> = FusorTensor::from_slice( - device, - shape, - &deterministic_values(batch * heads * seq_len * head_dim, 9, 0.01), - ); - let cos: FusorTensor<2, f32> = FusorTensor::from_slice(device, pos_shape, &cos_values); - let sin: FusorTensor<2, f32> = FusorTensor::from_slice(device, pos_shape, &sin_values); - materialize_inputs(&[&input]).await; - materialize_inputs(&[&cos, &sin]).await; - let samples = time_samples(config, || { - let output = input.rope_fused(&cos, &sin); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - Ok(sweep_report( - "webgpu", - case, - size, - config, - samples, - format!("{} fused rope", shape_label(&shape)), - )) + "attention_causal_small" => { + cases::attention_case(device, config, name, value, [31, 32, 33], true, "attention") + .await } + "rope_fused_decode" => cases::rope_fused_decode_case(device, config, name, value).await, _ => Err(format!("unknown WebGPU benchmark sweep: {case}").into()), } } -fn topk_values(input_len: usize) -> Vec { - (0..input_len) - .map(|index| { - let base = ((index * 67 + 29) % 10_007) as f32 * 0.001; - let bump = if index % 4099 == 0 { 20.0 } else { 0.0 }; - base + bump - (index % 13) as f32 * 0.0001 - }) - .collect() -} - -fn rope_values(shape: [usize; 2], head_dim: usize, cos: bool) -> Vec { - (0..shape[0]) - .flat_map(|i| { - (0..shape[1]).map(move |j| { - let value = (i as f32) / 10000f32.powf((2 * (j / 2)) as f32 / head_dim as f32); - if cos { value.cos() } else { value.sin() } - }) - }) - .collect() -} - -#[cfg(feature = "burn-bench")] -type BurnTensor = ::burn::tensor::Tensor<::burn::backend::Wgpu, R>; - -#[cfg(feature = "burn-bench")] -fn burn_tensor( - values: Vec, - shape: [usize; R], - device: &::burn::backend::wgpu::WgpuDevice, -) -> BurnTensor { - ::burn::tensor::Tensor::<::burn::backend::Wgpu, R>::from_data( - ::burn::tensor::TensorData::new(values, shape), - device, - ) -} - -#[cfg(feature = "burn-bench")] -async fn burn_materialize(tensor: BurnTensor) -> BenchmarkResult<()> { - let _ = tensor.into_data_async().await; - Ok(()) -} - -#[cfg(feature = "burn-bench")] -async fn burn_materialize_inputs(inputs: &[BurnTensor]) -> BenchmarkResult<()> { - for input in inputs { - burn_materialize(input.clone()).await?; - } - Ok(()) -} - #[cfg(feature = "burn-bench")] async fn run_burn_case( case: &str, size: BenchmarkSweepSize, config: BenchmarkConfig, ) -> BenchmarkResult { - use ::burn::{ - backend::Wgpu, - nn::{LayerNormConfig, RmsNormConfig, RotaryEncodingConfig}, - tensor::{ - activation, module, - ops::{AttentionModuleOptions, ConvOptions}, - }, - }; - - let device = crate::bench::burn::initialized_device().await; + use crate::bench::{burn as cases, webgpu}; + let name = format!("burn::{case}@{}", size.label); + let value = size.value; match case { - "elementwise_add_square" => { - let shape = [size.value, size.value]; - let lhs = burn_tensor( - deterministic_values(elements(&shape), 1, 0.01), - shape, - &device, - ); - let rhs = burn_tensor( - deterministic_values(elements(&shape), 2, 0.008), - shape, - &device, - ); - burn_materialize_inputs(&[lhs.clone(), rhs.clone()]).await?; - let samples = time_samples(config, || { - let output = lhs.clone() + rhs.clone(); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} f32 add", shape_label(&shape)), - )) - } + "elementwise_add_square" => cases::elementwise_add_square_case(config, name, value).await, "elementwise_mul_rank4" => { - let shape = rank4_shape(size.value); - let lhs = burn_tensor( - deterministic_values(elements(&shape), 3, 0.012), - shape, - &device, - ); - let rhs = burn_tensor( - deterministic_values(elements(&shape), 4, 0.009), - shape, - &device, - ); - burn_materialize_inputs(&[lhs.clone(), rhs.clone()]).await?; - let samples = time_samples(config, || { - let output = lhs.clone() * rhs.clone(); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} f32 mul", shape_label(&shape)), - )) - } - "unary_trig_chain" => { - let shape = [size.value, size.value]; - let input = burn_tensor( - deterministic_values(elements(&shape), 10, 0.01), - shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = input.clone().sin() + input.clone().cos(); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} sin+cos", shape_label(&shape)), - )) - } - "activation_gelu" => { - let shape = [size.value, size.value]; - let input = burn_tensor( - deterministic_values(elements(&shape), 11, 0.015), - shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = activation::gelu(input.clone()); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} gelu", shape_label(&shape)), - )) - } - "broadcast_add" => { - let matrix_shape = [size.value, 512usize]; - let vector_shape = [512usize]; - let matrix = burn_tensor( - deterministic_values(elements(&matrix_shape), 12, 0.006), - matrix_shape, - &device, - ); - let vector = burn_tensor( - deterministic_values(elements(&vector_shape), 13, 0.01), - vector_shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&matrix)).await?; - burn_materialize_inputs(std::slice::from_ref(&vector)).await?; - let samples = time_samples(config, || { - let output = matrix.clone() + vector.clone().reshape([1, vector_shape[0]]); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} + broadcast 512", shape_label(&matrix_shape)), - )) + cases::elementwise_mul_rank4_case(config, name, rank4_shape(value)).await } + "unary_trig_chain" => cases::unary_trig_chain_case(config, name, value).await, + "activation_gelu" => cases::activation_gelu_case(config, name, [value, value]).await, + "broadcast_add" => cases::broadcast_add_case(config, name, value).await, "transpose_then_elementwise" => { - let shape = [size.value, size.value + size.value / 2]; - let input = burn_tensor( - deterministic_values(elements(&shape), 14, 0.01), - shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let transposed = input.clone().transpose(); - let output = transposed.clone() * transposed; - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} transpose, square", shape_label(&shape)), - )) - } - "reduction_sum_last_dim" => { - let shape = [size.value, 512usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 15, 0.004), - shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = input.clone().sum_dim(1); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} sum axis 1", shape_label(&shape)), - )) + cases::transpose_then_elementwise_case(config, name, [value, value + value / 2]).await } + "reduction_sum_last_dim" => cases::reduction_sum_last_dim_case(config, name, value).await, "reduction_max_middle_axis" => { - let shape = middle_shape(size.value); - let input = burn_tensor( - deterministic_values(elements(&shape), 16, 0.004), - shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = input.clone().max_dim(1); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} max axis 1", shape_label(&shape)), - )) - } - "softmax_last_dim" => { - let shape = [size.value, 256usize]; - let input = burn_tensor( - deterministic_values(elements(&shape), 5, 0.006), - shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = activation::softmax(input.clone(), 1); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} last-axis softmax", shape_label(&shape)), - )) + cases::reduction_max_middle_axis_case(config, name, middle_shape(value)).await } + "softmax_last_dim" => cases::softmax_last_dim_case(config, name, value).await, "softmax_middle_axis" => { - let shape = middle_shape(size.value); - let input = burn_tensor( - deterministic_values(elements(&shape), 17, 0.004), - shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = activation::softmax(input.clone(), 1); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} softmax axis 1", shape_label(&shape)), - )) + cases::softmax_middle_axis_case(config, name, middle_shape(value)).await } "layer_norm_last_dim" => { - let shape = [4usize, size.value, 512usize]; - let last_dim = shape[2]; - let input = burn_tensor( - deterministic_values(elements(&shape), 18, 0.01), - shape, - &device, - ); - let layer = LayerNormConfig::new(last_dim) - .with_epsilon(1.0e-5) - .init::(&device); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = layer.clone().forward(input.clone()); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} layer norm", shape_label(&shape)), - )) - } - "rms_norm_fused" => { - let shape = [4usize, size.value, 512usize]; - let last_dim = shape[2]; - let input = burn_tensor( - deterministic_values(elements(&shape), 21, 0.01), - shape, - &device, - ); - let rms = RmsNormConfig::new(last_dim) - .with_epsilon(1.0e-5) - .init::(&device); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = rms.clone().forward(input.clone()); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} rms norm", shape_label(&shape)), - )) - } - "dense_matmul_square" => { - let lhs_shape = [size.value, size.value]; - let rhs_shape = [size.value, size.value]; - let lhs = burn_tensor( - deterministic_values(elements(&lhs_shape), 6, 0.004), - lhs_shape, - &device, - ); - let rhs = burn_tensor( - deterministic_values(elements(&rhs_shape), 7, 0.004), - rhs_shape, - &device, - ); - burn_materialize_inputs(&[lhs.clone(), rhs.clone()]).await?; - let samples = time_samples(config, || { - let output = lhs.clone().matmul(rhs.clone()); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!( - "{} @ {} f32", - shape_label(&lhs_shape), - shape_label(&rhs_shape) - ), - )) + cases::layer_norm_last_dim_case(config, name, [4, value, 512]).await } + "rms_norm_fused" => cases::rms_norm_fused_case(config, name, [4, value, 512]).await, + "dense_matmul_square" => cases::dense_matmul_square_case(config, name, value).await, "dense_batched_matmul" => { - let batch = 4usize; - let k = size.value + 32; - let lhs_shape = [batch, size.value, k]; - let rhs_shape = [batch, k, size.value]; - let lhs = burn_tensor( - deterministic_values(elements(&lhs_shape), 23, 0.004), - lhs_shape, - &device, - ); - let rhs = burn_tensor( - deterministic_values(elements(&rhs_shape), 24, 0.004), - rhs_shape, - &device, - ); - burn_materialize_inputs(&[lhs.clone(), rhs.clone()]).await?; - let samples = time_samples(config, || { - let output = lhs.clone().matmul(rhs.clone()); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!( - "{} @ {} f32", - shape_label(&lhs_shape), - shape_label(&rhs_shape) - ), - )) - } - "conv1d_small" => { - let input_shape = [4usize, 8usize, size.value]; - let weight_shape = [16usize, 8usize, 5usize]; - let bias_shape = [16usize]; - let input = burn_tensor( - deterministic_values(elements(&input_shape), 25, 0.01), - input_shape, - &device, - ); - let weight = burn_tensor( - deterministic_values(elements(&weight_shape), 26, 0.01), - weight_shape, - &device, - ); - let bias = burn_tensor( - deterministic_values(elements(&bias_shape), 27, 0.001), - bias_shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - burn_materialize_inputs(std::slice::from_ref(&weight)).await?; - burn_materialize_inputs(std::slice::from_ref(&bias)).await?; - let samples = time_samples(config, || { - let output = module::conv1d( - input.clone(), - weight.clone(), - Some(bias.clone()), - ConvOptions::new([2], [1], [1], 1), - ); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!( - "{} conv {}", - shape_label(&input_shape), - shape_label(&weight_shape) - ), - )) - } - "top_k_large" | "top_k_qwen_vocab" => { - let input_len = size.value; - let k = if case == "top_k_qwen_vocab" { 40 } else { 64 }; - let input = burn_tensor(topk_values(input_len), [input_len], &device); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = input.clone().topk(k, 0); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{input_len} logits, k={k}"), - )) + cases::dense_batched_matmul_case(config, name, 4, value, value + 32).await } - "q8_0_qgemv" | "q4k_qgemv" => { - let weight_shape = if case == "q8_0_qgemv" { - q8_shape(size.value) - } else { - q4_shape(size.value) - }; - let input_shape = [1usize, weight_shape[1]]; - let dense_weight_shape = [weight_shape[1], weight_shape[0]]; - let input = burn_tensor( - deterministic_values( - elements(&input_shape), - if case == "q8_0_qgemv" { 8 } else { 29 }, - 0.003, - ), - input_shape, - &device, - ); - let weights = burn_tensor( - deterministic_values( - elements(&dense_weight_shape), - if case == "q8_0_qgemv" { 80 } else { 81 }, - 0.003, - ), - dense_weight_shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - burn_materialize_inputs(std::slice::from_ref(&weights)).await?; - let samples = time_samples(config, || { - let output = input.clone().matmul(weights.clone()); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!( - "1x{} @ dense f32 {}", - weight_shape[1], - shape_label(&dense_weight_shape) - ), - )) + "conv1d_small" => cases::conv1d_small_case(config, name, value).await, + "top_k_large" => { + cases::top_k_case(config, name, value, 64, webgpu::topk_values(value)).await } - "q4k_paired_silu" => { - let weight_shape = [size.value * 2, 1024usize]; - let input_shape = [1usize, weight_shape[1]]; - let dense_weight_shape = [weight_shape[1], weight_shape[0]]; - let input = burn_tensor( - deterministic_values(elements(&input_shape), 30, 0.003), - input_shape, - &device, - ); - let weights = burn_tensor( - deterministic_values(elements(&dense_weight_shape), 82, 0.003), - dense_weight_shape, - &device, - ); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - burn_materialize_inputs(std::slice::from_ref(&weights)).await?; - let samples = time_samples(config, || { - let projected = input.clone().matmul(weights.clone()); - let gate = projected.clone().narrow(1, 0, size.value); - let up = projected.narrow(1, size.value, size.value); - let output = activation::silu(gate) * up; - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!( - "1x1024 @ dense f32 {} + paired SiLU", - shape_label(&dense_weight_shape) - ), - )) + "top_k_qwen_vocab" => { + cases::top_k_case(config, name, value, 40, webgpu::topk_values(value)).await } - "flash_attention_small" | "flash_attention_causal_small" => { - let seq_len = size.value; - let shape = [1usize, 4usize, seq_len, 64usize]; - let q = burn_tensor( - deterministic_values(elements(&shape), 31, 0.003), - shape, - &device, - ); - let k = burn_tensor( - deterministic_values(elements(&shape), 32, 0.003), - shape, - &device, - ); - let v = burn_tensor( - deterministic_values(elements(&shape), 33, 0.003), - shape, - &device, - ); - burn_materialize_inputs(&[q.clone(), k.clone(), v.clone()]).await?; - let samples = time_samples(config, || { - let output = module::attention( - q.clone(), - k.clone(), - v.clone(), - None, - None, - AttentionModuleOptions { - scale: Some(1.0 / (64.0f64).sqrt()), - softcap: None, - is_causal: case == "flash_attention_causal_small", - }, - ); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} attention", shape_label(&shape)), - )) + "q8_0_qgemv" => cases::qgemv_dense_case(config, name, q8_shape(value), 8, 80, "").await, + "q4k_qgemv" => cases::qgemv_dense_case(config, name, q4_shape(value), 29, 81, "").await, + "q4k_paired_silu" => cases::q4k_paired_silu_case(config, name, [value * 2, 1024], "").await, + "attention_small" => { + cases::attention_case(config, name, value, [31, 32, 33], false, "attention").await } - "rope_fused_decode" => { - let seq_len = size.value; - let shape = [1usize, 8usize, seq_len, 64usize]; - let [batch, heads, _, head_dim] = shape; - let input = burn_tensor( - deterministic_values(batch * heads * seq_len * head_dim, 9, 0.01), - shape, - &device, - ); - let rope = RotaryEncodingConfig::new(seq_len * 2, head_dim).init::(&device); - burn_materialize_inputs(std::slice::from_ref(&input)).await?; - let samples = time_samples(config, || { - let output = rope.clone().forward(input.clone()); - async move { burn_materialize(output).await } - }) - .await?; - Ok(sweep_report( - "burn", - case, - size, - config, - samples, - format!("{} rotary encoding", shape_label(&shape)), - )) + "attention_causal_small" => { + cases::attention_case(config, name, value, [31, 32, 33], true, "attention").await } + "rope_fused_decode" => cases::rope_fused_decode_case(config, name, value).await, _ => Err(format!("unknown Burn benchmark sweep: {case}").into()), } } diff --git a/fusor-ml/conformance/src/bench/webgpu.rs b/fusor-ml/conformance/src/bench/webgpu.rs index 2e36f23a5..82bd04c15 100644 --- a/fusor-ml/conformance/src/bench/webgpu.rs +++ b/fusor-ml/conformance/src/bench/webgpu.rs @@ -1,4 +1,9 @@ //! Concrete WebGPU benchmark cases. +//! +//! Each case body is parameterized by the sizes that vary between the fixed +//! registry entry and the per-size sweep, so both are thin callers into one +//! implementation. The `fixed_cases!` invocation at the bottom pins the sizes +//! every registry entry runs at. use fusor::{Device, GgmlType, MaskKind, QMatrix, Tensor}; @@ -22,6 +27,27 @@ fn deterministic_values(len: usize, seed: usize, scale: f32) -> Vec { .collect() } +pub(super) fn topk_values(input_len: usize) -> Vec { + (0..input_len) + .map(|index| { + let base = ((index * 67 + 29) % 10_007) as f32 * 0.001; + let bump = if index % 4099 == 0 { 20.0 } else { 0.0 }; + base + bump - (index % 13) as f32 * 0.0001 + }) + .collect() +} + +fn rope_values(shape: [usize; 2], head_dim: usize, cos: bool) -> Vec { + (0..shape[0]) + .flat_map(|i| { + (0..shape[1]).map(move |j| { + let value = (i as f32) / 10000f32.powf((2 * (j / 2)) as f32 / head_dim as f32); + if cos { value.cos() } else { value.sin() } + }) + }) + .collect() +} + fn shape_label(shape: &[usize]) -> String { shape .iter() @@ -34,17 +60,37 @@ fn elements(shape: &[usize]) -> usize { shape.iter().product() } -async fn materialize_inputs(inputs: &[&Tensor]) { - for input in inputs { - input.materialize().await; - } +async fn values_input( + device: &Device, + shape: [usize; R], + values: &[f32], +) -> Tensor { + let tensor: Tensor = Tensor::from_slice(device, shape, values); + tensor.materialize().await; + tensor +} + +async fn input_tensor( + device: &Device, + shape: [usize; R], + seed: usize, + scale: f32, +) -> Tensor { + values_input( + device, + shape, + &deterministic_values(elements(&shape), seed, scale), + ) + .await } fn bench_case( name: &'static str, - run: impl for<'a> FnOnce(&'a Device, BenchmarkConfig) -> super::CaseFuture<'a> + 'static, + run: impl for<'a> FnOnce(&'a Device, BenchmarkConfig, String) -> super::CaseFuture<'a> + 'static, ) -> BenchmarkCase { - BenchmarkCase::new(name, run) + BenchmarkCase::new(name, move |device, config| { + run(device, config, name.to_string()) + }) } pub async fn run_webgpu_bench_suite(device: &Device) -> BenchmarkResult> { @@ -60,729 +106,618 @@ pub async fn run_webgpu_bench_suite_with_progress( .await } -pub fn elementwise_add_square() -> BenchmarkCase { - bench_case("webgpu::elementwise_add_square", |device, config| { - Box::pin(async move { - let shape = [512usize, 512usize]; - let lhs_values = deterministic_values(elements(&shape), 1, 0.01); - let rhs_values = deterministic_values(elements(&shape), 2, 0.008); - let lhs: Tensor<2, f32> = Tensor::from_slice(device, shape, &lhs_values); - let rhs: Tensor<2, f32> = Tensor::from_slice(device, shape, &rhs_values); - materialize_inputs(&[&lhs, &rhs]).await; - - let samples = time_samples(config, || { - let output = (&lhs + &rhs).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::elementwise_add_square", - config, - samples, - format!("{} f32 add", shape_label(&shape)), - )) - }) +pub(super) async fn elementwise_add_square_case( + device: &Device, + config: BenchmarkConfig, + name: String, + size: usize, +) -> BenchmarkResult { + let shape = [size, size]; + let lhs = input_tensor(device, shape, 1, 0.01).await; + let rhs = input_tensor(device, shape, 2, 0.008).await; + let samples = time_samples(config, || { + let output = (&lhs + &rhs).to_concrete(); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} f32 add", shape_label(&shape)), + )) } -pub fn elementwise_mul_rank4() -> BenchmarkCase { - bench_case("webgpu::elementwise_mul_rank4", |device, config| { - Box::pin(async move { - let shape = [9usize, 11usize, 32usize, 16usize]; - let lhs_values = deterministic_values(elements(&shape), 3, 0.012); - let rhs_values = deterministic_values(elements(&shape), 4, 0.009); - let lhs: Tensor<4, f32> = Tensor::from_slice(device, shape, &lhs_values); - let rhs: Tensor<4, f32> = Tensor::from_slice(device, shape, &rhs_values); - materialize_inputs(&[&lhs, &rhs]).await; - - let samples = time_samples(config, || { - let output = (&lhs * &rhs).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::elementwise_mul_rank4", - config, - samples, - format!("{} f32 mul", shape_label(&shape)), - )) - }) +pub(super) async fn elementwise_mul_rank4_case( + device: &Device, + config: BenchmarkConfig, + name: String, + shape: [usize; 4], +) -> BenchmarkResult { + let lhs = input_tensor(device, shape, 3, 0.012).await; + let rhs = input_tensor(device, shape, 4, 0.009).await; + let samples = time_samples(config, || { + let output = (&lhs * &rhs).to_concrete(); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} f32 mul", shape_label(&shape)), + )) } -pub fn unary_trig_chain() -> BenchmarkCase { - bench_case("webgpu::unary_trig_chain", |device, config| { - Box::pin(async move { - let shape = [384usize, 384usize]; - let values = deterministic_values(elements(&shape), 10, 0.01); - let input: Tensor<2, f32> = Tensor::from_slice(device, shape, &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let output = (input.sin() + input.cos()).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::unary_trig_chain", - config, - samples, - format!("{} sin+cos", shape_label(&shape)), - )) - }) +pub(super) async fn unary_trig_chain_case( + device: &Device, + config: BenchmarkConfig, + name: String, + size: usize, +) -> BenchmarkResult { + let shape = [size, size]; + let input = input_tensor(device, shape, 10, 0.01).await; + let samples = time_samples(config, || { + let output = (input.sin() + input.cos()).to_concrete(); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} sin+cos", shape_label(&shape)), + )) } -pub fn activation_gelu() -> BenchmarkCase { - bench_case("webgpu::activation_gelu", |device, config| { - Box::pin(async move { - let shape = [512usize, 256usize]; - let values = deterministic_values(elements(&shape), 11, 0.015); - let input: Tensor<2, f32> = Tensor::from_slice(device, shape, &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let output = input.gelu(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::activation_gelu", - config, - samples, - format!("{} gelu", shape_label(&shape)), - )) - }) +pub(super) async fn activation_gelu_case( + device: &Device, + config: BenchmarkConfig, + name: String, + shape: [usize; 2], +) -> BenchmarkResult { + let input = input_tensor(device, shape, 11, 0.015).await; + let samples = time_samples(config, || { + let output = input.gelu(); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} gelu", shape_label(&shape)), + )) } -pub fn broadcast_add() -> BenchmarkCase { - bench_case("webgpu::broadcast_add", |device, config| { - Box::pin(async move { - let matrix_shape = [256usize, 512usize]; - let vector_shape = [512usize]; - let matrix_values = deterministic_values(elements(&matrix_shape), 12, 0.006); - let vector_values = deterministic_values(elements(&vector_shape), 13, 0.01); - let matrix: Tensor<2, f32> = Tensor::from_slice(device, matrix_shape, &matrix_values); - let vector: Tensor<1, f32> = Tensor::from_slice(device, vector_shape, &vector_values); - materialize_inputs(&[&matrix]).await; - materialize_inputs(&[&vector]).await; - - let samples = time_samples(config, || { - let vector_row = vector.reshape([1, vector_shape[0]]); - let vector_broadcast = vector_row.broadcast_as(matrix_shape); - let output = (&matrix + vector_broadcast).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::broadcast_add", - config, - samples, - "256x512 + broadcast 512", - )) - }) +pub(super) async fn broadcast_add_case( + device: &Device, + config: BenchmarkConfig, + name: String, + rows: usize, +) -> BenchmarkResult { + let matrix_shape = [rows, 512usize]; + let vector_shape = [512usize]; + let matrix = input_tensor(device, matrix_shape, 12, 0.006).await; + let vector = input_tensor(device, vector_shape, 13, 0.01).await; + let samples = time_samples(config, || { + let vector_row = vector.reshape([1, vector_shape[0]]); + let output = (&matrix + vector_row.broadcast_as(matrix_shape)).to_concrete(); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} + broadcast 512", shape_label(&matrix_shape)), + )) } -pub fn transpose_then_elementwise() -> BenchmarkCase { - bench_case("webgpu::transpose_then_elementwise", |device, config| { - Box::pin(async move { - let shape = [256usize, 384usize]; - let values = deterministic_values(elements(&shape), 14, 0.01); - let input: Tensor<2, f32> = Tensor::from_slice(device, shape, &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let transposed = input.transpose(0, 1); - let output = (transposed.clone() * transposed).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::transpose_then_elementwise", - config, - samples, - "256x384 transpose, square", - )) - }) +pub(super) async fn transpose_then_elementwise_case( + device: &Device, + config: BenchmarkConfig, + name: String, + shape: [usize; 2], +) -> BenchmarkResult { + let input = input_tensor(device, shape, 14, 0.01).await; + let samples = time_samples(config, || { + let transposed = input.transpose(0, 1); + let output = (transposed.clone() * transposed).to_concrete(); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} transpose, square", shape_label(&shape)), + )) } -pub fn reduction_sum_last_dim() -> BenchmarkCase { - bench_case("webgpu::reduction_sum_last_dim", |device, config| { - Box::pin(async move { - let shape = [256usize, 512usize]; - let values = deterministic_values(elements(&shape), 15, 0.004); - let input: Tensor<2, f32> = Tensor::from_slice(device, shape, &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let output = input.sum::<1>(1); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::reduction_sum_last_dim", - config, - samples, - format!("{} sum axis 1", shape_label(&shape)), - )) - }) +pub(super) async fn reduction_sum_last_dim_case( + device: &Device, + config: BenchmarkConfig, + name: String, + rows: usize, +) -> BenchmarkResult { + let shape = [rows, 512usize]; + let input = input_tensor(device, shape, 15, 0.004).await; + let samples = time_samples(config, || { + let output = input.sum::<1>(1); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} sum axis 1", shape_label(&shape)), + )) } -pub fn reduction_max_middle_axis() -> BenchmarkCase { - bench_case("webgpu::reduction_max_middle_axis", |device, config| { - Box::pin(async move { - let shape = [64usize, 128usize, 64usize]; - let values = deterministic_values(elements(&shape), 16, 0.004); - let input: Tensor<3, f32> = Tensor::from_slice(device, shape, &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let output = input.max::<2>(1); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::reduction_max_middle_axis", - config, - samples, - format!("{} max axis 1", shape_label(&shape)), - )) - }) +pub(super) async fn reduction_max_middle_axis_case( + device: &Device, + config: BenchmarkConfig, + name: String, + shape: [usize; 3], +) -> BenchmarkResult { + let input = input_tensor(device, shape, 16, 0.004).await; + let samples = time_samples(config, || { + let output = input.max::<2>(1); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} max axis 1", shape_label(&shape)), + )) } -pub fn softmax_last_dim() -> BenchmarkCase { - bench_case("webgpu::softmax_last_dim", |device, config| { - Box::pin(async move { - let shape = [512usize, 256usize]; - let values = deterministic_values(elements(&shape), 5, 0.006); - let input: Tensor<2, f32> = Tensor::from_slice(device, shape, &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let output = input.softmax_last_dim::<1>(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::softmax_last_dim", - config, - samples, - format!("{} last-axis softmax", shape_label(&shape)), - )) - }) +pub(super) async fn softmax_last_dim_case( + device: &Device, + config: BenchmarkConfig, + name: String, + rows: usize, +) -> BenchmarkResult { + let shape = [rows, 256usize]; + let input = input_tensor(device, shape, 5, 0.006).await; + let samples = time_samples(config, || { + let output = input.softmax_last_dim::<1>(); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} last-axis softmax", shape_label(&shape)), + )) } -pub fn softmax_middle_axis() -> BenchmarkCase { - bench_case("webgpu::softmax_middle_axis", |device, config| { - Box::pin(async move { - let shape = [32usize, 128usize, 64usize]; - let values = deterministic_values(elements(&shape), 17, 0.004); - let input: Tensor<3, f32> = Tensor::from_slice(device, shape, &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let output = input.softmax::<2>(1); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::softmax_middle_axis", - config, - samples, - format!("{} softmax axis 1", shape_label(&shape)), - )) - }) +pub(super) async fn softmax_middle_axis_case( + device: &Device, + config: BenchmarkConfig, + name: String, + shape: [usize; 3], +) -> BenchmarkResult { + let input = input_tensor(device, shape, 17, 0.004).await; + let samples = time_samples(config, || { + let output = input.softmax::<2>(1); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} softmax axis 1", shape_label(&shape)), + )) } -pub fn layer_norm_last_dim() -> BenchmarkCase { - bench_case("webgpu::layer_norm_last_dim", |device, config| { - Box::pin(async move { - let shape = [8usize, 128usize, 512usize]; - let last_dim = shape[2]; - let values = deterministic_values(elements(&shape), 18, 0.01); - let weight_values = deterministic_values(last_dim, 19, 0.002) - .into_iter() - .map(|value| value + 1.0) - .collect::>(); - let bias_values = deterministic_values(last_dim, 20, 0.001); - let input: Tensor<3, f32> = Tensor::from_slice(device, shape, &values); - let weight: Tensor<1, f32> = Tensor::from_slice(device, [last_dim], &weight_values); - let bias: Tensor<1, f32> = Tensor::from_slice(device, [last_dim], &bias_values); - materialize_inputs(&[&input]).await; - materialize_inputs(&[&weight, &bias]).await; - - let samples = time_samples(config, || { - let output = - input.layer_norm_last_dim_fused::<2, 1, _, _>(&weight, Some(&bias), 1.0e-5); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::layer_norm_last_dim", - config, - samples, - format!("{} layer norm", shape_label(&shape)), - )) - }) +pub(super) async fn layer_norm_last_dim_case( + device: &Device, + config: BenchmarkConfig, + name: String, + shape: [usize; 3], +) -> BenchmarkResult { + let last_dim = shape[2]; + let input = input_tensor(device, shape, 18, 0.01).await; + let weight_values = deterministic_values(last_dim, 19, 0.002) + .into_iter() + .map(|value| value + 1.0) + .collect::>(); + let weight = values_input(device, [last_dim], &weight_values).await; + let bias = input_tensor(device, [last_dim], 20, 0.001).await; + let samples = time_samples(config, || { + let output = input.layer_norm_last_dim_fused::<2, 1, _, _>(&weight, Some(&bias), 1.0e-5); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} layer norm", shape_label(&shape)), + )) } -pub fn rms_norm_fused() -> BenchmarkCase { - bench_case("webgpu::rms_norm_fused", |device, config| { - Box::pin(async move { - let shape = [8usize, 128usize, 512usize]; - let last_dim = shape[2]; - let values = deterministic_values(elements(&shape), 21, 0.01); - let weight_values = deterministic_values(last_dim, 22, 0.002) - .into_iter() - .map(|value| value + 1.0) - .collect::>(); - let input: Tensor<3, f32> = Tensor::from_slice(device, shape, &values); - let weight: Tensor<1, f32> = Tensor::from_slice(device, [last_dim], &weight_values); - materialize_inputs(&[&input]).await; - materialize_inputs(&[&weight]).await; - - let samples = time_samples(config, || { - let output = input.rms_norm_fused_no_bias::<1, 2>(&weight, 1.0e-5); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::rms_norm_fused", - config, - samples, - format!("{} rms norm", shape_label(&shape)), - )) - }) +pub(super) async fn rms_norm_fused_case( + device: &Device, + config: BenchmarkConfig, + name: String, + shape: [usize; 3], +) -> BenchmarkResult { + let last_dim = shape[2]; + let input = input_tensor(device, shape, 21, 0.01).await; + let weight_values = deterministic_values(last_dim, 22, 0.002) + .into_iter() + .map(|value| value + 1.0) + .collect::>(); + let weight = values_input(device, [last_dim], &weight_values).await; + let samples = time_samples(config, || { + let output = input.rms_norm_fused_no_bias::<1, 2>(&weight, 1.0e-5); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} rms norm", shape_label(&shape)), + )) } -pub fn dense_matmul_square() -> BenchmarkCase { - bench_case("webgpu::dense_matmul_square", |device, config| { - Box::pin(async move { - let lhs_shape = [256usize, 256usize]; - let rhs_shape = [256usize, 256usize]; - let lhs_values = deterministic_values(elements(&lhs_shape), 6, 0.004); - let rhs_values = deterministic_values(elements(&rhs_shape), 7, 0.004); - let lhs: Tensor<2, f32> = Tensor::from_slice(device, lhs_shape, &lhs_values); - let rhs: Tensor<2, f32> = Tensor::from_slice(device, rhs_shape, &rhs_values); - materialize_inputs(&[&lhs, &rhs]).await; - - let samples = time_samples(config, || { - let output = lhs.matmul(&rhs); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::dense_matmul_square", - config, - samples, - "256x256 @ 256x256 f32", - )) - }) +pub(super) async fn dense_matmul_square_case( + device: &Device, + config: BenchmarkConfig, + name: String, + size: usize, +) -> BenchmarkResult { + let lhs_shape = [size, size]; + let rhs_shape = [size, size]; + let lhs = input_tensor(device, lhs_shape, 6, 0.004).await; + let rhs = input_tensor(device, rhs_shape, 7, 0.004).await; + let samples = time_samples(config, || { + let output = lhs.matmul(&rhs); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "{} @ {} f32", + shape_label(&lhs_shape), + shape_label(&rhs_shape) + ), + )) } -pub fn dense_batched_matmul() -> BenchmarkCase { - bench_case("webgpu::dense_batched_matmul", |device, config| { - Box::pin(async move { - let lhs_shape = [8usize, 64usize, 96usize]; - let rhs_shape = [8usize, 96usize, 64usize]; - let lhs_values = deterministic_values(elements(&lhs_shape), 23, 0.004); - let rhs_values = deterministic_values(elements(&rhs_shape), 24, 0.004); - let lhs: Tensor<3, f32> = Tensor::from_slice(device, lhs_shape, &lhs_values); - let rhs: Tensor<3, f32> = Tensor::from_slice(device, rhs_shape, &rhs_values); - materialize_inputs(&[&lhs, &rhs]).await; - - let samples = time_samples(config, || { - let output = lhs.matmul(&rhs); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::dense_batched_matmul", - config, - samples, - "8x64x96 @ 8x96x64 f32", - )) - }) +pub(super) async fn dense_batched_matmul_case( + device: &Device, + config: BenchmarkConfig, + name: String, + batch: usize, + m: usize, + k: usize, +) -> BenchmarkResult { + let lhs_shape = [batch, m, k]; + let rhs_shape = [batch, k, m]; + let lhs = input_tensor(device, lhs_shape, 23, 0.004).await; + let rhs = input_tensor(device, rhs_shape, 24, 0.004).await; + let samples = time_samples(config, || { + let output = lhs.matmul(&rhs); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "{} @ {} f32", + shape_label(&lhs_shape), + shape_label(&rhs_shape) + ), + )) } -pub fn conv1d_small() -> BenchmarkCase { - bench_case("webgpu::conv1d_small", |device, config| { - Box::pin(async move { - let input_shape = [4usize, 8usize, 256usize]; - let weight_shape = [16usize, 8usize, 5usize]; - let bias_shape = [16usize]; - let input_values = deterministic_values(elements(&input_shape), 25, 0.01); - let weight_values = deterministic_values(elements(&weight_shape), 26, 0.01); - let bias_values = deterministic_values(elements(&bias_shape), 27, 0.001); - let input: Tensor<3, f32> = Tensor::from_slice(device, input_shape, &input_values); - let weight: Tensor<3, f32> = Tensor::from_slice(device, weight_shape, &weight_values); - let bias: Tensor<1, f32> = Tensor::from_slice(device, bias_shape, &bias_values); - materialize_inputs(&[&input]).await; - materialize_inputs(&[&weight]).await; - materialize_inputs(&[&bias]).await; - - let samples = time_samples(config, || { - let output = input.conv(&weight, Some(&bias), [2], [1]); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::conv1d_small", - config, - samples, - "4x8x256 conv 16x8x5", - )) - }) +pub(super) async fn conv1d_small_case( + device: &Device, + config: BenchmarkConfig, + name: String, + len: usize, +) -> BenchmarkResult { + let input_shape = [4usize, 8usize, len]; + let weight_shape = [16usize, 8usize, 5usize]; + let bias_shape = [16usize]; + let input = input_tensor(device, input_shape, 25, 0.01).await; + let weight = input_tensor(device, weight_shape, 26, 0.01).await; + let bias = input_tensor(device, bias_shape, 27, 0.001).await; + let samples = time_samples(config, || { + let output = input.conv(&weight, Some(&bias), [2], [1]); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "{} conv {}", + shape_label(&input_shape), + shape_label(&weight_shape) + ), + )) } -pub fn top_k_large() -> BenchmarkCase { - bench_case("webgpu::top_k_large", |device, config| { - Box::pin(async move { - let input_len = 65_537usize; - let k = 64usize; - let values = (0..input_len) - .map(|index| { - let base = ((index * 67 + 29) % 10_007) as f32 * 0.001; - let bump = if index % 4099 == 0 { 20.0 } else { 0.0 }; - base + bump - (index % 13) as f32 * 0.0001 - }) - .collect::>(); - let input: Tensor<1, f32> = Tensor::from_slice(device, [input_len], &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || async { - let top = input.top_k_pairs(k).await?; - if top.len() != k { - return Err(format!("top_k returned {} pairs, expected {k}", top.len()).into()); - } - Ok(()) - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::top_k_large", - config, - samples, - format!("{input_len} logits, k={k}"), - )) - }) +pub(super) async fn top_k_case( + device: &Device, + config: BenchmarkConfig, + name: String, + input_len: usize, + k: usize, + values: Vec, +) -> BenchmarkResult { + let input = values_input(device, [input_len], &values).await; + let samples = time_samples(config, || async { + let top = input.top_k_pairs(k).await?; + if top.len() != k { + return Err(format!("top_k returned {} pairs, expected {k}", top.len()).into()); + } + Ok(()) }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{input_len} logits, k={k}"), + )) } -pub fn top_k_qwen_vocab() -> BenchmarkCase { - bench_case("webgpu::top_k_qwen_vocab", |device, config| { - Box::pin(async move { - let input_len = 151_936usize; - let k = 40usize; - let values = deterministic_values(input_len, 28, 0.01); - let input: Tensor<1, f32> = Tensor::from_slice(device, [input_len], &values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || async { - let top = input.top_k_pairs(k).await?; - if top.len() != k { - return Err(format!("top_k returned {} pairs, expected {k}", top.len()).into()); - } - Ok(()) - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::top_k_qwen_vocab", - config, - samples, - format!("{input_len} logits, k={k}"), - )) - }) +pub(super) async fn q8_0_qgemv_case( + device: &Device, + config: BenchmarkConfig, + name: String, + weight_shape: [usize; 2], +) -> BenchmarkResult { + let input_shape = [1usize, weight_shape[1]]; + let raw_bytes = q8_0_raw_bytes(weight_shape); + let matrix: QMatrix = qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q8_0); + let input = input_tensor(device, input_shape, 8, 0.003).await; + let samples = time_samples(config, || { + let output = input.q_mat_mul(&matrix); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "1x{} @ Q8_0 {}", + weight_shape[1], + shape_label(&weight_shape) + ), + )) } -pub fn q8_0_qgemv() -> BenchmarkCase { - bench_case("webgpu::q8_0_qgemv", |device, config| { - Box::pin(async move { - let weight_shape = [4096usize, 896usize]; - let input_shape = [1usize, weight_shape[1]]; - let raw_bytes = q8_0_raw_bytes(weight_shape); - let matrix: QMatrix = - qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q8_0); - let input_values = deterministic_values(elements(&input_shape), 8, 0.003); - let input: Tensor<2, f32> = Tensor::from_slice(device, input_shape, &input_values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let output = input.q_mat_mul(&matrix); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::q8_0_qgemv", - config, - samples, - "1x896 @ Q8_0 4096x896", - )) - }) +pub(super) async fn q4k_qgemv_case( + device: &Device, + config: BenchmarkConfig, + name: String, + weight_shape: [usize; 2], +) -> BenchmarkResult { + let input_shape = [1usize, weight_shape[1]]; + let raw_bytes = q4k_raw_bytes(weight_shape); + let matrix: QMatrix = qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q4K); + let input = input_tensor(device, input_shape, 29, 0.003).await; + let samples = time_samples(config, || { + let output = input.q_mat_mul(&matrix); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("1x{} @ Q4K {}", weight_shape[1], shape_label(&weight_shape)), + )) } -pub fn q4k_qgemv() -> BenchmarkCase { - bench_case("webgpu::q4k_qgemv", |device, config| { - Box::pin(async move { - let weight_shape = [2048usize, 1024usize]; - let input_shape = [1usize, weight_shape[1]]; - let raw_bytes = q4k_raw_bytes(weight_shape); - let matrix: QMatrix = - qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q4K); - let input_values = deterministic_values(elements(&input_shape), 29, 0.003); - let input: Tensor<2, f32> = Tensor::from_slice(device, input_shape, &input_values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let output = input.q_mat_mul(&matrix); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::q4k_qgemv", - config, - samples, - "1x1024 @ Q4K 2048x1024", - )) - }) +pub(super) async fn q4k_paired_silu_case( + device: &Device, + config: BenchmarkConfig, + name: String, + weight_shape: [usize; 2], +) -> BenchmarkResult { + let input_shape = [1usize, weight_shape[1]]; + let raw_bytes = q4k_raw_bytes(weight_shape); + let matrix: QMatrix = qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q4K); + let input = input_tensor(device, input_shape, 30, 0.003).await; + let samples = time_samples(config, || { + let pair_len = weight_shape[0] / 2; + let projected = input.q_mat_mul(&matrix); + let gate = projected + .narrow(fusor::D::Minus1, 0, pair_len) + .to_concrete(); + let up = projected + .narrow(fusor::D::Minus1, pair_len, pair_len) + .to_concrete(); + let output = (gate.silu() * up).to_concrete(); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!( + "1x{} @ paired Q4K {}", + weight_shape[1], + shape_label(&weight_shape) + ), + )) } -pub fn q4k_paired_silu() -> BenchmarkCase { - bench_case("webgpu::q4k_paired_silu", |device, config| { - Box::pin(async move { - let weight_shape = [2048usize, 1024usize]; - let input_shape = [1usize, weight_shape[1]]; - let raw_bytes = q4k_raw_bytes(weight_shape); - let matrix: QMatrix = - qmatrix_from_raw_bytes(device, weight_shape, &raw_bytes, GgmlType::Q4K); - let input_values = deterministic_values(elements(&input_shape), 30, 0.003); - let input: Tensor<2, f32> = Tensor::from_slice(device, input_shape, &input_values); - materialize_inputs(&[&input]).await; - - let samples = time_samples(config, || { - let pair_len = weight_shape[0] / 2; - let projected = input.q_mat_mul(&matrix); - let gate = projected - .narrow(fusor::D::Minus1, 0, pair_len) - .to_concrete(); - let up = projected - .narrow(fusor::D::Minus1, pair_len, pair_len) - .to_concrete(); - let output = (gate.silu() * up).to_concrete(); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::q4k_paired_silu", - config, - samples, - "1x1024 @ paired Q4K 2048x1024", - )) - }) +pub(super) async fn attention_case( + device: &Device, + config: BenchmarkConfig, + name: String, + seq_len: usize, + seeds: [usize; 3], + causal: bool, + detail_op: &'static str, +) -> BenchmarkResult { + let shape = [1usize, 4usize, seq_len, 64usize]; + let q = input_tensor(device, shape, seeds[0], 0.003).await; + let k = input_tensor(device, shape, seeds[1], 0.003).await; + let v = input_tensor(device, shape, seeds[2], 0.003).await; + let mask_shape = [seq_len, seq_len]; + let mask = if causal { + Some(values_input(device, mask_shape, &vec![0.0f32; elements(&mask_shape)]).await) + } else { + None + }; + let samples = time_samples(config, || { + let mask_arg = mask.as_ref().map(|mask| (mask, MaskKind::Causal)); + let output = q.attention(&k, &v, 1.0 / (64.0f32).sqrt(), mask_arg); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} {detail_op}", shape_label(&shape)), + )) } -pub fn flash_attention_small() -> BenchmarkCase { - bench_case("webgpu::flash_attention_small", |device, config| { - Box::pin(async move { - let shape = [1usize, 4usize, 128usize, 64usize]; - let q_values = deterministic_values(elements(&shape), 31, 0.003); - let k_values = deterministic_values(elements(&shape), 32, 0.003); - let v_values = deterministic_values(elements(&shape), 33, 0.003); - let q: Tensor<4, f32> = Tensor::from_slice(device, shape, &q_values); - let k: Tensor<4, f32> = Tensor::from_slice(device, shape, &k_values); - let v: Tensor<4, f32> = Tensor::from_slice(device, shape, &v_values); - materialize_inputs(&[&q, &k, &v]).await; - - let samples = time_samples(config, || { - let output = q.flash_attention(&k, &v, 1.0 / (64.0f32).sqrt(), None); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::flash_attention_small", - config, - samples, - format!("{} flash attention", shape_label(&shape)), - )) - }) +pub(super) async fn rope_fused_decode_case( + device: &Device, + config: BenchmarkConfig, + name: String, + seq_len: usize, +) -> BenchmarkResult { + let shape = [1usize, 8usize, seq_len, 64usize]; + let [_, _, _, head_dim] = shape; + let pos_shape = [seq_len * 2, head_dim / 2]; + let input = input_tensor(device, shape, 9, 0.01).await; + let cos = values_input(device, pos_shape, &rope_values(pos_shape, head_dim, true)).await; + let sin = values_input(device, pos_shape, &rope_values(pos_shape, head_dim, false)).await; + let samples = time_samples(config, || { + let output = input.rope_fused(&cos, &sin); + async move { + output.materialize().await; + Ok(()) + } }) + .await?; + Ok(BenchmarkReport::new( + name, + config, + samples, + format!("{} fused rope", shape_label(&shape)), + )) } -pub fn flash_attention_causal_small() -> BenchmarkCase { - bench_case("webgpu::flash_attention_causal_small", |device, config| { - Box::pin(async move { - let shape = [1usize, 4usize, 128usize, 64usize]; - let mask_shape = [128usize, 128usize]; - let q_values = deterministic_values(elements(&shape), 34, 0.003); - let k_values = deterministic_values(elements(&shape), 35, 0.003); - let v_values = deterministic_values(elements(&shape), 36, 0.003); - let mask_values = vec![0.0f32; elements(&mask_shape)]; - let q: Tensor<4, f32> = Tensor::from_slice(device, shape, &q_values); - let k: Tensor<4, f32> = Tensor::from_slice(device, shape, &k_values); - let v: Tensor<4, f32> = Tensor::from_slice(device, shape, &v_values); - let mask: Tensor<2, f32> = Tensor::from_slice(device, mask_shape, &mask_values); - materialize_inputs(&[&q, &k, &v]).await; - materialize_inputs(&[&mask]).await; - - let samples = time_samples(config, || { - let output = q.flash_attention( - &k, - &v, - 1.0 / (64.0f32).sqrt(), - Some((&mask, MaskKind::Causal)), - ); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::flash_attention_causal_small", - config, - samples, - format!("{} causal flash attention", shape_label(&shape)), - )) - }) - }) +macro_rules! fixed_cases { + ($($case:ident => $body:ident($($arg:expr),* $(,)?);)*) => { + $( + pub fn $case() -> BenchmarkCase { + bench_case( + concat!("webgpu::", stringify!($case)), + |device, config, name| Box::pin($body(device, config, name, $($arg),*)), + ) + } + )* + }; } -pub fn rope_fused_decode() -> BenchmarkCase { - bench_case("webgpu::rope_fused_decode", |device, config| { - Box::pin(async move { - let shape = [1usize, 8usize, 256usize, 64usize]; - let [batch, heads, seq_len, head_dim] = shape; - let pos_shape = [seq_len * 2, head_dim / 2]; - let cos_values = (0..pos_shape[0]) - .flat_map(|i| { - (0..pos_shape[1]).map(move |j| { - ((i as f32) / 10000f32.powf((2 * (j / 2)) as f32 / head_dim as f32)).cos() - }) - }) - .collect::>(); - let sin_values = (0..pos_shape[0]) - .flat_map(|i| { - (0..pos_shape[1]).map(move |j| { - ((i as f32) / 10000f32.powf((2 * (j / 2)) as f32 / head_dim as f32)).sin() - }) - }) - .collect::>(); - let input_values = deterministic_values(batch * heads * seq_len * head_dim, 9, 0.01); - let input: Tensor<4, f32> = Tensor::from_slice(device, shape, &input_values); - let cos: Tensor<2, f32> = Tensor::from_slice(device, pos_shape, &cos_values); - let sin: Tensor<2, f32> = Tensor::from_slice(device, pos_shape, &sin_values); - materialize_inputs(&[&input]).await; - materialize_inputs(&[&cos, &sin]).await; - - let samples = time_samples(config, || { - let output = input.rope_fused(&cos, &sin); - async move { - output.materialize().await; - Ok(()) - } - }) - .await?; - - Ok(BenchmarkReport::new( - "webgpu::rope_fused_decode", - config, - samples, - format!("{} fused rope", shape_label(&shape)), - )) - }) - }) +fixed_cases! { + elementwise_add_square => elementwise_add_square_case(512); + elementwise_mul_rank4 => elementwise_mul_rank4_case([9, 11, 32, 16]); + unary_trig_chain => unary_trig_chain_case(384); + activation_gelu => activation_gelu_case([512, 256]); + broadcast_add => broadcast_add_case(256); + transpose_then_elementwise => transpose_then_elementwise_case([256, 384]); + reduction_sum_last_dim => reduction_sum_last_dim_case(256); + reduction_max_middle_axis => reduction_max_middle_axis_case([64, 128, 64]); + softmax_last_dim => softmax_last_dim_case(512); + softmax_middle_axis => softmax_middle_axis_case([32, 128, 64]); + layer_norm_last_dim => layer_norm_last_dim_case([8, 128, 512]); + rms_norm_fused => rms_norm_fused_case([8, 128, 512]); + dense_matmul_square => dense_matmul_square_case(256); + dense_batched_matmul => dense_batched_matmul_case(8, 64, 96); + conv1d_small => conv1d_small_case(256); + top_k_large => top_k_case(65_537, 64, topk_values(65_537)); + top_k_qwen_vocab => top_k_case(151_936, 40, deterministic_values(151_936, 28, 0.01)); + q8_0_qgemv => q8_0_qgemv_case([4096, 896]); + q4k_qgemv => q4k_qgemv_case([2048, 1024]); + q4k_paired_silu => q4k_paired_silu_case([2048, 1024]); + attention_small => attention_case(128, [31, 32, 33], false, "attention"); + attention_causal_small => attention_case(128, [34, 35, 36], true, "causal attention"); + rope_fused_decode => rope_fused_decode_case(256); } diff --git a/fusor-ml/conformance/src/common/quantized.rs b/fusor-ml/conformance/src/common/quantized.rs index 13caad635..a851068bd 100644 --- a/fusor-ml/conformance/src/common/quantized.rs +++ b/fusor-ml/conformance/src/common/quantized.rs @@ -121,7 +121,7 @@ pub fn assert_q_mat_mul_matches_host_reference( fixture: &QuantizedFixture, fuzz: QMatMulFuzz, ) -> AssertionCase { - use fusor::ToVec2; + use fusor::ToVec; let ty = fixture.ty; let weight_shape = fixture.weight_shape; @@ -148,7 +148,7 @@ pub fn assert_q_mat_mul_matches_host_reference( let expected_weights = expected_weights.clone(); async move { let device = input.device(); - let input_values = input.as_slice().await.unwrap().to_vec2(); + let input_values = input.as_slice().await.unwrap().to_vec(); Tensor::new( &device, &super::matmul2(&input_values, &super::transpose2(&expected_weights)), diff --git a/fusor-ml/conformance/src/goldens.rs b/fusor-ml/conformance/src/goldens.rs new file mode 100644 index 000000000..74cdb8a6b --- /dev/null +++ b/fusor-ml/conformance/src/goldens.rs @@ -0,0 +1,214 @@ +//! Machine-pinned output-hash goldens. +//! +//! Each test replays one deterministic GPU trace and compares an FNV-1a hash +//! of the exact output bytes against a golden captured on the baseline +//! machine (`goldens/`). The traces are bit-reproducible for a fixed +//! device/driver, so any hash drift means a numeric change: refactors that +//! claim "no numeric change" must keep these green, and intentional numeric +//! changes must re-capture the goldens from the failure output. + +use fusor::autograd::Graph; +use fusor::{Device, GgmlType, MaskKind, Tensor}; + +use crate::common::quantized::{q4k_raw_bytes, q6k_raw_bytes, qmatrix_from_raw_bytes}; + +fn fnv1a(bytes: &[u8]) -> u64 { + let mut hash = 0xcbf29ce484222325u64; + for &byte in bytes { + hash ^= byte as u64; + hash = hash.wrapping_mul(0x100000001b3); + } + hash +} + +fn fill(seed: u32, len: usize) -> Vec { + let mut state = seed as u64 | 1; + (0..len) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5 + }) + .collect() +} + +async fn tensor_hash(tensor: &Tensor) -> u64 { + let len = tensor.shape().iter().product(); + let flat: Tensor<1, f32> = tensor.reshape([len]).to_concrete(); + let values = flat.as_slice().await.unwrap().as_slice().to_vec(); + let mut bytes = Vec::with_capacity(values.len() * 4); + for value in &values { + bytes.extend_from_slice(&value.to_le_bytes()); + } + fnv1a(&bytes) +} + +fn assert_matches_golden(name: &str, golden: &str, actual: &str) { + assert!( + golden.trim() == actual.trim(), + "{name} golden mismatch; measured values:\n{actual}" + ); +} + +/// One grouped-query attention config with a causal additive mask, forward +/// and backward: 8 query heads over 2 key/value heads takes the composite +/// replay backward, so the hashes pin both the fused forward kernel and the +/// recomputed probability chain. +#[allow(clippy::await_holding_lock)] +#[tokio::test] +async fn attention_gqa_causal_fwd_bwd_matches_golden() { + let _gpu_guard = crate::suite::registry::gpu_test_guard(); + let Ok(device) = Device::gpu().await else { + return; + }; + const BATCH: usize = 2; + const Q_HEADS: usize = 8; + const KV_HEADS: usize = 2; + const SEQ: usize = 32; + const HEAD_DIM: usize = 64; + let scale = 1.0 / (HEAD_DIM as f32).sqrt(); + + let mask_values: Vec = (0..SEQ * SEQ) + .map(|i| if i % SEQ <= i / SEQ { 0.0 } else { -1e9 }) + .collect(); + let mask: Tensor<2, f32> = Tensor::from_slice(&device, [SEQ, SEQ], &mask_values); + + let graph = Graph::new(); + let q = graph.leaf(Tensor::from_slice( + &device, + [BATCH, Q_HEADS, SEQ, HEAD_DIM], + &fill(11, BATCH * Q_HEADS * SEQ * HEAD_DIM), + )); + let k = graph.leaf(Tensor::from_slice( + &device, + [BATCH, KV_HEADS, SEQ, HEAD_DIM], + &fill(23, BATCH * KV_HEADS * SEQ * HEAD_DIM), + )); + let v = graph.leaf(Tensor::from_slice( + &device, + [BATCH, KV_HEADS, SEQ, HEAD_DIM], + &fill(37, BATCH * KV_HEADS * SEQ * HEAD_DIM), + )); + + let out = q.attention(&k, &v, scale, Some((&mask, MaskKind::Causal))); + let gradients = out + .reshape([BATCH * Q_HEADS * SEQ * HEAD_DIM]) + .sum() + .backward() + .unwrap(); + let dq = gradients.get(&q).expect("missing q gradient"); + let dk = gradients.get(&k).expect("missing k gradient"); + let dv = gradients.get(&v).expect("missing v gradient"); + + let actual = [ + format!("out {:#018x}", tensor_hash(out.raw()).await), + format!("dq {:#018x}", tensor_hash(&dq).await), + format!("dk {:#018x}", tensor_hash(&dk).await), + format!("dv {:#018x}", tensor_hash(&dv).await), + ] + .join("\n"); + assert_matches_golden( + "attention_gqa_causal", + include_str!("../goldens/attention_gqa_causal.txt"), + &actual, + ); +} + +/// A bidirectional recurrent trace in the bilstm style: per-timestep narrows +/// off one input, a matmul + recurrent matmul + tanh per step, and the step +/// outputs reassembled with `cat` (per direction, then across directions). +#[allow(clippy::await_holding_lock)] +#[tokio::test] +async fn bilstm_split_op_cat_trace_matches_golden() { + let _gpu_guard = crate::suite::registry::gpu_test_guard(); + let Ok(device) = Device::gpu().await else { + return; + }; + const STEPS: usize = 8; + const BATCH: usize = 4; + const FEATURES: usize = 32; + const HIDDEN: usize = 16; + + let x: Tensor<3, f32> = Tensor::from_slice( + &device, + [STEPS, BATCH, FEATURES], + &fill(5, STEPS * BATCH * FEATURES), + ); + let direction = |input_seed: u32, recurrent_seed: u32, reverse: bool| { + let input_weight: Tensor<2, f32> = Tensor::from_slice( + &device, + [FEATURES, HIDDEN], + &fill(input_seed, FEATURES * HIDDEN), + ); + let recurrent_weight: Tensor<2, f32> = Tensor::from_slice( + &device, + [HIDDEN, HIDDEN], + &fill(recurrent_seed, HIDDEN * HIDDEN), + ); + let mut hidden: Tensor<2, f32> = Tensor::zeros(&device, [BATCH, HIDDEN]); + let mut outputs = vec![Tensor::zeros(&device, [1, BATCH, HIDDEN]); STEPS]; + for step in 0..STEPS { + let step = if reverse { STEPS - 1 - step } else { step }; + let input = x + .narrow(0, step, 1) + .reshape([BATCH, FEATURES]) + .to_concrete(); + hidden = (input.mat_mul(&input_weight) + hidden.mat_mul(&recurrent_weight)) + .tanh() + .to_concrete(); + outputs[step] = hidden.reshape([1, BATCH, HIDDEN]).to_concrete(); + } + Tensor::cat(outputs, 0) + }; + let forward = direction(41, 43, false); + let backward = direction(47, 53, true); + let out = Tensor::cat(vec![forward, backward], 2); + + let actual = format!("out {:#018x}", tensor_hash(&out).await); + assert_matches_golden( + "bilstm_trace", + include_str!("../goldens/bilstm_trace.txt"), + &actual, + ); +} + +/// Every Q4K/Q6K decode (M=1 qgemv) shape exercised by the ggml qgemv +/// lowering suite (`tile-ir-kernels/tests/lowering.rs`), run end-to-end with +/// deterministic weights: the main 4096x8192 shape, the tail-column and mid +/// variants, and the Q6K shape. +#[allow(clippy::await_holding_lock)] +#[tokio::test] +async fn qgemv_decode_ggml_shapes_match_golden() { + let _gpu_guard = crate::suite::registry::gpu_test_guard(); + let Ok(device) = Device::gpu().await else { + return; + }; + let cases: [(GgmlType, &str, usize, usize); 4] = [ + (GgmlType::Q4K, "q4k", 4096, 8192), + (GgmlType::Q4K, "q4k", 4096, 8193), + (GgmlType::Q4K, "q4k", 4096, 5120), + (GgmlType::Q6K, "q6k", 4096, 8192), + ]; + let mut actual = Vec::new(); + for (ty, label, rows, cols) in cases { + let weight_shape = [cols, rows]; + let raw_bytes = match ty { + GgmlType::Q4K => q4k_raw_bytes(weight_shape), + GgmlType::Q6K => q6k_raw_bytes(weight_shape), + _ => unreachable!(), + }; + let weights = qmatrix_from_raw_bytes(&device, weight_shape, &raw_bytes, ty); + let input: Tensor<2, f32> = Tensor::from_slice(&device, [1, rows], &fill(3, rows)); + let out = input.q_mat_mul(&weights).to_concrete(); + actual.push(format!( + "{label}_{rows}x{cols} {:#018x}", + tensor_hash(&out).await + )); + } + assert_matches_golden( + "qgemv_decode_ggml", + include_str!("../goldens/qgemv_decode_ggml.txt"), + &actual.join("\n"), + ); +} diff --git a/fusor-ml/conformance/src/lib.rs b/fusor-ml/conformance/src/lib.rs index 3f9a267c5..2be396a82 100644 --- a/fusor-ml/conformance/src/lib.rs +++ b/fusor-ml/conformance/src/lib.rs @@ -6,6 +6,8 @@ mod builder; mod comparison; mod fuzz; +#[cfg(test)] +mod goldens; mod table; mod tuple_macros; diff --git a/fusor-ml/conformance/src/suite/native/flash_attention_ops.rs b/fusor-ml/conformance/src/suite/native/attention_ops.rs similarity index 83% rename from fusor-ml/conformance/src/suite/native/flash_attention_ops.rs rename to fusor-ml/conformance/src/suite/native/attention_ops.rs index 388ee206d..3557e0955 100644 --- a/fusor-ml/conformance/src/suite/native/flash_attention_ops.rs +++ b/fusor-ml/conformance/src/suite/native/attention_ops.rs @@ -1,4 +1,4 @@ -//! Flash attention conformance cases. +//! Attention conformance cases. use fusor::{Device, MaskKind, Tensor}; use fusor_conformance::{ @@ -8,7 +8,7 @@ use fusor_conformance::{ use half::f16; #[derive(Clone, Copy)] -struct FlashCase { +struct AttentionCase { batch: usize, num_heads: usize, num_kv_heads: usize, @@ -45,9 +45,9 @@ fn batch_key_mask_data(batch: usize, kv_seq_len: usize) -> Vec { data } -fn assert_flash_attention_case_f16( +fn assert_attention_case_f16( name: impl Into, - case: FlashCase, + case: AttentionCase, mask: Option<(Vec, MaskKind, [usize; 2])>, tol: f16, ) -> AssertionCase { @@ -108,10 +108,10 @@ fn assert_flash_attention_case_f16( if let Some((mask_data, kind, shape)) = mask.as_ref() { let mask_f16: Vec = mask_data.iter().copied().map(f16::from_f32).collect(); let device_mask: Tensor<2, f16> = Tensor::from_slice(&device, *shape, &mask_f16); - q.flash_attention(&k, &v, scale, Some((&device_mask, *kind))) + q.attention(&k, &v, scale, Some((&device_mask, *kind))) .to_concrete() } else { - q.flash_attention(&k, &v, scale, None).to_concrete() + q.attention(&k, &v, scale, None).to_concrete() } } }) @@ -121,16 +121,16 @@ fn assert_flash_attention_case_f16( .runs(1); // Chrome/SwiftShader browser WebGPU has been observed to hang/crash in - // flash-attention GPU kernels. Native builds still run GPU variants. + // attention GPU kernels. Native builds still run GPU variants. #[cfg(target_arch = "wasm32")] let assertion = assertion.devices([Device::Cpu]); assertion.into_case(name) } -pub fn flash_attention_f16_matches_cpu_reference_on_varied_shapes() -> AssertionCases { +pub fn attention_f16_matches_cpu_reference_on_varied_shapes() -> AssertionCases { [ - FlashCase { + AttentionCase { batch: 1, num_heads: 1, num_kv_heads: 1, @@ -138,7 +138,7 @@ pub fn flash_attention_f16_matches_cpu_reference_on_varied_shapes() -> Assertion kv_seq_len: 2, head_dim: 2, }, - FlashCase { + AttentionCase { batch: 2, num_heads: 2, num_kv_heads: 2, @@ -146,7 +146,7 @@ pub fn flash_attention_f16_matches_cpu_reference_on_varied_shapes() -> Assertion kv_seq_len: 5, head_dim: 3, }, - FlashCase { + AttentionCase { batch: 1, num_heads: 2, num_kv_heads: 1, @@ -157,9 +157,9 @@ pub fn flash_attention_f16_matches_cpu_reference_on_varied_shapes() -> Assertion ] .into_iter() .map(|case| { - assert_flash_attention_case_f16( + assert_attention_case_f16( format!( - "flash_attention_ops::flash_attention_f16_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", + "attention_ops::attention_f16_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", case.batch, case.num_heads, case.num_kv_heads, @@ -176,9 +176,9 @@ pub fn flash_attention_f16_matches_cpu_reference_on_varied_shapes() -> Assertion .into() } -pub fn flash_attention_f16_with_qk_mask_matches_cpu_reference() -> AssertionCases { +pub fn attention_f16_with_qk_mask_matches_cpu_reference() -> AssertionCases { [ - FlashCase { + AttentionCase { batch: 1, num_heads: 1, num_kv_heads: 1, @@ -186,7 +186,7 @@ pub fn flash_attention_f16_with_qk_mask_matches_cpu_reference() -> AssertionCase kv_seq_len: 2, head_dim: 2, }, - FlashCase { + AttentionCase { batch: 1, num_heads: 2, num_kv_heads: 2, @@ -198,9 +198,9 @@ pub fn flash_attention_f16_with_qk_mask_matches_cpu_reference() -> AssertionCase .into_iter() .map(|case| { let shape = [case.q_seq_len, case.kv_seq_len]; - assert_flash_attention_case_f16( + assert_attention_case_f16( format!( - "flash_attention_ops::flash_attention_f16_with_qk_mask_matches_cpu_reference::b{}_h{}_kvh{}_q{}_kv{}_d{}", + "attention_ops::attention_f16_with_qk_mask_matches_cpu_reference::b{}_h{}_kvh{}_q{}_kv{}_d{}", case.batch, case.num_heads, case.num_kv_heads, @@ -221,9 +221,9 @@ pub fn flash_attention_f16_with_qk_mask_matches_cpu_reference() -> AssertionCase .into() } -fn assert_flash_attention_case( +fn assert_attention_case( name: impl Into, - case: FlashCase, + case: AttentionCase, mask: Option<(Vec, MaskKind, [usize; 2])>, tol: f32, runs: usize, @@ -275,10 +275,10 @@ fn assert_flash_attention_case( ); if let Some((mask_data, kind, shape)) = mask.as_ref() { let device_mask = Tensor::from_slice(&device, *shape, mask_data); - q.flash_attention(&k, &v, scale, Some((&device_mask, *kind))) + q.attention(&k, &v, scale, Some((&device_mask, *kind))) .to_concrete() } else { - q.flash_attention(&k, &v, scale, None).to_concrete() + q.attention(&k, &v, scale, None).to_concrete() } } }) @@ -287,7 +287,7 @@ fn assert_flash_attention_case( .runs(runs); // Chrome/SwiftShader browser WebGPU has been observed to hang/crash in - // flash-attention GPU kernels. Native builds still run GPU variants. + // attention GPU kernels. Native builds still run GPU variants. #[cfg(target_arch = "wasm32")] let assertion = assertion.devices([Device::Cpu]); @@ -302,7 +302,7 @@ fn assert_flash_attention_case( /// /// Each shape is run multiple times because earlier decode failures were /// non-deterministic workgroup-memory races. -pub fn flash_attention_decode_tiled_matches_cpu_reference() -> AssertionCases { +pub fn attention_decode_tiled_matches_cpu_reference() -> AssertionCases { // (num_heads, num_kv_heads, kv_seq_len) // Shapes specifically chosen to stress decode tile boundaries in the // attention row program: kv_seq_len spans ragged single tiles and @@ -323,7 +323,7 @@ pub fn flash_attention_decode_tiled_matches_cpu_reference() -> AssertionCases { shapes .into_iter() .map(|(num_heads, num_kv_heads, kv_seq_len)| { - let case = FlashCase { + let case = AttentionCase { batch: 1, num_heads, num_kv_heads, @@ -332,9 +332,9 @@ pub fn flash_attention_decode_tiled_matches_cpu_reference() -> AssertionCases { head_dim: 128, }; let tol = if kv_seq_len > 1024 { 5e-3 } else { 1e-3 }; - assert_flash_attention_case( + assert_attention_case( format!( - "flash_attention_ops::flash_attention_decode_tiled_matches_cpu_reference::h{num_heads}_kvh{num_kv_heads}_kv{kv_seq_len}" + "attention_ops::attention_decode_tiled_matches_cpu_reference::h{num_heads}_kvh{num_kv_heads}_kv{kv_seq_len}" ), case, None, @@ -351,7 +351,7 @@ pub fn flash_attention_decode_tiled_matches_cpu_reference() -> AssertionCases { /// in `models/kalosm-llama/src/raw/attention_layer.rs`. The row program /// reads Q through its layout strides, so different strides hit different /// memory addresses. -pub fn flash_attention_decode_tiled_with_transposed_q_matches_cpu_reference() -> AssertionCases { +pub fn attention_decode_tiled_with_transposed_q_matches_cpu_reference() -> AssertionCases { let shapes = [(16, 2, 129), (16, 2, 257), (16, 2, 384), (16, 2, 569)]; shapes @@ -389,7 +389,7 @@ pub fn flash_attention_decode_tiled_with_transposed_q_matches_cpu_reference() -> Tensor::from_slice(&device, [batch, num_kv_heads, kv_seq_len, head_dim], &k_data); let v: Tensor<4, f32> = Tensor::from_slice(&device, [batch, num_kv_heads, kv_seq_len, head_dim], &v_data); - q.flash_attention(&k, &v, scale, None).to_concrete() + q.attention(&k, &v, scale, None).to_concrete() } }) .arg(|device: &Device| device.clone()) @@ -405,20 +405,20 @@ pub fn flash_attention_decode_tiled_with_transposed_q_matches_cpu_reference() -> Tensor::from_slice(&device, [batch, num_kv_heads, kv_seq_len, head_dim], &k_data); let v: Tensor<4, f32> = Tensor::from_slice(&device, [batch, num_kv_heads, kv_seq_len, head_dim], &v_data); - q.flash_attention(&k, &v, scale, None).to_concrete() + q.attention(&k, &v, scale, None).to_concrete() } }) .compare_with(approx_compare::<4, f32>(1e-3)) .runs(4) .into_case(format!( - "flash_attention_ops::flash_attention_decode_tiled_with_transposed_q_matches_cpu_reference::h{num_heads}_kvh{num_kv_heads}_kv{kv_seq_len}" + "attention_ops::attention_decode_tiled_with_transposed_q_matches_cpu_reference::h{num_heads}_kvh{num_kv_heads}_kv{kv_seq_len}" )) }) .collect::>() .into() } -pub fn flash_attention_subgroup_fallback_preserves_gpu_backend() -> AssertionCase { +pub fn attention_without_subgroups_preserves_gpu_backend() -> AssertionCase { let q_shape = [1, 1, 2, 4]; let kv_shape = [1, 1, 3, 4]; let q_data = attention_data(q_shape.iter().product(), 0.1); @@ -434,7 +434,7 @@ pub fn flash_attention_subgroup_fallback_preserves_gpu_backend() -> AssertionCas let q = Tensor::from_slice(&device, q_shape, &q_data); let k = Tensor::from_slice(&device, kv_shape, &k_data); let v = Tensor::from_slice(&device, kv_shape, &v_data); - q.flash_attention(&k, &v, scale, None).is_gpu() + q.attention(&k, &v, scale, None).is_gpu() } }) .arg(|device: &Device| device.clone()) @@ -453,12 +453,12 @@ pub fn flash_attention_subgroup_fallback_preserves_gpu_backend() -> AssertionCas .collect() }) .runs(1) - .into_case("flash_attention_ops::flash_attention_subgroup_fallback_preserves_gpu_backend") + .into_case("attention_ops::attention_without_subgroups_preserves_gpu_backend") } -pub fn flash_attention_matches_cpu_reference_on_varied_shapes() -> AssertionCases { +pub fn attention_matches_cpu_reference_on_varied_shapes() -> AssertionCases { [ - FlashCase { + AttentionCase { batch: 1, num_heads: 1, num_kv_heads: 1, @@ -466,7 +466,7 @@ pub fn flash_attention_matches_cpu_reference_on_varied_shapes() -> AssertionCase kv_seq_len: 2, head_dim: 2, }, - FlashCase { + AttentionCase { batch: 2, num_heads: 2, num_kv_heads: 2, @@ -474,7 +474,7 @@ pub fn flash_attention_matches_cpu_reference_on_varied_shapes() -> AssertionCase kv_seq_len: 5, head_dim: 3, }, - FlashCase { + AttentionCase { batch: 1, num_heads: 3, num_kv_heads: 3, @@ -482,7 +482,7 @@ pub fn flash_attention_matches_cpu_reference_on_varied_shapes() -> AssertionCase kv_seq_len: 3, head_dim: 4, }, - FlashCase { + AttentionCase { batch: 1, num_heads: 2, num_kv_heads: 1, @@ -493,9 +493,9 @@ pub fn flash_attention_matches_cpu_reference_on_varied_shapes() -> AssertionCase ] .into_iter() .map(|case| { - assert_flash_attention_case( + assert_attention_case( format!( - "flash_attention_ops::flash_attention_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", + "attention_ops::attention_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", case.batch, case.num_heads, case.num_kv_heads, @@ -513,9 +513,9 @@ pub fn flash_attention_matches_cpu_reference_on_varied_shapes() -> AssertionCase .into() } -pub fn flash_attention_with_qk_mask_matches_cpu_reference_on_varied_shapes() -> AssertionCases { +pub fn attention_with_qk_mask_matches_cpu_reference_on_varied_shapes() -> AssertionCases { [ - FlashCase { + AttentionCase { batch: 1, num_heads: 1, num_kv_heads: 1, @@ -523,7 +523,7 @@ pub fn flash_attention_with_qk_mask_matches_cpu_reference_on_varied_shapes() -> kv_seq_len: 2, head_dim: 2, }, - FlashCase { + AttentionCase { batch: 2, num_heads: 3, num_kv_heads: 3, @@ -531,7 +531,7 @@ pub fn flash_attention_with_qk_mask_matches_cpu_reference_on_varied_shapes() -> kv_seq_len: 6, head_dim: 3, }, - FlashCase { + AttentionCase { batch: 1, num_heads: 2, num_kv_heads: 2, @@ -543,9 +543,9 @@ pub fn flash_attention_with_qk_mask_matches_cpu_reference_on_varied_shapes() -> .into_iter() .map(|case| { let shape = [case.q_seq_len, case.kv_seq_len]; - assert_flash_attention_case( + assert_attention_case( format!( - "flash_attention_ops::flash_attention_with_qk_mask_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", + "attention_ops::attention_with_qk_mask_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", case.batch, case.num_heads, case.num_kv_heads, @@ -567,9 +567,9 @@ pub fn flash_attention_with_qk_mask_matches_cpu_reference_on_varied_shapes() -> .into() } -pub fn flash_attention_gqa_matches_cpu_reference_on_varied_shapes() -> AssertionCases { +pub fn attention_gqa_matches_cpu_reference_on_varied_shapes() -> AssertionCases { [ - FlashCase { + AttentionCase { batch: 1, num_heads: 4, num_kv_heads: 2, @@ -577,7 +577,7 @@ pub fn flash_attention_gqa_matches_cpu_reference_on_varied_shapes() -> Assertion kv_seq_len: 2, head_dim: 2, }, - FlashCase { + AttentionCase { batch: 2, num_heads: 6, num_kv_heads: 2, @@ -585,7 +585,7 @@ pub fn flash_attention_gqa_matches_cpu_reference_on_varied_shapes() -> Assertion kv_seq_len: 5, head_dim: 3, }, - FlashCase { + AttentionCase { batch: 1, num_heads: 8, num_kv_heads: 4, @@ -596,9 +596,9 @@ pub fn flash_attention_gqa_matches_cpu_reference_on_varied_shapes() -> Assertion ] .into_iter() .map(|case| { - assert_flash_attention_case( + assert_attention_case( format!( - "flash_attention_ops::flash_attention_gqa_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", + "attention_ops::attention_gqa_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", case.batch, case.num_heads, case.num_kv_heads, @@ -616,12 +616,12 @@ pub fn flash_attention_gqa_matches_cpu_reference_on_varied_shapes() -> Assertion .into() } -pub fn flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes() -> AssertionCases { +pub fn attention_with_kv_cache_matches_cpu_reference_on_varied_shapes() -> AssertionCases { // KV-cache regression: short Q sequence with a longer K/V sequence — the // typical autoregressive decode shape after appending to a KvCache. - // Replaces the deleted `core/src/composite/flash_attention.rs::test_flash_attention_kv_cache_fuzz`. + // Replaces the deleted `core/src/composite/attention.rs::test_attention_kv_cache_fuzz`. [ - FlashCase { + AttentionCase { batch: 1, num_heads: 1, num_kv_heads: 1, @@ -629,7 +629,7 @@ pub fn flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes() -> kv_seq_len: 5, head_dim: 4, }, - FlashCase { + AttentionCase { batch: 2, num_heads: 4, num_kv_heads: 4, @@ -637,7 +637,7 @@ pub fn flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes() -> kv_seq_len: 16, head_dim: 8, }, - FlashCase { + AttentionCase { batch: 1, num_heads: 32, num_kv_heads: 8, @@ -645,7 +645,7 @@ pub fn flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes() -> kv_seq_len: 10, head_dim: 128, }, - FlashCase { + AttentionCase { batch: 2, num_heads: 8, num_kv_heads: 8, @@ -653,7 +653,7 @@ pub fn flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes() -> kv_seq_len: 17, head_dim: 16, }, - FlashCase { + AttentionCase { batch: 1, num_heads: 6, num_kv_heads: 2, @@ -668,9 +668,9 @@ pub fn flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes() -> // an all-zero mask. let mask = vec![0.0f32; case.q_seq_len * case.kv_seq_len]; let shape = [case.q_seq_len, case.kv_seq_len]; - assert_flash_attention_case( + assert_attention_case( format!( - "flash_attention_ops::flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", + "attention_ops::attention_with_kv_cache_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", case.batch, case.num_heads, case.num_kv_heads, @@ -688,10 +688,9 @@ pub fn flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes() -> .into() } -pub fn flash_attention_with_batch_key_mask_matches_cpu_reference_on_varied_shapes() -> AssertionCases -{ +pub fn attention_with_batch_key_mask_matches_cpu_reference_on_varied_shapes() -> AssertionCases { [ - FlashCase { + AttentionCase { batch: 2, num_heads: 1, num_kv_heads: 1, @@ -699,7 +698,7 @@ pub fn flash_attention_with_batch_key_mask_matches_cpu_reference_on_varied_shape kv_seq_len: 3, head_dim: 2, }, - FlashCase { + AttentionCase { batch: 3, num_heads: 2, num_kv_heads: 2, @@ -707,7 +706,7 @@ pub fn flash_attention_with_batch_key_mask_matches_cpu_reference_on_varied_shape kv_seq_len: 5, head_dim: 3, }, - FlashCase { + AttentionCase { batch: 2, num_heads: 4, num_kv_heads: 4, @@ -719,9 +718,9 @@ pub fn flash_attention_with_batch_key_mask_matches_cpu_reference_on_varied_shape .into_iter() .map(|case| { let shape = [case.batch, case.kv_seq_len]; - assert_flash_attention_case( + assert_attention_case( format!( - "flash_attention_ops::flash_attention_with_batch_key_mask_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", + "attention_ops::attention_with_batch_key_mask_matches_cpu_reference_on_varied_shapes::b{}_h{}_kvh{}_q{}_kv{}_d{}", case.batch, case.num_heads, case.num_kv_heads, @@ -743,7 +742,7 @@ pub fn flash_attention_with_batch_key_mask_matches_cpu_reference_on_varied_shape .into() } -/// Exercises the tiled (Q-batched) streaming flash attention kernel. The +/// Exercises the tiled (Q-batched) streaming attention kernel. The /// selector switches to that variant when `q_seq_len >= 64` and `head_dim` is /// a multiple of 8. Shapes are chosen to span: /// - exact Q-block alignment (q_seq_len = 64 = 8*8), @@ -752,15 +751,15 @@ pub fn flash_attention_with_batch_key_mask_matches_cpu_reference_on_varied_shape /// /// The trailing `masked` sub-cases additionally drive the tiled kernel through /// the additive-causal QKMask path (previously a separate -/// `flash_attention_tiled_with_mask_matches_cpu_reference` test), keeping the +/// `attention_tiled_with_mask_matches_cpu_reference` test), keeping the /// tighter 1e-3 tolerance. -pub fn flash_attention_tiled_matches_cpu_reference_on_varied_shapes() -> AssertionCases { +pub fn attention_tiled_matches_cpu_reference_on_varied_shapes() -> AssertionCases { // (case, masked): unmasked shapes span the Q-block alignment cases above; // masked shapes exercise the additive-causal QKMask path through the same // tiled kernel. [ ( - FlashCase { + AttentionCase { batch: 1, num_heads: 2, num_kv_heads: 2, @@ -771,7 +770,7 @@ pub fn flash_attention_tiled_matches_cpu_reference_on_varied_shapes() -> Asserti false, ), ( - FlashCase { + AttentionCase { batch: 1, num_heads: 2, num_kv_heads: 1, @@ -782,7 +781,7 @@ pub fn flash_attention_tiled_matches_cpu_reference_on_varied_shapes() -> Asserti false, ), ( - FlashCase { + AttentionCase { batch: 2, num_heads: 4, num_kv_heads: 2, @@ -793,7 +792,7 @@ pub fn flash_attention_tiled_matches_cpu_reference_on_varied_shapes() -> Asserti false, ), ( - FlashCase { + AttentionCase { batch: 1, num_heads: 4, num_kv_heads: 4, @@ -804,7 +803,7 @@ pub fn flash_attention_tiled_matches_cpu_reference_on_varied_shapes() -> Asserti false, ), ( - FlashCase { + AttentionCase { batch: 1, num_heads: 2, num_kv_heads: 2, @@ -815,7 +814,7 @@ pub fn flash_attention_tiled_matches_cpu_reference_on_varied_shapes() -> Asserti true, ), ( - FlashCase { + AttentionCase { batch: 2, num_heads: 2, num_kv_heads: 2, @@ -830,7 +829,7 @@ pub fn flash_attention_tiled_matches_cpu_reference_on_varied_shapes() -> Asserti .map(|(case, masked)| { let prefix = if masked { "masked::" } else { "" }; let name = format!( - "flash_attention_ops::flash_attention_tiled_matches_cpu_reference_on_varied_shapes::{prefix}b{}_h{}_kvh{}_q{}_kv{}_d{}", + "attention_ops::attention_tiled_matches_cpu_reference_on_varied_shapes::{prefix}b{}_h{}_kvh{}_q{}_kv{}_d{}", case.batch, case.num_heads, case.num_kv_heads, @@ -845,7 +844,7 @@ pub fn flash_attention_tiled_matches_cpu_reference_on_varied_shapes() -> Asserti [case.q_seq_len, case.kv_seq_len], ) }); - assert_flash_attention_case(name, case, mask, 1e-3, 1) + assert_attention_case(name, case, mask, 1e-3, 1) }) .collect::>() .into() diff --git a/fusor-ml/conformance/src/suite/native/fusion_behavior.rs b/fusor-ml/conformance/src/suite/native/fusion_behavior.rs index 713b8f947..d5688cf51 100644 --- a/fusor-ml/conformance/src/suite/native/fusion_behavior.rs +++ b/fusor-ml/conformance/src/suite/native/fusion_behavior.rs @@ -2,7 +2,8 @@ use fusor::{Device, Tensor}; use fusor_conformance::{ - AssertionCase, AssertionCases, approx_compare, available_devices, exact_value_compare, + AssertionCase, AssertionCases, approx_compare, approx_or_relative_compare, available_devices, + exact_value_compare, }; async fn gpu_devices() -> Vec { @@ -127,9 +128,7 @@ pub fn gpu_nary_triple_add_fuses_into_one_kernel() -> AssertionCases { let c = Tensor::from_slice(&device, shape, &kernel_c_data); let sum = &a + &b; let result = &sum + &c; - result - .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + result.as_gpu().is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); } @@ -164,9 +163,7 @@ pub fn gpu_nary_unary_chain_fuses_into_one_kernel() -> AssertionCases { let b = Tensor::from_slice(&device, shape, &kernel_b_data); let sum = (-a.clone()) + b.sin(); let result = sum.cos() + 1.0; - result - .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + result.as_gpu().is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); } @@ -195,7 +192,7 @@ pub fn gpu_nary_same_input_multiple_times_deduplicates_bindings() -> AssertionCa let result = &sum + &a; result .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + .is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); } @@ -232,16 +229,14 @@ pub fn gpu_nary_where_cond_fuses_into_one_kernel() -> AssertionCases { let on_true = Tensor::from_slice(&device, shape, &kernel_on_true_data); let on_false = Tensor::from_slice(&device, shape, &kernel_on_false_data); let result = condition.where_cond(&on_true, &on_false); - result - .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + result.as_gpu().is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); } assertions } -pub fn gpu_flash_attention_fuses_into_one_kernel() -> AssertionCases { +pub fn gpu_attention_fuses_into_one_kernel() -> AssertionCases { let q_shape = [1, 2, 3, 4]; let kv_shape = [1, 2, 5, 4]; let q_data = attention_data(q_shape.iter().product(), 0.1); @@ -254,17 +249,17 @@ pub fn gpu_flash_attention_fuses_into_one_kernel() -> AssertionCases { let mut assertions = AssertionCases::new(); assertions.push(assert_gpu_tensor_case( - "fusion_behavior::gpu_flash_attention_fuses_into_one_kernel::correctness", + "fusion_behavior::gpu_attention_fuses_into_one_kernel::correctness", move |device| { let q = Tensor::from_slice(&device, q_shape, &q_data); let k = Tensor::from_slice(&device, kv_shape, &k_data); let v = Tensor::from_slice(&device, kv_shape, &v_data); - q.flash_attention(&k, &v, scale, None).to_concrete() + q.attention(&k, &v, scale, None).to_concrete() }, 1e-4, )); assertions.push(assert_gpu_kernel_property( - "fusion_behavior::gpu_flash_attention_fuses_into_one_kernel::kernels", + "fusion_behavior::gpu_attention_fuses_into_one_kernel::kernels", move |device| { let Some(gpu) = device.as_gpu() else { return true; @@ -275,9 +270,9 @@ pub fn gpu_flash_attention_fuses_into_one_kernel() -> AssertionCases { let q = Tensor::from_slice(&device, q_shape, &kernel_q_data); let k = Tensor::from_slice(&device, kv_shape, &kernel_k_data); let v = Tensor::from_slice(&device, kv_shape, &kernel_v_data); - q.flash_attention(&k, &v, scale, None) + q.attention(&k, &v, scale, None) .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + .is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); assertions @@ -316,7 +311,7 @@ pub fn gpu_residual_rms_norm_fuses_into_one_kernel() -> AssertionCases { input .rms_norm_residual_fused::<1, 2, _>(&residual, &weight, None, 1e-5) .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + .is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); assertions @@ -343,7 +338,12 @@ pub fn gpu_nary_fusion_respects_binding_limit() -> AssertionCases { let num_tensors = nary_binding_limit_stress_input_count(&device).unwrap_or(1); Tensor::from_slice(&device, shape, &binding_limit_sum_data(shape, num_tensors)) }) - .compare_with(approx_compare::<2, f32>(5e-6)) + // Exceeding the binding budget is the point of the case, so the sum + // splits into differently associated partial sums than the reference + // accumulates. The element magnitude grows with the adapter's budget + // (one f32 ulp already exceeds an absolute 5e-6 by the hundredth + // term), so the accumulation term has to be relative. + .compare_with(approx_or_relative_compare::<2>(5e-6, 1e-5)) .baseline_on_test_device() .devices_async(gpu_devices()) .runs(1) @@ -389,7 +389,7 @@ pub fn gpu_gelu_lowers_to_one_kernel() -> AssertionCases { Tensor::from_slice(&device, shape, &kernel_data) .gelu() .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + .is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); } @@ -420,9 +420,7 @@ pub fn gpu_matmul_then_unary_chain_fuses_into_one_kernel() -> AssertionCases { let b = Tensor::from_slice(&device, b_shape, &kernel_b_data); let matmul = a.mat_mul(&b); let result = matmul.cos() + 1.0; - result - .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + result.as_gpu().is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); assertions @@ -453,7 +451,39 @@ pub fn gpu_unary_inputs_fuse_into_matmul_kernel() -> AssertionCases { (-a.clone()) .mat_mul(&b.sin()) .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + .is_some_and(|gpu| gpu.resolves_in::<1>()) + }, + )); + assertions +} + +pub fn gpu_coop_matmul_fuses_pre_and_post_unary_chains() -> AssertionCases { + // Large enough to select a cooperative tile on capable adapters. The + // expression exercises both operand staging hooks and the in-dispatch + // post-store map. + let shape = [64, 64]; + let a_data = matrix_data(shape, 0.07); + let b_data = matrix_data(shape, -0.03); + let kernel_a_data = a_data.clone(); + let kernel_b_data = b_data.clone(); + let mut assertions = AssertionCases::new(); + assertions.push(assert_gpu_tensor_case( + "fusion_behavior::gpu_coop_matmul_fuses_pre_and_post_unary_chains::correctness", + move |device| { + let a = Tensor::from_slice(&device, shape, &a_data); + let b = Tensor::from_slice(&device, shape, &b_data); + ((-a).mat_mul(&b.sin()).cos() + 1.0).to_concrete() + }, + 1e-3, + )); + assertions.push(assert_gpu_kernel_property( + "fusion_behavior::gpu_coop_matmul_fuses_pre_and_post_unary_chains::kernels", + move |device| { + let a = Tensor::from_slice(&device, shape, &kernel_a_data); + let b = Tensor::from_slice(&device, shape, &kernel_b_data); + let matmul = (-a).mat_mul(&b.sin()); + let out = matmul.cos() + 1.0; + out.as_gpu().is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); assertions @@ -478,9 +508,7 @@ pub fn gpu_reduce_then_unary_chain_fuses_into_one_kernel() -> AssertionCases { let tensor = Tensor::from_slice(&device, shape, &kernel_data); let reduced = tensor.sum::<1>(0); let result = reduced.cos() + 1.0; - result - .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 1) + result.as_gpu().is_some_and(|gpu| gpu.resolves_in::<1>()) }, )); assertions @@ -528,9 +556,7 @@ pub fn gpu_reduce_then_gelu_uses_two_kernels() -> AssertionCases { .sum_keepdim::<1>(0) .gelu(); // Resize between Reduce and Gelu prevents fusion of the two kernels. - result - .as_gpu() - .is_some_and(|gpu| gpu.count_kernels_to_resolve() == 2) + result.as_gpu().is_some_and(|gpu| gpu.resolves_in::<2>()) }, )); } diff --git a/fusor-ml/conformance/src/suite/native/layer_ops.rs b/fusor-ml/conformance/src/suite/native/layer_ops.rs index 9f3b4f14f..8ed7160dc 100644 --- a/fusor-ml/conformance/src/suite/native/layer_ops.rs +++ b/fusor-ml/conformance/src/suite/native/layer_ops.rs @@ -2,7 +2,7 @@ use fusor::{ Device, Tensor, - layers::{Conv1d, Conv1dConfig, Embedding, LayerNorm, RmsNorm}, + layers::{ConvNd, ConvNdConfig, Embedding, LayerNorm, RmsNorm}, }; use fusor_conformance::{ AssertionCase, AssertionCases, approx_compare, exact_compare, exact_value_compare, @@ -37,10 +37,10 @@ fn assert_conv1d_case(case: ConvCase) -> AssertionCase { -0.35, ); let bias_data = case.with_bias.then(|| layer_data(case.out_channels, 0.1)); - let config = Conv1dConfig { - padding: case.padding, - stride: case.stride, - ..Default::default() + let config = ConvNdConfig { + padding: [case.padding], + stride: [case.stride], + groups: 1, }; fusor_conformance::assert(move |device: Device| { @@ -61,7 +61,7 @@ fn assert_conv1d_case(case: ConvCase) -> AssertionCase { let bias = bias_data .as_ref() .map(|data| Tensor::from_slice(&device, [case.out_channels], data)); - Conv1d::new(weight, bias, config) + ConvNd::new(weight, bias, config) .forward(&input) .to_concrete() } @@ -214,7 +214,7 @@ fn assert_layer_norm_2d_case(batch: usize, features: usize, with_bias: bool) -> .map(|data| Tensor::from_slice(&device, [features], data)); let layer_norm = LayerNorm::new(weight, bias, 1e-5); let input = Tensor::from_slice(&device, [batch, features], &input_data); - layer_norm.forward_2d(&input).to_concrete() + layer_norm.forward(&input).to_concrete() } }) .arg(|device: &Device| device.clone()) @@ -273,7 +273,7 @@ fn assert_rms_norm_2d_case(batch: usize, features: usize, with_bias: bool) -> As .map(|data| Tensor::from_slice(&device, [features], data)); let rms_norm = RmsNorm::new(weight, bias, 1e-5); let input = Tensor::from_slice(&device, [batch, features], &input_data); - rms_norm.forward_2d(&input).to_concrete() + rms_norm.forward(&input).to_concrete() } }) .arg(|device: &Device| device.clone()) @@ -338,7 +338,7 @@ fn assert_rms_norm_4d_case( .map(|data| Tensor::from_slice(&device, [features], data)); let rms_norm = RmsNorm::new(weight, bias, 1e-5); let input = Tensor::from_slice(&device, [batch, heads, seq_len, features], &input_data); - rms_norm.forward_4d(&input).to_concrete() + rms_norm.forward(&input).to_concrete() } }) .arg(|device: &Device| device.clone()) @@ -399,22 +399,22 @@ pub fn conv1d_properties_match_configuration() -> AssertionCases { [out_channels, in_channels, kernel_size], &vec![0.0f32; out_channels * in_channels * kernel_size], ); - let conv = Conv1d::new( + let conv = ConvNd::new( weight, None, - Conv1dConfig { - padding, - stride, - ..Default::default() + ConvNdConfig { + padding: [padding], + stride: [stride], + groups: 1, }, ); ( conv.in_channels(), conv.out_channels(), - conv.kernel_size(), - conv.config().padding, - conv.config().stride, + conv.weight().shape()[2], + conv.config().padding[0], + conv.config().stride[0], ) }) .arg(|device: &Device| device.clone()) diff --git a/fusor-ml/conformance/src/suite/native/mod.rs b/fusor-ml/conformance/src/suite/native/mod.rs index bcb8c3c93..35190f648 100644 --- a/fusor-ml/conformance/src/suite/native/mod.rs +++ b/fusor-ml/conformance/src/suite/native/mod.rs @@ -5,10 +5,10 @@ //! one `#[tokio::test]` per case; the cases that aren't (`should_panic`, //! `vision_block_pattern`) carry their own `#[cfg(test)]` test module. +pub mod attention_ops; pub mod cache_ops; pub mod dtypes; pub mod elementwise_ops; -pub mod flash_attention_ops; pub mod fusion_behavior; pub mod fusion_correctness; pub mod layer_ops; diff --git a/fusor-ml/conformance/src/suite/native/quantized_matmul_batched.rs b/fusor-ml/conformance/src/suite/native/quantized_matmul_batched.rs index b4474e026..98403653f 100644 --- a/fusor-ml/conformance/src/suite/native/quantized_matmul_batched.rs +++ b/fusor-ml/conformance/src/suite/native/quantized_matmul_batched.rs @@ -4,7 +4,7 @@ use crate::common::quantized::{ deterministic_input, q4k_raw_bytes, q8_0_raw_bytes, qmatrix_from_raw_bytes, }; use crate::common::{matmul2, transpose2}; -use fusor::{BlockQ4K, Device, GgmlType, GgufBlock, Tensor, ToVec2}; +use fusor::{BlockQ4K, Device, GgmlType, GgufBlock, Tensor, ToVec}; use fusor_conformance::{AssertionCase, AssertionCases, approx_compare, available_devices}; use std::mem::size_of; @@ -53,7 +53,7 @@ fn assert_q_mat_mul_3d_contiguous(input_rows: usize, batch: usize) -> AssertionC .as_slice() .await .unwrap() - .to_vec2(); + .to_vec(); let weights_t = transpose2(&dequantized_rows); let mut expected_rows = Vec::with_capacity(batch); for b in 0..batch { @@ -114,7 +114,7 @@ fn assert_q_mat_mul_3d_transposed(input_rows: usize, batch: usize) -> AssertionC .as_slice() .await .unwrap() - .to_vec2(); + .to_vec(); let weights_t = transpose2(&dequantized_rows); let mut expected_rows = Vec::with_capacity(batch); for b in 0..batch { @@ -329,7 +329,7 @@ pub fn q_mat_mul_batched_matches_unbatched_property() -> AssertionCase { let unbatched: Tensor<2, f32> = Tensor::from_slice(&device, [input_rows, weight_shape[1]], &slice_data); let result = unbatched.q_mat_mul(&weights).to_concrete(); - expected.push(result.as_slice().await.unwrap().to_vec2()); + expected.push(result.as_slice().await.unwrap().to_vec()); } Tensor::new(&device, &expected) } diff --git a/fusor-ml/conformance/src/suite/native/quantized_matmul_fusion.rs b/fusor-ml/conformance/src/suite/native/quantized_matmul_fusion.rs index 7feba210e..86201e260 100644 --- a/fusor-ml/conformance/src/suite/native/quantized_matmul_fusion.rs +++ b/fusor-ml/conformance/src/suite/native/quantized_matmul_fusion.rs @@ -96,7 +96,7 @@ pub fn rmsnorm_post_relu_resolves_to_single_kernel() -> AssertionCase { .relu() .to_concrete() .as_gpu() - .is_some_and(|gpu_out| gpu_out.count_kernels_to_resolve() == 1) + .is_some_and(|gpu_out| gpu_out.resolves_in::<1>()) }, ) } @@ -137,7 +137,7 @@ pub fn q4k_qmatmul_fusion_kernels() -> AssertionCases { .q_mat_mul(&weights) .to_concrete() .as_gpu() - .is_some_and(|gpu_out| gpu_out.count_kernels_to_resolve() == 1) + .is_some_and(|gpu_out| gpu_out.resolves_in::<1>()) }, ), // `q_mat_mul(...).relu()` post-activation fusion (all GPU backends). @@ -152,7 +152,7 @@ pub fn q4k_qmatmul_fusion_kernels() -> AssertionCases { .relu() .to_concrete() .as_gpu() - .is_some_and(|gpu_out| gpu_out.count_kernels_to_resolve() == 1) + .is_some_and(|gpu_out| gpu_out.resolves_in::<1>()) }, ), ]) diff --git a/fusor-ml/conformance/src/suite/native/quantized_matmul_paired.rs b/fusor-ml/conformance/src/suite/native/quantized_matmul_paired.rs index dfa8db99c..e0727b80d 100644 --- a/fusor-ml/conformance/src/suite/native/quantized_matmul_paired.rs +++ b/fusor-ml/conformance/src/suite/native/quantized_matmul_paired.rs @@ -4,7 +4,7 @@ use crate::common::quantized::{ concrete_to_rows, q_mat_mul_input_fuzz, q4k_raw_bytes, qmatrix_from_raw_bytes, }; use crate::common::{matmul2, transpose2}; -use fusor::{BlockQ4K, Device, GgmlType, GgufBlock, QuantizedTensor, Tensor, ToVec2}; +use fusor::{BlockQ4K, Device, GgmlType, GgufBlock, QuantizedTensor, Tensor, ToVec}; use fusor_conformance::{ AssertionCase, AssertionCases, approx_compare, approx_or_relative_compare, available_devices, cases_from_rows, @@ -361,7 +361,7 @@ fn gated_matches_cpu_for_rows(input_row_count: usize, kind: GatedKind) -> Assert let expected_weights = expected_weights.clone(); async move { let device = input.device(); - let input_values = input.as_slice().await.unwrap().to_vec2(); + let input_values = input.as_slice().await.unwrap().to_vec(); let projected = matmul2(&input_values, &transpose2(&expected_weights)); let expected = projected .iter() diff --git a/fusor-ml/conformance/src/suite/native/rope_ops.rs b/fusor-ml/conformance/src/suite/native/rope_ops.rs index 19acf3ebd..81e8f9c41 100644 --- a/fusor-ml/conformance/src/suite/native/rope_ops.rs +++ b/fusor-ml/conformance/src/suite/native/rope_ops.rs @@ -1,7 +1,7 @@ //! RoPE op conformance cases. use crate::common::{reshape4, rope_interleaved_4d, rope_normal_4d}; -use fusor::{Device, RopeCache, Tensor, ToVec1, base_inverse_frequency}; +use fusor::{Device, RopeCache, Tensor, ToVec, base_inverse_frequency}; use fusor_conformance::{AssertionCases, FuzzGenerator, approx_compare, exact_value_compare}; use rand::distr::Uniform; @@ -81,7 +81,7 @@ async fn rope_reference_tensor( .as_slice() .await .unwrap(); - let flat: Vec = slice.to_vec1(); + let flat: Vec = slice.to_vec(); let host = reshape4(&flat, [1, 2, 3, 4]); let out = match reference { RopeReference::Normal => rope_normal_4d(&host, &cos, &sin), diff --git a/fusor-ml/conformance/src/suite/native/vision_block_pattern.rs b/fusor-ml/conformance/src/suite/native/vision_block_pattern.rs index 4ced08d9e..0f5257bf6 100644 --- a/fusor-ml/conformance/src/suite/native/vision_block_pattern.rs +++ b/fusor-ml/conformance/src/suite/native/vision_block_pattern.rs @@ -3,7 +3,7 @@ //! Regression test for the qwen-vision multi-block resolve pattern. //! //! Each "block" mimics the shape of `VisionBlock::forward`: narrow a fused QKV -//! tensor into Q/K/V, run windowed flash attention (multiple narrows + flash +//! tensor into Q/K/V, run windowed attention (multiple narrows + attention //! per block, then sum), add a residual, then an MLP-shaped elementwise pass. //! The point isn't numerical correctness against a reference, it's to exercise //! the same fusion+resolve interaction the qwen vision encoder hits — without @@ -30,7 +30,7 @@ const MLP_INTERMEDIATE: usize = 3420; fn ramp_data(len: usize, scale: f32) -> Vec { // Small magnitudes keep the values finite through many stacked blocks of - // matmul + flash + residual + mlp without depending on real + // matmul + attention + residual + mlp without depending on real // initialization. (0..len) .map(|i| (((i % 23) as f32) / 23.0 - 0.5) * scale + 0.001) @@ -80,7 +80,7 @@ async fn run_blocks(device: &fusor::Device, flush_every: Option) -> Vec = qw.flash_attention(&kw, &vw, scale, None); + let attn_w: Tensor<4, f32> = qw.attention(&kw, &vw, scale, None); window_outputs.push(attn_w); } let attn_out: Tensor<4, f32> = fusor::cat(window_outputs, 2).to_concrete(); diff --git a/fusor-ml/conformance/src/suite/registry.rs b/fusor-ml/conformance/src/suite/registry.rs index 400b16130..37dbf11b8 100644 --- a/fusor-ml/conformance/src/suite/registry.rs +++ b/fusor-ml/conformance/src/suite/registry.rs @@ -151,18 +151,19 @@ registry! { elementwise_ops::tanh_exact_saturation_at_large_magnitudes, elementwise_ops::unary_math_ops_match_host_reference, elementwise_ops::where_cond_fuzzed, - flash_attention_ops::flash_attention_decode_tiled_matches_cpu_reference, - flash_attention_ops::flash_attention_decode_tiled_with_transposed_q_matches_cpu_reference, - flash_attention_ops::flash_attention_f16_matches_cpu_reference_on_varied_shapes, - flash_attention_ops::flash_attention_f16_with_qk_mask_matches_cpu_reference, - flash_attention_ops::flash_attention_gqa_matches_cpu_reference_on_varied_shapes, - flash_attention_ops::flash_attention_matches_cpu_reference_on_varied_shapes, - flash_attention_ops::flash_attention_subgroup_fallback_preserves_gpu_backend, - flash_attention_ops::flash_attention_tiled_matches_cpu_reference_on_varied_shapes, - flash_attention_ops::flash_attention_with_batch_key_mask_matches_cpu_reference_on_varied_shapes, - flash_attention_ops::flash_attention_with_kv_cache_matches_cpu_reference_on_varied_shapes, - flash_attention_ops::flash_attention_with_qk_mask_matches_cpu_reference_on_varied_shapes, - fusion_behavior::gpu_flash_attention_fuses_into_one_kernel, + attention_ops::attention_decode_tiled_matches_cpu_reference, + attention_ops::attention_decode_tiled_with_transposed_q_matches_cpu_reference, + attention_ops::attention_f16_matches_cpu_reference_on_varied_shapes, + attention_ops::attention_f16_with_qk_mask_matches_cpu_reference, + attention_ops::attention_gqa_matches_cpu_reference_on_varied_shapes, + attention_ops::attention_matches_cpu_reference_on_varied_shapes, + attention_ops::attention_without_subgroups_preserves_gpu_backend, + attention_ops::attention_tiled_matches_cpu_reference_on_varied_shapes, + attention_ops::attention_with_batch_key_mask_matches_cpu_reference_on_varied_shapes, + attention_ops::attention_with_kv_cache_matches_cpu_reference_on_varied_shapes, + attention_ops::attention_with_qk_mask_matches_cpu_reference_on_varied_shapes, + fusion_behavior::gpu_attention_fuses_into_one_kernel, + fusion_behavior::gpu_coop_matmul_fuses_pre_and_post_unary_chains, fusion_behavior::gpu_gelu_lowers_to_one_kernel, fusion_behavior::gpu_indexing_then_arithmetic_matches_cpu, fusion_behavior::gpu_matmul_then_unary_chain_fuses_into_one_kernel, diff --git a/fusor-ml/conformance/src/suite/webgpu.rs b/fusor-ml/conformance/src/suite/webgpu.rs index 53071eb81..e14b8d49f 100644 --- a/fusor-ml/conformance/src/suite/webgpu.rs +++ b/fusor-ml/conformance/src/suite/webgpu.rs @@ -112,11 +112,11 @@ pub async fn run_webgpu_kernel_suite_with_progress( } fn skip_browser_registry_case(name: &str) -> bool { - name.starts_with("flash_attention_ops::flash_attention_decode_tiled_matches_cpu_reference::") + name.starts_with("attention_ops::attention_decode_tiled_matches_cpu_reference::") || name.starts_with( - "flash_attention_ops::flash_attention_decode_tiled_with_transposed_q_matches_cpu_reference::", + "attention_ops::attention_decode_tiled_with_transposed_q_matches_cpu_reference::", ) - || name == "flash_attention_ops::flash_attention_subgroup_fallback_preserves_gpu_backend" + || name == "attention_ops::attention_without_subgroups_preserves_gpu_backend" } // --------------------------------------------------------------------------- @@ -526,10 +526,6 @@ mod tests { /// browser. The suite internally expands each case across the /// {subgroups, no subgroups} × {cold pool, poisoned pool} device matrix, so /// the no-subgroup kernel fallbacks the web build takes are covered natively. - /// - /// To approximate the browser quantized storage layout on a native GPU, set - /// `FUSOR_Q_NATIVE=0` to force the `GpuF32Scales` layout (the web build - /// disables `SHADER_F16`, so it never uses the native f16-scale layout). #[allow(clippy::await_holding_lock)] #[tokio::test] async fn webgpu_kernel_suite_runs_on_gpu() { diff --git a/fusor-ml/core/Cargo.toml b/fusor-ml/core/Cargo.toml index adb3f0581..cf894498d 100644 --- a/fusor-ml/core/Cargo.toml +++ b/fusor-ml/core/Cargo.toml @@ -20,6 +20,7 @@ tracing = "0.1.41" lru = { version = "0.14.0", default-features = false } pollster = "0.4.0" petgraph = { version = "0.8.3", features = ["stable_graph", "std"], default-features = false } +egg = { version = "0.11", default-features = false } fusor-tile-ir = { workspace = true } fusor-tile-ir-kernels = { workspace = true } fusor-tile-ir-runtime = { workspace = true } diff --git a/fusor-ml/core/examples/bench_contraction.rs b/fusor-ml/core/examples/bench_contraction.rs deleted file mode 100644 index 92315ffc2..000000000 --- a/fusor-ml/core/examples/bench_contraction.rs +++ /dev/null @@ -1,284 +0,0 @@ -//! A/B kernel timing for the dense matmul routes and composed contractions. -//! Run on two builds and compare medians; wall-clock per resolve with a -//! device sync, warmup excluded. - -use fusor_core::{Device, QMatrix, StrideSpec, Tensor}; -use fusor_gguf::GgmlType; -use std::time::Instant; - -fn values(len: usize, scale: f32) -> Vec { - (0..len).map(|i| ((i as f32) * scale).sin()).collect() -} - -fn time_case(name: &str, warmup: usize, iters: usize, mut run: impl FnMut()) { - for _ in 0..warmup { - run(); - } - let mut samples = Vec::with_capacity(iters); - for _ in 0..iters { - let start = Instant::now(); - run(); - samples.push(start.elapsed().as_secs_f64() * 1e3); - } - samples.sort_by(|a, b| a.partial_cmp(b).unwrap()); - let median = samples[samples.len() / 2]; - let min = samples[0]; - println!("{name}: median {median:.3} ms, min {min:.3} ms over {iters} iters"); -} - -fn main() { - pollster::block_on(async { - let device = Device::new().await.expect("gpu device"); - - // Coop-ineligible dense matmul (1000 % 64 != 0): the workgroup-tiled - // route. - { - let a = Tensor::from_slice(&device, [1000, 1000], &values(1_000_000, 0.13)); - let b = Tensor::from_slice(&device, [1000, 1000], &values(1_000_000, 0.07)); - time_case("matmul_1000_tiled", 3, 15, || { - let out = a.mat_mul(&b); - out.materialize_sync(); - }); - } - - // Coop-eligible square (512 divisible by 64): should be unchanged. - { - let a = Tensor::from_slice(&device, [512, 512], &values(512 * 512, 0.13)); - let b = Tensor::from_slice(&device, [512, 512], &values(512 * 512, 0.07)); - time_case("matmul_512_coop", 3, 30, || { - let out = a.mat_mul(&b); - out.materialize_sync(); - }); - } - - // Low tile utilization (65/96/65): the register-tile route. - { - let a = Tensor::from_slice(&device, [16, 65, 96], &values(16 * 65 * 96, 0.13)); - let b = Tensor::from_slice(&device, [16, 96, 65], &values(16 * 96 * 65, 0.07)); - time_case("matmul_batched_65_register", 3, 30, || { - let out = a.mat_mul(&b); - out.materialize_sync(); - }); - } - - // Weighted sum: `sum_k w[k] * x[m, k]` — k-dependent inputs share no - // (row, col) pair, but `w` misses the row dim. - { - let (m, k) = (4096usize, 4096usize); - let x = Tensor::from_slice(&device, [m, k], &values(m * k, 0.13)); - let w = Tensor::from_slice(&device, [k], &values(k, 0.07)); - time_case("weighted_sum_4096", 3, 30, || { - let w2 = w.restride([StrideSpec::dim_with(0, m, 0), StrideSpec::dim(0, k)]); - let out = (&x * &w2).sum(1); - out.materialize_sync(); - }); - } - - // Broadcast scale: `x[m, n] * w[n]` — `w` is invariant along rows. - { - let (m, n) = (4096usize, 4096usize); - let x = Tensor::from_slice(&device, [m, n], &values(m * n, 0.13)); - let w = Tensor::from_slice(&device, [n], &values(n, 0.07)); - time_case("broadcast_scale_4096", 3, 30, || { - let w2 = w.restride([StrideSpec::dim_with(0, m, 0), StrideSpec::dim(0, n)]); - let out = &x * &w2; - out.materialize_sync(); - }); - } - - // Broadcast table apply: `x[b, s, h] * t[s, h]` with a table too - // large for cache — the invariant table loads hoist out of each - // thread's output run. - { - let (b, s, h) = (4usize, 2048usize, 2048usize); - let x = Tensor::from_slice(&device, [b, s, h], &values(b * s * h, 0.13)); - let t = Tensor::from_slice(&device, [s, h], &values(s * h, 0.07)); - time_case("broadcast_table_4x2048x2048", 3, 30, || { - let t3 = t.restride([ - StrideSpec::dim_with(0, b, 0), - StrideSpec::dim(0, s), - StrideSpec::dim(1, h), - ]); - let out = &x * &t3; - out.materialize_sync(); - }); - } - - // Dequantize a Q8_0 matrix to dense f32. - { - const Q8_BLOCK: usize = 32; - let (n, k) = (4096usize, 4096usize); - let scale = half::f16::from_f32(0.02); - let mut bytes = Vec::with_capacity(n * k / Q8_BLOCK * 34); - for block in 0..(n * k / Q8_BLOCK) { - bytes.extend_from_slice(&scale.to_le_bytes()); - for i in 0..Q8_BLOCK { - bytes.push((((block * 7 + i * 5) % 64) as i32 - 32) as i8 as u8); - } - } - let w = QMatrix::from_parts( - &device, - &bytes, - vec![n, k].into_boxed_slice(), - GgmlType::Q8_0, - ) - .unwrap(); - time_case("dequantize_q8_4096", 3, 30, || { - let out = w.dequantize::(); - out.materialize_sync(); - }); - } - - // Dense gemv: [m, k] x [k, 1]. - { - let (m, k) = (4096usize, 4096usize); - let a = Tensor::from_slice(&device, [m, k], &values(m * k, 0.13)); - let b = Tensor::from_slice(&device, [k, 1], &values(k, 0.07)); - time_case("gemv_4096", 3, 50, || { - let out = a.mat_mul(&b); - out.materialize_sync(); - }); - } - - // Norm shapes: decode-like single row and prefill-like many rows. - { - let hidden = 4096usize; - let x1 = Tensor::from_slice(&device, [1, hidden], &values(hidden, 0.13)); - let xs = Tensor::from_slice(&device, [512, hidden], &values(512 * hidden, 0.13)); - let w = Tensor::from_slice(&device, [hidden], &values(hidden, 0.07)); - time_case("rms_norm_1x4096", 3, 50, || { - let out = x1.rms_norm_fused(&w, None, 1e-5); - out.materialize_sync(); - }); - time_case("rms_norm_512x4096", 3, 30, || { - let out = xs.rms_norm_fused(&w, None, 1e-5); - out.materialize_sync(); - }); - time_case("softmax_512x4096", 3, 30, || { - let out = xs.softmax(1); - out.materialize_sync(); - }); - let small = Tensor::from_slice(&device, [32, 128], &values(32 * 128, 0.17)); - time_case("softmax_32x128", 3, 50, || { - let out = small.softmax(1); - out.materialize_sync(); - }); - } - - // Decode-shape attention (q_len = 1): the per-token hot path. kv - // lengths cover the single-dispatch decode buckets and (at 2048) the - // split two-dispatch route. - { - let (batch, heads, head_dim) = (1usize, 32usize, 128usize); - let scale = 1.0 / (head_dim as f32).sqrt(); - let q = Tensor::from_slice( - &device, - [batch, heads, 1, head_dim], - &values(batch * heads * head_dim, 0.13), - ); - for kv in [512usize, 1024, 2048, 4096, 8192] { - let k = Tensor::from_slice( - &device, - [batch, heads, kv, head_dim], - &values(batch * heads * kv * head_dim, 0.07), - ); - let v = Tensor::from_slice( - &device, - [batch, heads, kv, head_dim], - &values(batch * heads * kv * head_dim, 0.11), - ); - time_case(&format!("attn_decode_kv{kv}"), 3, 50, || { - let out = q.flash_attention(&k, &v, scale, None); - out.materialize_sync(); - }); - } - // Grouped-query variant: 32 query heads over 8 KV heads. - let kv_heads = 8usize; - let kv = 1024usize; - let k = Tensor::from_slice( - &device, - [batch, kv_heads, kv, head_dim], - &values(batch * kv_heads * kv * head_dim, 0.07), - ); - let v = Tensor::from_slice( - &device, - [batch, kv_heads, kv, head_dim], - &values(batch * kv_heads * kv * head_dim, 0.11), - ); - time_case("attn_decode_gqa_kv1024", 3, 50, || { - let out = q.flash_attention(&k, &v, scale, None); - out.materialize_sync(); - }); - } - - // Prefill-shape attention: causal self-attention at q == kv (the - // streaming-tiled regime) and a mid-length query block (the plain - // streaming regime). - { - let (batch, heads, head_dim) = (1usize, 32usize, 128usize); - let scale = 1.0 / (head_dim as f32).sqrt(); - let seq = 512usize; - let q = Tensor::from_slice( - &device, - [batch, heads, seq, head_dim], - &values(batch * heads * seq * head_dim, 0.13), - ); - let k = Tensor::from_slice( - &device, - [batch, heads, seq, head_dim], - &values(batch * heads * seq * head_dim, 0.07), - ); - let v = Tensor::from_slice( - &device, - [batch, heads, seq, head_dim], - &values(batch * heads * seq * head_dim, 0.11), - ); - time_case("attn_prefill_512_causal", 3, 20, || { - let out = q.flash_attention_causal(&k, &v, scale); - out.materialize_sync(); - }); - - let q64 = Tensor::from_slice( - &device, - [batch, heads, 64, head_dim], - &values(batch * heads * 64 * head_dim, 0.13), - ); - let kv = 1024usize; - let k = Tensor::from_slice( - &device, - [batch, heads, kv, head_dim], - &values(batch * heads * kv * head_dim, 0.07), - ); - let v = Tensor::from_slice( - &device, - [batch, heads, kv, head_dim], - &values(batch * heads * kv * head_dim, 0.11), - ); - time_case("attn_prefill_q64_kv1024", 3, 30, || { - let out = q64.flash_attention(&k, &v, scale, None); - out.materialize_sync(); - }); - } - - // Composed broadcast contraction: not recognized as a matmul. - { - let (m, n, k) = (256usize, 256usize, 256usize); - let a = Tensor::from_slice(&device, [m, k], &values(m * k, 0.13)); - let b = Tensor::from_slice(&device, [n, k], &values(n * k, 0.07)); - time_case("broadcast_contraction_256", 3, 30, || { - let a3 = a.restride([ - StrideSpec::dim(0, m), - StrideSpec::dim_with(0, n, 0), - StrideSpec::dim(1, k), - ]); - let b3 = b.restride([ - StrideSpec::dim_with(0, m, 0), - StrideSpec::dim(0, n), - StrideSpec::dim(1, k), - ]); - let out = (&a3 * &b3).sum(2); - out.materialize_sync(); - }); - } - }); -} diff --git a/fusor-ml/core/examples/bench_conv_implicit_gemm.rs b/fusor-ml/core/examples/bench_conv_implicit_gemm.rs deleted file mode 100644 index e53c0be47..000000000 --- a/fusor-ml/core/examples/bench_conv_implicit_gemm.rs +++ /dev/null @@ -1,190 +0,0 @@ -// A/B microbench: conv-as-im2col matmul with the gather materialization -// (the old pipeline) vs the implicit-GEMM unflatten (the new pipeline). -// -// Holding the flat [M, K] view tensor across materialization trips the -// resolver's live-reference guard, which declines the unflatten — giving -// exactly the old gather + matmul pipeline in the same binary. -// -// Run with: -// cargo run --package fusor-core --example bench_conv_implicit_gemm --release - -use std::time::{Duration, Instant}; - -use fusor_core::{Device, StrideSpec, Tensor}; - -const WARMUP: usize = 5; -const MEASURED: usize = 30; - -struct ConvCase { - name: &'static str, - b: usize, - c: usize, - h: usize, - w: usize, - n: usize, - kh: usize, - kw: usize, -} - -fn main() -> Result<(), Box> { - pollster::block_on(async { - let device = Device::new().await?; - println!("bench_conv_implicit_gemm (warmup {WARMUP}, measured {MEASURED})"); - println!(); - - let cases = [ - // Coop-tile-unaligned M: both variants run the generic reduce. - ConvCase { - name: "small_unaligned", - b: 2, - c: 8, - h: 16, - w: 16, - n: 16, - kh: 3, - kw: 3, - }, - ConvCase { - name: "mid_unaligned_n256", - b: 1, - c: 128, - h: 35, - w: 35, - n: 256, - kh: 3, - kw: 3, - }, - ConvCase { - name: "mid_unaligned_n128", - b: 1, - c: 128, - h: 35, - w: 35, - n: 128, - kh: 3, - kw: 3, - }, - ConvCase { - name: "mid_unaligned_n64", - b: 1, - c: 128, - h: 35, - w: 35, - n: 64, - kh: 3, - kw: 3, - }, - // Coop-tile-aligned M/K/N: the matmul runs the hardware kernel. - ConvCase { - name: "large_aligned", - b: 2, - c: 64, - h: 34, - w: 34, - n: 128, - kh: 3, - kw: 3, - }, - ConvCase { - name: "vision_aligned", - b: 1, - c: 256, - h: 66, - w: 66, - n: 256, - kh: 3, - kw: 3, - }, - ]; - - for case in &cases { - bench_case(&device, case); - } - Ok(()) - }) -} - -fn bench_case(device: &Device, case: &ConvCase) { - let &ConvCase { - name, - b, - c, - h, - w, - n, - kh, - kw, - } = case; - let (oh, ow) = (h - kh + 1, w - kw + 1); - let (m, k) = (b * oh * ow, c * kh * kw); - - let input_host: Vec = (0..b * c * h * w).map(|i| (i % 13) as f32 * 0.1).collect(); - let weight_host: Vec = (0..n * k).map(|i| (i % 7) as f32 * 0.01).collect(); - let input = Tensor::from_slice(device, [b, c, h, w], &input_host); - let weight = Tensor::from_slice(device, [n, k], &weight_host); - input.materialize_sync(); - weight.materialize_sync(); - - let build = |hold_flat: bool| -> (Tensor, Option) { - let windows = input.restride([ - StrideSpec::dim(0, b), - StrideSpec::dim_with(2, oh, 1), - StrideSpec::dim_with(3, ow, 1), - StrideSpec::dim(1, c), - StrideSpec::dim(2, kh), - StrideSpec::dim(3, kw), - ]); - let a = windows.reshape([m, k]); - let b_mat = weight.restride([StrideSpec::dim(1, k), StrideSpec::dim(0, n)]); - let out = a.mat_mul(&b_mat); - (out, hold_flat.then_some(a)) - }; - - let run = |hold_flat: bool| -> (Vec, usize) { - let mut kernels = 0; - for _ in 0..WARMUP { - let (out, held) = build(hold_flat); - kernels = out.count_kernels_to_resolve(); - device.poll_wait(); - drop(held); - drop(out); - } - let mut samples = Vec::with_capacity(MEASURED); - for _ in 0..MEASURED { - let (out, held) = build(hold_flat); - let start = Instant::now(); - let _ = out.count_kernels_to_resolve(); - device.poll_wait(); - samples.push(start.elapsed()); - drop(held); - drop(out); - } - samples.sort_unstable(); - (samples, kernels) - }; - - let (gather, gather_kernels) = run(true); - let (implicit, implicit_kernels) = run(false); - - let stats = |samples: &[Duration]| { - let mean = samples.iter().sum::() / samples.len() as u32; - let p50 = samples[samples.len() / 2]; - ( - mean.as_secs_f64() * 1000.0, - p50.as_secs_f64() * 1000.0, - samples[0].as_secs_f64() * 1000.0, - ) - }; - let (g_mean, g_p50, g_min) = stats(&gather); - let (i_mean, i_p50, i_min) = stats(&implicit); - - println!("{name}: conv {b}x{c}x{h}x{w} k{kh}x{kw} -> matmul {m}x{k} @ {k}x{n}"); - println!( - " gather+matmul ({gather_kernels} dispatches): mean {g_mean:.3} ms, p50 {g_p50:.3} ms, min {g_min:.3} ms" - ); - println!( - " implicit-GEMM ({implicit_kernels} dispatches): mean {i_mean:.3} ms, p50 {i_p50:.3} ms, min {i_min:.3} ms" - ); - println!(" speedup (p50): {:.2}x", g_p50 / i_p50); - println!(); -} diff --git a/fusor-ml/core/examples/bench_coop_tiles.rs b/fusor-ml/core/examples/bench_coop_tiles.rs new file mode 100644 index 000000000..8ffe932fa --- /dev/null +++ b/fusor-ml/core/examples/bench_coop_tiles.rs @@ -0,0 +1,346 @@ +//! Raw sweep: every cooperative tile over the training step's hot matmul +//! shapes, dispatched directly (no resolver, no selection). Establishes each +//! shape's best achievable rate with the current kernels against the +//! measured ~8.9 TF/s simdgroup ceiling. + +use fusor_core::Device; +use fusor_tile_ir::tile; +use fusor_tile_ir_kernels::{ + coop_tile_entries, DenseCoopMatmulConfig, DenseMatmulShape, DenseMatmulTensors, + SubgroupConfig, +}; + +const SHAPES: [(u32, u32, u32); 10] = [ + (16384, 384, 384), + (16384, 384, 1536), + (16384, 1536, 384), + (384, 16384, 1536), + (1536, 16384, 384), + (4096, 4096, 4096), + (16384, 3072, 1536), + (1024, 1024, 1024), + (1000, 1024, 1024), + (1944, 1280, 3840), +]; + +fn main() { + pollster::block_on(async { + let device = Device::new().await.unwrap(); + let wgpu_device = device.wgpu_device(); + let queue = device.wgpu_queue(); + wgpu_device.on_uncaptured_error(std::sync::Arc::new(|error: wgpu::Error| { + eprintln!("wgpu error: {error}"); + })); + let subgroups = SubgroupConfig::fixed(fusor_tile_ir::SubgroupToken::new_unchecked(), 32); + let coop = fusor_tile_ir::CoopMatrixToken::new_unchecked(); + + use wgpu::util::DeviceExt; + let host = |len: usize| -> Vec { + (0..len).map(|i| ((i % 61) as f32) * 0.01 - 0.3).collect() + }; + // DTYPE=f16 benches the native f16-storage kernels (mixed MMA with + // f32 accumulation); verification stays on the f32 path. + let f16_storage = std::env::var("DTYPE").as_deref() == Ok("f16"); + // SHAPES=m,k,n;m,k,n sweeps caller-supplied contractions instead of + // the built-in training-step set. + let shapes: Vec<(u32, u32, u32)> = std::env::var("SHAPES") + .map(|list| { + list.split(';') + .map(|entry| { + let dims: Vec = entry + .split(',') + .map(|v| v.trim().parse().expect("SHAPES=m,k,n;..")) + .collect(); + (dims[0], dims[1], dims[2]) + }) + .collect() + }) + .unwrap_or_else(|_| SHAPES.to_vec()); + // SWIZZLE=1,2,4,8,16 sweeps traversal-order groups per tile entry + // (1 = plain row-major decomposition). + let swizzle_groups: Vec = std::env::var("SWIZZLE") + .map(|list| list.split(',').filter_map(|v| v.parse().ok()).collect()) + .unwrap_or_else(|_| vec![fusor_tile_ir_kernels::DEFAULT_SWIZZLE_GROUP_M]); + let make = |data: &[f32]| { + if f16_storage { + let halves: Vec = + data.iter().map(|&value| half::f16::from_f32(value)).collect(); + wgpu_device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(&halves), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + }) + } else { + wgpu_device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: bytemuck::cast_slice(data), + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + }) + } + }; + let check = std::env::var_os("CHECK").is_some(); + let staging_f16 = std::env::var("STAGING").as_deref() == Ok("f16"); + // BUFFERS=1 stages every entry from one tile pair (half the + // workgroup footprint, no load/MMA overlap) — the raw-sweep side of + // the staging-depth trade the planner scores. + let stage_buffers: u32 = std::env::var("BUFFERS") + .ok() + .and_then(|v| v.trim().parse().ok()) + .unwrap_or(2); + + // Optional shape filter for tight alternation windows: SHAPE=MxKxN. + let shape_filter = std::env::var("SHAPE").ok(); + for (m, k, n) in shapes { + if let Some(filter) = &shape_filter { + if *filter != format!("{m}x{k}x{n}") { + continue; + } + } + let flops = 2.0 * m as f64 * k as f64 * n as f64; + let a_host = host((m * k) as usize); + let b_host = host((k * n) as usize); + let a_buf = make(&a_host); + let b_buf = make(&b_host); + println!("--- {m}x{k}x{n} ({:.1} GFLOP)", flops / 1e9); + // Candidates built up-front, then measured round-robin: each + // rep times every entry once, and each entry keeps its minimum + // across reps. Under bursty GPU contention a contended burst + // poisons single entries per rep instead of whole entries, so + // per-entry minima converge toward uncontended times. + struct Candidate { + label: String, + pipeline: wgpu::ComputePipeline, + binds: Vec, + grid: [u32; 3], + best: f64, + } + let mut candidates: Vec = Vec::new(); + for (entry, swizzle_group_m) in coop_tile_entries() + .iter() + .flat_map(|entry| swizzle_groups.iter().map(move |&group| (entry, group))) + { + let tile = entry.tile; + let (bm, bn) = (tile.bm, tile.bn); + // `SUBGROUP_SPLIT=rg,cg` forces a factorization the closed + // form would not pick — the vehicle for isolating what a + // threadgroup fragment load actually costs. + let (row_groups, col_groups) = match std::env::var("SUBGROUP_SPLIT") { + Ok(forced) => { + let (rg, cg) = forced.split_once(',').expect("SUBGROUP_SPLIT=rg,cg"); + (rg.trim().parse().unwrap(), cg.trim().parse().unwrap()) + } + Err(_) => entry.subgroup_split(), + }; + if row_groups * col_groups != entry.subgroups { + continue; + } + let label = format!( + "{bm}x{bn} rg{row_groups} cg{col_groups} np{} sw{swizzle_group_m}", + entry.n_passes, + ); + let m_pad = m.div_ceil(bm) * bm; + let n_pad = n.div_ceil(bn) * bn; + let total_tiles = (m_pad / bm) * (n_pad / bn); + if total_tiles > 65535 { + continue; + } + let element_size = if f16_storage { 2 } else { 4 }; + let y_bufs: Vec<_> = (0..4) + .map(|_| { + wgpu_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: (m_pad as u64) * (n_pad as u64) * element_size, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }) + }) + .collect(); + let elem = if f16_storage { + fusor_tile_ir::ScalarElement::F16.element() + } else { + fusor_tile_ir::ScalarElement::F32.element() + }; + let mut ok = true; + let ir = tile::build(|phase| { + let a = phase.storage_read(elem, fusor_tile_ir::Shape::new([m, k])); + let b = phase.storage_read(elem, fusor_tile_ir::Shape::new([k, n])); + let y = phase.storage_write(elem, fusor_tile_ir::Shape::new([m_pad, n_pad])); + ok = fusor_tile_ir_kernels::try_batched_coop_matmul( + phase, + DenseMatmulTensors { + a: &a, + b: &b, + y: &y, + }, + DenseMatmulShape { batch: 1, m, k, n }, + &fusor_tile_ir_kernels::DenseMatmulEpilogues::empty(), + 65535, + DenseCoopMatmulConfig { + coop, + subgroups, + tile, + row_groups, + col_groups, + staging: staging_f16.then_some(fusor_tile_ir::ScalarElement::F16), + stage_buffers, + swizzle_group_m, + }, + ); + }); + if !ok { + continue; + } + let grid = ir.grid; + let Ok(kernel) = ir.lower_to_naga() else { + println!(" {label}: lowering failed"); + continue; + }; + let module = unsafe { + wgpu_device.create_shader_module_trusted( + wgpu::ShaderModuleDescriptor { + label: None, + source: wgpu::ShaderSource::Naga(std::borrow::Cow::Owned( + kernel.module().clone(), + )), + }, + wgpu::ShaderRuntimeChecks::unchecked(), + ) + }; + let pipeline = + wgpu_device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: None, + module: &module, + entry_point: Some("main"), + compilation_options: wgpu::PipelineCompilationOptions { + zero_initialize_workgroup_memory: false, + ..Default::default() + }, + cache: None, + }); + let layout = pipeline.get_bind_group_layout(0); + let binds: Vec<_> = y_bufs + .iter() + .map(|y| { + wgpu_device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: a_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: b_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: y.as_entire_binding(), + }, + ], + }) + }) + .collect(); + let run = |iters: u32| { + let mut encoder = wgpu_device.create_command_encoder(&Default::default()); + for i in 0..iters { + let mut pass = encoder.begin_compute_pass(&Default::default()); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &binds[(i as usize) % binds.len()], &[]); + pass.dispatch_workgroups(grid[0], grid[1], grid[2]); + } + queue.submit([encoder.finish()]); + }; + run(3); + device.poll_wait(); + if check && !f16_storage { + let staging = wgpu_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: y_bufs[2].size(), + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut encoder = + wgpu_device.create_command_encoder(&Default::default()); + encoder.copy_buffer_to_buffer(&y_bufs[2], 0, &staging, 0, staging.size()); + queue.submit([encoder.finish()]); + let slice = staging.slice(..); + let (tx, rx) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |r| tx.send(r).unwrap()); + device.poll_wait(); + rx.recv().unwrap().unwrap(); + let view = slice.get_mapped_range(); + let y_host: &[f32] = bytemuck::cast_slice(&view); + let mut worst = 0f64; + for sample in 0..997u64 { + let i = ((sample * 7919) % m as u64) as usize; + let j = ((sample * 104729) % n as u64) as usize; + let mut acc = 0f64; + for kk in 0..k as usize { + acc += a_host[i * k as usize + kk] as f64 + * b_host[kk * n as usize + j] as f64; + } + let got = y_host[i * n_pad as usize + j] as f64; + let err = (got - acc).abs() / acc.abs().max(1.0); + worst = worst.max(err); + if err >= 5e-3 && std::env::var_os("LENIENT").is_none() { + panic!("{label} mismatch at ({i},{j}): got {got}, want {acc}"); + } + } + println!(" {label} verified, worst rel err {worst:.2e}"); + } + candidates.push(Candidate { + label, + pipeline, + binds, + grid, + best: f64::MAX, + }); + } + let iters: u32 = std::env::var("ITERS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(10); + let reps: u32 = std::env::var("REPS") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(6); + for _ in 0..reps { + for candidate in candidates.iter_mut() { + let start = std::time::Instant::now(); + let mut encoder = wgpu_device.create_command_encoder(&Default::default()); + for i in 0..iters { + let mut pass = encoder.begin_compute_pass(&Default::default()); + pass.set_pipeline(&candidate.pipeline); + pass.set_bind_group( + 0, + &candidate.binds[(i as usize) % candidate.binds.len()], + &[], + ); + pass.dispatch_workgroups( + candidate.grid[0], + candidate.grid[1], + candidate.grid[2], + ); + } + queue.submit([encoder.finish()]); + device.poll_wait(); + candidate.best = candidate + .best + .min(start.elapsed().as_secs_f64() / iters as f64); + } + } + for candidate in &candidates { + let best = candidate.best; + println!( + " {}: {:.3} ms, {:.2} TF/s ({:.0}%)", + candidate.label, + best * 1e3, + flops / best / 1e12, + flops / best / 1e12 / 8.86 * 100.0 + ); + } + } + }); +} diff --git a/fusor-ml/core/examples/bench_specialized_kernels.rs b/fusor-ml/core/examples/bench_specialized_kernels.rs deleted file mode 100644 index 0bf70ef6d..000000000 --- a/fusor-ml/core/examples/bench_specialized_kernels.rs +++ /dev/null @@ -1,288 +0,0 @@ -use std::time::Duration; - -use fusor_core::{Device, GpuMirostat2Sampler, GpuMirostat2SamplerParams, Tensor}; - -const DEFAULT_WARMUP_BATCHES: usize = 3; -const DEFAULT_MEASURED_BATCHES: usize = 20; -const DEFAULT_DISPATCHES_PER_BATCH: usize = 16; - -fn env_usize(name: &str, default: usize) -> usize { - std::env::var(name) - .ok() - .and_then(|value| value.parse().ok()) - .unwrap_or(default) -} - -fn main() -> Result<(), Box> { - pollster::block_on(async { - let warmup_batches = env_usize( - "FUSOR_SPECIALIZED_BENCH_WARMUP_BATCHES", - DEFAULT_WARMUP_BATCHES, - ); - let measured_batches = env_usize( - "FUSOR_SPECIALIZED_BENCH_MEASURED_BATCHES", - DEFAULT_MEASURED_BATCHES, - ); - let dispatches_per_batch = env_usize( - "FUSOR_SPECIALIZED_BENCH_DISPATCHES_PER_BATCH", - DEFAULT_DISPATCHES_PER_BATCH, - ) - .max(1); - - let device = Device::new().await?; - println!("bench_specialized_kernels"); - println!("warmup_batches: {warmup_batches}"); - println!("measured_batches: {measured_batches}"); - println!("dispatches_per_batch: {dispatches_per_batch}"); - - bench_rms_norm_vec4( - &device, - warmup_batches, - measured_batches, - dispatches_per_batch, - ); - bench_flash_attention_streaming( - &device, - warmup_batches, - measured_batches, - dispatches_per_batch, - ); - bench_flash_attention_decode( - &device, - warmup_batches, - measured_batches, - dispatches_per_batch, - ); - bench_top_k_pairs(&device, warmup_batches, measured_batches).await?; - bench_mirostat2(&device, warmup_batches, measured_batches).await?; - - Ok(()) - }) -} - -fn bench_rms_norm_vec4( - device: &Device, - warmup_batches: usize, - measured_batches: usize, - dispatches_per_batch: usize, -) { - let input = Tensor::splat(device, 0.25f32, [64, 4096]); - let weight = Tensor::splat(device, 1.0f32, [4096]); - input.materialize_sync(); - weight.materialize_sync(); - - bench_tensor_case( - "rms_norm_vec4_64x4096", - warmup_batches, - measured_batches, - dispatches_per_batch, - || input.rms_norm_fused(&weight, None, 1e-5), - device, - ); -} - -fn bench_flash_attention_streaming( - device: &Device, - warmup_batches: usize, - measured_batches: usize, - dispatches_per_batch: usize, -) { - let q = Tensor::splat(device, 0.125f32, [1, 32, 48, 128]); - let k = Tensor::splat(device, 0.25f32, [1, 8, 48, 128]); - let v = Tensor::splat(device, 0.5f32, [1, 8, 48, 128]); - q.materialize_sync(); - k.materialize_sync(); - v.materialize_sync(); - - bench_tensor_case( - "flash_attention_streaming_1x32x48x128_by_8x48", - warmup_batches, - measured_batches, - dispatches_per_batch, - || q.flash_attention(&k, &v, 1.0 / f32::sqrt(128.0), None), - device, - ); -} - -fn bench_flash_attention_decode( - device: &Device, - warmup_batches: usize, - measured_batches: usize, - dispatches_per_batch: usize, -) { - let q = Tensor::splat(device, 0.125f32, [1, 32, 1, 128]); - let k = Tensor::splat(device, 0.25f32, [1, 8, 512, 128]); - let v = Tensor::splat(device, 0.5f32, [1, 8, 512, 128]); - q.materialize_sync(); - k.materialize_sync(); - v.materialize_sync(); - - bench_tensor_case( - "flash_attention_decode_1x32x1x128_by_8x512", - warmup_batches, - measured_batches, - dispatches_per_batch, - || q.flash_attention(&k, &v, 1.0 / f32::sqrt(128.0), None), - device, - ); -} - -fn bench_tensor_case( - name: &str, - warmup_batches: usize, - measured_batches: usize, - dispatches_per_batch: usize, - mut make_output: F, - device: &Device, -) where - F: FnMut() -> Tensor, -{ - for _ in 0..warmup_batches { - run_tensor_batch(dispatches_per_batch, &mut make_output, device); - } - - let mut samples = Vec::with_capacity(measured_batches); - let mut kernels = 0usize; - for _ in 0..measured_batches { - let (elapsed, batch_kernels) = - run_tensor_batch(dispatches_per_batch, &mut make_output, device); - samples.push(elapsed / dispatches_per_batch as u32); - kernels = batch_kernels; - } - - print_summary( - name, - &samples, - Some(kernels as f64 / dispatches_per_batch as f64), - ); -} - -fn run_tensor_batch(dispatches: usize, make_output: &mut F, device: &Device) -> (Duration, usize) -where - F: FnMut() -> Tensor, -{ - let mut outputs = Vec::with_capacity(dispatches); - for _ in 0..dispatches { - let output = make_output(); - outputs.push(output); - } - - let start = std::time::Instant::now(); - let kernels = outputs.iter().map(Tensor::count_kernels_to_resolve).sum(); - device.poll_wait(); - let elapsed = start.elapsed(); - drop(outputs); - - (elapsed, kernels) -} - -async fn bench_top_k_pairs( - device: &Device, - warmup_batches: usize, - measured_batches: usize, -) -> Result<(), wgpu::BufferAsyncError> { - let logits_data = bench_logits(8192); - let logits = Tensor::new(device, &logits_data); - logits.materialize_sync(); - - for _ in 0..warmup_batches { - let _ = logits.top_k_pairs(512).await?; - } - - let mut samples = Vec::with_capacity(measured_batches); - for _ in 0..measured_batches { - let start = std::time::Instant::now(); - let _ = logits.top_k_pairs(512).await?; - samples.push(start.elapsed()); - } - - print_summary("top_k_pairs_8192_k512", &samples, None); - Ok(()) -} - -async fn bench_mirostat2( - device: &Device, - warmup_batches: usize, - measured_batches: usize, -) -> Result<(), wgpu::BufferAsyncError> { - let logits_data = bench_logits(8192); - let logits = Tensor::new(device, &logits_data); - logits.materialize_sync(); - let params = GpuMirostat2SamplerParams { - top_k: 512, - temperature: 0.8, - repetition_penalty: 1.05, - tau: 5.0, - eta: 0.1, - random: 0.35, - }; - - for _ in 0..warmup_batches { - let mut sampler = GpuMirostat2Sampler::new(device, 2.0 * params.tau); - let _ = logits - .sample_mirostat2_token(&mut sampler, &[], params) - .await?; - } - - let mut samples = Vec::with_capacity(measured_batches); - for _ in 0..measured_batches { - let mut sampler = GpuMirostat2Sampler::new(device, 2.0 * params.tau); - let start = std::time::Instant::now(); - let _ = logits - .sample_mirostat2_token(&mut sampler, &[], params) - .await?; - samples.push(start.elapsed()); - } - - print_summary("mirostat2_8192_k512", &samples, None); - Ok(()) -} - -fn bench_logits(len: usize) -> Vec { - (0..len) - .map(|index| { - let chunk_rank = index % 256; - let chunk = index / 256; - 10.0 - chunk_rank as f32 * 0.01 - chunk as f32 * 0.0001 - }) - .collect() -} - -fn print_summary(name: &str, samples: &[Duration], kernels_per_dispatch: Option) { - let mut samples = samples.to_vec(); - samples.sort_unstable(); - let mean = mean_duration(&samples); - let p50 = percentile_duration(&samples, 50); - let p90 = percentile_duration(&samples, 90); - let min = samples.first().copied().unwrap_or_default(); - let max = samples.last().copied().unwrap_or_default(); - - println!("{name}:"); - if let Some(kernels_per_dispatch) = kernels_per_dispatch { - println!(" kernels_per_dispatch: {kernels_per_dispatch:.3}"); - } - println!(" mean_us: {:.3}", duration_us(mean)); - println!(" p50_us: {:.3}", duration_us(p50)); - println!(" p90_us: {:.3}", duration_us(p90)); - println!(" min_us: {:.3}", duration_us(min)); - println!(" max_us: {:.3}", duration_us(max)); -} - -fn mean_duration(samples: &[Duration]) -> Duration { - if samples.is_empty() { - return Duration::ZERO; - } - samples.iter().copied().sum::() / samples.len() as u32 -} - -fn percentile_duration(samples: &[Duration], percentile: usize) -> Duration { - if samples.is_empty() { - return Duration::ZERO; - } - let index = ((samples.len() - 1) * percentile).div_ceil(100); - samples[index] -} - -fn duration_us(duration: Duration) -> f64 { - duration.as_secs_f64() * 1.0e6 -} diff --git a/fusor-ml/core/examples/bench_vision_matmul.rs b/fusor-ml/core/examples/bench_vision_matmul.rs deleted file mode 100644 index 442014572..000000000 --- a/fusor-ml/core/examples/bench_vision_matmul.rs +++ /dev/null @@ -1,259 +0,0 @@ -// Microbench for the vision-encoder matmul shapes hit during Qwen2.5-VL -// prefill. The vision encoder runs ~133 dense f32 matmuls per image; this -// binary isolates one shape and reports per-call ms so we can compare the -// shared-tile fallback against any new kernel variant in isolation. -// -// Run with: -// cargo run --package fusor-core --example bench_vision_matmul --release - -use std::time::{Duration, Instant}; - -use fusor_core::{Device, Tensor}; - -const WARMUP_BATCHES: usize = 3; -const MEASURED_BATCHES: usize = 10; -const DISPATCHES_PER_BATCH: usize = 4; - -fn main() -> Result<(), Box> { - pollster::block_on(async { - let device = Device::new().await?; - println!("bench_vision_matmul"); - println!("warmup_batches: {WARMUP_BATCHES}"); - println!("measured_batches: {MEASURED_BATCHES}"); - println!("dispatches_per_batch: {DISPATCHES_PER_BATCH}"); - println!(); - - // Shapes that show up in the Qwen2.5-VL vision encoder prefill for the - // standard demo image (M=1944 = 1944 patches after merge). N varies by - // projection; K is 1280 (embed dim) or 3420 (mlp inner / 2). - let cases: &[(&str, usize, usize, usize)] = &[ - ("vision_qkv", 1944, 1280, 3840), // fused qkv projection - ("vision_o", 1944, 1280, 1280), // attention output proj - ("vision_mlp_gate_up", 1944, 1280, 6840), // gate+up fused - ("vision_mlp_down", 1944, 3420, 1280), // down proj - // Aligned reference shape (M%128=0) so we can see how much faster the - // coop tile path is for the same K/N when M is friendly. - ("aligned_ref_1920_qkv", 1920, 1280, 3840), - ("aligned_ref_2048_qkv", 2048, 1280, 3840), - ]; - - for &(name, m, k, n) in cases { - bench_matmul(&device, name, m, k, n); - } - - bench_flash_attention_vision(&device); - - Ok(()) - }) -} - -fn bench_flash_attention_vision(device: &Device) { - // Mirror the vision attention shape: 16 heads, seq=1944, head_dim=80, - // unmasked self-attention (the per-window mask is dense and irrelevant - // to throughput). - let q = Tensor::splat(device, 0.1f32, [1, 16, 1944, 80]); - let k = Tensor::splat(device, 0.1f32, [1, 16, 1944, 80]); - let v = Tensor::splat(device, 0.1f32, [1, 16, 1944, 80]); - q.materialize_sync(); - k.materialize_sync(); - v.materialize_sync(); - - // Cold-start - let cold = Instant::now(); - { - let y = q.flash_attention(&k, &v, 1.0 / (80f32).sqrt(), None); - let _ = y.count_kernels_to_resolve(); - device.poll_wait(); - drop(y); - } - let cold_elapsed = cold.elapsed(); - - for _ in 0..WARMUP_BATCHES { - let y = q.flash_attention(&k, &v, 1.0 / (80f32).sqrt(), None); - let _ = y.count_kernels_to_resolve(); - device.poll_wait(); - drop(y); - } - - let mut samples = Vec::with_capacity(MEASURED_BATCHES); - for _ in 0..MEASURED_BATCHES { - let start = Instant::now(); - let y = q.flash_attention(&k, &v, 1.0 / (80f32).sqrt(), None); - let _ = y.count_kernels_to_resolve(); - device.poll_wait(); - samples.push(start.elapsed()); - drop(y); - } - - let mean = mean_duration(&samples); - let mut sorted = samples.clone(); - sorted.sort_unstable(); - let p50 = percentile_duration(&sorted, 50); - let p90 = percentile_duration(&sorted, 90); - let min = sorted.first().copied().unwrap_or_default(); - println!(); - println!("vision_flash_attention NO MASK (Q=K=V [1,16,1944,80]):"); - println!(" cold ms: {:.3}", cold_elapsed.as_secs_f64() * 1000.0); - println!(" mean_ms: {:.3}", mean.as_secs_f64() * 1000.0); - println!(" p50_ms: {:.3}", p50.as_secs_f64() * 1000.0); - println!(" p90_ms: {:.3}", p90.as_secs_f64() * 1000.0); - println!(" min_ms: {:.3}", min.as_secs_f64() * 1000.0); - - // Now bench WITH a mask (matches vision attention call site). - let mask: Tensor = Tensor::splat(device, 0.0f32, [1944, 1944]); - mask.materialize_sync(); - - // Bench with TRANSPOSED Q/K/V layout — the model produces Q via - // `xs.transpose(0,1).unsqueeze(0)` so the underlying tensor has - // non-contiguous strides. If that defeats a fast path in - // `try_flash_attention_direct`, this case will be much slower. - // Build a [1, 16, 1944, 80] tensor whose underlying layout is the - // [1944, 16, 80] memory order — same as `q.transpose(0, 1).unsqueeze(0)` - // in the model. We do this by allocating [1944, 16, 80] and then using - // restride to expose it as [1, 16, 1944, 80] with the seq/head strides - // swapped (head: 80, seq: 16*80=1280). - // (the transposed-Q/K/V bench was removed — see kernel-level analysis - // in qwen_vision_block.rs: the issue is V's non-contiguous layout - // defeats coalesced GPU loads in the streaming flash kernel.) - let _ = mask; - - for _ in 0..WARMUP_BATCHES { - let y = q.flash_attention(&k, &v, 1.0 / (80f32).sqrt(), Some(&mask)); - let _ = y.count_kernels_to_resolve(); - device.poll_wait(); - drop(y); - } - - let mut masked_samples = Vec::with_capacity(MEASURED_BATCHES); - for _ in 0..MEASURED_BATCHES { - let start = Instant::now(); - let y = q.flash_attention(&k, &v, 1.0 / (80f32).sqrt(), Some(&mask)); - let _ = y.count_kernels_to_resolve(); - device.poll_wait(); - masked_samples.push(start.elapsed()); - drop(y); - } - let mut sorted_m = masked_samples.clone(); - sorted_m.sort_unstable(); - println!(); - println!("vision_flash_attention WITH MASK (1944x1944):"); - println!( - " mean_ms: {:.3}", - mean_duration(&masked_samples).as_secs_f64() * 1000.0 - ); - println!( - " p50_ms: {:.3}", - percentile_duration(&sorted_m, 50).as_secs_f64() * 1000.0 - ); - println!( - " p90_ms: {:.3}", - percentile_duration(&sorted_m, 90).as_secs_f64() * 1000.0 - ); - println!( - " min_ms: {:.3}", - sorted_m.first().copied().unwrap_or_default().as_secs_f64() * 1000.0 - ); -} - -fn bench_matmul(device: &Device, name: &str, m: usize, k: usize, n: usize) { - let a = Tensor::splat(device, 0.001f32, [1, m, k]); - let b = Tensor::splat(device, 0.001f32, [1, k, n]); - a.materialize_sync(); - b.materialize_sync(); - - // Cold-start: measure the first dispatch separately. Shader pipeline - // creation (WGSL -> MSL -> Metal pipeline) happens lazily on first use - // — if this is large relative to warm runs, shader compile is the - // dominant cost in real prefill, not matmul math. - let cold_start = Instant::now(); - { - let y = a.mat_mul(&b); - let _ = y.count_kernels_to_resolve(); - device.poll_wait(); - drop(y); - } - let cold = cold_start.elapsed(); - println!(" COLD first-call ms: {:.3}", cold.as_secs_f64() * 1000.0); - - for _ in 0..WARMUP_BATCHES { - run_batch(device, &a, &b); - } - - let mut samples = Vec::with_capacity(MEASURED_BATCHES); - let mut kernels = 0usize; - for _ in 0..MEASURED_BATCHES { - let (elapsed, k_count) = run_batch(device, &a, &b); - samples.push(elapsed / DISPATCHES_PER_BATCH as u32); - kernels = k_count; - } - - let gflops = { - // M * N * K * 2 (mul + add) per matmul - let ops_per_call = (m as f64) * (n as f64) * (k as f64) * 2.0; - let mean_secs = mean_duration(&samples).as_secs_f64(); - if mean_secs > 0.0 { - ops_per_call / mean_secs / 1.0e9 - } else { - 0.0 - } - }; - - print_summary(name, m, k, n, &samples, kernels, gflops); -} - -fn run_batch(device: &Device, a: &Tensor, b: &Tensor) -> (Duration, usize) { - let mut outputs = Vec::with_capacity(DISPATCHES_PER_BATCH); - for _ in 0..DISPATCHES_PER_BATCH { - let y = a.mat_mul(b); - outputs.push(y); - } - let start = Instant::now(); - let kernels = outputs.iter().map(Tensor::count_kernels_to_resolve).sum(); - device.poll_wait(); - let elapsed = start.elapsed(); - drop(outputs); - (elapsed, kernels) -} - -fn print_summary( - name: &str, - m: usize, - k: usize, - n: usize, - samples: &[Duration], - kernels: usize, - gflops: f64, -) { - let mut sorted = samples.to_vec(); - sorted.sort_unstable(); - let mean = mean_duration(samples); - let p50 = percentile_duration(&sorted, 50); - let p90 = percentile_duration(&sorted, 90); - let min = sorted.first().copied().unwrap_or_default(); - let max = sorted.last().copied().unwrap_or_default(); - println!(); - println!("{name} (M={m}, K={k}, N={n}):"); - println!(" kernels_per_dispatch: {kernels}"); - println!(" mean_ms: {:.3}", mean.as_secs_f64() * 1000.0); - println!(" p50_ms: {:.3}", p50.as_secs_f64() * 1000.0); - println!(" p90_ms: {:.3}", p90.as_secs_f64() * 1000.0); - println!(" min_ms: {:.3}", min.as_secs_f64() * 1000.0); - println!(" max_ms: {:.3}", max.as_secs_f64() * 1000.0); - println!(" gflops_mean: {gflops:.1}"); - println!(); -} - -fn mean_duration(samples: &[Duration]) -> Duration { - if samples.is_empty() { - return Duration::ZERO; - } - samples.iter().copied().sum::() / samples.len() as u32 -} - -fn percentile_duration(sorted: &[Duration], percentile: usize) -> Duration { - if sorted.is_empty() { - return Duration::ZERO; - } - let index = ((sorted.len() - 1) * percentile) / 100; - sorted[index] -} diff --git a/fusor-ml/core/examples/wgpu_overhead_probe.rs b/fusor-ml/core/examples/wgpu_overhead_probe.rs deleted file mode 100644 index a4f094102..000000000 --- a/fusor-ml/core/examples/wgpu_overhead_probe.rs +++ /dev/null @@ -1,126 +0,0 @@ -use std::time::{Duration, Instant}; - -use fusor_core::{Device, Tensor}; - -async fn wait_for_submitted_work(device: &Device) { - let (sender, receiver) = futures_channel::oneshot::channel(); - device.wgpu_queue().on_submitted_work_done(|| { - _ = sender.send(()); - }); - let _ = receiver.await; -} - -async fn measure_async_empty_submit(device: &Device) -> Duration { - let encoder = device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("empty async overhead probe"), - }); - let start = Instant::now(); - device.wgpu_queue().submit(Some(encoder.finish())); - wait_for_submitted_work(device).await; - start.elapsed() -} - -fn measure_sync_empty_submit(device: &Device) -> Duration { - let encoder = device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("empty sync overhead probe"), - }); - let start = Instant::now(); - device.wgpu_queue().submit(Some(encoder.finish())); - device.poll_wait(); - start.elapsed() -} - -fn poll_until_empty(device: &Device) { - loop { - let status = device - .wgpu_device() - .poll(wgpu::PollType::Poll) - .expect("failed to poll GPU device"); - if status.is_queue_empty() { - break; - } - std::thread::yield_now(); - } -} - -fn measure_poll_loop_empty_submit(device: &Device) -> Duration { - let encoder = device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("empty poll-loop overhead probe"), - }); - let start = Instant::now(); - device.wgpu_queue().submit(Some(encoder.finish())); - poll_until_empty(device); - start.elapsed() -} - -async fn measure_callback_only(device: &Device) -> Duration { - let start = Instant::now(); - wait_for_submitted_work(device).await; - start.elapsed() -} - -fn print_samples(label: &str, samples: &[Duration]) { - let mut sorted = samples.to_vec(); - sorted.sort(); - let median = sorted[sorted.len() / 2]; - let min = sorted[0]; - let max = sorted[sorted.len() - 1]; - tracing::info!("{label}: median={median:?} min={min:?} max={max:?} samples={samples:?}"); -} - -fn main() { - let _ = tracing_subscriber::fmt::try_init(); - - pollster::block_on(async { - let repeats = std::env::args() - .nth(1) - .and_then(|arg| arg.parse::().ok()) - .unwrap_or(20); - - let device = Device::new().await.unwrap(); - tracing::info!("device={:?}", device.wgpu_adapter().get_info()); - - let data = vec![vec![1.0f32; 100]; 100]; - let tensor = Tensor::new(&device, &data); - _ = tensor.as_slice::<2, f32>().await.unwrap(); - - let add = tensor.clone() + 1.0; - add.materialize().await; - - let mut callback_only = Vec::with_capacity(repeats); - let mut empty_async = Vec::with_capacity(repeats); - let mut empty_sync = Vec::with_capacity(repeats); - let mut empty_poll_loop = Vec::with_capacity(repeats); - let mut leaf_materialize = Vec::with_capacity(repeats); - let mut add_materialize = Vec::with_capacity(repeats); - - for _ in 0..repeats { - callback_only.push(measure_callback_only(&device).await); - empty_async.push(measure_async_empty_submit(&device).await); - empty_sync.push(measure_sync_empty_submit(&device)); - empty_poll_loop.push(measure_poll_loop_empty_submit(&device)); - - let start = Instant::now(); - tensor.materialize().await; - leaf_materialize.push(start.elapsed()); - - let add = tensor.clone() + 1.0; - let start = Instant::now(); - add.materialize().await; - add_materialize.push(start.elapsed()); - } - - print_samples("callback-only", &callback_only); - print_samples("empty-submit-async", &empty_async); - print_samples("empty-submit-sync", &empty_sync); - print_samples("empty-submit-poll-loop", &empty_poll_loop); - print_samples("leaf-materialize", &leaf_materialize); - print_samples("add-materialize", &add_materialize); - }); -} diff --git a/fusor-ml/core/src/composite/attention.rs b/fusor-ml/core/src/composite/attention.rs new file mode 100644 index 000000000..dd75e07cc --- /dev/null +++ b/fusor-ml/core/src/composite/attention.rs @@ -0,0 +1,352 @@ +use crate::{ + DataTypeEnum, Layout, Tensor, + nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction, NaryOp, NaryScalar}, + view::ViewOperation, +}; + +/// Finite stand-in for -inf on masked score lanes. Naga validation rejects +/// non-finite float literals, so the composed causal select cannot lower with +/// a real -inf; this matches the flash kernels' `MASKED_SCORE` so the fused +/// and composed forms agree (exp of a masked lane underflows to zero either +/// way, and causal rows always keep at least one live lane for the max). +pub(crate) const MASKED_SCORE_F32: f32 = -3.0e38; +/// The f16 counterpart: the most negative finite half. +pub(crate) const MASKED_SCORE_F16: half::f16 = half::f16::MIN; + +impl Tensor { + /// A view layered directly on this tensor's node, without collapsing + /// into any underlying view chain. Composed-attention clusters use these + /// so recognition can peel the exact GQA-expand / transpose / mask + /// layouts back to the original q/k/v/mask nodes. + fn attached_view(&self, layout: Layout) -> Tensor { + Tensor::from_parts(self.data().view(ViewOperation::fully_defined( + self.key(), + layout, + self.shape(), + self.datatype(), + ))) + } + + /// Causal attention in its composed form: scores at kv positions + /// beyond the query position are replaced with `-inf` via an + /// index-comparison select (`kv_pos <= q_pos`), so causality is pure + /// index arithmetic — no mask tensor. The resolver recognizes the + /// cluster and routes it to the attention row program, whose axis bound + /// skips the masked upper-triangle tiles entirely. + pub fn attention_causal(&self, k: &Self, v: &Self, scale: f32) -> Self { + assert_eq!( + self.shape()[2], + k.shape()[2], + "causal attention requires q_seq_len == kv_seq_len \ + (self-attention prefill); use an explicit mask otherwise" + ); + self.compose_attention(k, v, scale, None, true) + } + + /// Scaled dot-product attention in its composed form: + /// `softmax(q · kᵀ · scale [+ mask]) · v`, with K/V expanded across query + /// heads for grouped-query attention. The resolver recognizes the + /// canonical cluster and routes it to the fused attention row program; + /// ineligible shapes lower through the recognized matmul + softmax + /// kernels (the same math). + pub fn attention(&self, k: &Self, v: &Self, scale: f32, mask: Option<&Tensor>) -> Self { + self.compose_attention(k, v, scale, mask, false) + } + + fn compose_attention( + &self, + k: &Self, + v: &Self, + scale: f32, + mask: Option<&Tensor>, + causal: bool, + ) -> Self { + self.assert_rank::<4>(); + k.assert_rank::<4>(); + v.assert_rank::<4>(); + assert_eq!(self.datatype(), k.datatype()); + assert_eq!(self.datatype(), v.datatype()); + if let Some(mask) = mask { + mask.assert_rank::<2>(); + assert_eq!(self.datatype(), mask.datatype()); + } + + let q_shape = self.shape(); + let k_shape = k.shape(); + + let batch = q_shape[0]; + let num_heads = q_shape[1]; + let q_seq_len = q_shape[2]; + let head_dim = q_shape[3]; + let num_kv_heads = k_shape[1]; + let kv_seq_len = k_shape[2]; + + assert!( + num_heads.is_multiple_of(num_kv_heads), + "Number of Q heads ({}) must be divisible by number of K/V heads ({})", + num_heads, + num_kv_heads + ); + + let groups = num_heads / num_kv_heads; + let expanded_shape = [batch, num_heads, kv_seq_len, head_dim]; + let expand = |tensor: &Tensor| -> Tensor { + if groups == 1 { + return tensor.clone(); + } + // Two attached views: a stride-0 broadcast across the group dim, + // then a flat reinterpret down to rank 4. + let grouped = tensor.attached_view(Layout::from_parts( + 0, + [batch, num_kv_heads, groups, kv_seq_len, head_dim].into(), + [ + num_kv_heads * kv_seq_len * head_dim, + kv_seq_len * head_dim, + 0, + head_dim, + 1, + ] + .into(), + )); + grouped.attached_view(Layout::contiguous(&expanded_shape)) + }; + let (k_expanded, v_expanded) = (expand(k), expand(v)); + + let k_t = k_expanded.attached_view(Layout::contiguous(&expanded_shape).transpose(2, 3)); + let scores = match self.datatype() { + DataTypeEnum::F32 => self.mat_mul(&k_t) * scale, + DataTypeEnum::F16 => self.mat_mul(&k_t) * half::f16::from_f32(scale), + DataTypeEnum::U32 => panic!("attention requires f32/f16 tensors"), + }; + let scores = if causal { + // Keep kv positions at or before the query position; everything + // later contributes exp(-inf) = 0 to the softmax. + let condition = NaryExpr::Op { + children: vec![NaryExpr::DimIndex(3), NaryExpr::DimIndex(2)], + function: NaryFunction::binary( + Some("causal_bound".to_string()), + NaryOp::LessEqual, + DataTypeEnum::U32, + DataTypeEnum::U32, + DataTypeEnum::U32, + ), + }; + let datatype = self.datatype(); + let neg_inf = match datatype { + DataTypeEnum::F32 => NaryScalar::F32(MASKED_SCORE_F32), + DataTypeEnum::F16 => NaryScalar::F16(MASKED_SCORE_F16), + DataTypeEnum::U32 => unreachable!("attention requires f32/f16"), + }; + let expression = NaryExpr::select( + condition, + NaryExpr::input(0, 4), + NaryExpr::scalar(neg_inf), + DataTypeEnum::U32, + datatype, + ); + Tensor::from_parts(scores.data().nary(ElementwiseOperation { + inputs: vec![scores.key()], + expression, + shape: [batch, num_heads, q_seq_len, kv_seq_len].into(), + output_datatype: datatype, + })) + } else if let Some(mask) = mask { + let mask_shape = mask.shape(); + assert_eq!( + mask_shape, + &[q_seq_len, kv_seq_len], + "attention mask shape {:?} does not match expected [{}, {}]", + mask_shape, + q_seq_len, + kv_seq_len + ); + let mask_view = mask.attached_view(Layout::from_parts( + 0, + [batch, num_heads, q_seq_len, kv_seq_len].into(), + [0, 0, kv_seq_len, 1].into(), + )); + scores + mask_view + } else { + scores + }; + + let weights = scores.softmax(3); + weights.mat_mul(&v_expanded) + } + + /// The canonical scaled-masked score cluster + /// (`q · kᵀ · scale [+ mask | causal-select]`) shared by the attention + /// composites, with K expanded across query heads for grouped-query + /// attention. Pattern recognition matches this exact structure. + fn attention_scores(&self, k: &Self, scale: f32, mask: Option<&Tensor>, causal: bool) -> Self { + self.assert_rank::<4>(); + k.assert_rank::<4>(); + assert_eq!(self.datatype(), k.datatype()); + + let q_shape = self.shape(); + let k_shape = k.shape(); + let batch = q_shape[0]; + let num_heads = q_shape[1]; + let q_seq_len = q_shape[2]; + let head_dim = q_shape[3]; + let num_kv_heads = k_shape[1]; + let kv_seq_len = k_shape[2]; + assert!(num_heads.is_multiple_of(num_kv_heads)); + let groups = num_heads / num_kv_heads; + let expanded_shape = [batch, num_heads, kv_seq_len, head_dim]; + let k_expanded = if groups == 1 { + k.clone() + } else { + let grouped = k.attached_view(Layout::from_parts( + 0, + [batch, num_kv_heads, groups, kv_seq_len, head_dim].into(), + [ + num_kv_heads * kv_seq_len * head_dim, + kv_seq_len * head_dim, + 0, + head_dim, + 1, + ] + .into(), + )); + grouped.attached_view(Layout::contiguous(&expanded_shape)) + }; + let k_t = k_expanded.attached_view(Layout::contiguous(&expanded_shape).transpose(2, 3)); + let scores = match self.datatype() { + DataTypeEnum::F32 => self.mat_mul(&k_t) * scale, + DataTypeEnum::F16 => self.mat_mul(&k_t) * half::f16::from_f32(scale), + DataTypeEnum::U32 => panic!("attention requires f32/f16 tensors"), + }; + if causal { + let condition = NaryExpr::Op { + children: vec![NaryExpr::DimIndex(3), NaryExpr::DimIndex(2)], + function: NaryFunction::binary( + Some("causal_bound".to_string()), + NaryOp::LessEqual, + DataTypeEnum::U32, + DataTypeEnum::U32, + DataTypeEnum::U32, + ), + }; + let datatype = self.datatype(); + let neg_inf = match datatype { + DataTypeEnum::F32 => NaryScalar::F32(MASKED_SCORE_F32), + DataTypeEnum::F16 => NaryScalar::F16(MASKED_SCORE_F16), + DataTypeEnum::U32 => unreachable!("attention requires f32/f16"), + }; + let expression = NaryExpr::select( + condition, + NaryExpr::input(0, 4), + NaryExpr::scalar(neg_inf), + DataTypeEnum::U32, + datatype, + ); + Tensor::from_parts(scores.data().nary(ElementwiseOperation { + inputs: vec![scores.key()], + expression, + shape: [batch, num_heads, q_seq_len, kv_seq_len].into(), + output_datatype: datatype, + })) + } else if let Some(mask) = mask { + assert_eq!(mask.shape(), &[q_seq_len, kv_seq_len]); + let mask_view = mask.attached_view(Layout::from_parts( + 0, + [batch, num_heads, q_seq_len, kv_seq_len].into(), + [0, 0, kv_seq_len, 1].into(), + )); + scores + mask_view + } else { + scores + } + } + + /// Row log-sum-exp of the attention scores over the KV axis: + /// `m + ln Σ exp(s − m)`, shape `[batch, heads, q_len]` — the statistic + /// that reconstructs attention probabilities as `exp(s − lse)`. + pub fn attention_lse(&self, k: &Self, scale: f32, mask: Option<&Tensor>, causal: bool) -> Self { + let scores = self.attention_scores(k, scale, mask, causal); + let shape = scores.shape().to_vec(); + let kept = [shape[0], shape[1], shape[2], 1]; + let m = scores.max(3); + let m_b = m.reshape(&kept).broadcast_as(&shape); + let sum = (&scores - &m_b).exp().sum(3); + &m + &sum.log() + } + + /// Gradients of [`Tensor::attention`] with respect to q, k, and v, + /// composed from the forward output and its row log-sum-exp: + /// probabilities are recomputed as `p = exp(s − lse)`, then + /// `dq = (p ∘ (dO·vᵀ − Σ dO∘o) · scale) · k` and its transposed + /// contractions produce dk and dv. Grouped-query K/V is not supported — + /// callers with fewer KV heads expand before differentiating. + #[allow(clippy::too_many_arguments)] + pub fn attention_grads( + &self, + k: &Self, + v: &Self, + o: &Self, + grad_o: &Self, + lse: &Self, + scale: f32, + mask: Option<&Tensor>, + causal: bool, + ) -> (Self, Self, Self) { + assert_eq!( + self.shape()[1], + k.shape()[1], + "attention_grads requires matching q/kv head counts" + ); + let scores = self.attention_scores(k, scale, mask, causal); + let shape = scores.shape().to_vec(); + let kept = [shape[0], shape[1], shape[2], 1]; + let row_broadcast = |row: &Tensor| row.reshape(&kept).broadcast_as(&shape); + let p = (&scores - &row_broadcast(lse)).exp(); + let dsum = (grad_o * o).sum(3); + let expanded = [shape[0], shape[1], shape[3], self.shape()[3]]; + let v_t = v.attached_view(Layout::contiguous(&expanded).transpose(2, 3)); + let dp = grad_o.mat_mul(&v_t); + let weighted = &p * &(&dp - &row_broadcast(&dsum)); + let ds = match self.datatype() { + DataTypeEnum::F16 => &weighted * half::f16::from_f32(scale), + _ => &weighted * scale, + }; + let dq = ds.mat_mul(k); + let score_transpose = + Layout::contiguous([shape[0], shape[1], shape[2], shape[3]].as_slice()).transpose(2, 3); + let dk = ds.attached_view(score_transpose.clone()).mat_mul(self); + let dv = p.attached_view(score_transpose).mat_mul(grad_o); + // Land both KV-side contractions in one combined tensor (dk rows, + // then dv rows, along the sequence axis): the paired streaming + // kernel claims the whole chain in one dispatch and shares the + // probability recomputation between them; the halves read back as + // zero-cost views. The composed fallback pays two slice copies. + let (batch, heads, kv_len) = (shape[0], shape[1], shape[3]); + let head_dim = self.shape()[3]; + let combined_shape = [batch, heads, 2 * kv_len, head_dim]; + let dkv = match self.datatype() { + DataTypeEnum::F16 => Tensor::splat(self.device(), half::f16::ZERO, combined_shape), + _ => Tensor::splat(self.device(), 0.0f32, combined_shape), + } + .slice_assign([0..batch, 0..heads, 0..kv_len, 0..head_dim], &dk) + .slice_assign([0..batch, 0..heads, kv_len..2 * kv_len, 0..head_dim], &dv); + let half_strides: Box<[usize]> = [ + heads * 2 * kv_len * head_dim, + 2 * kv_len * head_dim, + head_dim, + 1, + ] + .into(); + let half_shape: Box<[usize]> = [batch, heads, kv_len, head_dim].into(); + let dk_view = dkv.attached_view(Layout::from_parts( + 0, + half_shape.clone(), + half_strides.clone(), + )); + let dv_view = dkv.attached_view(Layout::from_parts( + kv_len * head_dim, + half_shape, + half_strides, + )); + (dq, dk_view, dv_view) + } +} diff --git a/fusor-ml/core/src/composite/flash_attention.rs b/fusor-ml/core/src/composite/flash_attention.rs deleted file mode 100644 index e36a4890e..000000000 --- a/fusor-ml/core/src/composite/flash_attention.rs +++ /dev/null @@ -1,167 +0,0 @@ -use crate::{ - DataTypeEnum, Layout, Tensor, - nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction, NaryOp, NaryScalar}, - view::ViewOperation, -}; - -impl Tensor { - /// A view layered directly on this tensor's node, without collapsing - /// into any underlying view chain. Composed-attention clusters use these - /// so recognition can peel the exact GQA-expand / transpose / mask - /// layouts back to the original q/k/v/mask nodes. - fn attached_view(&self, layout: Layout) -> Tensor { - Tensor::from_parts(self.data().view(ViewOperation::fully_defined( - self.key(), - layout, - self.shape(), - self.datatype(), - ))) - } - - /// Causal flash attention in its composed form: scores at kv positions - /// beyond the query position are replaced with `-inf` via an - /// index-comparison select (`kv_pos <= q_pos`), so causality is pure - /// index arithmetic — no mask tensor. The resolver recognizes the - /// cluster and routes it to the attention row program, whose axis bound - /// skips the masked upper-triangle tiles entirely. - pub fn flash_attention_causal(&self, k: &Self, v: &Self, scale: f32) -> Self { - assert_eq!( - self.shape()[2], - k.shape()[2], - "causal flash attention requires q_seq_len == kv_seq_len \ - (self-attention prefill); use an explicit mask otherwise" - ); - self.compose_attention(k, v, scale, None, true) - } - - /// Scaled dot-product attention in its composed form: - /// `softmax(q · kᵀ · scale [+ mask]) · v`, with K/V expanded across query - /// heads for grouped-query attention. The resolver recognizes the - /// canonical cluster and routes it to the fused attention row program; - /// ineligible shapes lower through the recognized matmul + softmax - /// kernels (the same math). - pub fn flash_attention(&self, k: &Self, v: &Self, scale: f32, mask: Option<&Tensor>) -> Self { - self.compose_attention(k, v, scale, mask, false) - } - - fn compose_attention( - &self, - k: &Self, - v: &Self, - scale: f32, - mask: Option<&Tensor>, - causal: bool, - ) -> Self { - self.assert_rank::<4>(); - k.assert_rank::<4>(); - v.assert_rank::<4>(); - assert_eq!(self.datatype(), k.datatype()); - assert_eq!(self.datatype(), v.datatype()); - if let Some(mask) = mask { - mask.assert_rank::<2>(); - assert_eq!(self.datatype(), mask.datatype()); - } - - let q_shape = self.shape(); - let k_shape = k.shape(); - - let batch = q_shape[0]; - let num_heads = q_shape[1]; - let q_seq_len = q_shape[2]; - let head_dim = q_shape[3]; - let num_kv_heads = k_shape[1]; - let kv_seq_len = k_shape[2]; - - assert!( - num_heads.is_multiple_of(num_kv_heads), - "Number of Q heads ({}) must be divisible by number of K/V heads ({})", - num_heads, - num_kv_heads - ); - - let groups = num_heads / num_kv_heads; - let expanded_shape = [batch, num_heads, kv_seq_len, head_dim]; - let expand = |tensor: &Tensor| -> Tensor { - if groups == 1 { - return tensor.clone(); - } - // Two attached views: a stride-0 broadcast across the group dim, - // then a flat reinterpret down to rank 4. - let grouped = tensor.attached_view(Layout::from_parts( - 0, - [batch, num_kv_heads, groups, kv_seq_len, head_dim].into(), - [ - num_kv_heads * kv_seq_len * head_dim, - kv_seq_len * head_dim, - 0, - head_dim, - 1, - ] - .into(), - )); - grouped.attached_view(Layout::contiguous(&expanded_shape)) - }; - let (k_expanded, v_expanded) = (expand(k), expand(v)); - - let k_t = k_expanded.attached_view(Layout::contiguous(&expanded_shape).transpose(2, 3)); - let scores = match self.datatype() { - DataTypeEnum::F32 => self.mat_mul(&k_t) * scale, - DataTypeEnum::F16 => self.mat_mul(&k_t) * half::f16::from_f32(scale), - DataTypeEnum::U32 => panic!("flash_attention requires f32/f16 tensors"), - }; - let scores = if causal { - // Keep kv positions at or before the query position; everything - // later contributes exp(-inf) = 0 to the softmax. - let condition = NaryExpr::Op { - children: vec![NaryExpr::DimIndex(3), NaryExpr::DimIndex(2)], - function: NaryFunction::binary( - Some("causal_bound".to_string()), - NaryOp::LessEqual, - DataTypeEnum::U32, - DataTypeEnum::U32, - DataTypeEnum::U32, - ), - }; - let datatype = self.datatype(); - let neg_inf = match datatype { - DataTypeEnum::F32 => NaryScalar::F32(f32::NEG_INFINITY), - DataTypeEnum::F16 => NaryScalar::F16(half::f16::NEG_INFINITY), - DataTypeEnum::U32 => unreachable!("attention requires f32/f16"), - }; - let expression = NaryExpr::select( - condition, - NaryExpr::input(0, 4), - NaryExpr::scalar(neg_inf), - DataTypeEnum::U32, - datatype, - ); - Tensor::from_parts(scores.data().nary(ElementwiseOperation { - inputs: vec![scores.key()], - expression, - shape: [batch, num_heads, q_seq_len, kv_seq_len].into(), - output_datatype: datatype, - })) - } else if let Some(mask) = mask { - let mask_shape = mask.shape(); - assert_eq!( - mask_shape, - &[q_seq_len, kv_seq_len], - "attention mask shape {:?} does not match expected [{}, {}]", - mask_shape, - q_seq_len, - kv_seq_len - ); - let mask_view = mask.attached_view(Layout::from_parts( - 0, - [batch, num_heads, q_seq_len, kv_seq_len].into(), - [0, 0, kv_seq_len, 1].into(), - )); - scores + mask_view - } else { - scores - }; - - let weights = scores.softmax(3); - weights.mat_mul(&v_expanded) - } -} diff --git a/fusor-ml/core/src/composite/mod.rs b/fusor-ml/core/src/composite/mod.rs index 55957c3cb..9a84d6453 100644 --- a/fusor-ml/core/src/composite/mod.rs +++ b/fusor-ml/core/src/composite/mod.rs @@ -1,4 +1,4 @@ -mod flash_attention; +pub(crate) mod attention; mod rms_norm_fused; mod rope_fused; mod softmax; diff --git a/fusor-ml/core/src/composite/rope_fused.rs b/fusor-ml/core/src/composite/rope_fused.rs index d6a3fa074..0601ee008 100644 --- a/fusor-ml/core/src/composite/rope_fused.rs +++ b/fusor-ml/core/src/composite/rope_fused.rs @@ -497,21 +497,29 @@ fn row_major_strides(shape: &[usize; 4]) -> Vec { } fn row_major_indices_from_flat(flat: NaryExpr, shape: &[usize; 4]) -> Vec { - let mut indices = Vec::with_capacity(shape.len()); - for axis in 0..shape.len() { - let divisor = shape[axis + 1..].iter().product::(); - let quotient = if divisor == 1 { - flat.clone() - } else { - NaryExpr::unary_op( - flat.clone(), + // Peel innermost-out with a running quotient — Metal miscompiles u32 + // div/mod chains with large non-power-of-two constants (see + // nary_direct::output_dims_from_flat). + let mut indices = vec![NaryExpr::scalar(NaryScalar::U32(0)); shape.len()]; + let mut rest = flat; + for axis in (0..shape.len()).rev() { + let dim = shape[axis]; + if dim == 1 { + continue; + } + if shape[..axis].iter().any(|&outer| outer != 1) { + indices[axis] = rem_const(rest.clone(), dim, "dim_index"); + rest = NaryExpr::unary_op( + rest, "div_stride", - NaryOp::DivConst(NaryScalar::U32(divisor as u32)), + NaryOp::DivConst(NaryScalar::U32(dim as u32)), DataTypeEnum::U32, DataTypeEnum::U32, - ) - }; - indices.push(rem_const(quotient, shape[axis], "dim_index")); + ); + } else { + indices[axis] = rest; + break; + } } indices } diff --git a/fusor-ml/core/src/compute_graph/mod.rs b/fusor-ml/core/src/compute_graph/mod.rs index f28a0c735..6174d50c7 100644 --- a/fusor-ml/core/src/compute_graph/mod.rs +++ b/fusor-ml/core/src/compute_graph/mod.rs @@ -4,27 +4,26 @@ use parking_lot::RwLock; pub use petgraph::graph::NodeIndex; use petgraph::prelude::StableGraph; use resolve::Resolver; +use resolve::flush_replay::{self, ReplayAction}; +use rustc_hash::FxHashMap; #[cfg(feature = "graphvis")] use tabbycat::Graph; +pub(crate) use resolve::FusionPlanStore; +pub(crate) use resolve::flush_replay::FlushPlanCache; + mod layout_pass; mod queue; -mod resolve; +pub(crate) mod resolve; #[cfg(test)] mod tests; #[cfg(feature = "graphvis")] mod visualize; use crate::{ - DataTypeEnum, Device, QMatrix, ReduceOperation, - compute_graph::resolve::ResolverResult, - dequantize::DequantizeOperation, - mir::{inputs::MirValue, operation::Operation}, - nary_wise::ElementwiseOperation, - quantized::matmul::QMatMulOperation, - slice_assign::SliceAssignOperation, - tensor::TensorData, - view::ViewOperation, + DataTypeEnum, Device, QMatrix, ReduceOperation, compute_graph::resolve::ResolverResult, + dequantize::DequantizeOperation, nary_wise::ElementwiseOperation, + slice_assign::SliceAssignOperation, tensor::TensorData, view::ViewOperation, visit_tiled::MaybeQData, }; @@ -65,13 +64,6 @@ impl ComputeGraph { self.create_node(ComputeGraphNodeVariant::View(op)) } - /// Build and submit one operation immediately against already-cached - /// inputs, bypassing the graph. Used by tuning APIs whose kernel - /// parameters cannot round-trip through the composed vocabulary. - pub(crate) fn execute_eager(&self, operation: &dyn Operation) -> Option { - self.with_mut(|inner| inner.execute_eager(operation)) - } - /// Clone the view at `key` if that node is an unresolved view. Used to /// collapse view chains at construction time. Cached views are excluded: /// a resolved view no longer keeps its base alive, so the base node may @@ -104,6 +96,23 @@ impl ComputeGraph { ))) } + /// Resolve every pending lazy output now, submitting the work to the + /// GPU without waiting for it or downloading anything. Keeps the pending + /// graph small in iteration-heavy workloads like training loops. + pub(crate) fn flush(&self) { + let mut removed = Vec::new(); + { + let mut inner = self.inner.write(); + inner.flush_all_pending(&mut removed); + inner.prune_deferred_dead(&mut removed); + #[cfg(feature = "extra_assertions")] + { + inner.verify_integrity() + } + } + drop(removed); + } + pub(crate) fn resolve(&self, key: NodeIndex) -> ResolverResult { if let Some(data) = { let inner = self.inner.read(); @@ -115,24 +124,12 @@ impl ComputeGraph { }; } - if let Some(data) = { - let mut inner = self.inner.write(); - let data = inner.try_resolve_direct_qmatmul(key); - #[cfg(feature = "extra_assertions")] - { - inner.verify_integrity() - } - data - } { - return data; - } - let (data, removed) = { let mut inner = self.inner.write(); let mut removed = Vec::new(); - let mut resolver = Resolver::new(&mut inner, key); - let data = resolver.run(&mut inner, &mut removed); + let (data, ()) = inner.resolve_target_with_replay(key, &mut removed, |_, _| ()); inner.try_auto_flush(&mut removed); + inner.prune_deferred_dead(&mut removed); #[cfg(feature = "extra_assertions")] { inner.verify_integrity() @@ -176,9 +173,9 @@ impl ComputeGraph { let (data, removed, tail_result) = { let mut inner = self.inner.write(); let mut removed = Vec::new(); - let mut resolver = Resolver::new(&mut inner, key); - let (data, tail_result) = resolver.run_with_tail(&mut inner, &mut removed, tail); + let (data, tail_result) = inner.resolve_target_with_replay(key, &mut removed, tail); inner.try_auto_flush(&mut removed); + inner.prune_deferred_dead(&mut removed); #[cfg(feature = "extra_assertions")] { inner.verify_integrity() @@ -204,6 +201,7 @@ impl ComputeGraph { let mut inner = self.inner.write(); let mut removed = Vec::new(); inner.remove_reference(key, &mut removed); + inner.prune_deferred_dead(&mut removed); #[cfg(feature = "extra_assertions")] { inner.verify_integrity() @@ -302,10 +300,6 @@ impl ComputeGraphNode { } } -pub(crate) trait GraphOperation: Operation + Send + Sync { - fn category(&self) -> &'static str; -} - /// The graph vocabulary. Exactly three core operations — elementwise /// visitation, reduction, and zero-dispatch views — over tensor and /// quantized-matrix leaves, plus the in-place region write (pure data @@ -354,23 +348,39 @@ pub(crate) struct ComputeGraphInner { // where the user would otherwise need to sprinkle explicit `resolve()` // calls. 0 disables. flush_threshold: usize, + // Incremental pending-sink set: every node with `reference_count > 0 && + // cached.is_none()`, tagged with a monotonically increasing insertion + // sequence number. Replaces the O(all-nodes) scan in `flush_all_pending` + // and makes sink enumeration deterministic in tape-construction order + // (StableGraph recycles indices, so `node_indices()` order is not stable + // across isomorphic steps) — which is what lets flush fingerprints of + // isomorphic steps collide. + pending_sinks: FxHashMap, + pending_seq: u64, + // Nodes whose `should_keep_cached()` flipped false during a resolve + // (their last alive descendant was cached). They cannot be removed at + // that point — the in-flight execution still reads their buffers by + // index — so removal is deferred to `prune_deferred_dead` at the end of + // the public operation. Without this, every cached-over node lingers as + // a permanent husk: `check_life` only runs on reference drops, and a + // dead node's references are already gone. + deferred_dead: Vec, } const DEFAULT_FLUSH_THRESHOLD: usize = 8192; -fn read_flush_threshold() -> usize { - std::env::var("FUSOR_GRAPH_FLUSH_THRESHOLD") - .ok() - .and_then(|v| v.parse::().ok()) - .unwrap_or(DEFAULT_FLUSH_THRESHOLD) -} - impl ComputeGraphInner { fn new(device: &Device) -> Self { Self { device: device.downgrade(), nodes: ComputeGraphNodes::default(), - flush_threshold: read_flush_threshold(), + flush_threshold: device + .config() + .graph_flush_threshold + .unwrap_or(DEFAULT_FLUSH_THRESHOLD), + pending_sinks: FxHashMap::default(), + pending_seq: 0, + deferred_dead: Vec::new(), } } @@ -380,6 +390,9 @@ impl ComputeGraphInner { device, nodes: ComputeGraphNodes::default(), flush_threshold: 0, + pending_sinks: FxHashMap::default(), + pending_seq: 0, + deferred_dead: Vec::new(), } } @@ -396,24 +409,87 @@ impl ComputeGraphInner { if self.nodes.nodes.node_count() < self.flush_threshold { return; } - let pending: Vec = self - .nodes - .nodes - .node_indices() - .filter(|&k| { - let Some(n) = self.nodes.nodes.node_weight(k) else { - return false; - }; - n.reference_count > 0 && n.cached.is_none() + self.flush_all_pending(removed); + } + + /// Materialize every pending lazy output in a single batched resolve. + /// + /// Consecutive structurally identical pending subgraphs go through the + /// flush-plan replay cache: the first occurrence records the full resolve + /// and later occurrences replay it, skipping execution-graph building, + /// optimization, lowering, and kernel building. + fn flush_all_pending(&mut self, removed: &mut Vec) { + // Enumerate pending sinks in insertion (tape-construction) order so + // fingerprints of isomorphic steps are deterministic for every graph. + let mut pending: Vec<(u64, NodeIndex)> = self + .pending_sinks + .iter() + .filter(|&(&key, _)| { + self.nodes + .nodes + .node_weight(key) + .map(|n| n.reference_count > 0 && n.cached.is_none()) + .unwrap_or(false) }) + .map(|(&key, &seq)| (seq, key)) .collect(); + pending.sort_unstable(); + let pending: Vec = pending.into_iter().map(|(_, key)| key).collect(); if pending.is_empty() { return; } + + match flush_replay::prepare_replay(self, &pending) { + ReplayAction::Replay { plan, fingerprint } => { + let _ = + flush_replay::execute_replay_with_tail(self, &plan, &fingerprint, |_, _| ()); + return; + } + ReplayAction::Record { key, fingerprint } => { + let mut resolver = Resolver::new_batch_with_recording(self, pending, fingerprint); + let _ = resolver.run(self, removed); + if let Some(plan) = resolver.take_recorded_plan() { + self.device().flush_plan_cache().insert(key, Arc::new(plan)); + } + return; + } + ReplayAction::Resolve => {} + } + let mut resolver = Resolver::new_batch(self, pending); let _ = resolver.run(self, removed); } + /// Resolve one target, recording or replaying the same bufferless plan + /// format used by batched flushes. This is the hot materialization path + /// for repeated isomorphic inference graphs such as `as_slice()` during + /// autoregressive generation. + fn resolve_target_with_replay( + &mut self, + target: NodeIndex, + removed: &mut Vec, + tail: impl FnOnce(&TensorData, &mut wgpu::CommandEncoder) -> T, + ) -> (ResolverResult, T) { + match flush_replay::prepare_replay(self, &[target]) { + ReplayAction::Replay { plan, fingerprint } => { + flush_replay::execute_replay_with_tail(self, &plan, &fingerprint, tail) + } + ReplayAction::Record { key, fingerprint } => { + let mut resolver = + Resolver::new_batch_with_recording(self, vec![target], fingerprint); + let result = resolver.run_with_tail(self, removed, tail); + if let Some(plan) = resolver.take_recorded_plan() { + self.device().flush_plan_cache().insert(key, Arc::new(plan)); + } + result + } + ReplayAction::Resolve => { + let mut resolver = Resolver::new(self, target); + resolver.run_with_tail(self, removed, tail) + } + } + } + /// Upgrade the weak device reference to a strong one. /// Panics if the device has been dropped (should not happen during normal operation). pub(crate) fn device(&self) -> Device { @@ -429,19 +505,31 @@ impl ComputeGraphInner { live_descendant_count: 0, cached: None, }); + // New node has `reference_count = 1` and no cached result: pending. + self.mark_pending(node); // New node has `reference_count = 1`, so it is alive. Adding edges // below propagates that liveness up to each dependency. self.add_dependency_edges(node); node } + /// Track `key` in the pending-sink set (referenced and uncached). + fn mark_pending(&mut self, key: NodeIndex) { + let seq = self.pending_seq; + self.pending_seq += 1; + self.pending_sinks.entry(key).or_insert(seq); + } + fn add_reference(&mut self, key: NodeIndex) { - let transitioned_alive = { + let (transitioned_alive, now_pending) = { let node = self.nodes.nodes.node_weight_mut(key).unwrap(); let prev_alive = node.alive_uncached(); node.reference_count += 1; - !prev_alive && node.alive_uncached() + (!prev_alive && node.alive_uncached(), node.cached.is_none()) }; + if now_pending { + self.mark_pending(key); + } if transitioned_alive { self.propagate_alive_change(key, true); } @@ -497,11 +585,12 @@ impl ComputeGraphInner { .neighbors_directed(child, petgraph::Direction::Incoming) .collect(); for parent in parents { - let parent_transitioned = { + let (parent_transitioned, parent_now_dead) = { let Some(parent_node) = self.nodes.nodes.node_weight_mut(parent) else { continue; }; let prev_parent_alive = parent_node.alive_uncached(); + let prev_parent_kept = parent_node.should_keep_cached(); if now_alive { parent_node.live_descendant_count = parent_node .live_descendant_count @@ -511,8 +600,21 @@ impl ComputeGraphInner { parent_node.live_descendant_count = parent_node.live_descendant_count.saturating_sub(1); } - prev_parent_alive != parent_node.alive_uncached() + ( + prev_parent_alive != parent_node.alive_uncached(), + prev_parent_kept && !parent_node.should_keep_cached(), + ) }; + // A parent whose last live descendant just went away is now + // unreachable by any future resolve and must eventually be + // removed. `check_life` cannot run here — during a resolve + // the execution graph still reads this node — so record it + // for `prune_deferred_dead`. Note this also catches CACHED + // parents, which never flip `alive_uncached` and so are + // invisible to the transition propagation below. + if parent_now_dead { + self.deferred_dead.push(parent); + } if parent_transitioned { stack.push(parent); } @@ -526,93 +628,19 @@ impl ComputeGraphInner { } } - fn ensure_tensor_cached(&mut self, key: NodeIndex) -> Option<()> { - if self.get_cached_result(key).is_some() { - return Some(()); - } - - let data = match self.nodes.nodes.node_weight(key)?.variant.clone() { - ComputeGraphNodeVariant::Tensor(data) => data, - _ => return None, - }; - self.set_cached_result(key, data); - Some(()) - } - - fn try_submit_direct_qmatmul( - &mut self, - operation: &QMatMulOperation, - ) -> Option<(TensorData, usize)> { - self.ensure_tensor_cached(operation.input)?; - - let device = self.device(); - let workgroup_shape = crate::mir::workgroup_shape::WorkgroupShape::new(1, 1, 1); - let inputs = operation.inputs(self); - let direct_kernel_plan = operation - .build_direct_kernels(self, &workgroup_shape, &inputs) - .ok()?; - let MirValue::Tensor(output) = operation.output(self, &inputs) else { - return None; - }; - - let mut command_encoder = - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("QMatMul Direct Encoder"), - }); - let total_kernels = direct_kernel_plan.dispatch_count(); - for direct_kernel in direct_kernel_plan.into_kernels() { - direct_kernel.run(device.kernel_cache(), &mut command_encoder); - } - if total_kernels > 0 { - device.wgpu_queue().submit(Some(command_encoder.finish())); - device.reset_initialized_buffers(); - } - - Some((output, total_kernels)) - } - - fn execute_eager(&mut self, operation: &dyn Operation) -> Option { - let device = self.device(); - let inputs = operation.inputs(self); - let workgroup_shape = operation - .workgroup_shape_constraints(&device) - .solve(device.max_subgroup_size(), &device.limits())?; - let kernel = operation.build_direct_kernel(self, &workgroup_shape, &inputs)?; - let MirValue::Tensor(output) = operation.output(self, &inputs) else { - return None; - }; - - let mut command_encoder = - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("Eager Operation Encoder"), - }); - kernel.run(device.kernel_cache(), &mut command_encoder); - device.wgpu_queue().submit(Some(command_encoder.finish())); - device.reset_initialized_buffers(); - Some(output) - } - - fn try_resolve_direct_qmatmul(&mut self, key: NodeIndex) -> Option { - let operation = self.match_direct_qmatmul(key)?; - let (output, total_kernels) = self.try_submit_direct_qmatmul(&operation)?; - self.set_cached_result(key, output.clone()); - Some(ResolverResult { - data: output, - total_kernels, - }) - } - fn remove_reference(&mut self, key: NodeIndex, removed: &mut Vec) { - let transitioned_dead = { + let (transitioned_dead, still_referenced) = { let node = self.nodes.nodes.node_weight_mut(key).unwrap(); let prev_alive = node.alive_uncached(); node.reference_count = node.reference_count.saturating_sub(1); - prev_alive && !node.alive_uncached() + ( + prev_alive && !node.alive_uncached(), + node.reference_count > 0, + ) }; + if !still_referenced { + self.pending_sinks.remove(&key); + } if transitioned_dead { self.propagate_alive_change(key, false); } @@ -620,39 +648,55 @@ impl ComputeGraphInner { } fn check_life(&mut self, key: NodeIndex, removed: &mut Vec) { - // The node is needed iff it has external references OR some - // uncached live descendant. `live_descendant_count` is maintained - // eagerly, so this is O(1). - match self - .nodes - .nodes - .node_weight(key) - .map(|n| n.should_keep_cached()) - { - Some(true) | None => return, - Some(false) => {} - } - - let mut dependencies = Vec::new(); - self.visit_dependencies(key, &mut |dependency| { - dependencies.push(dependency); - }); + // Iterative worklist, NOT recursion: teardown cascades one frame per + // node, and dropping the last handle to a long-lived chain (e.g. an + // optimizer moment at the end of a training run) must not overflow + // the stack. + let mut worklist = vec![key]; + while let Some(key) = worklist.pop() { + // The node is needed iff it has external references OR some + // uncached live descendant. `live_descendant_count` is maintained + // eagerly, so this is O(1). + match self + .nodes + .nodes + .node_weight(key) + .map(|n| n.should_keep_cached()) + { + Some(true) | None => continue, + Some(false) => {} + } - // Not needed — remove it. Per the invariant above, the node's - // `alive_uncached` was already false (cached.is_some() or - // ref==luc==0), so its contribution to each parent's - // `live_descendant_count` is already 0; no further bookkeeping is - // needed when the edges go away with the node. - self.remove_key(key, removed); + // Not needed — remove it. Per the invariant above, the node's + // `alive_uncached` was already false (cached.is_some() or + // ref==luc==0), so its contribution to each parent's + // `live_descendant_count` is already 0; no further bookkeeping is + // needed when the edges go away with the node. + self.visit_dependencies(key, &mut |dependency| { + worklist.push(dependency); + }); + self.remove_key(key, removed); + } + } - for dependency in dependencies { - self.check_life(dependency, removed); + /// Remove nodes whose liveness died inside a resolve (recorded in + /// `deferred_dead` by `propagate_alive_change`). Runs at the end of the + /// public graph operations, once the execution that was still reading + /// those nodes' buffers has been submitted. Entries removed by an earlier + /// cascade are skipped by `check_life`'s existence check. + fn prune_deferred_dead(&mut self, removed: &mut Vec) { + while let Some(key) = self.deferred_dead.pop() { + self.check_life(key, removed); } } fn remove_key(&mut self, key: NodeIndex, removed: &mut Vec) { // Remove the node from the graph (this also removes all edges) if let Some(node) = self.nodes.nodes.remove_node(key) { + // A removable node has `reference_count == 0`, so it should + // already be out of the pending set; defensive removal keeps the + // set exact even if that invariant ever slips. + self.pending_sinks.remove(&key); removed.push(node); } } @@ -674,16 +718,24 @@ impl ComputeGraphInner { } pub(crate) fn set_cached_result(&mut self, key: NodeIndex, data: TensorData) { + // A cached node is no longer a pending sink. + self.pending_sinks.remove(&key); // Setting `cached` flips `alive_uncached` false: a cached node no // longer needs to be recomputed, so its parents can free their own // cached buffers once no other uncached descendant remains. Propagate // the transition so ancestor counters reflect the new state. - let transitioned_dead = { + let (transitioned_dead, now_dead) = { let node = self.nodes.nodes.node_weight_mut(key).unwrap(); let prev_alive = node.alive_uncached(); node.cached = Some(data); - prev_alive && !node.alive_uncached() + ( + prev_alive && !node.alive_uncached(), + !node.should_keep_cached(), + ) }; + if now_dead { + self.deferred_dead.push(key); + } if transitioned_dead { self.propagate_alive_change(key, false); } @@ -724,6 +776,24 @@ impl ComputeGraphInner { #[cfg(feature = "extra_assertions")] fn verify_integrity(&self) { + // Dead nodes (no references, no live uncached descendant) are pruned + // eagerly — by the `check_life` cascade on reference drops and by + // `prune_deferred_dead` after resolves — so none may survive past the + // end of a public operation. A node lingering here is a husk: it + // would accumulate once per training step and make final teardown + // O(steps). + assert!( + self.deferred_dead.is_empty(), + "deferred dead set not drained" + ); + for key in self.nodes.nodes.node_indices() { + let node = self.nodes.nodes.node_weight(key).unwrap(); + assert!( + node.should_keep_cached(), + "dead node {key:?} survived pruning" + ); + } + // Check that all edges point to existing nodes for key in self.nodes.nodes.node_indices() { for neighbor in self.nodes.nodes.neighbors(key) { @@ -787,5 +857,27 @@ impl ComputeGraphInner { "live_descendant_count mismatch at {key:?}: expected {expected}, got {actual}" ); } + + // Check that the incremental pending-sink set exactly matches the + // predicate it caches (`reference_count > 0 && cached.is_none()`). + for key in self.nodes.nodes.node_indices() { + let pending = self + .nodes + .nodes + .node_weight(key) + .map(|n| n.reference_count > 0 && n.cached.is_none()) + .unwrap_or(false); + assert_eq!( + self.pending_sinks.contains_key(&key), + pending, + "pending_sinks mismatch at {key:?}: expected pending={pending}" + ); + } + for key in self.pending_sinks.keys() { + assert!( + self.nodes.nodes.contains_node(*key), + "pending_sinks contains removed node {key:?}" + ); + } } } diff --git a/fusor-ml/core/src/compute_graph/resolve/alloc_reuse.rs b/fusor-ml/core/src/compute_graph/resolve/alloc_reuse.rs new file mode 100644 index 000000000..74f4fa685 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/alloc_reuse.rs @@ -0,0 +1,296 @@ +//! Liveness-based allocation reuse: an operation output claims the buffer +//! of a dead intermediate instead of allocating fresh. +//! +//! Claims are decided serially in queue order while the compute-graph +//! write lock is held for the whole resolve (no user thread can clone or +//! drop handles), so `Arc::strong_count` reads are stable and the claim +//! decision is a pure function of graph structure. A replayed plan +//! re-creates the same claims through `OutputSource::Alias`. +//! +//! Claim sources are restricted to buffers allocated during this resolve +//! (intermediates), and an operation never claims a buffer it reads: +//! in-place read-write claims need the kernel to fold the pair into one +//! read-write binding (wgpu rejects one buffer bound read-only and +//! read-write in the same dispatch), so they are declined here. +//! +//! The accounting is fail-safe by construction: every strong-reference +//! holder the resolver knows about is enumerated, and any *unaccounted* +//! holder makes the observed `Arc::strong_count` exceed the expectation, so +//! the claim is declined and a fresh buffer is allocated instead — a missed +//! registration can cause a missed optimization, never a wrong result. + +use std::collections::VecDeque; + +use super::*; + +pub(super) struct BufferLedger { + enabled: bool, + /// Set while claims may still be made: release events feed the free + /// list only until the claim window closes. + accepting: bool, + device: crate::Device, + /// Shadow of the remaining-consumer accounting, advanced at gather time + /// when releases run later than allocations (the batched dense queue, + /// where every allocation happens before any release). Empty when the + /// real release runs in step with allocation and feeds the free list + /// directly via [`Self::note_released`]. + shadow: FxHashMap, + /// Dead intermediate buffers by (size, usage), in death order. One entry + /// per buffer pointer (the entry holds its own strong clone). + free: FxHashMap<(u64, wgpu::BufferUsages), VecDeque<(NodeIndex, Arc)>>, + /// Buffer pointers currently queued in `free`. + queued: FxHashSet, + /// ptr -> lingering `cached` clones of dead graph nodes (phase 3 has not + /// released them yet). + dead_cached: FxHashMap, + /// ptr -> phase-1 resolver clones registered so far (work-item inputs and + /// outputs, view results, copy records). + transient: FxHashMap, + /// Buffers allocated during this resolve's phase 1. + allocated_here: FxHashSet, + /// Claimer node -> source node: the recorder's sole authority for + /// classifying an output-provenance hit as a chosen alias. + chosen: FxHashMap, + /// ptr -> strong clones held by the flush-plan recorder (boundary pins + /// taken at recorder construction). + recorder_pins: FxHashMap, + pub(super) claims: usize, +} + +impl BufferLedger { + pub(super) fn new( + device: &crate::Device, + shadow_consumers: Option<&FxHashMap>, + ) -> Self { + let enabled = !device.poisons_allocations(); + Self { + enabled, + accepting: enabled, + device: device.clone(), + shadow: match shadow_consumers { + Some(consumers) if enabled => consumers.clone(), + _ => FxHashMap::default(), + }, + free: FxHashMap::default(), + queued: FxHashSet::default(), + dead_cached: FxHashMap::default(), + transient: FxHashMap::default(), + allocated_here: FxHashSet::default(), + chosen: FxHashMap::default(), + recorder_pins: FxHashMap::default(), + claims: 0, + } + } + + /// Register strong clones an armed flush-plan recorder holds, so claimed + /// buffers it pinned still account exactly. + pub(super) fn register_recorder_pins(&mut self, pins: impl Iterator) { + if !self.enabled { + return; + } + for ptr in pins { + *self.recorder_pins.entry(ptr).or_insert(0) += 1; + } + } + + pub(super) fn enabled(&self) -> bool { + self.enabled + } + + /// Close the claim window: later release events no longer feed the free + /// list (their buffers can no longer be claimed by anything). + pub(super) fn freeze(&mut self) { + self.accepting = false; + } + + /// A dead node's cached buffer is about to be released: make it + /// claimable if it was allocated during this resolve. + pub(super) fn note_released(&mut self, source: NodeIndex, cached: &TensorData) { + if !self.enabled || !self.accepting { + return; + } + let buffer = cached.buffer(); + let ptr = Arc::as_ptr(buffer) as usize; + if !self.allocated_here.contains(&ptr) || !self.queued.insert(ptr) { + return; + } + self.free + .entry((buffer.size(), buffer.usage())) + .or_default() + .push_back((source, buffer.clone())); + } + + /// The source node this claimer's output aliases, if any. + pub(super) fn chosen_source(&self, claimer: NodeIndex) -> Option { + self.chosen.get(&claimer).copied() + } + + /// Register a buffer allocated during this phase 1. + pub(super) fn note_alloc(&mut self, data: &TensorData) { + if self.enabled { + self.allocated_here + .insert(Arc::as_ptr(data.buffer()) as usize); + } + } + + /// Register one resolver-held clone (work-item input/output, view + /// result, copy record) of `buffer`. + pub(super) fn note_transient(&mut self, buffer: &Arc) { + if self.enabled { + *self + .transient + .entry(Arc::as_ptr(buffer) as usize) + .or_insert(0) += 1; + } + } + + /// Advance the shadow release accounting for one produced node's + /// dependencies (mirror of `release_consumed`, without the + /// release): nodes whose last consumer this is enter the free list when + /// their buffer was allocated this resolve. + pub(super) fn consume( + &mut self, + graph: &ComputeGraphInner, + deps: &[NodeIndex], + targets: &FxHashSet, + ) { + if !self.enabled || !self.accepting { + return; + } + for &dep in deps { + let Some(count) = self.shadow.get_mut(&dep) else { + continue; + }; + *count = count.saturating_sub(1); + if *count != 0 || targets.contains(&dep) || graph.has_live_lazy_descendant(dep) { + continue; + } + let Some(cached) = graph.get_cached_result(dep) else { + continue; + }; + let buffer = cached.buffer(); + let ptr = Arc::as_ptr(buffer) as usize; + // Every dead node's lingering `cached` clone counts, even when + // several dead nodes (a view and its base) share one buffer. + *self.dead_cached.entry(ptr).or_insert(0) += 1; + if !self.allocated_here.contains(&ptr) || !self.queued.insert(ptr) { + continue; + } + self.free + .entry((buffer.size(), buffer.usage())) + .or_default() + .push_back((dep, buffer.clone())); + } + } + + /// Claim a dead intermediate's buffer for `claimer`'s output, or `None` + /// to allocate fresh. `forbidden` is the set of buffer pointers the + /// claiming operation reads (in-place claims need binding folding in + /// the kernel, so they are declined here). + pub(super) fn try_claim( + &mut self, + claimer: NodeIndex, + output: &TensorData, + forbidden: &FxHashSet, + ) -> Option { + if !self.enabled { + return None; + } + let key = (output.buffer().size(), output.buffer().usage()); + let candidates = self.free.get_mut(&key)?; + let mut picked = None; + for (index, (source, buffer)) in candidates.iter().enumerate() { + let ptr = Arc::as_ptr(buffer) as usize; + if forbidden.contains(&ptr) { + continue; + } + // Expected holders: dead graph clones + this free-list entry + + // registered phase-1 transients + the pool's tracked clone. + let expected = self.dead_cached.get(&ptr).copied().unwrap_or(0) + + 1 + + self.transient.get(&ptr).copied().unwrap_or(0) + + self.recorder_pins.get(&ptr).copied().unwrap_or(0) + + u32::from(self.device.buffer_pool_is_tracked(key.0, key.1, buffer)); + if Arc::strong_count(buffer) as u32 != expected { + continue; + } + picked = Some((index, *source)); + break; + } + let (index, source) = picked?; + let (_, buffer) = candidates.remove(index).expect("index in range"); + self.queued.remove(&(Arc::as_ptr(&buffer) as usize)); + self.chosen.insert(claimer, source); + self.claims += 1; + Some(TensorData::new_from_parts( + &self.device, + buffer, + output.layout().clone(), + output.datatype(), + )) + } + + /// Claim a specific dead node's buffer for an output that will be + /// written by the same dispatch that reads it. The caller guarantees + /// the kernel-level safety conditions (identity-indexed reads, one + /// read-write binding, no later reader of the source within the + /// dispatch); this checks liveness and exact reference accounting. + pub(super) fn try_claim_in_place( + &mut self, + claimer: NodeIndex, + output: &TensorData, + source: NodeIndex, + graph: &ComputeGraphInner, + targets: &FxHashSet, + ) -> Option { + if !self.enabled { + return None; + } + // The source must be dead at this queue position: the shadow + // accounting has already consumed the claiming operation's reads. + if self.shadow.get(&source).copied().unwrap_or(usize::MAX) != 0 + || targets.contains(&source) + || graph.has_live_lazy_descendant(source) + { + return None; + } + let cached = graph.get_cached_result(source)?; + if cached.datatype() != output.datatype() + || cached.layout() != output.layout() + || cached.buffer().size() != output.buffer().size() + || cached.buffer().usage() != output.buffer().usage() + { + return None; + } + let buffer = cached.buffer().clone(); + let ptr = Arc::as_ptr(&buffer) as usize; + // One extra holder for the `buffer` clone taken just above; free-list + // membership adds another. + let expected = self.dead_cached.get(&ptr).copied().unwrap_or(0) + + 1 + + u32::from(self.queued.contains(&ptr)) + + self.transient.get(&ptr).copied().unwrap_or(0) + + self.recorder_pins.get(&ptr).copied().unwrap_or(0) + + u32::from( + self.device + .buffer_pool_is_tracked(buffer.size(), buffer.usage(), &buffer), + ); + if Arc::strong_count(&buffer) as u32 != expected { + return None; + } + // Retire any free-list entry so nothing else claims this buffer. + if self.queued.remove(&ptr) { + if let Some(entries) = self.free.get_mut(&(buffer.size(), buffer.usage())) { + entries.retain(|(_, entry)| Arc::as_ptr(entry) as usize != ptr); + } + } + self.chosen.insert(claimer, source); + self.claims += 1; + Some(TensorData::new_from_parts( + &self.device, + buffer, + output.layout().clone(), + output.datatype(), + )) + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/HOISTING_SPIKE.md b/fusor-ml/core/src/compute_graph/resolve/egraph/HOISTING_SPIKE.md new file mode 100644 index 000000000..2c41e010c --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/HOISTING_SPIKE.md @@ -0,0 +1,136 @@ +# Recognition-hoisting measurement spike + +`recognize.rs` keeps recognition outside equality saturation because ingesting +the un-preshrunk graph for every generated token was believed to be expensive, +and because a generator that re-derives a matmul needs a structural window +deeper than the two-step horizon plan sharing cuts at +(`structural_memo.rs`, `WINDOW_STUB_DEPTH`). Both costs were predictions. This +spike measures them. + +The knobs are measurement-only and default off: + +| variable | effect | +|---|---| +| `FUSOR_SPIKE_HOISTING` | emit the `hoisting_spike_ingest` / `hoisting_spike_windows` ledgers; changes no decision | +| `FUSOR_SPIKE_NO_RECOGNITION=` | skip the pre-ingest recognition sweep for resolves with at most `n` execution nodes | +| `FUSOR_SPIKE_WINDOW_DEPTH=` | widen the structural window horizon from the built-in stub depth | + +## Method + +One process, the real decode trace from the verification baseline: + +``` +FUSOR_TRACE_RESOLVE=1 FUSOR_TRACE_RESOLVE_HOST=1 FUSOR_SPIKE_HOISTING=1 [knob] \ + cargo run --release -p fusor --example transformer -- \ + --steps 2 --progress-every 0 --trace-resolve --trace-host +``` + +Two training steps, an eval pass, then 400 generated tokens. The 399 identical +steady-state decode resolves (`node_count=667` before recognition, 77 kernels +per token) are the population; every number below is the p50 over those 399 +unless stated otherwise. Every configuration was run twice: all counts are +byte-identical run to run and the p50 timings agree within 1%. A run with +every spike knob off reproduces the golden `resolve_egg_plans` line and +decode dispatch histogram, and all four configurations generate byte-identical +text. + +`FUSOR_SPIKE_NO_RECOGNITION` is scoped by graph size because the un-preshrunk +form of a *training* step does not run at all: with recognition disabled +everywhere, the first training resolve dies allocating 8.59 GB more (20.57 GB +already live under a 22.91 GB in-flight cap) — the contraction's `[.., N, K]` +product materialized instead of contracted. Budget `1000` keeps recognition +for the two training resolves (3498 and 3265 execution nodes) and drops it for +every eval, prefill and decode resolve (693, 663 and 667 nodes), all of which +run to completion un-preshrunk. + +## (a) Ingest cost of the un-preshrunk decode graph + +| per decode token | recognized | un-preshrunk | delta | +|---|---|---|---| +| execution nodes ingested | 387 | 667 | +72% | +| provenances | 466 | 746 | +60% | +| e-nodes / e-classes | 454 / 454 | 734 / 734 | +62% | +| interned payloads / specs | 47 / 47 | 72 / 72 | +53% | +| recognition sweep | 320 µs | 0 µs | −320 µs | +| ingest | 379 µs | 606 µs | +227 µs | +| window capture | 480 µs (311 windows, 1.55 µs each) | 911 µs (724 windows, 1.26 µs each) | +431 µs | +| `optimize` phase total | 1879 µs | 3019 µs | +1140 µs | +| resolver host total | 4144 µs | 5494 µs | +1350 µs | + +Ingest itself is not the bill, and the never-measured cost turns out to be the +smaller half: the sweep costs 320 µs and saves 227 µs of ingest, so the +pre-shrink is a 93 µs *net win* on ingest alone. The rest lands in extraction +(1415 µs → 2765 µs): +431 µs capturing 724 windows instead of 311, and ++692 µs generating, looking up and costing candidates over a graph 60% larger. + +The extraction *result* is much worse, which is the finding that matters: + +| per decode token | recognized | un-preshrunk | +|---|---|---| +| dispatch categories | `flash_attention 6, matmul_f32 19, merged_matmul 6, nary_direct 33, row_program 13` | `merged_row 6, nary_direct 70, row_program 69` | +| kernels | 77 | 145 | +| extractor cost `work` | 9.86e9 | 4.90e10 | + +No matmul and no attention kernel survives. Today's generators do not +re-derive a contraction from `Elementwise(Mul) + Reduce(Sum)`, so removing the +sweep does not move recognition into the e-graph — it deletes it, and the +generic row/nary lowering picks up the pieces at 5x the arithmetic work and +1.9x the dispatches. (It is still correct: the 400 generated tokens are +byte-identical to the recognized run, which is the `recognize.rs` "slower, but +correct" claim, measured.) + +A recognizer-to-generator port therefore has to pay all of the above *and* +carry the matcher logic into the generator set; the un-preshrunk ingest is the +floor of its cost, not its cost. + +## (b) Window-depth widening + +Recognition on, only the structural horizon changes: + +| stub depth | unique windows | intra-resolve hits | capture | capture per window | `optimize` phase | +|---|---|---|---|---|---| +| 2 (shipped) | 64 | 247 | 480 µs | 1.55 µs | 1879 µs | +| 4 | 81 | 230 | 823 µs | 2.65 µs | 2265 µs | +| 6 | 101 | 210 | 1316 µs | 4.23 µs | 2846 µs | + +Windows captured stays 311 at every depth — widening does not plan more, it +plans the same windows less shareably: unique windows grow 64 → 81 → 101 and +per-window capture grows 1.7x / 2.7x, for +386 µs and +967 µs on the per-token +`optimize` budget (+21% / +51%) and +8% / +22% on the resolver's host total. + +Plan-store miss rate is unaffected: after the first decode resolve of a +process warms the device store (`store_misses` 6 / 14 / 29 at depth 2 / 4 / 6), +every one of the remaining 398 resolves reports `store_misses=0` with +`store_hits` equal to its miss count (64 / 81 / 101). The store is +device-scoped and in memory, so a new process re-warms it on its first token; +across two process runs the ledgers are byte-identical. + +Extraction output is unchanged at every depth: `dispatches=177 +bytes=165875972 work=9863870292` and the same dispatch histogram. The depth-2 +horizon is not costing fusion quality on this workload, so widening is pure +loss here — its only justification would be a generator that needs to see +further, which is exactly the recognizer-to-generator case. + +## Go / no-go contract + +A recognizer-to-generator campaign may proceed only if, measured on this +trace with these ledgers: + +1. the un-preshrunk ingest growth plus the widened capture cost fits inside + the decode `optimize_phases` budget the ported recognizers vacate — i.e. + per-token `optimize` does not exceed the shipped 1879 µs p50 (recognition + 320 µs + extraction 1415 µs, 45% of a 4.1 ms host resolve), and +2. the second and later decode resolves of a process report + `store_misses=0` — a port that makes windows layer-unique loses the + device store and pays generation on every token, and +3. the extracted plan is unchanged: `dispatches`, `bytes`, `work` and the + dispatch histogram must still match the goldens. + +As measured, none of the three holds today. The minimum viable port — drop the +sweep and widen the horizon to 4 so a generator can see a contraction's +factors through their views — costs 3019 µs measured un-preshrunk plus the +386 µs depth-4 capture delta measured separately, roughly 3.4 ms against a +1879 µs budget, and its generators lose every matmul and every attention +kernel. Until a spike shows otherwise, pre-ingest destructive placement is the +accepted end state, and this document is the evidence that it is a measurement +and not a preference. diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/analysis.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/analysis.rs new file mode 100644 index 000000000..5dbfd14d3 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/analysis.rs @@ -0,0 +1,56 @@ +//! Per-class analysis data and the driver context rules read through the +//! e-graph. +//! +//! Per-observation facts live in [`FusorAnalysis::facts`]. E-class membership +//! is maintained by the driver because hash-consing may attach several graph +//! observations to one class without invoking `Analysis::merge`. + +use egg::{Analysis, DidMerge, EGraph, Id}; +use rustc_hash::FxHashMap; + +use super::interner::PayloadTable; +use super::lang::{FusorLang, Prov}; +use crate::compute_graph::NodeIndex; + +/// Facts about one execution node (or cached-boundary leaf), snapshotted at +/// ingestion. Indexed by `Prov`. +#[derive(Debug, Clone)] +pub(super) struct NodeFacts { + pub(super) inner: NodeIndex, + /// The execution-graph node, `None` for cached-boundary leaves (which + /// are excluded from the execution graph). + pub(super) exec: Option, + /// `reference_count > 0` at ingestion: a user handle exists. Blocks + /// recognition cluster claims; does NOT block nary fusion (matching the + /// destructive optimizer's gates). + pub(super) externally_live: bool, + /// A resolve target: must materialize, may never be killed. + pub(super) is_target: bool, +} + +#[derive(Default)] +pub(super) struct FusorAnalysis { + /// Indexed by `Prov`. + pub(super) facts: Vec, + pub(super) payloads: PayloadTable, + /// Inner node -> the e-class assigned during ingestion. Ids may become + /// non-canonical after unions; callers canonicalize with `EGraph::find`. + pub(super) class_of_inner: FxHashMap, +} + +impl FusorAnalysis { + pub(super) fn facts_of(&self, prov: Prov) -> &NodeFacts { + &self.facts[prov.0 as usize] + } +} + +impl Analysis for FusorAnalysis { + /// No per-class data: every fact this optimizer needs is per observation. + type Data = (); + + fn make(_egraph: &mut EGraph, _enode: &FusorLang, _id: Id) -> Self::Data {} + + fn merge(&mut self, _a: &mut Self::Data, _b: Self::Data) -> DidMerge { + DidMerge(false, false) + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/apply.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/apply.rs new file mode 100644 index 000000000..ac5b9fb4f --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/apply.rs @@ -0,0 +1,106 @@ +//! Delta application: write extraction's non-identity selections back onto +//! the execution graph. +//! +//! [`Resolver::commit_recognized`] is the one commit surgery — the pre-ingest +//! recognizers land their clusters through it too — so the two invariants +//! every rewrite maintains hold by construction: execution-graph edges match +//! the new payload's `inputs`, and `add_physical_dependencies` records +//! persistent inner-graph edges (firing the flush-replay recording hook +//! identically). Killed producers fall out through `remove_node_if_dead`, +//! mirroring extraction's kill cascade. + +use egg::Language; + +use super::super::{ExecutionNodeIndex, ExecutionVariant, Resolver}; +use super::EGraphDriver; +use super::extract::Extraction; +use super::interner::{rebind_variant_dependencies, variant_dependencies}; +use crate::compute_graph::{ComputeGraphInner, NodeIndex}; + +impl Resolver { + pub(super) fn apply_egraph_deltas( + &mut self, + graph: &mut ComputeGraphInner, + driver: &EGraphDriver, + extraction: &Extraction, + ) -> usize { + let mut applied = 0; + for (prov, enode) in extraction.deltas() { + let facts = driver.egraph.analysis.facts_of(prov); + let exec_idx = facts + .exec + .expect("deltas only select alternatives for execution nodes"); + // An earlier delta's commit can kill this delta's target: when + // one recognized cluster's root is another cluster's + // intermediate (semantic identity lets both carry deltas), the + // outer commit rewires past the inner root and its kill cascade + // removes it. The removed node is unconsumed and not a target, + // so its rewrite is vacuous — both application orders converge + // to the same final graph. + if !self.execution_graph.contains_node(exec_idx) { + continue; + } + let payload = enode + .payload() + .expect("non-identity selections carry a payload"); + let mut variant = driver.egraph.analysis.payloads.get(payload).clone(); + // The payload may have been interned by a different + // structurally-identical instance; its concrete inputs belong to + // that instance. Rebind them to this e-node's actual children, + // resolved through the same class-representative mapping + // extraction used for its read/kill accounting, so the graph + // edges agree with what extraction kept alive. + let child_inners: Vec = enode + .children() + .iter() + .map(|&child| { + let child_prov = driver.prov_of_class(child, &extraction.needed); + driver.egraph.analysis.facts_of(child_prov).inner + }) + .collect(); + rebind_variant_dependencies(&mut variant, &child_inners); + let dependencies = variant_dependencies(&variant); + debug_assert_eq!( + dependencies, child_inners, + "rebinding must place every child in a dependency slot" + ); + self.commit_recognized(graph, exec_idx, &dependencies, variant); + applied += 1; + } + applied + } + + /// Replace a rewritten node's variant: drop every edge from the form it + /// replaces, wire the operation's dependencies directly, and let the + /// now-unconsumed producers fall out of the execution graph. + pub(in super::super) fn commit_recognized( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + dependencies: &[NodeIndex], + variant: ExecutionVariant, + ) { + self.execution_graph[node_idx].variant = variant; + + let previous: Vec = self + .execution_graph + .neighbors_directed(node_idx, petgraph::Direction::Incoming) + .collect(); + for &prev in &previous { + if let Some(edge) = self.execution_graph.find_edge(prev, node_idx) { + self.execution_graph.remove_edge(edge); + } + } + for &dependency in dependencies { + if let Some(exec) = self.get_input_node_in_exec_graph(dependency) + && self.execution_graph.find_edge(exec, node_idx).is_none() + { + self.execution_graph.add_edge(exec, node_idx, ()); + } + } + self.add_physical_dependencies(graph, node_idx, dependencies); + for prev in previous { + self.remove_node_if_dead(prev); + } + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/compose.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/compose.rs new file mode 100644 index 000000000..400ab573a --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/compose.rs @@ -0,0 +1,390 @@ +//! One expression-composition engine. +//! +//! Every fusion rewrite walks an n-ary expression the same way: a load's +//! index expressions are rewritten first, then the load itself is replaced, +//! folded into a producer's expression or left in place, while dimension +//! references go through a coordinate map. [`rewrite`] owns that walk and +//! each rule below is only its per-load decision, so an inlinability gate is +//! stated once and every rewrite composes coordinates identically. + +use rustc_hash::{FxHashMap, FxHashSet}; + +use super::super::ExecutionVariant; +use crate::Layout; +use crate::compute_graph::NodeIndex; +use crate::nary_wise::{ElementwiseOperation, ExtractedUnaryChain, NaryExpr, NaryFunction}; +use crate::view::ViewOperation; + +/// Rewrite `expr` bottom-up: `coords` maps every dimension reference, and +/// `load` replaces every input read, receiving the read's slot, its original +/// index expressions and those same expressions after rewriting. `None` from +/// either hook aborts the whole rewrite. +fn rewrite( + expr: &NaryExpr, + coords: &mut impl FnMut(usize) -> Option, + load: &mut impl FnMut(usize, &[NaryExpr], Vec) -> Option, +) -> Option { + Some(match expr { + NaryExpr::Op { children, function } => NaryExpr::Op { + children: children + .iter() + .map(|child| rewrite(child, &mut *coords, &mut *load)) + .collect::>>()?, + function: function.clone(), + }, + NaryExpr::IndexedInput { input_idx, indices } => { + let mapped = indices + .iter() + .map(|index| rewrite(index, &mut *coords, &mut *load)) + .collect::>>()?; + load(*input_idx, indices, mapped)? + } + NaryExpr::DimIndex(dim) => coords(*dim)?, + NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), + }) +} + +/// [`rewrite`] over loads alone, leaving the index space untouched. +pub(super) fn rewrite_loads( + expr: &NaryExpr, + load: &mut impl FnMut(usize, &[NaryExpr], Vec) -> Option, +) -> Option { + rewrite(expr, &mut |dim| Some(NaryExpr::DimIndex(dim)), load) +} + +/// [`rewrite_loads`] for rules that cannot decline, so neither can the walk. +pub(in super::super) fn map_loads( + expr: &NaryExpr, + load: &mut impl FnMut(usize, &[NaryExpr], Vec) -> NaryExpr, +) -> NaryExpr { + rewrite_loads(expr, &mut |input_idx, indices, mapped| { + Some(load(input_idx, indices, mapped)) + }) + .expect("a total load rule never aborts the walk") +} + +/// Evaluate `expr` (written in its own index space) at the coordinates +/// given by `indices`: every `DimIndex(d)` becomes `indices[d]`. `None` +/// when `expr` references a dimension `indices` does not provide. +pub(super) fn compose_expr_with_indices(expr: &NaryExpr, indices: &[NaryExpr]) -> Option { + rewrite( + expr, + &mut |dim| indices.get(dim).cloned(), + &mut |input_idx, _, mapped| { + Some(NaryExpr::IndexedInput { + input_idx, + indices: mapped, + }) + }, + ) +} + +/// Renumber every input slot through `mapping`. +pub(super) fn remap_inputs(expr: &NaryExpr, mapping: &[usize]) -> NaryExpr { + map_loads(expr, &mut |input_idx, _, indices| NaryExpr::IndexedInput { + input_idx: mapping[input_idx], + indices, + }) +} + +/// Add offset to all input indices in an expression. +pub(super) fn offset_input_indices(expr: &NaryExpr, offset: usize) -> NaryExpr { + map_loads(expr, &mut |input_idx, _, indices| NaryExpr::IndexedInput { + input_idx: input_idx + offset, + indices, + }) +} + +/// Substitute IndexedInput(target_idx) with element-wise access with the +/// replacement expression. Returns (new_expression, success) where success is +/// true if all references to target_idx were successfully substituted. If +/// false, the input should NOT be removed from the graph. +pub(super) fn substitute_input_in_expr( + expr: &NaryExpr, + target_idx: usize, + replacement: &NaryExpr, +) -> (NaryExpr, bool) { + let elementwise_replacement = match replacement { + NaryExpr::IndexedInput { input_idx, indices } + if NaryExpr::is_elementwise_indices(indices) => + { + Some(*input_idx) + } + _ => None, + }; + let mut success = true; + let expr = map_loads(expr, &mut |input_idx, indices, mapped| { + if input_idx != target_idx { + return NaryExpr::IndexedInput { + input_idx, + indices: mapped, + }; + } + if NaryExpr::is_elementwise_indices(indices) { + // Element-wise can be fully replaced with any expression + return replacement.clone(); + } + // Custom indexing can only substitute if replacement is also + // element-wise; anything else cannot fuse into the indexed load. + let Some(input_idx) = elementwise_replacement else { + success = false; + return NaryExpr::IndexedInput { + input_idx, + indices: mapped, + }; + }; + NaryExpr::IndexedInput { + input_idx, + indices: mapped, + } + }); + (expr, success) +} + +/// Substitute every read of input `target_idx` — elementwise *or* +/// custom-indexed — with `replacement` evaluated at the read's index +/// expressions. Where [`substitute_input_in_expr`] declines custom-indexed +/// reads unless the replacement is a bare input, this composes the +/// replacement expression with the index list instead: `input_t[i0, i1]` +/// becomes `replacement` with `DimIndex(d)` rewritten to `i_d`. Returns +/// `None` when the composition is impossible (an index list shorter than the +/// replacement's rank). +pub(super) fn substitute_input_composed( + expr: &NaryExpr, + target_idx: usize, + replacement: &NaryExpr, +) -> Option { + rewrite_loads(expr, &mut |input_idx, _, mapped| { + if input_idx == target_idx { + compose_expr_with_indices(replacement, &mapped) + } else { + Some(NaryExpr::IndexedInput { + input_idx, + indices: mapped, + }) + } + }) +} + +/// Replace every element-wise read of an input carrying a replacement with +/// that expression. `None` when such an input is read through custom +/// indexing, which the replacement's index space cannot serve. +pub(super) fn replace_inputs_in_expr( + expr: &NaryExpr, + replacements: &[Option], +) -> Option { + rewrite_loads(expr, &mut |input_idx, indices, mapped| match replacements + .get(input_idx) + .and_then(|r| r.as_ref()) + { + Some(replacement) => NaryExpr::is_elementwise_indices(indices).then(|| replacement.clone()), + None => Some(NaryExpr::IndexedInput { + input_idx, + indices: mapped, + }), + }) +} + +/// Rewrite every access to input `target_idx` through `view`'s +/// coordinate map: the original index expressions (the view's output +/// coordinates) walk down the stage stack to base coordinates, with +/// fill selects and in-bounds clamps where stages are partially defined +/// (both select branches evaluate). +pub(super) fn rewrite_view_input( + expr: &NaryExpr, + target_idx: usize, + view: &ViewOperation, +) -> Option { + rewrite_loads(expr, &mut |input_idx, _, mapped| { + if input_idx == target_idx { + Some(view.value_expression(input_idx, &mapped)?.0) + } else { + Some(NaryExpr::IndexedInput { + input_idx, + indices: mapped, + }) + } + }) +} + +/// Remove unused inputs and deduplicate, returning new inputs and remapped expression. +pub(in super::super) fn deduplicate_inputs( + inputs: Vec, + expr: NaryExpr, +) -> (Vec, NaryExpr) { + // Collect which input indices are actually used + let mut seen_indices = FxHashSet::default(); + let mut used_indices = Vec::new(); + collect_used_inputs(&expr, &mut seen_indices, &mut used_indices); + + // Build the input-index remap, collecting only used inputs. + let mut new_inputs = Vec::new(); + let mut old_to_new = FxHashMap::default(); + + for old_idx in used_indices { + let node = inputs[old_idx]; + // Check if this node already exists in new_inputs (deduplication) + let new_idx = if let Some(existing) = new_inputs.iter().position(|&n| n == node) { + existing + } else { + let idx = new_inputs.len(); + new_inputs.push(node); + idx + }; + old_to_new.insert(old_idx, new_idx); + } + + let new_expr = map_loads(&expr, &mut |input_idx, _, indices| NaryExpr::IndexedInput { + input_idx: old_to_new[&input_idx], + indices, + }); + (new_inputs, new_expr) +} + +fn collect_used_inputs(expr: &NaryExpr, seen: &mut FxHashSet, used: &mut Vec) { + match expr { + NaryExpr::Op { children, .. } => { + for child in children { + collect_used_inputs(child, seen, used); + } + } + NaryExpr::IndexedInput { input_idx, indices } => { + if seen.insert(*input_idx) { + used.push(*input_idx); + } + for c in indices { + collect_used_inputs(c, seen, used); + } + } + NaryExpr::DimIndex(_) => {} + NaryExpr::Scalar(_) => {} + } +} + +/// The worst re-read factor across this slot's loads: the product of +/// index-space dims a load's coordinates never reference — each such dim +/// re-reads the same element once per step. +pub(super) fn input_reread_factor(expr: &NaryExpr, shape: &[usize], slot: usize) -> usize { + fn collect_dims(expr: &NaryExpr, referenced: &mut [bool]) { + match expr { + NaryExpr::Op { children, .. } => { + for child in children { + collect_dims(child, referenced); + } + } + NaryExpr::IndexedInput { indices, .. } => { + for index in indices { + collect_dims(index, referenced); + } + } + NaryExpr::DimIndex(dim) => referenced[*dim] = true, + NaryExpr::Scalar(_) => {} + } + } + fn visit_loads(expr: &NaryExpr, shape: &[usize], slot: usize, worst: &mut usize) { + match expr { + NaryExpr::Op { children, .. } => { + for child in children { + visit_loads(child, shape, slot, worst); + } + } + NaryExpr::IndexedInput { input_idx, indices } => { + for index in indices { + visit_loads(index, shape, slot, worst); + } + if *input_idx == slot { + let mut referenced = vec![false; shape.len()]; + for index in indices { + collect_dims(index, &mut referenced); + } + let factor: usize = shape + .iter() + .zip(&referenced) + .filter(|(_, referenced)| !**referenced) + .map(|(size, _)| *size) + .product(); + *worst = (*worst).max(factor); + } + } + NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => {} + } + } + let mut worst = 1; + visit_loads(expr, shape, slot, &mut worst); + worst +} + +/// Try to extract a unary function chain from a node variant. +/// Only Nary ops with a single input and element-wise access can be converted. +pub(super) fn try_get_unary_chain(variant: &ExecutionVariant) -> Option { + match variant { + ExecutionVariant::Elementwise(nary) => nary.try_extract_unary_chain(), + _ => None, + } +} + +/// Extract a (possibly empty) unary function chain over exactly one +/// read of input 0 with arbitrary index expressions, innermost function +/// first. The index expressions must not read any input themselves. +pub(super) fn extract_unary_chain_indexed( + nary: &ElementwiseOperation, +) -> Option<(Vec, Vec)> { + fn contains_input(expr: &NaryExpr) -> bool { + match expr { + NaryExpr::Op { children, .. } => children.iter().any(contains_input), + NaryExpr::IndexedInput { .. } => true, + NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => false, + } + } + let mut functions = Vec::new(); + let mut expr = &nary.expression; + loop { + match expr { + NaryExpr::Op { children, function } + if children.len() == 1 && function.input_types.len() == 1 => + { + functions.push(function.clone()); + expr = &children[0]; + } + NaryExpr::IndexedInput { + input_idx: 0, + indices, + } if !indices.iter().any(contains_input) => { + functions.reverse(); + return Some((functions, indices.clone())); + } + _ => return None, + } + } +} + +/// Walk through view nodes from `inner` down to the first non-view +/// node, composing each view's collapsed stage stack. Public tensor ops +/// collapse into single view nodes at construction, but composed +/// clusters (attention's attached GQA/transpose views) still layer view +/// nodes deliberately. `view_of` selects the chain — the collapsed layout and +/// input of the node's current view form, `None` where the walk must stop. +/// Returns the base node and the composed layout over the base's logical +/// value space; the layout is `None` when `inner` is not a view (identity). +/// Views that don't collapse or compose (or carry a fill region) act as chain +/// breaks: the walk stops without seeing through them. +pub(in super::super) fn walk_view_chain( + mut inner: NodeIndex, + mut view_of: impl FnMut(NodeIndex) -> Option<(Layout, NodeIndex)>, +) -> (NodeIndex, Option) { + let mut composed: Option = None; + loop { + let Some((collapsed, input)) = view_of(inner) else { + return (inner, composed); + }; + let next = match &composed { + None => collapsed, + Some(outer) => match crate::view::compose_layouts(outer, &collapsed) { + Some(layout) => layout, + None => return (inner, composed), + }, + }; + composed = Some(next); + inner = input; + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/cost.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/cost.rs new file mode 100644 index 000000000..7978d04c7 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/cost.rs @@ -0,0 +1,462 @@ +//! GPU-oriented extraction cost. +//! +//! Dispatches dominate small and medium fusion decisions, materialized bytes +//! are the next-order cost, and approximate arithmetic work breaks remaining +//! ties. The model is deliberately monotone and deterministic; legality is +//! still owned by the rewrite generators. + +use std::ops::{Add, AddAssign}; + +use super::super::ExecutionVariant; +use super::EGraphDriver; +use super::extract::{ExtractState, Selection}; +use super::interner::variant_dependencies; +use super::lang::Prov; +use crate::DataTypeEnum; +use crate::nary_wise::NaryExpr; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord)] +pub(super) struct GpuCost { + pub(super) dispatches: u64, + pub(super) materialized_bytes: u128, + pub(super) work: u128, +} + +impl Add for GpuCost { + type Output = Self; + + fn add(self, rhs: Self) -> Self::Output { + Self { + dispatches: self.dispatches.saturating_add(rhs.dispatches), + materialized_bytes: self + .materialized_bytes + .saturating_add(rhs.materialized_bytes), + work: self.work.saturating_add(rhs.work), + } + } +} + +impl AddAssign for GpuCost { + fn add_assign(&mut self, rhs: Self) { + *self = *self + rhs; + } +} + +/// Roofline constants for the scalarized model, measured on this machine +/// through `kernel_bench`'s anchors (`roof_bw`, `roof_flops`) and the +/// inter-kernel gap of a training step. They convert the three terms to one +/// clock so a dispatch, a byte and a flop can actually outweigh each other. +/// +/// Under the lexicographic tuple `dispatches` is effectively infinite: the +/// model will pay unbounded traffic and unbounded arithmetic to remove one +/// launch. On the measured step, dispatches are 0.2% of modeled time. +const DISPATCH_NS: u128 = 1_000; +/// ~340 GB/s achievable (402 MB stream add in ~1.17 ms). +const BYTES_PER_NS: u128 = 340; +/// ~5.5 TFLOP/s achievable (68.7 GFLOP merged matmul in ~12.6 ms). +const FLOPS_PER_NS: u128 = 5_468; + +/// The three terms on one clock, scaled by both rates so the comparison +/// stays exact in integers rather than truncating two divisions. +fn scaled_nanos(dispatches: i128, materialized_bytes: i128, work: i128) -> i128 { + let dispatch_scale = (DISPATCH_NS * BYTES_PER_NS * FLOPS_PER_NS) as i128; + dispatches.saturating_mul(dispatch_scale) + + materialized_bytes.saturating_mul(FLOPS_PER_NS as i128) + + work.saturating_mul(BYTES_PER_NS as i128) +} + +#[derive(Debug, Clone, Copy, Default)] +pub(super) struct CostDelta { + dispatches: i128, + materialized_bytes: i128, + work: i128, + /// Compare on one clock instead of the lexicographic tuple + /// (`FUSOR_SPIKE_SCALAR_COST`). Uniform across a run. + scalar: bool, +} + +impl CostDelta { + /// Fusion rules are one-way structural simplifications. An equal-cost + /// canonicalization (notably unit-reduce -> elementwise) is useful because + /// it exposes the next fusion and cannot cycle back to its old form. + pub(super) fn non_worse(self) -> bool { + self <= Self::default() + } + + /// The ordering key: one clock when scalarized, otherwise the + /// lexicographic tuple this model shipped with. + fn key(self) -> (i128, i128, i128) { + if self.scalar { + ( + scaled_nanos(self.dispatches, self.materialized_bytes, self.work), + 0, + 0, + ) + } else { + (self.dispatches, self.materialized_bytes, self.work) + } + } +} + +// Ordering is by `key`, so equality must be too: `min_by_key` and `non_worse` +// both rely on `Ord` agreeing with `Eq`. +impl PartialEq for CostDelta { + fn eq(&self, other: &Self) -> bool { + self.key() == other.key() + } +} + +impl Eq for CostDelta {} + +impl PartialOrd for CostDelta { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for CostDelta { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.key().cmp(&other.key()) + } +} + +impl EGraphDriver { + /// The write footprint of one variant. + fn output_bytes(&self, variant: &ExecutionVariant) -> u128 { + match variant { + // Leaves cost no dispatch, but reading one still moves its bytes. + ExecutionVariant::Tensor(data) => bytes(data.info.layout.shape(), data.info.datatype), + // Attention writes a Q-shaped result; resolve it through Q. + ExecutionVariant::Attention(operation) => self + .prov_of + .get(&operation.q) + .map(|&prov| self.dependency_bytes(prov)) + .unwrap_or_default(), + other => variant_cost(other).materialized_bytes, + } + } + + /// Bytes a consumer moves to read `prov`'s value: that producer's own + /// output footprint. Shape and datatype are invariant across a node's + /// selections, so the identity form answers for every form. + fn dependency_bytes(&self, prov: Prov) -> u128 { + self.identity_variant(prov) + .map(|variant| self.output_bytes(variant)) + .unwrap_or_default() + } + + /// [`variant_cost`] plus the bytes the variant reads from its inputs. + /// + /// The stock model counts output writes only, which makes producer + /// duplication invisible: inlining a producer deletes one write and adds + /// a read of each of that producer's own inputs, and only the write was + /// ever scored. Under `FUSOR_SPIKE_READ_TRAFFIC` the byte term becomes + /// total traffic, so a fusion that trades one write for two reads is + /// priced as the loss it is. + fn traffic_cost(&self, state: &ExtractState, variant: &ExecutionVariant) -> GpuCost { + let mut cost = variant_cost(variant); + if !state.read_traffic { + return cost; + } + cost.materialized_bytes = self.output_bytes(variant); + let reads = variant_dependencies(variant) + .into_iter() + .filter_map(|inner| self.prov_of.get(&inner).copied()) + .fold(0u128, |total, prov| { + total.saturating_add(self.dependency_bytes(prov)) + }); + cost.materialized_bytes = cost.materialized_bytes.saturating_add(reads); + cost + } + + pub(super) fn switch_cost_delta( + &self, + state: &ExtractState, + prov: Prov, + candidate: &ExecutionVariant, + kills: &[u32], + ) -> CostDelta { + let current = self + .selected_variant(state, prov) + .map(|variant| self.traffic_cost(state, variant)) + .unwrap_or_default(); + let removed = kills.iter().fold(GpuCost::default(), |mut cost, &dead| { + if let Some(variant) = self.selected_variant(state, Prov(dead)) { + cost += self.traffic_cost(state, variant); + } + cost + }); + let candidate = self.traffic_cost(state, candidate); + CostDelta { + dispatches: i128::from(candidate.dispatches) + - i128::from(current.dispatches) + - i128::from(removed.dispatches), + materialized_bytes: as_i128(candidate.materialized_bytes) + - as_i128(current.materialized_bytes) + - as_i128(removed.materialized_bytes), + work: as_i128(candidate.work) - as_i128(current.work) - as_i128(removed.work), + scalar: state.scalar_cost, + } + } + + pub(super) fn extraction_cost(&self, state: &ExtractState) -> GpuCost { + let mut cost = GpuCost::default(); + for prov in 0..state.sel.len() as u32 { + if state.needed[prov as usize] + && let Some(variant) = self.selected_variant(state, Prov(prov)) + { + cost += self.traffic_cost(state, variant); + } + } + cost + } + + pub(super) fn selection_cost(&self, state: &ExtractState, prov: Prov) -> GpuCost { + self.selected_variant(state, prov) + .map(|variant| self.traffic_cost(state, variant)) + .unwrap_or_default() + } + + fn selected_variant<'a>( + &'a self, + state: &'a ExtractState, + prov: Prov, + ) -> Option<&'a ExecutionVariant> { + match &state.sel[prov.0 as usize] { + Selection::Identity => self.identity_variant(prov), + Selection::Alt(enode) => enode + .payload() + .map(|payload| self.egraph.analysis.payloads.get(payload)), + } + } +} + +fn as_i128(value: u128) -> i128 { + value.min(i128::MAX as u128) as i128 +} + +fn elements(shape: &[usize]) -> u128 { + shape + .iter() + .fold(1u128, |size, &dim| size.saturating_mul(dim as u128)) +} + +fn bytes(shape: &[usize], datatype: DataTypeEnum) -> u128 { + elements(shape).saturating_mul(datatype.element_size() as u128) +} + +fn expr_work(expression: &NaryExpr) -> u128 { + match expression { + NaryExpr::Op { children, .. } => 1 + children.iter().map(expr_work).sum::(), + NaryExpr::IndexedInput { indices, .. } => 1 + indices.iter().map(expr_work).sum::(), + NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => 1, + } +} + +pub(super) fn variant_cost(variant: &ExecutionVariant) -> GpuCost { + match variant { + ExecutionVariant::Tensor(_) => GpuCost::default(), + ExecutionVariant::QMatrix(operation) => { + let output = bytes(operation.matrix.shape(), operation.datatype); + GpuCost { + dispatches: 1, + materialized_bytes: output, + work: elements(operation.matrix.shape()) + .saturating_mul(1 + operation.post_dequantize.functions.len() as u128), + } + } + ExecutionVariant::Elementwise(operation) => { + let output_elements = elements(&operation.shape); + GpuCost { + dispatches: 1, + materialized_bytes: bytes(&operation.shape, operation.output_datatype), + work: output_elements.saturating_mul(expr_work(&operation.expression)), + } + } + // A fold writes every output and evaluates its step at every + // coordinate of the folded index space. Blocking multiplies the + // dispatch by the number of blocks, which is what makes a split + // visible to the cost model at all. + ExecutionVariant::Fold(operation) => { + let out_shape = operation.out_shape(); + let materialized_bytes = operation + .outputs + .iter() + .map(|output| bytes(&out_shape, output.datatype)) + .fold(0u128, |total, output| total.saturating_add(output)); + GpuCost { + dispatches: 1, + materialized_bytes, + work: elements(&operation.shape).saturating_mul(operation.step_work()), + } + } + ExecutionVariant::Reduce(operation) => { + let output_shape = operation.out_shape(); + GpuCost { + dispatches: 1, + materialized_bytes: bytes(&output_shape, operation.out_datatype()), + work: elements(&operation.shape).saturating_mul( + expr_work(&operation.expression) + + 1 + + operation.post_element_wise.functions.len() as u128, + ), + } + } + ExecutionVariant::View(operation) => { + // Some views become zero-dispatch aliases and some materialize a + // gather. Treat one dispatch as the safe upper bound: folding a + // view can only make this estimate more accurate, never add work. + GpuCost { + dispatches: 1, + materialized_bytes: bytes(operation.shape(), operation.datatype), + work: elements(operation.shape()), + } + } + ExecutionVariant::Assign(_) => GpuCost { + dispatches: 1, + materialized_bytes: 0, + work: 1, + }, + ExecutionVariant::Region(operation) => { + let output_count = operation.output_count() as u128; + GpuCost { + dispatches: 1, + materialized_bytes: elements(&operation.shape) + .saturating_mul(4) + .saturating_mul(output_count), + work: elements(&operation.shape).saturating_mul(operation.statements.len() as u128), + } + } + ExecutionVariant::MatMul(operation) => { + let batch = elements(operation.a.batch_shape()); + let m = operation.a.rows() as u128; + let n = operation.b.cols() as u128; + let k = operation.a.cols() as u128; + let epilogue = operation.pre_element_wise[0].functions.len() + + operation.pre_element_wise[1].functions.len() + + operation.post_element_wise.functions.len(); + GpuCost { + dispatches: 1, + materialized_bytes: bytes(&operation.out_shape, operation.datatype), + work: batch + .saturating_mul(m) + .saturating_mul(n) + .saturating_mul(k.saturating_mul(2).saturating_add(epilogue as u128)), + } + } + ExecutionVariant::QMatMul(operation) => { + let k = operation.matrix.shape()[1] as u128; + let n = operation.matrix.shape()[0] as u128; + let rows = elements(&operation.in_shape[..operation.in_shape.len() - 1]); + let epilogue = operation + .pre_element_wise_expr + .as_ref() + .map(|epilogue| expr_work(&epilogue.expression)) + .unwrap_or_default() + + operation + .post_element_wise_expr + .as_ref() + .map(|epilogue| expr_work(&epilogue.expression)) + .unwrap_or_default(); + GpuCost { + dispatches: 1, + materialized_bytes: bytes(&operation.out_shape, operation.input_datatype), + work: rows + .saturating_mul(n) + .saturating_mul(k.saturating_mul(2).saturating_add(epilogue)), + } + } + ExecutionVariant::QEmbedding(operation) => GpuCost { + dispatches: 1, + materialized_bytes: bytes(&operation.out_shape, operation.datatype), + work: elements(&operation.out_shape), + }, + // A row program's `shape` is its operating footprint; for a reducing + // program the write is smaller, so this is an upper bound. `work` + // stays a placeholder — the per-step cost is not derivable here. + ExecutionVariant::RowProgram(operation) => GpuCost { + dispatches: 1, + materialized_bytes: bytes(&operation.shape, operation.output_datatype), + work: elements(&operation.shape).saturating_mul(operation.steps.len() as u128), + }, + // Attention cannot size its own output: the shape lives on its Q + // operand, so `traffic_cost` fills the write in from the dependency + // table. `work` remains a placeholder, so the model under-counts + // attention compute. + ExecutionVariant::Attention(_) => GpuCost { + dispatches: 1, + materialized_bytes: 0, + work: 1, + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn lexicographic(dispatches: i128, materialized_bytes: i128, work: i128) -> CostDelta { + CostDelta { + dispatches, + materialized_bytes, + work, + scalar: false, + } + } + + fn scalar(dispatches: i128, materialized_bytes: i128, work: i128) -> CostDelta { + CostDelta { + dispatches, + materialized_bytes, + work, + scalar: true, + } + } + + #[test] + fn dispatch_reduction_dominates_secondary_costs() { + assert!(lexicographic(-1, 1_000_000, 1_000_000) < CostDelta::default()); + } + + #[test] + fn bytes_then_work_break_dispatch_ties() { + assert!(lexicographic(0, -1, 1_000_000) < CostDelta::default()); + assert!(lexicographic(0, 0, -1) < CostDelta::default()); + assert!(CostDelta::default().non_worse()); + assert!(!lexicographic(0, 0, 1).non_worse()); + } + + #[test] + fn scalar_cost_lets_traffic_outweigh_a_dispatch() { + // One saved launch is worth 1 us. 1 MB of extra traffic costs about + // 3 us at the measured roof, so the trade is a loss — a verdict the + // lexicographic tuple cannot reach. + assert!(lexicographic(-1, 1_000_000, 0).non_worse()); + assert!(!scalar(-1, 1_000_000, 0).non_worse()); + } + + #[test] + fn scalar_cost_still_takes_a_cheap_dispatch_saving() { + // Same saved launch, but only 100 KB of added traffic: still a win. + assert!(scalar(-1, 100_000, 0).non_worse()); + } + + #[test] + fn scalar_cost_prices_recompute_against_a_saved_write() { + // Dropping a 1 MB write in exchange for recomputing 1M flops wins: + // the write is ~2.9 us of bandwidth, the arithmetic ~0.2 us. + assert!(scalar(0, -1_000_000, 1_000_000).non_worse()); + // Ten times the arithmetic for the same saved write does not. + assert!(!scalar(0, -1_000_000, 20_000_000).non_worse()); + } + + #[test] + fn ordering_agrees_with_equality_in_both_modes() { + // `min_by_key` and `non_worse` both require Ord to agree with Eq. + let a = scalar(1, -(BYTES_PER_NS as i128) * 1_000, 0); + let b = scalar(0, 0, 0); + assert_eq!(a == b, a.cmp(&b) == std::cmp::Ordering::Equal); + let c = lexicographic(0, 0, 5); + assert_eq!(c == c, c.cmp(&c) == std::cmp::Ordering::Equal); + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/extract.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/extract.rs new file mode 100644 index 000000000..e7f434b58 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/extract.rs @@ -0,0 +1,599 @@ +//! Cost-guided extraction: choose one legal form per execution node, +//! minimizing estimated GPU runtime. +//! +//! Selection state remains per observation so liveness constraints are +//! explicit, while equivalent observations may reference the same e-class. +//! Every observation starts at its identity selection. Fusion generators +//! drive every switch: legal alternatives are planned once per +//! allocation-independent local window and reused by repeated layers. +//! Successful generations are recorded into the value e-graph and committed +//! as switches. +//! +//! Consumer counts are multisets over the current selection, initialized +//! from the identity selections (one entry per read occurrence). A switch's +//! kill set is the transitive closure of producers whose counts drop to +//! zero; targets are never killed (they materialize regardless, even when a +//! consumer also inlined their expression). +//! +//! The cost tuple is lexicographic: dispatch count, materialized bytes, then +//! estimated arithmetic work. Determinism comes from provenance-order +//! worklists, fixed generator order and generation-order tie breaks; no +//! decision consults hash-map iteration order. + +use egg::Language; + +use super::EGraphDriver; +use super::interner::variant_dependencies; +use super::lang::{FusorLang, Prov}; +use super::rules_fuse::{FusionCtx, FusionView}; +use super::structural_memo::{FusionPlanMemo, PlanLookup}; + +/// What extraction chose for one execution node. +#[derive(Debug, Clone)] +pub(super) enum Selection { + Identity, + Alt(FusorLang), +} + +pub(super) struct Extraction { + /// Indexed by `Prov`. + pub(super) sel: Vec, + /// Indexed by `Prov`; false = killed (no longer materialized). + pub(super) needed: Vec, +} + +impl Extraction { + /// The non-identity selections of needed nodes, in provenance order. + pub(super) fn deltas(&self) -> impl Iterator { + self.sel + .iter() + .enumerate() + .filter(|(prov, _)| self.needed[*prov]) + .filter_map(|(prov, selection)| match selection { + Selection::Identity => None, + Selection::Alt(enode) => Some((Prov(prov as u32), enode)), + }) + } +} + +/// Live extraction state, indexed by `Prov`. +pub(super) struct ExtractState { + pub(super) sel: Vec, + pub(super) needed: Vec, + /// Multiset consumer count under the current selection. + pub(super) reads: Vec, + /// Reverse index: consumers (by prov) of each node under the current + /// selection, one entry per read occurrence. + pub(super) consumers: Vec>, + /// Spike: price producer duplication instead of forbidding it + /// (`FUSOR_SPIKE_NO_DUP_GATE`). Default false keeps the gate. + pub(super) no_dup_gate: bool, + /// Spike: count input reads in the byte term (`FUSOR_SPIKE_READ_TRAFFIC`). + pub(super) read_traffic: bool, + /// Spike: compare costs on one clock instead of the lexicographic tuple + /// (`FUSOR_SPIKE_SCALAR_COST`). Implies `read_traffic` — scalarizing a + /// write-only byte term would weigh a number that is not bandwidth. + pub(super) scalar_cost: bool, + /// Spike ledger: distinct producers the gate refused to inline, and how + /// many candidates it rejected. Recorded whether or not the gate is + /// enforcing, so the surface is measurable from an unmodified run. + /// Off by default; the bookkeeping allocates and the gate is hot. + pub(super) dup_ledger: bool, + pub(super) dup_blocked: std::cell::RefCell>, + pub(super) dup_firings: std::cell::Cell, +} + +impl ExtractState { + /// Every ingested observation starts needed and at its identity form. + pub(super) fn new(driver: &EGraphDriver) -> Self { + let count = driver.egraph.analysis.facts.len(); + let mut state = ExtractState { + sel: vec![Selection::Identity; count], + needed: vec![true; count], + reads: vec![0; count], + consumers: vec![Vec::new(); count], + no_dup_gate: false, + read_traffic: false, + scalar_cost: false, + dup_ledger: false, + dup_blocked: Default::default(), + dup_firings: Default::default(), + }; + for prov in 0..count as u32 { + for child in state.selected_child_provs(driver, Prov(prov)) { + state.reads[child as usize] += 1; + state.consumers[child as usize].push(prov); + } + } + state + } + + /// Child provenances of the current selection, one entry per read + /// occurrence. + pub(super) fn selected_child_provs(&self, driver: &EGraphDriver, prov: Prov) -> Vec { + self.selected_enode(driver, prov) + .children() + .iter() + .map(|&child| driver.prov_of_class(child, &self.needed).0) + .collect() + } + + /// Child provenances of a node the cascade is killing. A child class + /// whose observations are all dead was killed earlier in the same + /// cascade — a chain-folding generator kills a producer and that + /// producer's own inputs — and its counts no longer reach any live + /// selection, so it drops out instead of resolving. + fn killed_child_provs(&self, driver: &EGraphDriver, prov: Prov) -> Vec { + self.selected_enode(driver, prov) + .children() + .iter() + .filter_map(|&child| driver.live_prov_of_class(child, &self.needed)) + .map(|prov| prov.0) + .collect() + } + + pub(super) fn selected_enode<'d>( + &'d self, + driver: &'d EGraphDriver, + prov: Prov, + ) -> &'d FusorLang { + match &self.sel[prov.0 as usize] { + Selection::Identity => driver.identity_enode(prov), + Selection::Alt(enode) => enode, + } + } + + /// The transitive kill set of switching `prov` to `candidate`. Targets + /// and leaves are never killed (they stay needed; a target producer that + /// stops being read still materializes, duplicating compute exactly as + /// the destructive optimizer does). + fn kills(&self, driver: &EGraphDriver, prov: Prov, candidate: &FusorLang) -> Vec { + self.kills_from_child_provs( + driver, + prov, + candidate + .children() + .iter() + .map(|&child| driver.prov_of_class(child, &self.needed).0), + ) + } + + fn kills_for_variant( + &self, + driver: &EGraphDriver, + prov: Prov, + candidate: &super::super::ExecutionVariant, + ) -> Vec { + self.kills_from_child_provs( + driver, + prov, + variant_dependencies(candidate) + .into_iter() + .filter_map(|inner| driver.prov_of.get(&inner).map(|prov| prov.0)), + ) + } + + /// A rewrite that inlines a materializing producer but cannot kill it + /// (another consumer or a target still needs it) duplicates GPU work. + /// Treat that as a hard extraction constraint rather than hoping an + /// approximate arithmetic tie-break notices it. + fn variant_duplicates_required_producer( + &self, + driver: &EGraphDriver, + prov: Prov, + candidate: &super::super::ExecutionVariant, + kills: &[u32], + ) -> bool { + self.duplicates_required_producer( + driver, + prov, + variant_dependencies(candidate) + .into_iter() + .filter_map(|inner| driver.prov_of.get(&inner).map(|prov| prov.0)), + kills, + ) + } + + fn duplicates_required_producer( + &self, + driver: &EGraphDriver, + prov: Prov, + candidate_children: impl IntoIterator, + kills: &[u32], + ) -> bool { + let candidate_children: rustc_hash::FxHashSet = + candidate_children.into_iter().collect(); + let mut duplicated = self + .selected_child_provs(driver, prov) + .into_iter() + .filter(|child| { + !candidate_children.contains(child) + && !kills.contains(child) + && self.needed[*child as usize] + && driver.selection_cost(self, Prov(*child)).dispatches > 0 + }); + if !self.dup_ledger { + // Default path: the first offender decides, no bookkeeping and no + // allocation, exactly as before the spike was added. + return duplicated.next().is_some() && !self.no_dup_gate; + } + let duplicated: Vec = duplicated.collect(); + if duplicated.is_empty() { + return false; + } + // The ledger records whether or not the gate enforces, so the surface + // is measurable from an otherwise-stock run. + self.dup_firings.set(self.dup_firings.get() + 1); + self.dup_blocked.borrow_mut().extend(duplicated); + !self.no_dup_gate + } + + fn kills_from_child_provs( + &self, + driver: &EGraphDriver, + prov: Prov, + candidate_children: impl IntoIterator, + ) -> Vec { + let mut overlay: rustc_hash::FxHashMap = Default::default(); + for child in self.selected_child_provs(driver, prov) { + *overlay.entry(child).or_default() -= 1; + } + for child in candidate_children { + *overlay.entry(child).or_default() += 1; + } + let mut kills = Vec::new(); + let mut frontier: Vec = overlay + .iter() + .filter(|&(&p, &delta)| { + delta < 0 && (self.reads[p as usize] as i64 + delta) <= 0 && self.needed[p as usize] + }) + .map(|(&p, _)| p) + .collect(); + frontier.sort_unstable(); + while let Some(dead) = frontier.pop() { + if kills.contains(&dead) { + continue; + } + let facts = &driver.egraph.analysis.facts[dead as usize]; + if facts.is_target || facts.exec.is_none() { + // Unkillable: targets materialize regardless; leaves have no + // dispatch to save. + continue; + } + kills.push(dead); + for child in self.selected_child_provs(driver, Prov(dead)) { + let entry = overlay.entry(child).or_default(); + *entry -= 1; + if (self.reads[child as usize] as i64 + *entry) <= 0 + && self.needed[child as usize] + && !kills.contains(&child) + { + frontier.push(child); + } + } + } + kills.sort_unstable(); + kills + } + + /// Commit the switch. Returns provs whose situation changed: the + /// switched node, everything whose read count changed, and the killed + /// nodes' surviving children. + fn commit( + &mut self, + driver: &EGraphDriver, + prov: Prov, + candidate: FusorLang, + kills: &[u32], + ) -> Vec { + let mut touched = vec![prov.0]; + for child in self.selected_child_provs(driver, prov) { + self.reads[child as usize] -= 1; + remove_one(&mut self.consumers[child as usize], prov.0); + touched.push(child); + } + for &child in candidate.children() { + let child_prov = driver.prov_of_class(child, &self.needed).0; + self.reads[child_prov as usize] += 1; + self.consumers[child_prov as usize].push(prov.0); + touched.push(child_prov); + } + self.sel[prov.0 as usize] = Selection::Alt(candidate); + for &dead in kills { + if !self.needed[dead as usize] { + continue; + } + self.needed[dead as usize] = false; + for child in self.killed_child_provs(driver, Prov(dead)) { + self.reads[child as usize] = self.reads[child as usize].saturating_sub(1); + remove_one(&mut self.consumers[child as usize], dead); + touched.push(child); + } + } + touched + } +} + +fn remove_one(consumers: &mut Vec, value: u32) { + if let Some(position) = consumers.iter().position(|&c| c == value) { + consumers.swap_remove(position); + } +} + +impl EGraphDriver { + pub(super) fn prov_of_class(&self, id: egg::Id, needed: &[bool]) -> Prov { + self.live_prov_of_class(id, needed) + .expect("a selected e-class child must retain a needed provenance") + } + + /// The live provenance of `id`, or `None` once every observation in the + /// class has been killed. + pub(super) fn live_prov_of_class(&self, id: egg::Id, needed: &[bool]) -> Option { + self.provs_of_class[&self.egraph.find(id)] + .iter() + .copied() + .find(|prov| needed[prov.0 as usize]) + } + + /// The identity e-node of a provenance: the one ingested for the + /// execution node itself (unique per class by construction). + pub(super) fn identity_enode(&self, prov: Prov) -> &FusorLang { + &self.identity_enodes[prov.0 as usize] + } + + /// The cheapest legal generator candidate for `prov` under live counts. + fn best_fusion_candidate( + &self, + state: &ExtractState, + view: &FusionView<'_>, + prov: Prov, + ) -> Option { + view.generate_candidates(prov) + .into_iter() + .enumerate() + .filter_map(|(order, variant)| { + let kills = state.kills_for_variant(self, prov, &variant); + if state.variant_duplicates_required_producer(self, prov, &variant, &kills) { + return None; + } + let delta = self.switch_cost_delta(state, prov, &variant, &kills); + delta.non_worse().then_some((delta, order, variant)) + }) + .min_by_key(|(delta, order, _)| (*delta, *order)) + .map(|(_, _, variant)| variant) + } + + fn verify_shared_plan( + &self, + state: &ExtractState, + view: &FusionView<'_>, + prov: Prov, + shared: Option<&super::super::ExecutionVariant>, + ) { + let fresh = self.best_fusion_candidate(state, view, prov); + match (shared, &fresh) { + (None, None) => {} + (Some(shared), Some(generated)) => { + assert!( + super::interner::planning_payload_eq(shared, generated), + "shared plan diverges from regeneration: window horizon misses generator input (prov {})", + prov.0 + ); + assert_eq!( + super::interner::variant_dependencies(shared), + super::interner::variant_dependencies(generated), + "shared plan dependencies diverge (prov {})", + prov.0 + ); + } + _ => panic!( + "shared plan presence diverges from regeneration (prov {}, shared={}, fresh={})", + prov.0, + shared.is_some(), + fresh.is_some() + ), + } + } + + /// Fusion extraction worklist, seeded with every + /// fusion-eligible node in provenance order; after each committed + /// switch, everything whose situation changed — the node itself, old and + /// new producers, killed nodes' children, and consumers reachable + /// through views — re-enters the worklist. Counts only decrease and the + /// generators' gates are antitone in them, so the loop converges to the + /// greatest fixpoint regardless of order: maximal legal fusion. + pub(super) fn extract_with_fusion(&mut self, ctx: &FusionCtx<'_>) -> Extraction { + let mut state = ExtractState::new(self); + let device = ctx.graph.device(); + state.no_dup_gate = device.config().spike_no_dup_gate; + state.dup_ledger = device.config().spike_dup_ledger; + state.scalar_cost = device.config().spike_scalar_cost; + state.read_traffic = device.config().spike_read_traffic || state.scalar_cost; + let mut plans = FusionPlanMemo::for_config(device.config()); + // The window horizon must cover everything a generator observes; + // this tripwire proves it by regenerating and comparing on every + // hit, per-resolve and device-store alike. + let verify_sharing = device.config().verify_plan_sharing; + let store = device.fusion_plan_store(); + let count = state.sel.len() as u32; + let mut worklist: std::collections::VecDeque = (0..count) + .filter(|&prov| { + let view = FusionView::new(self, &state, ctx); + view.is_seed_candidate(Prov(prov)) + }) + .collect(); + let mut queued = vec![false; count as usize]; + for &prov in &worklist { + queued[prov as usize] = true; + } + + while let Some(prov) = worklist.pop_front() { + queued[prov as usize] = false; + if !state.needed[prov as usize] { + continue; + } + if self.egraph.analysis.facts[prov as usize].exec.is_none() { + continue; + } + let pre_consumers: Vec = state.consumers[prov as usize].clone(); + + let (generated, plan_root, known_spec) = { + let view = FusionView::new(self, &state, ctx); + let instance = plans.capture(self, &state, &view, Prov(prov)); + match plans.lookup(&instance, &view) { + PlanLookup::Hit(result) => { + if verify_sharing { + self.verify_shared_plan(&state, &view, Prov(prov), result.as_ref()); + } + let spec = plans.known_spec(&instance); + (result, instance.root, spec) + } + PlanLookup::Miss => { + let key = plans.window_key(instance.root, self); + let result = match store.instantiate(key, &instance, &view) { + Some(result) => { + plans.note_store_hit(); + if verify_sharing { + self.verify_shared_plan( + &state, + &view, + Prov(prov), + result.as_ref(), + ); + } + plans.record(&instance, &view, result.as_ref()); + result + } + None => { + plans.note_store_miss(); + let result = self.best_fusion_candidate(&state, &view, Prov(prov)); + if let Some(decision) = + plans.record(&instance, &view, result.as_ref()) + { + store.record(key, decision); + } + result + } + }; + (result, instance.root, None) + } + } + }; + let Some(variant) = generated else { + continue; + }; + let planned_kills = state.kills_for_variant(self, Prov(prov), &variant); + if state.variant_duplicates_required_producer( + self, + Prov(prov), + &variant, + &planned_kills, + ) || !self + .switch_cost_delta(&state, Prov(prov), &variant, &planned_kills) + .non_worse() + { + continue; + } + // Record the fused form as an alternative of this node's class, + // then commit the switch with live counts. + let enode = self.mint_alternative_unique(Prov(prov), variant, known_spec); + let actual_spec = self + .egraph + .analysis + .payloads + .spec_of(enode.payload().expect("fused alternative has payload")); + plans.record_spec(plan_root, actual_spec); + let kills = state.kills(self, Prov(prov), &enode); + debug_assert_eq!(planned_kills, kills); + let mut touched = state.commit(self, Prov(prov), enode, &kills); + touched.extend(kills.iter().copied()); + + let view = FusionView::new(self, &state, ctx); + // The node itself (chains continue on the fused form), consumers + // from before and after the switch, and everything the kill + // cascade touched: a dead producer's siblings can become + // sole-consumed and newly fusible. + let new_consumers = state.consumers[prov as usize].clone(); + let seeds = std::iter::once(prov) + .chain(pre_consumers) + .chain(new_consumers) + .chain(touched); + view.enqueue_downstream(&state, seeds, &mut worklist, &mut queued); + } + if device.config().spike_dup_ledger { + // The remat surface: producers a candidate wanted to inline while + // they still materialize for another consumer. `write_bytes` is + // what dies if every consumer inlines; `recompute_work` is what + // each extra consumer then pays. Note the cost model counts only + // output writes, so neither side includes input re-reads. + let blocked = state.dup_blocked.borrow(); + let (mut live, mut write_bytes, mut recompute_work, mut extra_consumers) = + (0u64, 0u128, 0u128, 0u64); + for &prov in blocked.iter() { + if !state.needed[prov as usize] { + continue; + } + let cost = self.selection_cost(&state, Prov(prov)); + live += 1; + write_bytes = write_bytes.saturating_add(cost.materialized_bytes); + recompute_work = recompute_work.saturating_add(cost.work); + extra_consumers += u64::from(state.reads[prov as usize].saturating_sub(1)); + } + tracing::info!( + "dup_gate_ledger enforcing={} firings={} producers={} live={} write_bytes={} recompute_work={} extra_consumers={}", + !state.no_dup_gate, + state.dup_firings.get(), + blocked.len(), + live, + write_bytes, + recompute_work, + extra_consumers, + ); + } + if device.config().spike_hoisting { + let sharing = plans.stats(); + tracing::info!( + "hoisting_spike_windows stub_depth={} windows={} unique={} hits={} misses={} store_hits={} store_misses={} capture_us={} capture_ns_per_window={}", + plans.stub_depth(), + sharing.windows, + sharing.unique_windows, + sharing.hits, + sharing.misses, + sharing.store_hits, + sharing.store_misses, + plans.capture_time().as_micros(), + plans.capture_time().as_nanos() / sharing.windows.max(1) as u128, + ); + } + if device.config().trace_resolve_host { + let sharing = plans.stats(); + let cost = self.extraction_cost(&state); + tracing::info!( + "resolve_egg_plans windows={} unique={} hits={} misses={} store_hits={} store_misses={} templates={} negative={} payloads={} specs={} dispatches={} bytes={} work={}", + sharing.windows, + sharing.unique_windows, + sharing.hits, + sharing.misses, + sharing.store_hits, + sharing.store_misses, + sharing.templates, + sharing.negative_templates, + self.egraph.analysis.payloads.payload_count(), + self.egraph.analysis.payloads.spec_count(), + cost.dispatches, + cost.materialized_bytes, + cost.work, + ); + if sharing.unshareable > 0 { + tracing::info!( + "resolve_egg_plans_unshareable rewrites={}", + sharing.unshareable + ); + } + } + Extraction { + sel: state.sel, + needed: state.needed, + } + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/ingest.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/ingest.rs new file mode 100644 index 000000000..f3f222c58 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/ingest.rs @@ -0,0 +1,560 @@ +//! Deterministic ingestion: execution graph → e-graph. +//! +//! Traversal mirrors `build_execution_graph` (and the flush-replay +//! fingerprint): targets in order, dependencies in `visit_dependencies` +//! order, provenance assigned at first visit (pre-order). Payloads are +//! interned byte-faithful — no dedup, no normalization — because recognition +//! matchers depend on the exact API-emitted forms. +//! +//! Inputs outside the execution graph (cached before this resolve, the +//! `resolved_set`) become opaque [`FusorLang::Boundary`] leaves: the +//! structural equivalent of every `check_cached` guard — fusion rules only +//! match operation e-nodes, and a cached producer simply is not one. + +use egg::{EGraph, Id}; +use rustc_hash::FxHashSet; + +use super::super::{ExecutionVariant, Resolver}; +use super::EGraphDriver; +use super::analysis::{FusorAnalysis, NodeFacts}; +use super::interner::{SpecId, variant_dependencies}; +use super::lang::{AllocationId, FusorLang, Prov}; +use crate::compute_graph::{ComputeGraphInner, NodeIndex}; + +/// Build the e-node for `variant` with the given children (in +/// `visit_dependencies` order). Shared by ingestion and rule appliers so +/// the payload/children lockstep has one owner. +pub(super) fn enode_for( + analysis: &mut FusorAnalysis, + variant: &ExecutionVariant, + prov: Prov, + children: Vec, + dedup: bool, + known_spec: Option, +) -> FusorLang { + let intern = |analysis: &mut FusorAnalysis, variant: &ExecutionVariant| { + if let Some(spec) = known_spec { + analysis + .payloads + .push_unique_with_spec(variant.clone(), spec) + } else if dedup { + analysis.payloads.intern(variant.clone()) + } else { + analysis.payloads.push_unique(variant.clone()) + } + }; + match variant { + // Folds share the `Reduce` e-node kind: same child order, same + // structural role. The payload tag distinguishes them, exactly as + // attention shares `RowProgram`. + ExecutionVariant::Fold(_) => { + FusorLang::Reduce(intern(analysis, variant), children.into_boxed_slice()) + } + ExecutionVariant::Tensor(data) => { + debug_assert!(children.is_empty()); + FusorLang::TensorLeaf(AllocationId(std::sync::Arc::as_ptr(data.buffer()) as usize)) + } + ExecutionVariant::QMatrix(op) => { + debug_assert!(children.is_empty()); + FusorLang::QMatrixLeaf( + AllocationId(std::sync::Arc::as_ptr(op.matrix.buffer()) as usize), + intern(analysis, variant), + ) + } + ExecutionVariant::Elementwise(_) => { + FusorLang::Elementwise(intern(analysis, variant), children.into_boxed_slice()) + } + ExecutionVariant::Reduce(_) => { + FusorLang::Reduce(intern(analysis, variant), children.into_boxed_slice()) + } + ExecutionVariant::View(_) => FusorLang::View(intern(analysis, variant), [children[0]]), + ExecutionVariant::Assign(_) => { + FusorLang::Assign(prov, intern(analysis, variant), [children[0], children[1]]) + } + ExecutionVariant::Region(_) => { + FusorLang::Region(prov, intern(analysis, variant), children.into_boxed_slice()) + } + ExecutionVariant::MatMul(_) => { + FusorLang::MatMul(intern(analysis, variant), [children[0], children[1]]) + } + ExecutionVariant::QMatMul(_) => { + FusorLang::QMatMul(intern(analysis, variant), children.into_boxed_slice()) + } + ExecutionVariant::QEmbedding(_) => { + FusorLang::QEmbedding(intern(analysis, variant), [children[0]]) + } + ExecutionVariant::RowProgram(_) | ExecutionVariant::Attention(_) => { + FusorLang::RowProgram(intern(analysis, variant), children.into_boxed_slice()) + } + } +} + +impl EGraphDriver { + /// Ingest the resolver's execution graph reachable from its targets. + pub(super) fn ingest(resolver: &Resolver, graph: &ComputeGraphInner) -> Self { + let mut driver = EGraphDriver { + egraph: EGraph::new(FusorAnalysis { + facts: Vec::new(), + payloads: Default::default(), + class_of_inner: Default::default(), + }), + class_of: Vec::new(), + identity_enodes: Vec::new(), + identity_variants: Vec::new(), + prov_of: Default::default(), + provs_of_class: Default::default(), + }; + let target_set: FxHashSet = resolver.targets.iter().copied().collect(); + + enum Frame { + Enter(NodeIndex), + Exit { inner: NodeIndex, prov: Prov }, + } + let mut stack = Vec::new(); + // Nodes whose Exit frame has not run yet. Re-entering one means the + // execution graph has a dependency cycle, which the placeholder + // class slots downstream would turn into an opaque unionfind panic. + #[cfg(debug_assertions)] + let mut open: FxHashSet = FxHashSet::default(); + for &target in resolver.targets.iter().rev() { + stack.push(Frame::Enter(target)); + } + while let Some(frame) = stack.pop() { + match frame { + Frame::Enter(inner) => { + if driver.prov_of.contains_key(&inner) { + #[cfg(debug_assertions)] + if open.contains(&inner) { + let describe = |node: NodeIndex| { + resolver + .node_mapping + .get(&node) + .map(|&exec| { + format!("{:?}", resolver.execution_graph[exec].variant) + }) + .unwrap_or_else(|| "".into()) + }; + let variant = describe(inner); + let deps = resolver + .node_mapping + .get(&inner) + .map(|&exec| { + variant_dependencies(&resolver.execution_graph[exec].variant) + .into_iter() + .map(|dep| format!("{dep:?}: {}", describe(dep))) + .collect::>() + }) + .unwrap_or_default(); + let spine: Vec = stack + .iter() + .filter_map(|frame| match frame { + Frame::Exit { inner, .. } => { + let mut text = describe(*inner); + text.truncate(120); + Some(format!("{inner:?}: {text}")) + } + _ => None, + }) + .collect(); + panic!( + "dependency cycle through inner node {inner:?} ({variant}) reached \ + egraph ingest; dependencies: {deps:#?}; open spine: {spine:#?}" + ); + } + continue; + } + let exec = resolver.node_mapping.get(&inner).copied(); + let Some(exec_idx) = exec else { + // Cached before this resolve (or otherwise excluded + // from the execution graph): opaque boundary leaf. + let prov = driver.alloc_prov( + inner, + NodeFacts { + inner, + exec: None, + externally_live: graph.has_live_reference(inner), + is_target: target_set.contains(&inner), + }, + ); + let allocation = graph + .get_cached_result(inner) + .map( + |data| AllocationId(std::sync::Arc::as_ptr(data.buffer()) as usize), + ) + .unwrap_or(AllocationId(inner.index())); + let boundary = FusorLang::Boundary(allocation); + let id = driver.egraph.add(boundary.clone()); + driver.egraph.analysis.class_of_inner.insert(inner, id); + driver.class_of.push(id); + driver.identity_enodes.push(boundary); + driver.identity_variants.push(None); + debug_assert_eq!(driver.class_of.len(), prov.0 as usize + 1); + continue; + }; + let prov = driver.alloc_prov( + inner, + NodeFacts { + inner, + exec: Some(exec_idx), + externally_live: graph.has_live_reference(inner), + is_target: target_set.contains(&inner), + }, + ); + // Reserve the class slot now (pre-order prov => index + // into class_of); filled on exit. Safe placeholder: a + // consumer only reads a dependency's slot after that + // dependency's Exit frame ran (DAG + stack ordering), + // and the validation pass below re-checks every slot. + driver.class_of.push(Id::from(0usize)); + driver + .identity_enodes + .push(FusorLang::Boundary(AllocationId(inner.index()))); + driver.identity_variants.push(None); + #[cfg(debug_assertions)] + open.insert(inner); + stack.push(Frame::Exit { inner, prov }); + let deps = variant_dependencies(&resolver.execution_graph[exec_idx].variant); + for &dep in deps.iter().rev() { + stack.push(Frame::Enter(dep)); + } + } + Frame::Exit { inner, prov } => { + #[cfg(debug_assertions)] + open.remove(&inner); + let exec_idx = resolver.node_mapping[&inner]; + let variant = resolver.execution_graph[exec_idx].variant.clone(); + let children: Vec = variant_dependencies(&variant) + .into_iter() + .map(|dep| driver.class_of[driver.prov_of[&dep].0 as usize]) + .collect(); + let enode = enode_for( + &mut driver.egraph.analysis, + &variant, + prov, + children, + true, + None, + ); + driver.identity_enodes[prov.0 as usize] = enode.clone(); + driver.identity_variants[prov.0 as usize] = Some(variant); + let id = driver.egraph.add(enode); + driver.class_of[prov.0 as usize] = id; + driver.egraph.analysis.class_of_inner.insert(inner, id); + } + } + } + driver.egraph.rebuild(); + driver.refresh_prov_classes(); + debug_assert_eq!( + driver.class_of.len(), + driver.egraph.analysis.facts.len(), + "one class per provenance" + ); + #[cfg(debug_assertions)] + for (index, &id) in driver.class_of.iter().enumerate() { + let class = driver.egraph.find(id); + assert!( + driver.provs_of_class[&class].contains(&Prov(index as u32)), + "class slot must contain its own provenance" + ); + } + driver + } + + fn alloc_prov(&mut self, inner: NodeIndex, facts: NodeFacts) -> Prov { + let prov = Prov(self.egraph.analysis.facts.len() as u32); + self.egraph.analysis.facts.push(facts); + self.prov_of.insert(inner, prov); + prov + } + + /// The (canonical) e-class of an ingested inner node. + pub(super) fn class_for(&self, inner: NodeIndex) -> Option { + self.prov_of + .get(&inner) + .map(|prov| self.egraph.find(self.class_of[prov.0 as usize])) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{Device, Tensor}; + + /// Build the execution graph for `targets` exactly like resolve pass 1, + /// ingest it, and hand everything to the assertion closure. + fn with_ingested( + device: &Device, + targets: &[&Tensor], + f: impl FnOnce(&Resolver, &ComputeGraphInner, EGraphDriver) -> R, + ) -> R { + let keys: Vec = targets.iter().map(|t| t.data().key).collect(); + device.compute_graph().with_mut(|inner| { + let mut resolver = Resolver::new_batch(inner, keys.clone()); + for &key in &keys { + resolver.build_execution_graph(inner, key); + } + let driver = EGraphDriver::ingest(&resolver, inner); + f(&resolver, inner, driver) + }) + } + + fn interesting_graph(device: &Device) -> (Tensor, Tensor) { + let rows = vec![vec![1.0f32, 2.0, 3.0, 4.0]; 8]; + let input = Tensor::new::(device, &rows); + let weight_rows = vec![vec![0.25f32, 0.5, 0.75, 1.0]; 4]; + let weight = Tensor::new::(device, &weight_rows); + // Elementwise -> composed matmul (views + multiply + sum) -> reduce: + // covers Tensor leaves, View, Elementwise, and Reduce variants. + let x = (&input * 2.0) + 1.0; + let m = x.mat_mul(&weight); + let s = m.sum(1); + (s, m) + } + + fn dump(driver: &EGraphDriver) -> String { + use std::fmt::Write; + let mut out = String::new(); + for (index, &id) in driver.class_of.iter().enumerate() { + let facts = &driver.egraph.analysis.facts[index]; + let class = &driver.egraph[driver.egraph.find(id)]; + let mut nodes: Vec = + class.nodes.iter().map(|node| format!("{node:?}")).collect(); + nodes.sort(); + writeln!( + out, + "prov={index} inner={} exec={:?} live={} target={} nodes={nodes:?}", + facts.inner.index(), + facts.exec.map(|e| e.index()), + facts.externally_live, + facts.is_target, + ) + .unwrap(); + } + out + } + + #[test] + fn ingestion_covers_execution_graph() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (target, _mid) = interesting_graph(&device); + with_ingested(&device, &[&target], |resolver, _inner, driver| { + let non_boundary = driver + .egraph + .analysis + .facts + .iter() + .filter(|facts| facts.exec.is_some()) + .count(); + assert_eq!( + non_boundary, + resolver.execution_graph.node_count(), + "every execution node ingests exactly once" + ); + // Identity-only ingestion: one e-node per provenance. + assert_eq!( + driver.egraph.total_number_of_nodes(), + driver.egraph.analysis.facts.len(), + ); + assert!(driver.class_for(resolver.targets[0]).is_some()); + let target_prov = driver.prov_of[&resolver.targets[0]]; + assert!(driver.egraph.analysis.facts_of(target_prov).is_target); + }); + }); + } + + #[test] + fn ingestion_is_deterministic() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (target, _mid) = interesting_graph(&device); + let first = with_ingested(&device, &[&target], |_, _, driver| dump(&driver)); + let second = with_ingested(&device, &[&target], |_, _, driver| dump(&driver)); + assert_eq!( + first, second, + "ingestion must be a pure function of the graph" + ); + }); + } + + #[test] + fn cached_inputs_become_boundaries() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let rows = vec![vec![1.0f32, 2.0, 3.0, 4.0]; 8]; + let input = Tensor::new::(&device, &rows); + let x = (&input * 2.0) + 1.0; + // Materialize `x` so it is cached before the next resolve starts. + let _ = x.data().materialize(); + let y = x.sin(); + with_ingested(&device, &[&y], |_resolver, _inner, driver| { + let x_prov = driver.prov_of[&x.data().key]; + let facts = driver.egraph.analysis.facts_of(x_prov); + assert!( + facts.exec.is_none(), + "cached producer must ingest as an opaque boundary leaf" + ); + let class = &driver.egraph[driver.egraph.find(driver.class_of[x_prov.0 as usize])]; + assert!(matches!(class.nodes.as_slice(), [FusorLang::Boundary(..)])); + }); + }); + } + + #[test] + fn equivalent_pure_nodes_share_an_eclass() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let input = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let left = &input * 2.0; + let right = &input * 2.0; + with_ingested(&device, &[&left, &right], |_, _, driver| { + let left_class = driver.class_for(left.data().key).unwrap(); + let right_class = driver.class_for(right.data().key).unwrap(); + assert_eq!(left_class, right_class); + assert_eq!(driver.provs_of_class[&left_class].len(), 2); + }); + }); + } + + #[test] + fn class_resolution_selects_a_needed_equivalent_observation() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let input = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let left = &input * 2.0; + let right = &input * 2.0; + with_ingested(&device, &[&left, &right], |_, _, driver| { + let class = driver.class_for(left.data().key).unwrap(); + let provenances = &driver.provs_of_class[&class]; + assert_eq!(provenances.len(), 2); + let mut needed = vec![true; driver.egraph.analysis.facts.len()]; + needed[provenances[0].0 as usize] = false; + + let selected = driver.prov_of_class(class, &needed); + assert_eq!(selected, provenances[1]); + }); + }); + } + + #[test] + fn allocation_identity_distinguishes_equal_contents() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let a = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let b = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let left = &a * 2.0; + let right = &b * 2.0; + with_ingested(&device, &[&left, &right], |_, _, driver| { + assert_ne!( + driver.class_for(left.data().key), + driver.class_for(right.data().key) + ); + }); + }); + } + + #[test] + fn congruence_shares_equivalent_nested_subgraphs() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let input = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let left = (&input * 2.0) + 1.0; + let right = (&input * 2.0) + 1.0; + with_ingested(&device, &[&left, &right], |_, _, driver| { + assert_eq!( + driver.class_for(left.data().key), + driver.class_for(right.data().key) + ); + }); + }); + } + + #[test] + fn shared_eclass_coalesces_to_one_execution_node() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let input = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let left = &input * 2.0; + let right = &input * 2.0; + let targets = vec![left.data().key, right.data().key]; + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.clone()); + for &target in &targets { + resolver.build_execution_graph(graph, target); + } + let driver = EGraphDriver::ingest(&resolver, graph); + resolver.coalesce_equivalent_eclasses(graph, &driver); + assert_eq!(resolver.execution_graph.node_count(), 2); + assert_eq!( + resolver + .shared_outputs + .values() + .flatten() + .copied() + .collect::>(), + vec![targets[1]] + ); + }); + }); + } + + #[test] + fn nested_shared_eclasses_coalesce_without_stale_execution_nodes() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let input = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let mut left = &input * 2.0; + let mut right = &input * 2.0; + const DEPTH: usize = 2; + for _ in 0..DEPTH { + left = &left + 1.0; + right = &right + 1.0; + } + let targets = vec![left.data().key, right.data().key]; + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.clone()); + for &target in &targets { + resolver.build_execution_graph(graph, target); + } + + let driver = EGraphDriver::ingest(&resolver, graph); + resolver.coalesce_equivalent_eclasses(graph, &driver); + + assert_eq!(resolver.execution_graph.node_count(), DEPTH + 2); + assert!( + resolver + .node_mapping + .values() + .all(|&execution| { resolver.execution_graph.contains_node(execution) }) + ); + assert_eq!( + resolver + .shared_outputs + .values() + .map(Vec::len) + .sum::(), + DEPTH + 1 + ); + }); + }); + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/interner.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/interner.rs new file mode 100644 index 000000000..d6a3b9e66 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/interner.rs @@ -0,0 +1,423 @@ +//! Payload interning: complete [`ExecutionVariant`]s keyed by a two-lane +//! 128-bit structural hash. +//! +//! The key reuses each operation's `hash_kernel_fields` — the kernel-cache +//! surface, which by definition covers every field that changes generated +//! source, plus the variant tag. Dependencies are e-node children, not +//! payload identity; keeping physical NodeIndexes here defeats congruence. +//! Hashes select a bucket only; [`semantic_payload_eq`] performs exact +//! equality, so collisions cannot conflate operators. +//! +//! Interning full `ExecutionVariant` clones means delta application reads +//! complete, builder-produced operations straight from the table — identical +//! shader source, kernel-cache keys, and dispatch names to the destructive +//! optimizer. + +use std::hash::Hash; + +use fusor_tile_ir_runtime::TwoLaneHasher; +use rustc_hash::{FxHashMap, FxHasher}; + +use super::super::ExecutionVariant; +use super::lang::PayloadId; +use crate::compute_graph::NodeIndex; + +/// Allocation-independent identity of an operation shape used by the +/// planner. A payload remains the concrete, executable operation; a spec is +/// the reusable part of that operation (kernel fields, tensor shapes and +/// expression structure) with dependency and buffer identities erased. +/// +/// Keeping this separate from [`PayloadId`] is essential: two transformer +/// layers may share a spec while still referring to different weights and +/// activations, and therefore must never become value-equal e-nodes. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(super) struct SpecId(pub(super) u32); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(super) struct PayloadKey(pub(super) [u64; 2]); + +/// Accumulating payload/window hasher on the canonical [`TwoLaneHasher`] +/// (see `fusor_tile_ir_runtime::two_lane` for the collision model). +pub(super) struct TwoLane(TwoLaneHasher); + +impl TwoLane { + pub(super) fn new() -> Self { + Self(TwoLaneHasher::new()) + } + + pub(super) fn write_u64(&mut self, value: u64) { + self.0.write_u64(value); + } + + pub(super) fn finish(self) -> PayloadKey { + PayloadKey(self.0.finish()) + } +} + +pub(super) use fusor_tile_ir_runtime::single_lane as local_hash; + +/// The payload's dependencies in `visit_dependencies` order — the order the +/// e-node's children mirror. +pub(super) fn variant_dependencies(variant: &ExecutionVariant) -> Vec { + let mut deps = Vec::new(); + variant.visit_dependencies(&mut |dep| deps.push(dep)); + deps +} + +/// Rewrite `variant`'s dependency slots — in `visit_dependencies` order — to +/// `new`. Interned payloads deduplicate with their inputs ignored +/// (`semantic_payload_eq`), so a payload fetched through the table may carry +/// the concrete input indices of a *different* structurally-identical +/// instance (another layer's matmul, say). Every materialization of a +/// payload back into the execution graph must rebind its inputs to the +/// e-node's actual children or it computes with the wrong operands. +pub(super) fn rebind_variant_dependencies(variant: &mut ExecutionVariant, new: &[NodeIndex]) { + let mut slots = new.iter().copied(); + variant.visit_dependencies_mut(&mut |slot| { + *slot = slots.next().expect("rebind arity"); + }); + debug_assert!( + slots.next().is_none(), + "rebind received more children than the variant has dependency slots" + ); +} + +fn variant_tag(variant: &ExecutionVariant) -> u8 { + match variant { + ExecutionVariant::Tensor(_) => 0, + ExecutionVariant::QMatrix(_) => 1, + ExecutionVariant::Elementwise(_) => 2, + ExecutionVariant::Reduce(_) => 3, + ExecutionVariant::View(_) => 4, + ExecutionVariant::Assign(_) => 5, + ExecutionVariant::Region(_) => 6, + ExecutionVariant::MatMul(_) => 7, + ExecutionVariant::QMatMul(_) => 8, + ExecutionVariant::QEmbedding(_) => 9, + ExecutionVariant::RowProgram(_) => 10, + ExecutionVariant::Attention(_) => 11, + ExecutionVariant::Fold(_) => 12, + } +} + +fn hash_variant_fields(variant: &ExecutionVariant, hasher: &mut FxHasher) { + use crate::mir::operation::Operation; + match variant { + // Tensor leaves are never interned (identified by provenance alone); + // hash nothing beyond the tag if one ever reaches here. + ExecutionVariant::Tensor(_) => {} + ExecutionVariant::QMatrix(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::Elementwise(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::Reduce(op) => op.hash_kernel_fields(hasher), + // A fold has no `Operation` impl: it is an algebra term, and only its + // reduce form is lowerable. Hash the carrier structure directly. + ExecutionVariant::Fold(op) => op.hash_carrier_fields(hasher), + ExecutionVariant::View(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::Assign(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::Region(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::MatMul(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::QMatMul(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::QEmbedding(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::RowProgram(op) => op.hash_kernel_fields(hasher), + ExecutionVariant::Attention(op) => op.hash_kernel_fields(hasher), + } +} + +fn payload_key(variant: &ExecutionVariant) -> PayloadKey { + let mut lanes = TwoLane::new(); + lanes.write_u64(variant_tag(variant) as u64); + lanes.write_u64(local_hash(|hasher| hash_variant_fields(variant, hasher))); + lanes.finish() +} + +/// Append-only table of interned payloads. Ids are assigned in first-intern +/// order, which is deterministic because ingestion and rule application are. +#[derive(Default)] +pub(super) struct PayloadTable { + payloads: Vec, + by_key: FxHashMap>, + specs: Vec, + specs_by_key: FxHashMap>, + /// The spec's structural hash. `SpecId`s are per-resolve; this key is the + /// resolve-independent identity the persistent plan store hashes instead. + spec_keys: Vec, + spec_of_payload: Vec, +} + +impl PayloadTable { + pub(super) fn intern(&mut self, variant: ExecutionVariant) -> PayloadId { + let key = payload_key(&variant); + if let Some(id) = self.by_key.get(&key).and_then(|ids| { + ids.iter() + .copied() + .find(|id| semantic_payload_eq(&self.payloads[id.0 as usize], &variant)) + }) { + return id; + } + let id = PayloadId(self.payloads.len() as u32); + let spec = self.intern_spec(key, &variant); + self.payloads.push(variant); + self.by_key.entry(key).or_default().push(id); + self.spec_of_payload.push(spec); + id + } + + /// Append without semantic payload dedup. Fusion generates each concrete + /// occurrence once, so idempotence lookup buys nothing; the first planning + /// occurrence still establishes its allocation-independent spec. Repeated + /// occurrences use [`Self::push_unique_with_spec`] and skip that structural + /// hash too. Ingestion keeps using [`Self::intern`], whose semantic dedup + /// is what places equivalent observations in one e-class. + pub(super) fn push_unique(&mut self, variant: ExecutionVariant) -> PayloadId { + let key = payload_key(&variant); + let spec = self.intern_spec(key, &variant); + self.push_unique_with_spec(variant, spec) + } + + /// Append a concrete occurrence whose allocation-independent spec was + /// already established by a shared planning template. This is the hot + /// repeated-layer path: it avoids re-hashing a potentially large fused + /// expression for every layer. + pub(super) fn push_unique_with_spec( + &mut self, + variant: ExecutionVariant, + spec: SpecId, + ) -> PayloadId { + debug_assert!((spec.0 as usize) < self.specs.len()); + debug_assert!(planning_payload_eq(&self.specs[spec.0 as usize], &variant)); + let id = PayloadId(self.payloads.len() as u32); + self.payloads.push(variant); + self.spec_of_payload.push(spec); + id + } + + pub(super) fn get(&self, id: PayloadId) -> &ExecutionVariant { + &self.payloads[id.0 as usize] + } + + pub(super) fn spec_of(&self, id: PayloadId) -> SpecId { + self.spec_of_payload[id.0 as usize] + } + + pub(super) fn spec_key(&self, id: SpecId) -> PayloadKey { + self.spec_keys[id.0 as usize] + } + + pub(super) fn spec_count(&self) -> usize { + self.specs.len() + } + + pub(super) fn payload_count(&self) -> usize { + self.payloads.len() + } + + fn intern_spec(&mut self, key: PayloadKey, variant: &ExecutionVariant) -> SpecId { + if let Some(id) = self.specs_by_key.get(&key).and_then(|ids| { + ids.iter() + .copied() + .find(|id| planning_payload_eq(&self.specs[id.0 as usize], variant)) + }) { + return id; + } + let id = SpecId(self.specs.len() as u32); + self.specs.push(variant.clone()); + self.specs_by_key.entry(key).or_default().push(id); + self.spec_keys.push(key); + id + } +} + +fn same_qmatrix_spec(a: &crate::quantized::QMatrix, b: &crate::quantized::QMatrix) -> bool { + a.datatype() == b.datatype() + && a.storage_layout() == b.storage_layout() + && a.shape() == b.shape() +} + +fn same_epilogue_spec( + a: &Option, + b: &Option, +) -> bool { + match (a, b) { + (None, None) => true, + (Some(a), Some(b)) => { + a.expression == b.expression + && a.extras.len() == b.extras.len() + && a.input_datatype == b.input_datatype + && a.output_datatype == b.output_datatype + } + _ => false, + } +} + +/// Exact equality for the allocation-independent planning surface. +/// `payload_key` is only a bucket selector, so this comparison deliberately +/// checks every relevant field and makes hash collisions harmless. +pub(super) fn planning_payload_eq(a: &ExecutionVariant, b: &ExecutionVariant) -> bool { + let zero = NodeIndex::new(0); + match (a, b) { + (ExecutionVariant::Tensor(a), ExecutionVariant::Tensor(b)) => { + a.datatype() == b.datatype() && a.layout() == b.layout() + } + (ExecutionVariant::QMatrix(a), ExecutionVariant::QMatrix(b)) => { + same_qmatrix_spec(&a.matrix, &b.matrix) + && a.datatype == b.datatype + && a.post_dequantize == b.post_dequantize + } + (ExecutionVariant::Elementwise(a), ExecutionVariant::Elementwise(b)) => { + a.expression == b.expression + && a.shape == b.shape + && a.output_datatype == b.output_datatype + && a.inputs.len() == b.inputs.len() + } + (ExecutionVariant::Reduce(a), ExecutionVariant::Reduce(b)) => { + a.expression == b.expression + && a.shape == b.shape + && a.function == b.function + && a.post_element_wise == b.post_element_wise + && a.axis == b.axis + && a.inputs.len() == b.inputs.len() + } + (ExecutionVariant::View(a), ExecutionVariant::View(b)) => { + a.stages == b.stages && a.datatype == b.datatype + } + (ExecutionVariant::MatMul(a), ExecutionVariant::MatMul(b)) => { + let mut a = a.clone(); + let mut b = b.clone(); + a.first = zero; + a.second = zero; + b.first = zero; + b.second = zero; + a == b + } + (ExecutionVariant::QMatMul(a), ExecutionVariant::QMatMul(b)) => { + a.input_datatype == b.input_datatype + && same_qmatrix_spec(&a.matrix, &b.matrix) + && a.in_shape == b.in_shape + && a.out_shape == b.out_shape + && same_epilogue_spec(&a.pre_element_wise_expr, &b.pre_element_wise_expr) + && same_epilogue_spec(&a.post_element_wise_expr, &b.post_element_wise_expr) + && a.post_accumulator_offsets == b.post_accumulator_offsets + } + (ExecutionVariant::QEmbedding(a), ExecutionVariant::QEmbedding(b)) => { + same_qmatrix_spec(&a.matrix, &b.matrix) + && a.out_shape == b.out_shape + && a.datatype == b.datatype + } + (ExecutionVariant::RowProgram(a), ExecutionVariant::RowProgram(b)) => { + let mut a = a.clone(); + let mut b = b.clone(); + a.inputs.clear(); + b.inputs.clear(); + a == b + } + (ExecutionVariant::Attention(a), ExecutionVariant::Attention(b)) => { + let zero_deps = |op: &crate::flash_attention::FlashAttentionOperation| { + let mut op = op.clone(); + op.q = zero; + op.k = zero; + op.v = op.v.map(|_| zero); + op.grad_o = op.grad_o.map(|_| zero); + op.lse = op.lse.map(|_| zero); + op.dsum = op.dsum.map(|_| zero); + op.mask = op.mask.map(|_| zero); + op + }; + zero_deps(a) == zero_deps(b) + } + // Effects and multi-output regions are observation-specific. + (ExecutionVariant::Assign(_), ExecutionVariant::Assign(_)) + | (ExecutionVariant::Region(_), ExecutionVariant::Region(_)) => false, + _ => false, + } +} + +fn semantic_payload_eq(a: &ExecutionVariant, b: &ExecutionVariant) -> bool { + let zero = NodeIndex::new(0); + match (a, b) { + (ExecutionVariant::Tensor(a), ExecutionVariant::Tensor(b)) => { + std::sync::Arc::ptr_eq(a.buffer(), b.buffer()) + } + (ExecutionVariant::QMatrix(a), ExecutionVariant::QMatrix(b)) => a == b, + (ExecutionVariant::Elementwise(a), ExecutionVariant::Elementwise(b)) => { + let mut a = a.clone(); + let mut b = b.clone(); + a.inputs.clear(); + b.inputs.clear(); + a == b + } + (ExecutionVariant::Reduce(a), ExecutionVariant::Reduce(b)) => { + let mut a = a.clone(); + let mut b = b.clone(); + a.inputs.clear(); + b.inputs.clear(); + a == b + } + (ExecutionVariant::View(a), ExecutionVariant::View(b)) => { + let mut a = a.clone(); + let mut b = b.clone(); + a.input = zero; + b.input = zero; + a == b + } + (ExecutionVariant::MatMul(a), ExecutionVariant::MatMul(b)) => { + let mut a = a.clone(); + let mut b = b.clone(); + a.first = zero; + a.second = zero; + b.first = zero; + b.second = zero; + a == b + } + (ExecutionVariant::QMatMul(a), ExecutionVariant::QMatMul(b)) => { + let mut a = a.as_ref().clone(); + let mut b = b.as_ref().clone(); + a.input = zero; + b.input = zero; + for epilogue in [&mut a.pre_element_wise_expr, &mut a.post_element_wise_expr] + .into_iter() + .flatten() + { + epilogue.extras.fill(zero); + } + for epilogue in [&mut b.pre_element_wise_expr, &mut b.post_element_wise_expr] + .into_iter() + .flatten() + { + epilogue.extras.fill(zero); + } + a == b + } + (ExecutionVariant::QEmbedding(a), ExecutionVariant::QEmbedding(b)) => { + let mut a = a.clone(); + let mut b = b.clone(); + a.indexes = zero; + b.indexes = zero; + a == b + } + (ExecutionVariant::RowProgram(a), ExecutionVariant::RowProgram(b)) => { + let mut a = a.clone(); + let mut b = b.clone(); + a.inputs.clear(); + b.inputs.clear(); + a == b + } + (ExecutionVariant::Attention(a), ExecutionVariant::Attention(b)) => { + let zero_deps = |op: &crate::flash_attention::FlashAttentionOperation| { + let mut op = op.clone(); + op.q = zero; + op.k = zero; + op.v = op.v.map(|_| zero); + op.grad_o = op.grad_o.map(|_| zero); + op.lse = op.lse.map(|_| zero); + op.dsum = op.dsum.map(|_| zero); + op.mask = op.mask.map(|_| zero); + op + }; + zero_deps(a) == zero_deps(b) + } + // Effects and multi-output regions are observation-specific. + (ExecutionVariant::Assign(_), ExecutionVariant::Assign(_)) + | (ExecutionVariant::Region(_), ExecutionVariant::Region(_)) => false, + _ => false, + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/lang.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/lang.rs new file mode 100644 index 000000000..070378f04 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/lang.rs @@ -0,0 +1,244 @@ +//! The e-graph term language over execution-graph nodes. +//! +//! Pure operators compare by semantic payload and child e-classes and carry +//! no observation metadata. Concrete and cached leaves compare by +//! [`AllocationId`]. Effectful assignments and multi-output regions retain +//! observation identity through the [`Prov`] they were ingested for. +//! +//! Payloads are complete [`ExecutionVariant`]s held in the driver's +//! [`super::interner::PayloadTable`], referenced by [`PayloadId`]. Children +//! are the operand e-classes in the payload's dependency order +//! (`visit_dependencies` order), kept in lockstep with the payload's +//! `inputs` vector by construction. +//! +//! [`ExecutionVariant`]: super::super::ExecutionVariant + +use egg::{Id, Language}; +use std::hash::{Hash, Hasher}; + +/// Dense observation id assigned to an execution `NodeIndex`. This indexes +/// liveness and target facts; it is not pure value identity. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(super) struct Prov(pub(super) u32); + +/// Index into the driver's payload table. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(super) struct PayloadId(pub(super) u32); + +/// Identity of an already-existing storage allocation. Allocation-backed +/// leaves are equal only when they name the same buffer. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(super) struct AllocationId(pub(super) usize); + +/// One alternative form of one execution node. +/// +/// The variants mirror `ExecutionVariant`, plus `Boundary` for inputs that +/// were already cached when the resolve started (the `resolved_set`): those +/// are opaque leaves exactly like `build_execution_graph` excluding them — +/// no rule may see through a cached boundary. +#[derive(Debug, Clone)] +pub(super) enum FusorLang { + /// Concrete tensor data already bound to a buffer. + TensorLeaf(AllocationId), + /// A node cached before this resolve began; contents opaque. + Boundary(AllocationId), + /// Quantized-matrix leaf (`DequantizeOperation` payload). + QMatrixLeaf(AllocationId, PayloadId), + Elementwise(PayloadId, Box<[Id]>), + Reduce(PayloadId, Box<[Id]>), + View(PayloadId, [Id; 1]), + Assign(Prov, PayloadId, [Id; 2]), + MatMul(PayloadId, [Id; 2]), + QMatMul(PayloadId, Box<[Id]>), + QEmbedding(PayloadId, [Id; 1]), + /// Structurally comparable fused row program, including attention. + RowProgram(PayloadId, Box<[Id]>), + /// Multi-output elementwise region (only after region formation). + Region(Prov, PayloadId, Box<[Id]>), +} + +impl PartialEq for FusorLang { + fn eq(&self, other: &Self) -> bool { + use FusorLang::*; + match (self, other) { + (TensorLeaf(a), TensorLeaf(b)) | (Boundary(a), Boundary(b)) => a == b, + (QMatrixLeaf(aa, a), QMatrixLeaf(ba, b)) => aa == ba && a == b, + (Elementwise(a, ac), Elementwise(b, bc)) + | (Reduce(a, ac), Reduce(b, bc)) + | (QMatMul(a, ac), QMatMul(b, bc)) + | (RowProgram(a, ac), RowProgram(b, bc)) => a == b && ac == bc, + (View(a, ac), View(b, bc)) | (QEmbedding(a, ac), QEmbedding(b, bc)) => { + a == b && ac == bc + } + (MatMul(a, ac), MatMul(b, bc)) => a == b && ac == bc, + (Assign(ap, a, ac), Assign(bp, b, bc)) => ap == bp && a == b && ac == bc, + (Region(ap, a, ac), Region(bp, b, bc)) => ap == bp && a == b && ac == bc, + _ => false, + } + } +} + +impl Eq for FusorLang {} + +impl PartialOrd for FusorLang { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for FusorLang { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + let tag = |node: &Self| match node { + Self::TensorLeaf(..) => 0u8, + Self::Boundary(..) => 1, + Self::QMatrixLeaf(..) => 2, + Self::Elementwise(..) => 3, + Self::Reduce(..) => 4, + Self::View(..) => 5, + Self::Assign(..) => 6, + Self::MatMul(..) => 7, + Self::QMatMul(..) => 8, + Self::QEmbedding(..) => 9, + Self::RowProgram(..) => 10, + Self::Region(..) => 11, + }; + tag(self) + .cmp(&tag(other)) + .then_with(|| match (self, other) { + (Self::TensorLeaf(a), Self::TensorLeaf(b)) + | (Self::Boundary(a), Self::Boundary(b)) => a.cmp(b), + (Self::QMatrixLeaf(aa, a), Self::QMatrixLeaf(ba, b)) => { + aa.cmp(ba).then_with(|| a.cmp(b)) + } + (Self::Assign(ap, a, ac), Self::Assign(bp, b, bc)) => { + ap.cmp(bp).then_with(|| a.cmp(b)).then_with(|| ac.cmp(bc)) + } + (Self::Region(ap, a, ac), Self::Region(bp, b, bc)) => { + ap.cmp(bp).then_with(|| a.cmp(b)).then_with(|| ac.cmp(bc)) + } + _ => self + .payload() + .cmp(&other.payload()) + .then_with(|| self.children().cmp(other.children())), + }) + } +} + +impl Hash for FusorLang { + fn hash(&self, state: &mut H) { + std::mem::discriminant(self).hash(state); + match self { + Self::TensorLeaf(allocation) | Self::Boundary(allocation) => allocation.hash(state), + Self::QMatrixLeaf(allocation, payload) => { + allocation.hash(state); + payload.hash(state); + } + Self::Assign(prov, payload, children) => { + prov.hash(state); + payload.hash(state); + children.hash(state); + } + Self::Region(prov, payload, children) => { + prov.hash(state); + payload.hash(state); + children.hash(state); + } + _ => { + self.payload().hash(state); + self.children().hash(state); + } + } + } +} + +impl FusorLang { + pub(super) fn payload(&self) -> Option { + match self { + Self::TensorLeaf(..) | Self::Boundary(..) => None, + Self::QMatrixLeaf(_, payload) + | Self::Elementwise(payload, _) + | Self::Reduce(payload, _) + | Self::View(payload, _) + | Self::Assign(_, payload, _) + | Self::MatMul(payload, _) + | Self::QMatMul(payload, _) + | Self::QEmbedding(payload, _) + | Self::RowProgram(payload, _) + | Self::Region(_, payload, _) => Some(*payload), + } + } +} + +impl Language for FusorLang { + /// Everything `matches` compares below the children: variant, leaf + /// allocation, payload, and observation identity for effectful nodes. + /// Discriminant equality must coincide with `matches`. + type Discriminant = ( + std::mem::Discriminant, + Option, + Option, + Option, + ); + + fn discriminant(&self) -> Self::Discriminant { + let allocation = match self { + Self::TensorLeaf(allocation) + | Self::Boundary(allocation) + | Self::QMatrixLeaf(allocation, _) => Some(*allocation), + _ => None, + }; + let prov = match self { + Self::Assign(prov, _, _) | Self::Region(prov, _, _) => Some(*prov), + _ => None, + }; + ( + std::mem::discriminant(self), + allocation, + self.payload(), + prov, + ) + } + + fn matches(&self, other: &Self) -> bool { + use FusorLang::*; + match (self, other) { + (TensorLeaf(a), TensorLeaf(b)) | (Boundary(a), Boundary(b)) => a == b, + (QMatrixLeaf(aa, a), QMatrixLeaf(ba, b)) => aa == ba && a == b, + (Assign(ap, a, _), Assign(bp, b, _)) | (Region(ap, a, _), Region(bp, b, _)) => { + ap == bp && a == b + } + _ => { + std::mem::discriminant(self) == std::mem::discriminant(other) + && self.payload() == other.payload() + } + } + } + + fn children(&self) -> &[Id] { + match self { + Self::TensorLeaf(..) | Self::Boundary(..) | Self::QMatrixLeaf(..) => &[], + Self::Elementwise(_, children) + | Self::Reduce(_, children) + | Self::QMatMul(_, children) + | Self::RowProgram(_, children) + | Self::Region(_, _, children) => children, + Self::View(_, children) | Self::QEmbedding(_, children) => children, + Self::MatMul(_, children) => children, + Self::Assign(_, _, children) => children, + } + } + + fn children_mut(&mut self) -> &mut [Id] { + match self { + Self::TensorLeaf(..) | Self::Boundary(..) | Self::QMatrixLeaf(..) => &mut [], + Self::Elementwise(_, children) + | Self::Reduce(_, children) + | Self::QMatMul(_, children) + | Self::RowProgram(_, children) + | Self::Region(_, _, children) => children, + Self::View(_, children) | Self::QEmbedding(_, children) => children, + Self::MatMul(_, children) => children, + Self::Assign(_, _, children) => children, + } + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/mod.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/mod.rs new file mode 100644 index 000000000..bd56b4c06 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/mod.rs @@ -0,0 +1,256 @@ +//! Value e-graph optimizer for operation fusion. +//! +//! Rewrites are strictly additive: a fusion generator's alternative e-node is +//! unioned into the root's class and nothing is ever removed. A GPU-oriented +//! extractor chooses among those alternatives, and one physical planner +//! applies the chosen terms back onto the execution graph. +//! +//! Pure e-node identity is semantic: operator payload plus child e-classes. +//! Allocation-backed leaves use allocation identity. An allocation-independent +//! structural interner lets isomorphic repeated layers share rewrite templates +//! without ever sharing their values. Multiple +//! execution-graph observations of one value e-class are materialized once +//! and cached under every observed `NodeIndex`. + +mod analysis; +mod apply; +pub(super) mod compose; +mod cost; +mod extract; +mod ingest; +mod interner; +mod lang; +mod rules_fuse; +mod rules_fuse_matmul; +mod structural_memo; +pub(crate) use structural_memo::FusionPlanStore; + +use egg::{EGraph, Id}; +use rustc_hash::FxHashMap; + +use self::analysis::FusorAnalysis; +use self::lang::FusorLang; +use super::{ExecutionVariant, Resolver}; +use crate::compute_graph::{ComputeGraphInner, NodeIndex}; + +/// Owns the resolve's value e-graph plus the provenance bookkeeping +/// connecting it to the resolver's execution graph. +pub(super) struct EGraphDriver { + egraph: EGraph, + /// Provenance -> e-class id (as returned at add time; canonicalize with + /// `egraph.find` after unions). + class_of: Vec, + identity_enodes: Vec, + identity_variants: Vec>, + /// Inner-graph node -> provenance. + prov_of: FxHashMap, + provs_of_class: FxHashMap>, +} + +impl EGraphDriver { + /// The API-emitted operation for one observation, before any egg + /// alternative is selected. + fn identity_variant(&self, prov: lang::Prov) -> Option<&ExecutionVariant> { + self.identity_variants[prov.0 as usize].as_ref() + } + + fn refresh_prov_classes(&mut self) { + self.provs_of_class.clear(); + for (index, &class) in self.class_of.iter().enumerate() { + self.provs_of_class + .entry(self.egraph.find(class)) + .or_default() + .push(lang::Prov(index as u32)); + } + } + /// Fusion mint: skips semantic payload dedup and can reuse a structural + /// spec learned from an isomorphic earlier occurrence. + fn mint_alternative_unique( + &mut self, + root: lang::Prov, + variant: ExecutionVariant, + known_spec: Option, + ) -> FusorLang { + let children: Vec = interner::variant_dependencies(&variant) + .into_iter() + .map(|dep| { + self.class_for(dep) + .expect("alternative dependencies must already be ingested") + }) + .collect(); + let enode = ingest::enode_for( + &mut self.egraph.analysis, + &variant, + root, + children, + false, + known_spec, + ); + let id = self.egraph.add(enode.clone()); + let root_id = self.class_of[root.0 as usize]; + self.egraph.union(root_id, id); + enode + } +} + +impl Resolver { + /// Recognize specialized operations, extend them with explicit cluster + /// builders, and extract fusion alternatives through one value e-graph. + /// Allocation-independent structural templates make repeated-layer + /// fusion planning proportional to unique local structure. + pub(super) fn optimize_operations(&mut self, graph: &mut ComputeGraphInner) { + let device = graph.device(); + let config = device.config(); + let recognized = !config + .spike_no_recognition + .is_some_and(|budget| self.execution_graph.node_count() <= budget); + let recognition_start = std::time::Instant::now(); + if recognized { + self.recognize_all(graph); + } + let recognition = recognition_start.elapsed(); + self.optimize_phases.recognition += recognition; + #[cfg(feature = "graphvis")] + if let Some(dir) = &config.dump_stages { + super::visualize::dump_stage( + dir, + &self.execution_graph, + super::visualize::Stage::Recognized, + ); + } + + let extraction_start = std::time::Instant::now(); + let mut driver = EGraphDriver::ingest(self, graph); + if config.spike_hoisting { + let analysis = &driver.egraph.analysis; + tracing::info!( + "hoisting_spike_ingest recognized={} exec_nodes={} recognition_us={} ingest_us={} provs={} enodes={} classes={} payloads={} specs={}", + recognized, + self.execution_graph.node_count(), + recognition.as_micros(), + extraction_start.elapsed().as_micros(), + analysis.facts.len(), + driver.egraph.total_number_of_nodes(), + driver.egraph.number_of_classes(), + analysis.payloads.payload_count(), + analysis.payloads.spec_count(), + ); + } + let extraction = { + let ctx = rules_fuse::FusionCtx { + graph, + layouts: std::cell::RefCell::new(Default::default()), + }; + driver.extract_with_fusion(&ctx) + }; + driver.egraph.rebuild(); + driver.refresh_prov_classes(); + self.apply_egraph_deltas(graph, &driver, &extraction); + self.coalesce_equivalent_eclasses(graph, &driver); + // After extraction, like coalescing: both leave observations behind, + // and a later rewrite that rewired past one would strand it. + self.sink_unary_chains_into_matmuls(graph); + self.optimize_phases.extraction += extraction_start.elapsed(); + #[cfg(feature = "graphvis")] + if let Some(dir) = &config.dump_stages { + super::visualize::dump_stage( + dir, + &self.execution_graph, + super::visualize::Stage::Extracted, + ); + } + } + + /// Every matcher that claims a composed cluster, in nesting order: + /// attention reads recognized contractions, and the assign chains close + /// over whatever the earlier passes left. + pub(super) fn recognize_all(&mut self, graph: &mut ComputeGraphInner) { + self.recognize_contractions(graph); + self.recognize_embeddings(graph); + self.recognize_attention(graph); + self.fuse_row_programs(graph); + self.recognize_assign_chains(graph); + } + + /// Collapse execution nodes that ingestion places in the same semantic + /// e-class. The representative performs the work; every other inner + /// `NodeIndex` remains an observation of that result in `shared_outputs`. + fn coalesce_equivalent_eclasses( + &mut self, + graph: &mut ComputeGraphInner, + driver: &EGraphDriver, + ) { + let groups: Vec> = driver.provs_of_class.values().cloned().collect(); + // Removing a duplicate can leave its dependencies dead. Do not prune + // those dependencies until every e-class from this ingestion snapshot + // has been coalesced: a dead dependency may itself be a duplicate in a + // later group, whose snapshotted `facts.exec` must remain valid long + // enough to register all of its shared output observations. + let mut potentially_dead = Vec::new(); + for group in groups { + let executions: Vec<_> = group + .into_iter() + .filter_map(|prov| { + let facts = driver.egraph.analysis.facts_of(prov); + let exec = facts.exec?; + let variant = &self.execution_graph.node_weight(exec)?.variant; + matches!( + variant, + ExecutionVariant::Elementwise(_) + | ExecutionVariant::Reduce(_) + | ExecutionVariant::View(_) + | ExecutionVariant::MatMul(_) + | ExecutionVariant::QMatMul(_) + | ExecutionVariant::QEmbedding(_) + | ExecutionVariant::RowProgram(_) + ) + .then_some(exec) + }) + .collect(); + let Some((&representative, duplicates)) = executions.split_first() else { + continue; + }; + let representative_inner = self.execution_graph[representative].inner_idx; + for &duplicate in duplicates { + if !self.execution_graph.contains_node(duplicate) { + continue; + } + let duplicate_inner = self.execution_graph[duplicate].inner_idx; + let consumers: Vec<_> = self + .execution_graph + .neighbors_directed(duplicate, petgraph::Direction::Outgoing) + .collect(); + let dependencies: Vec<_> = self + .execution_graph + .neighbors_directed(duplicate, petgraph::Direction::Incoming) + .collect(); + for consumer in consumers { + if consumer != representative + && self + .execution_graph + .find_edge(representative, consumer) + .is_none() + { + self.execution_graph.add_edge(representative, consumer, ()); + } + } + self.execution_graph.remove_node(duplicate); + self.node_mapping.remove(&duplicate_inner); + self.shared_outputs + .entry(representative_inner) + .or_default() + .push(duplicate_inner); + graph.add_dependency_edge(representative_inner, duplicate_inner); + if let Some(recorder) = &self.recorder { + recorder + .borrow_mut() + .record_physical_edge(representative_inner, duplicate_inner); + } + potentially_dead.extend(dependencies); + } + } + for dependency in potentially_dead { + self.remove_node_if_dead(dependency); + } + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/rules_fuse.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/rules_fuse.rs new file mode 100644 index 000000000..9ecee58bb --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/rules_fuse.rs @@ -0,0 +1,593 @@ +//! Fusion generators: per-node view folding, nary +//! inlining, the reduce-fusion family, matmul/qmatmul epilogues), consulted +//! by the extraction worklist with live consumer counts. +//! +//! Each generator is a pure function from the node's current form (and the +//! evolving selection state, read through [`FusionView`]) to a legal +//! alternative. The extractor compares alternatives with the GPU cost model, +//! commits the winner, and cascades the kills. Gates cover binding budgets, +//! dtype/device capabilities and duplication checks against live counts. + +use std::cell::RefCell; + +use rustc_hash::FxHashSet; + +use super::super::ExecutionVariant; +use super::EGraphDriver; +use super::compose; +use super::extract::{ExtractState, Selection}; +use super::lang::Prov; +use crate::compute_graph::layout_pass::LayoutPass; +use crate::compute_graph::{ComputeGraphInner, NodeIndex}; +use crate::nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction, UnaryFunctionChain}; +use crate::{DataTypeEnum, Layout}; + +/// Where the two producer-inlining rewrites disagree; everything else about +/// them is shared. +struct InlineGate { + /// A producer that materializes anyway is left alone: inlining it into an + /// elementwise consumer duplicates its compute. Reduce consumers never see + /// one, because region formation claims it first. + skip_externally_live: bool, + /// Substitute directly only when the producer spans the consumer's index + /// space. A reduce's index space includes the reduced axis, so a + /// differently-shaped producer must come in through the composed path. + require_same_index_space: bool, +} + +pub(super) struct FusionCtx<'a> { + pub(super) graph: &'a ComputeGraphInner, + /// Memoized layout inference over the (immutable) inner graph, used by + /// qmatmul extra normalization. Fresh per stage; recomputation is + /// correctness-neutral. + pub(super) layouts: RefCell, +} + +/// Read-only view of the evolving optimization state, mirroring exactly what +/// the destructive fixpoint reads from the execution graph mid-rewrite. +pub(super) struct FusionView<'a> { + driver: &'a EGraphDriver, + state: &'a ExtractState, + ctx: &'a FusionCtx<'a>, +} + +impl<'a> FusionView<'a> { + pub(super) fn new( + driver: &'a EGraphDriver, + state: &'a ExtractState, + ctx: &'a FusionCtx<'a>, + ) -> Self { + Self { driver, state, ctx } + } + + /// The current form of an inner node: its selection's payload. `None` + /// mirrors every case where the destructive code bails — not in the + /// execution graph (cached boundary / tensor input handled separately), + /// or already killed. + pub(super) fn variant_of(&self, inner: NodeIndex) -> Option<&ExecutionVariant> { + let prov = *self.driver.prov_of.get(&inner)?; + let facts = self.driver.egraph.analysis.facts_of(prov); + facts.exec?; + if !self.state.needed[prov.0 as usize] { + return None; + } + match &self.state.sel[prov.0 as usize] { + Selection::Identity => self.driver.identity_variant(prov), + Selection::Alt(enode) => enode + .payload() + .map(|payload| self.driver.egraph.analysis.payloads.get(payload)), + } + } + + /// `check_cached` equivalent: the node was cached when the resolve + /// started (ingested as an opaque boundary). + pub(super) fn is_cached(&self, inner: NodeIndex) -> bool { + self.driver + .prov_of + .get(&inner) + .is_some_and(|&prov| self.driver.egraph.analysis.facts_of(prov).exec.is_none()) + } + + /// Live consumer count (the destructive `neighbors_directed(..).count()`). + pub(super) fn consumer_count(&self, inner: NodeIndex) -> u32 { + self.driver + .prov_of + .get(&inner) + .map(|&prov| self.state.reads[prov.0 as usize]) + .unwrap_or(0) + } + + fn externally_live(&self, inner: NodeIndex) -> bool { + self.ctx + .graph + .nodes + .nodes + .node_weight(inner) + .is_some_and(|node| node.reference_count > 0) + } + + pub(super) fn device(&self) -> crate::Device { + self.ctx.graph.device() + } + + /// [`compose::walk_view_chain`] over current selections. + pub(super) fn walk_view_chain(&self, inner: NodeIndex) -> (NodeIndex, Option) { + compose::walk_view_chain(inner, |inner| { + let ExecutionVariant::View(view) = self.variant_of(inner)? else { + return None; + }; + Some((view.composed_layout()?, view.input)) + }) + } + + pub(super) fn layout_of(&self, inner: NodeIndex) -> Option { + let mut layouts = self.ctx.layouts.borrow_mut(); + layouts.visit(self.ctx.graph, inner); + layouts.output_layout.get(&inner).cloned() + } + + /// Look through a qmatmul epilogue operand's view chain to the + /// contiguous f32 producer the epilogue can index directly. + pub(super) fn normalize_qmatmul_post_extra( + &self, + extra_inner: NodeIndex, + output_shape: &[usize], + ) -> Option { + let last_dim = *output_shape.last()?; + let extra_info = self.layout_of(extra_inner)?; + if extra_info.datatype() != DataTypeEnum::F32 || extra_info.layout().shape() != output_shape + { + return None; + } + let layout = extra_info.layout(); + let is_column_broadcast = layout.offset() == 0 + && layout.strides().last().copied() == Some(1) + && layout.shape().last().copied() == Some(last_dim) + && layout.strides()[..layout.strides().len().saturating_sub(1)] + .iter() + .all(|stride| *stride == 0); + if !is_column_broadcast { + return Some(extra_inner); + } + let (base_inner, _) = self.walk_view_chain(extra_inner); + let base_info = self.layout_of(base_inner)?; + let base_layout = base_info.layout(); + if base_info.datatype() == DataTypeEnum::F32 + && base_layout.shape() == [last_dim] + && base_layout.is_contiguous() + && base_layout.offset() == 0 + { + Some(base_inner) + } else { + Some(extra_inner) + } + } + + /// Whether the selected operation can participate in any fusion family. + pub(super) fn is_seed_candidate(&self, prov: Prov) -> bool { + let facts = self.driver.egraph.analysis.facts_of(prov); + if facts.exec.is_none() || !self.state.needed[prov.0 as usize] { + return false; + } + let Some(variant) = self.variant_of(facts.inner) else { + return false; + }; + matches!( + variant, + ExecutionVariant::Elementwise(_) + | ExecutionVariant::MatMul(_) + | ExecutionVariant::QMatMul(_) + | ExecutionVariant::Reduce(_) + ) + } + + /// `enqueue_downstream_candidates` transcription: enqueue candidates + /// reachable from `seeds`, descending through view nodes. + pub(super) fn enqueue_downstream( + &self, + state: &ExtractState, + seeds: impl IntoIterator, + worklist: &mut std::collections::VecDeque, + queued: &mut [bool], + ) { + let mut stack: Vec = seeds.into_iter().collect(); + let mut visited = FxHashSet::default(); + while let Some(prov) = stack.pop() { + if !state.needed[prov as usize] || !visited.insert(prov) { + continue; + } + if self.is_seed_candidate(Prov(prov)) { + if !queued[prov as usize] { + queued[prov as usize] = true; + worklist.push_back(prov); + } + } else { + let facts = &self.driver.egraph.analysis.facts[prov as usize]; + if matches!( + self.variant_of(facts.inner), + Some(ExecutionVariant::View(_)) + ) { + stack.extend(state.consumers[prov as usize].iter().copied()); + } + } + } + } + + /// All immediately legal alternatives in deterministic tie-break order: + /// fold views, fuse naries, reduce-fusion family, then matmul fusion. + /// Extraction compares their GPU costs and re-enqueues the winner, so + /// chained rewrites happen on later pops after that form becomes current. + pub(super) fn generate_candidates(&self, prov: Prov) -> Vec { + let facts = self.driver.egraph.analysis.facts_of(prov); + let Some(current) = self.variant_of(facts.inner).cloned() else { + return Vec::new(); + }; + let mut candidates = Vec::new(); + match ¤t { + ExecutionVariant::Elementwise(nary) => { + if let Some(folded) = self.gen_fold_views_elementwise(nary) { + candidates.push(folded); + } + if let Some(fused) = self.gen_fuse_naries(nary) { + candidates.push(fused); + } + } + ExecutionVariant::Reduce(_) => {} + _ => {} + } + candidates.extend(self.gen_fuse_reduce_candidates(¤t)); + if let Some(matmul) = self.gen_fuse_into_matmul(¤t) { + candidates.push(matmul); + } + candidates + } + + /// Fold view producers of this nary's inputs into its index expressions. + fn gen_fold_views_elementwise(&self, nary: &ElementwiseOperation) -> Option { + let (final_inputs, final_expression) = + self.fold_view_inputs(&nary.inputs, &nary.expression, &nary.shape)?; + Some(ExecutionVariant::Elementwise(ElementwiseOperation { + inputs: final_inputs, + expression: final_expression, + shape: nary.shape.clone(), + output_datatype: nary.output_datatype, + })) + } + + /// Shared body of the two view-folding rewrites (elementwise + reduce). + fn fold_view_inputs( + &self, + inputs: &[NodeIndex], + expression: &NaryExpr, + shape: &[usize], + ) -> Option<(Vec, NaryExpr)> { + let mut expression = expression.clone(); + let mut inputs = inputs.to_vec(); + let mut folded = false; + for (slot, input_inner) in inputs.clone().into_iter().enumerate() { + if self.is_cached(input_inner) { + continue; + } + let Some(ExecutionVariant::View(view)) = self.variant_of(input_inner) else { + continue; + }; + let needs_delinearize = view.stages.iter().any(|stage| { + crate::view::affine_dim_indices(&stage.layout, &stage.input_shape).is_none() + }); + if needs_delinearize && compose::input_reread_factor(&expression, shape, slot) > 1 { + continue; + } + let view = view.clone(); + let Some(rewritten) = compose::rewrite_view_input(&expression, slot, &view) else { + continue; + }; + expression = rewritten; + inputs[slot] = view.input; + folded = true; + } + if !folded { + return None; + } + Some(compose::deduplicate_inputs(inputs, expression)) + } + + /// Inline every sole-consumed elementwise producer into this nary, + /// within the direct-input binding budget. + fn gen_fuse_naries(&self, nary: &ElementwiseOperation) -> Option { + let (final_inputs, final_expression) = self.inline_producers( + &nary.inputs, + &nary.expression, + &nary.shape, + InlineGate { + skip_externally_live: true, + require_same_index_space: false, + }, + )?; + Some(ExecutionVariant::Elementwise(ElementwiseOperation { + inputs: final_inputs, + expression: final_expression, + shape: nary.shape.clone(), + output_datatype: nary.output_datatype, + })) + } + + /// Shared body of the two producer-inlining rewrites: substitute every + /// sole-consumed elementwise producer into `expression`, directly where + /// the read is element-wise and by composing the producer with the read's + /// coordinates otherwise. `None` when nothing inlined. + fn inline_producers( + &self, + inputs: &[NodeIndex], + expression: &NaryExpr, + shape: &[usize], + gate: InlineGate, + ) -> Option<(Vec, NaryExpr)> { + let mut expression = expression.clone(); + let mut all_inputs = inputs.to_vec(); + let mut fused_any = false; + let max_fused_inputs = self.device().nary_direct_input_binding_budget(); + + for &input_inner in inputs.iter() { + if self.is_cached(input_inner) { + continue; + } + // An externally live producer materializes regardless, so + // inlining it here would duplicate its compute. Region formation + // fuses it with consumers and emits it as another output. + if gate.skip_externally_live && self.externally_live(input_inner) { + continue; + } + // Inlining duplicates the producer's work unless this node is + // its only consumer. A user-held reference alone doesn't block + // fusion — only another consumer in this resolve does. + if self.consumer_count(input_inner) != 1 { + continue; + } + let Some(ExecutionVariant::Elementwise(input_nary)) = self.variant_of(input_inner) + else { + continue; + }; + let target_slots: Vec = all_inputs + .iter() + .enumerate() + .filter_map(|(slot, value)| (*value == input_inner).then_some(slot)) + .collect(); + let offset = all_inputs.len(); + let inlined = compose::offset_input_indices(&input_nary.expression, offset); + let mut new_expression = expression.clone(); + let mut success = !gate.require_same_index_space || input_nary.shape.as_ref() == shape; + if success { + for slot in &target_slots { + let (next, s) = + compose::substitute_input_in_expr(&new_expression, *slot, &inlined); + new_expression = next; + success &= s; + } + } + if !success + && target_slots + .iter() + .all(|&slot| compose::input_reread_factor(&expression, shape, slot) == 1) + { + let mut composed = expression.clone(); + success = true; + for slot in &target_slots { + match compose::substitute_input_composed(&composed, *slot, &inlined) { + Some(next) => composed = next, + None => { + success = false; + break; + } + } + } + if success { + new_expression = composed; + } + } + + if success { + let unique_inputs: FxHashSet<_> = all_inputs + .iter() + .chain(input_nary.inputs.iter()) + .copied() + .collect(); + if unique_inputs.len() > max_fused_inputs { + continue; + } + expression = new_expression; + all_inputs.extend(input_nary.inputs.iter().copied()); + fused_any = true; + } + } + if !fused_any { + return None; + } + Some(compose::deduplicate_inputs(all_inputs, expression)) + } + + /// All legal reduce-fusion alternatives. Extraction's cost and + /// duplication constraints decide whether any candidate commits. + fn gen_fuse_reduce_candidates(&self, current: &ExecutionVariant) -> Vec { + let mut candidates = Vec::new(); + candidates.extend(self.gen_collapse_unit_reduce(current)); + candidates.extend(self.gen_fold_views_into_reduce(current)); + candidates.extend(self.gen_unary_into_reduce(current)); + candidates.extend(self.gen_indexed_unary_into_reduce(current)); + candidates.extend(self.gen_producer_into_reduce(current)); + candidates + } + + /// Append a consumer's unary chain to its reduce producer's epilogue. + fn gen_unary_into_reduce(&self, current: &ExecutionVariant) -> Option { + let el_op = compose::try_get_unary_chain(current)?; + let input_inner = el_op.value; + if self.is_cached(input_inner) { + return None; + } + let ExecutionVariant::Reduce(reduce_op) = self.variant_of(input_inner)? else { + return None; + }; + let mut new_reduce = reduce_op.clone(); + let mut existing_post = new_reduce.post_element_wise.functions.clone(); + existing_post.extend(el_op.functions.functions.iter().cloned()); + new_reduce.post_element_wise = + UnaryFunctionChain::new(existing_post, reduce_op.post_element_wise.input_datatype()); + Some(ExecutionVariant::Reduce(new_reduce)) + } + + /// Rewrite a reduce over a unit axis as the equivalent elementwise. + fn gen_collapse_unit_reduce(&self, current: &ExecutionVariant) -> Option { + let ExecutionVariant::Reduce(reduce) = current else { + return None; + }; + if reduce.shape[reduce.axis] != 1 + || !reduce.post_element_wise.functions.is_empty() + || reduce.function.datatype() != reduce.out_datatype() + { + return None; + } + let mut mapping = Vec::with_capacity(reduce.shape.len()); + let mut out_pos = 0; + for dim in 0..reduce.shape.len() { + if dim == reduce.axis { + mapping.push(NaryExpr::Scalar(crate::nary_wise::NaryScalar::U32(0))); + } else { + mapping.push(NaryExpr::DimIndex(out_pos)); + out_pos += 1; + } + } + let expression = compose::compose_expr_with_indices(&reduce.expression, &mapping)?; + use crate::reduce::ReduceOp; + let init = reduce.function.initial_value; + let fold_op = match reduce.function.op { + ReduceOp::Sum => crate::nary_wise::NaryOp::AddConst(init), + ReduceOp::Product => crate::nary_wise::NaryOp::MulConst(init), + ReduceOp::Max => crate::nary_wise::NaryOp::MaxConst(init), + ReduceOp::Min => crate::nary_wise::NaryOp::MinConst(init), + }; + let dtype = reduce.function.datatype(); + let expression = NaryExpr::Op { + children: vec![expression], + function: NaryFunction::unary( + Some(format!("unit_{}", reduce.function.name())), + fold_op, + dtype, + dtype, + ), + }; + Some(ExecutionVariant::Elementwise(ElementwiseOperation { + inputs: reduce.inputs.clone(), + expression, + shape: reduce.out_shape().into(), + output_datatype: reduce.out_datatype(), + })) + } + + /// Fold view producers of this reduce's inputs into its index + /// expressions. + fn gen_fold_views_into_reduce(&self, current: &ExecutionVariant) -> Option { + let ExecutionVariant::Reduce(reduce) = current else { + return None; + }; + let (final_inputs, final_expression) = + self.fold_view_inputs(&reduce.inputs, &reduce.expression, &reduce.shape)?; + let mut new_reduce = reduce.clone(); + new_reduce.inputs = final_inputs; + new_reduce.expression = final_expression; + Some(ExecutionVariant::Reduce(new_reduce)) + } + + /// Inline a reduce producer read through an index expression into the + /// consuming nary, turning it into a reduce over the outer axis. + fn gen_indexed_unary_into_reduce( + &self, + current: &ExecutionVariant, + ) -> Option { + let ExecutionVariant::Elementwise(nary) = current else { + return None; + }; + if nary.inputs.len() != 1 { + return None; + } + let (functions, indices) = compose::extract_unary_chain_indexed(nary)?; + let input_inner = nary.inputs[0]; + if self.is_cached(input_inner) { + return None; + } + let ExecutionVariant::Reduce(reduce) = self.variant_of(input_inner)? else { + return None; + }; + let rows: usize = reduce + .shape + .iter() + .enumerate() + .filter_map(|(dim, &size)| (dim != reduce.axis).then_some(size)) + .product(); + if nary.shape.iter().product::() != rows || indices.len() + 1 != reduce.shape.len() { + return None; + } + let mut current_dtype = reduce.out_datatype(); + for function in &functions { + if function.input_types.as_slice() != [current_dtype] { + return None; + } + current_dtype = function.output_type; + } + if current_dtype != nary.output_datatype { + return None; + } + let node_rank = nary.shape.len(); + let mut mapping = Vec::with_capacity(reduce.shape.len()); + let mut out_pos = 0; + for dim in 0..reduce.shape.len() { + if dim == reduce.axis { + mapping.push(NaryExpr::DimIndex(node_rank)); + } else { + mapping.push(indices[out_pos].clone()); + out_pos += 1; + } + } + let expression = compose::compose_expr_with_indices(&reduce.expression, &mapping)?; + let mut shape: Vec = nary.shape.to_vec(); + shape.push(reduce.shape[reduce.axis]); + let mut post = reduce.post_element_wise.functions.clone(); + post.extend(functions); + Some(ExecutionVariant::Reduce(crate::reduce::ReduceOperation { + inputs: reduce.inputs.clone(), + expression, + shape: shape.into(), + function: reduce.function.clone(), + post_element_wise: UnaryFunctionChain::new( + post, + reduce.post_element_wise.input_datatype(), + ), + axis: node_rank, + })) + } + + /// Inline a sole-consumed elementwise producer into this reduce's + /// expression. + fn gen_producer_into_reduce(&self, current: &ExecutionVariant) -> Option { + let ExecutionVariant::Reduce(reduce) = current else { + return None; + }; + let (final_inputs, final_expression) = self.inline_producers( + &reduce.inputs, + &reduce.expression, + &reduce.shape, + InlineGate { + skip_externally_live: false, + require_same_index_space: true, + }, + )?; + let mut new_reduce = reduce.clone(); + new_reduce.inputs = final_inputs; + new_reduce.expression = final_expression; + Some(ExecutionVariant::Reduce(new_reduce)) + } + + /// The matmul/qmatmul epilogue family; see `rules_fuse_matmul.rs`. + fn gen_fuse_into_matmul(&self, current: &ExecutionVariant) -> Option { + self.gen_matmul_family(current) + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/rules_fuse_matmul.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/rules_fuse_matmul.rs new file mode 100644 index 000000000..7c0334e5d --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/rules_fuse_matmul.rs @@ -0,0 +1,785 @@ +//! Matmul/qmatmul fusion generators, consulted by the extraction worklist. +//! +//! Each branch reads graph state through [`FusionView`] and returns the new +//! variant for the node being rewritten; the extractor's switch/kill +//! machinery is the commit. The accumulator-offset epilogue's own expression +//! rules ride on the shared [`compose`] walk at the end of this file. + +use rustc_hash::FxHashMap; + +use super::super::{ExecutionVariant, Resolver}; +use super::compose; +use super::rules_fuse::FusionView; +use crate::Layout; +use crate::compute_graph::NodeIndex; +use crate::nary_wise::{ElementwiseOperation, NaryExpr, NaryOp, NaryScalar, UnaryFunctionChain}; +use crate::quantized::matmul::{ElementwiseEpilogue, QMatMulOperation}; + +impl FusionView<'_> { + /// Dense matmul post unary chains, qmatmul narrow-accumulator, indexed + /// post, general elementwise post, qmatmul pre epilogues, and dense + /// matmul pre unary chains, in attempt order. Returns the new variant + /// for the node being rewritten (first success wins). + pub(super) fn gen_matmul_family(&self, current: &ExecutionVariant) -> Option { + // Post-op: fuse elementwise after matmul (dense or quantized). + if let Some(el_op) = compose::try_get_unary_chain(current) { + let input_inner = el_op.value; + if !self.is_cached(input_inner) + && let Some(input_variant) = self.variant_of(input_inner) + { + // Dtype-preserving unary chains are hosted after the + // cooperative store, independently of how A/B are mapped. + // Unsupported chains still lower through the generic fused + // reduction. + if let ExecutionVariant::MatMul(matmul_op) = input_variant { + let mut new_matmul = matmul_op.clone(); + let mut existing_post = new_matmul.post_element_wise.functions.clone(); + existing_post.extend(el_op.functions.functions.iter().cloned()); + new_matmul.post_element_wise = UnaryFunctionChain::new( + existing_post, + matmul_op.post_element_wise.input_datatype(), + ); + return Some(ExecutionVariant::MatMul(new_matmul)); + } + } + } + + // Post-op (QMatMul): fuse a general element-wise expression after + // qmatmul. This handles composite expressions like GELU and ordered + // extra inputs whose layouts match the output visitation shape. + if let ExecutionVariant::Elementwise(nary) = current { + // Split/gate expressions built from `narrow` views of a qmatmul + // output (e.g. SwiGLU's gate/up halves) reach the qmatmul through + // MapLayout chains with distinct last-dimension column offsets. + // Absorb them into the accumulator-offset post epilogue before the + // per-input scan below. + if let Some(fused) = self.gen_fuse_qmatmul_narrow_accumulators(nary) { + return Some(fused); + } + for (candidate_input_idx, &input_inner) in nary.inputs.iter().enumerate() { + if self.variant_of(input_inner).is_none() { + continue; + } + let (qmatmul_inner, map_chain) = self.walk_view_chain(input_inner); + let Some(ExecutionVariant::QMatMul(qmatmul_op)) = self.variant_of(qmatmul_inner) + else { + continue; + }; + let qmatmul_op = qmatmul_op.clone(); + if map_chain.is_none() + && !self.is_cached(input_inner) + && qmatmul_op.post_element_wise_expr.is_none() + && qmatmul_op.in_shape[..qmatmul_op.in_shape.len() - 1] + .iter() + .product::() + == 1 + && let Some((expression, accumulator_offsets, extras)) = self + .try_extract_indexed_qmatmul_post_expr( + nary, + candidate_input_idx, + &qmatmul_op.out_shape, + ) + { + let Some(input_datatype) = nary + .expression + .elementwise_input_datatype(candidate_input_idx) + else { + continue; + }; + if input_datatype != crate::DataTypeEnum::F32 + || nary.output_datatype != crate::DataTypeEnum::F32 + { + continue; + } + if !qmatmul_op.supports_indexed_post_accumulator_offsets( + &self.device(), + &nary.shape, + &accumulator_offsets, + ) { + continue; + } + + let post_element_wise_expr = ElementwiseEpilogue { + expression, + extras: extras.clone(), + input_datatype, + output_datatype: nary.output_datatype, + }; + + let mut new_q = qmatmul_op.clone(); + new_q.out_shape = nary.shape.clone(); + new_q.post_element_wise_expr = Some(post_element_wise_expr); + new_q.post_accumulator_offsets = accumulator_offsets.into_boxed_slice(); + + if !new_q.fits_binding_budget(&self.device()) { + continue; + } + + return Some(ExecutionVariant::QMatMul(new_q)); + } + let Some(mapped_layout) = Resolver::apply_view_chain( + &Layout::contiguous(&qmatmul_op.out_shape), + &map_chain, + ) else { + continue; + }; + if mapped_layout != Layout::contiguous(&nary.shape) { + continue; + } + if !nary.expression.uses_input(candidate_input_idx) + || nary + .expression + .uses_custom_indexing_for_input(candidate_input_idx) + { + continue; + }; + let Some(input_datatype) = nary + .expression + .elementwise_input_datatype(candidate_input_idx) + else { + continue; + }; + let mut extras = Vec::new(); + let mut replacements = vec![None; nary.inputs.len()]; + let mut valid_expression = true; + for (input_idx, &nary_input) in nary.inputs.iter().enumerate() { + let (base_inner, chain) = self.walk_view_chain(nary_input); + let base_qmatmul = match self.variant_of(base_inner) { + Some(ExecutionVariant::QMatMul(op)) => Some(op.clone()), + _ => None, + }; + if let Some(base_qmatmul) = base_qmatmul + && qmatmul_same_base(&qmatmul_op, &base_qmatmul) + { + let alias_layout = Resolver::apply_view_chain( + &Layout::contiguous(&base_qmatmul.out_shape), + &chain, + ); + if alias_layout == Some(Layout::contiguous(&nary.shape)) + && !nary.expression.uses_custom_indexing_for_input(input_idx) + { + replacements[input_idx] = + qmatmul_output_expr(&base_qmatmul, &mut extras, nary.shape.len()); + continue; + } + valid_expression = false; + break; + } + + let Some(extra) = self.normalize_qmatmul_post_extra(nary_input, &nary.shape) + else { + valid_expression = false; + break; + }; + replacements[input_idx] = + Some(NaryExpr::input(extras.len() + 1, nary.shape.len())); + extras.push(extra); + } + if !valid_expression { + continue; + } + let Some(expression) = + compose::replace_inputs_in_expr(&nary.expression, &replacements) + else { + continue; + }; + if self.is_cached(input_inner) + || input_datatype != crate::DataTypeEnum::F32 + || nary.output_datatype != crate::DataTypeEnum::F32 + || !qmatmul_op.supports_elementwise_epilogue_fusion(&self.device()) + { + continue; + } + + let post_element_wise_expr = ElementwiseEpilogue { + expression, + extras: extras.clone(), + input_datatype: qmatmul_op + .post_element_wise_expr + .as_ref() + .map(|existing| existing.input_datatype) + .unwrap_or(input_datatype), + output_datatype: nary.output_datatype, + }; + + let mut new_q = qmatmul_op.clone(); + new_q.post_element_wise_expr = Some(post_element_wise_expr); + + if !new_q.fits_binding_budget(&self.device()) { + continue; + } + + return Some(ExecutionVariant::QMatMul(new_q)); + } + } + + // Pre-op (QMatMul): fuse a general element-wise expression upstream + // of a single-row qmatmul input. For batched/tiled qmatmul, the + // transformed activation tile is reloaded for each output-column + // tile, so expensive expressions like GELU would be recomputed many + // times. Keep those chains materialized once instead. + if let ExecutionVariant::QMatMul(qmatmul_op) = current + && qmatmul_op.in_shape[..qmatmul_op.in_shape.len() - 1] + .iter() + .product::() + == 1 + && qmatmul_op.supports_elementwise_epilogue_fusion(&self.device()) + && !self.is_cached(qmatmul_op.input) + && self.variant_of(qmatmul_op.input).is_some() + { + let (nary_inner, nary_map_chain) = self.walk_view_chain(qmatmul_op.input); + let Some(ExecutionVariant::Elementwise(nary)) = self.variant_of(nary_inner) else { + return None; + }; + let nary = nary.clone(); + let mapped_layout = + Resolver::apply_view_chain(&Layout::contiguous(&nary.shape), &nary_map_chain); + if mapped_layout != Some(Layout::contiguous(&qmatmul_op.in_shape)) { + return None; + } + + for (candidate_input_idx, &primary_input) in nary.inputs.iter().enumerate() { + if !nary.expression.uses_input(candidate_input_idx) + || nary + .expression + .uses_custom_indexing_for_input(candidate_input_idx) + { + continue; + } + let Some(input_datatype) = nary + .expression + .elementwise_input_datatype(candidate_input_idx) + else { + continue; + }; + if input_datatype != crate::DataTypeEnum::F32 + || nary.output_datatype != crate::DataTypeEnum::F32 + { + continue; + } + + let (primary_inner, primary_chain) = self.walk_view_chain(primary_input); + let Some(primary_info) = self.layout_of(primary_inner) else { + continue; + }; + let Some(primary_layout) = + Resolver::apply_view_chain(primary_info.layout(), &primary_chain) + else { + continue; + }; + if primary_layout != Layout::contiguous(&nary.shape) { + continue; + } + + let mut mapping = vec![usize::MAX; nary.inputs.len()]; + let mut extras = Vec::new(); + let mut valid_expression = true; + for (input_idx, &nary_input) in nary.inputs.iter().enumerate() { + let (base_inner, chain) = self.walk_view_chain(nary_input); + if base_inner == primary_inner { + let alias_layout = + Resolver::apply_view_chain(primary_info.layout(), &chain); + if alias_layout == Some(Layout::contiguous(&nary.shape)) + && !nary.expression.uses_custom_indexing_for_input(input_idx) + { + mapping[input_idx] = 0; + continue; + } + valid_expression = false; + break; + } + + let Some(extra) = self.normalize_qmatmul_post_extra(nary_input, &nary.shape) + else { + valid_expression = false; + break; + }; + mapping[input_idx] = extras.len() + 1; + extras.push(extra); + } + if !valid_expression { + continue; + } + let expression = compose::remap_inputs(&nary.expression, &mapping); + + let pre_element_wise_expr = if let Some(existing) = + &qmatmul_op.pre_element_wise_expr + { + if existing.input_datatype != nary.output_datatype { + continue; + } + let mut mapping = Vec::with_capacity(1 + existing.extras.len()); + mapping.push(0); + mapping.extend((0..existing.extras.len()).map(|i| i + 1 + extras.len())); + let shifted_existing = compose::remap_inputs(&existing.expression, &mapping); + let (expression, success) = + compose::substitute_input_in_expr(&shifted_existing, 0, &expression); + if !success { + continue; + } + let mut combined_extras = extras.clone(); + combined_extras.extend(existing.extras.clone()); + ElementwiseEpilogue { + expression, + extras: combined_extras, + input_datatype, + output_datatype: existing.output_datatype, + } + } else { + ElementwiseEpilogue { + expression, + extras: extras.clone(), + input_datatype, + output_datatype: nary.output_datatype, + } + }; + + let mut new_q = qmatmul_op.clone(); + new_q.input = primary_inner; + new_q.pre_element_wise_expr = Some(pre_element_wise_expr); + + if !new_q.fits_binding_budget(&self.device()) { + continue; + } + + return Some(ExecutionVariant::QMatMul(new_q)); + } + } + + // Pre-op: fuse elementwise before plain matmul inputs. Cooperative + // matmuls apply dtype-preserving chains while staging A/B; other + // chains lower through the generic fused reduction. Un-flattened + // operands remain excluded because their producer mapping is already + // being absorbed by cooperative staging. + if let ExecutionVariant::MatMul(matmul_op) = current + && matmul_op.a.is_plain() + && matmul_op.b.is_plain() + { + let mut new_matmul = matmul_op.clone(); + let mut changed = false; + + // Check first input + if !self.is_cached(matmul_op.first) + && let Some(first_variant) = self.variant_of(matmul_op.first) + && let Some(el_op) = compose::try_get_unary_chain(first_variant) + { + new_matmul.first = el_op.value; + let mut functions = el_op.functions.functions.clone(); + functions.extend(new_matmul.pre_element_wise[0].functions.iter().cloned()); + new_matmul.pre_element_wise[0] = + UnaryFunctionChain::new(functions, el_op.functions.input_datatype()); + changed = true; + } + + // Check second input + if !self.is_cached(matmul_op.second) + && let Some(second_variant) = self.variant_of(matmul_op.second) + && let Some(el_op) = compose::try_get_unary_chain(second_variant) + { + new_matmul.second = el_op.value; + let mut functions = el_op.functions.functions.clone(); + functions.extend(new_matmul.pre_element_wise[1].functions.iter().cloned()); + new_matmul.pre_element_wise[1] = + UnaryFunctionChain::new(functions, el_op.functions.input_datatype()); + changed = true; + } + + if changed { + return Some(ExecutionVariant::MatMul(new_matmul)); + } + } + + None + } + + /// Absorb a split/gate n-ary whose inputs are + /// `narrow` (MapLayout) views of a single-row qmatmul output into that + /// qmatmul's accumulator-offset post epilogue. Each distinct + /// last-dimension column offset (e.g. the gate half at 0 and the up half + /// at `pair_len`) becomes one accumulator value, so a SwiGLU-style + /// `silu(gate) * up` resolves to a single dynamic qmatmul kernel where + /// the backend supports it. Returns `None` when the pattern, dtype, + /// layout, accumulator offsets, or binding budget are unsupported. + fn gen_fuse_qmatmul_narrow_accumulators( + &self, + nary: &ElementwiseOperation, + ) -> Option { + if nary.output_datatype != crate::DataTypeEnum::F32 { + return None; + } + + // Find the qmatmul reached through a narrow MapLayout view. A direct + // (chain-less) reference is the indexed-input form handled below. + let mut base = None; + for &input in &nary.inputs { + let (base_inner, chain) = self.walk_view_chain(input); + if chain.is_none() { + continue; + } + if let Some(ExecutionVariant::QMatMul(op)) = self.variant_of(base_inner) { + // A qmatmul that already carries a post epilogue isn't a clean + // accumulator-offset base; leave it to the general scan. + if op.post_element_wise_expr.is_some() { + continue; + } + base = Some((base_inner, op.clone())); + break; + } + } + let Some((qmatmul_inner, qmatmul_op)) = base else { + return None; + }; + if self.is_cached(qmatmul_inner) { + return None; + } + + let Some((expression, accumulator_offsets, extras)) = + self.try_extract_mapped_qmatmul_post_expr(nary, qmatmul_inner, &qmatmul_op.out_shape) + else { + return None; + }; + + if !qmatmul_op.supports_indexed_post_accumulator_offsets( + &self.device(), + &nary.shape, + &accumulator_offsets, + ) { + return None; + } + + let post_element_wise_expr = ElementwiseEpilogue { + expression, + extras, + input_datatype: crate::DataTypeEnum::F32, + output_datatype: nary.output_datatype, + }; + + let mut new_q = qmatmul_op; + new_q.out_shape = nary.shape.clone(); + new_q.post_element_wise_expr = Some(post_element_wise_expr); + new_q.post_accumulator_offsets = accumulator_offsets.into_boxed_slice(); + + if !new_q.fits_binding_budget(&self.device()) { + return None; + } + + Some(ExecutionVariant::QMatMul(new_q)) + } + + /// Build the post epilogue expression, + /// accumulator column offsets, and extra-tensor dependencies for an n-ary + /// whose inputs are last-dimension `narrow` views of `qmatmul_inner`. + /// Inputs that view the qmatmul become accumulator values (indices + /// `0..offsets.len()`, deduplicated by column offset); every other input + /// becomes a normalized extra tensor (indices after the accumulators). + /// Returns `None` when an input isn't a clean last-dimension narrow, uses + /// custom indexing, or can't be normalized. + fn try_extract_mapped_qmatmul_post_expr( + &self, + nary: &ElementwiseOperation, + qmatmul_inner: NodeIndex, + qmatmul_out_shape: &[usize], + ) -> Option<(NaryExpr, Vec, Vec)> { + if nary.shape.len() != qmatmul_out_shape.len() { + return None; + } + // The accumulator-offset epilogue is only lowered by the single-row + // qgemv path, so every leading dimension must collapse to one row. + if qmatmul_out_shape[..qmatmul_out_shape.len() - 1] + .iter() + .product::() + != 1 + { + return None; + } + let output_cols = nary.shape.last().copied()? as u32; + let matrix_cols = qmatmul_out_shape.last().copied()? as u32; + // A full-width (or wider) output isn't a split; the general scan owns + // that case. + if output_cols >= matrix_cols { + return None; + } + + let qmatmul_out_layout = Layout::contiguous(qmatmul_out_shape); + let rank = nary.shape.len(); + + enum MappedInput { + Accumulator(usize), + Extra(usize), + } + + let mut accumulator_offsets = Vec::new(); + let mut accumulator_map = FxHashMap::default(); + let mut extras = Vec::new(); + let mut mapped = Vec::with_capacity(nary.inputs.len()); + for (input_idx, &nary_input) in nary.inputs.iter().enumerate() { + if !nary.expression.uses_input(input_idx) { + mapped.push(None); + continue; + } + if nary.expression.uses_custom_indexing_for_input(input_idx) { + return None; + } + let (base_inner, chain) = self.walk_view_chain(nary_input); + if base_inner == qmatmul_inner { + let view = Resolver::apply_view_chain(&qmatmul_out_layout, &chain)?; + let offset = qmatmul_last_dim_view_offset(&view, &nary.shape, matrix_cols)?; + let value_idx = *accumulator_map.entry(offset).or_insert_with(|| { + let idx = accumulator_offsets.len(); + accumulator_offsets.push(offset); + idx + }); + mapped.push(Some(MappedInput::Accumulator(value_idx))); + } else { + let extra = self.normalize_qmatmul_post_extra(nary_input, &nary.shape)?; + let pos = extras.len(); + extras.push(extra); + mapped.push(Some(MappedInput::Extra(pos))); + } + } + + // Two distinct column offsets are the smallest split worth folding into + // the accumulator-offset path; a single offset is either the default + // full-width store or a partial column the qgemv path can't cover. + if accumulator_offsets.len() < 2 { + return None; + } + + let accumulator_count = accumulator_offsets.len(); + let mut replacements = vec![None; nary.inputs.len()]; + for (input_idx, kind) in mapped.into_iter().enumerate() { + match kind { + Some(MappedInput::Accumulator(value_idx)) => { + replacements[input_idx] = Some(NaryExpr::input(value_idx, rank)); + } + Some(MappedInput::Extra(pos)) => { + replacements[input_idx] = Some(NaryExpr::input(accumulator_count + pos, rank)); + } + None => {} + } + } + + let expression = compose::replace_inputs_in_expr(&nary.expression, &replacements)?; + Some((expression, accumulator_offsets, extras)) + } + + /// Build the post epilogue expression and extra-tensor dependencies for + /// an n-ary that reads `qmatmul_inner` through an index expression. + fn try_extract_indexed_qmatmul_post_expr( + &self, + nary: &ElementwiseOperation, + qmatmul_input_idx: usize, + qmatmul_out_shape: &[usize], + ) -> Option<(NaryExpr, Vec, Vec)> { + if nary.output_datatype != crate::DataTypeEnum::F32 + || nary.shape.len() != qmatmul_out_shape.len() + || nary.shape.as_ref() == qmatmul_out_shape + { + return None; + } + let output_cols = nary.shape.last().copied()? as u32; + let matrix_cols = qmatmul_out_shape.last().copied()? as u32; + if output_cols >= matrix_cols { + return None; + } + + let temp_input_base = nary.inputs.len(); + let mut accumulator_offsets = Vec::new(); + let mut accumulator_map = FxHashMap::default(); + let expression = replace_indexed_qmatmul_accumulators( + &nary.expression, + qmatmul_input_idx, + nary.shape.len(), + output_cols, + matrix_cols, + temp_input_base, + &mut accumulator_offsets, + &mut accumulator_map, + )?; + if accumulator_offsets.len() < 2 { + return None; + } + + let mut replacements = vec![None; nary.inputs.len()]; + let mut extras = Vec::new(); + for (input_idx, &input) in nary.inputs.iter().enumerate() { + if input_idx == qmatmul_input_idx || !nary.expression.uses_input(input_idx) { + continue; + } + if nary.expression.uses_custom_indexing_for_input(input_idx) { + return None; + } + let extra = self.normalize_qmatmul_post_extra(input, &nary.shape)?; + replacements[input_idx] = Some(NaryExpr::input( + accumulator_offsets.len() + extras.len(), + nary.shape.len(), + )); + extras.push(extra); + } + + let expression = compose::replace_inputs_in_expr(&expression, &replacements)?; + let expression = + remap_temp_accumulator_inputs(&expression, temp_input_base, accumulator_offsets.len()); + Some((expression, accumulator_offsets, extras)) + } +} + +/// Whether two qmatmuls compute the same accumulators, so a view of one can +/// alias the other's output. +fn qmatmul_same_base(first: &QMatMulOperation, second: &QMatMulOperation) -> bool { + first.input_datatype == second.input_datatype + && first.input == second.input + && first.matrix == second.matrix + && first.in_shape == second.in_shape + && first.out_shape == second.out_shape + && first.pre_element_wise_expr == second.pre_element_wise_expr + && first.post_accumulator_offsets == second.post_accumulator_offsets +} + +/// The expression a qmatmul's output presents to a consumer: its existing +/// post epilogue with the epilogue's own extras appended to `extras`, or a +/// bare read of the accumulator. +fn qmatmul_output_expr( + qmatmul: &QMatMulOperation, + extras: &mut Vec, + rank: usize, +) -> Option { + if let Some(epilogue) = &qmatmul.post_element_wise_expr { + let value_arity = qmatmul.post_accumulator_offsets.len().max(1); + let mut mapping = Vec::with_capacity(value_arity + epilogue.extras.len()); + mapping.extend(0..value_arity); + mapping.extend((0..epilogue.extras.len()).map(|i| extras.len() + value_arity + i)); + extras.extend(epilogue.extras.iter().copied()); + Some(compose::remap_inputs(&epilogue.expression, &mapping)) + } else { + Some(NaryExpr::input(0, rank)) + } +} + +/// If `view` is a contiguous last-dimension narrow of a single-row qmatmul +/// output whose shape matches `output_shape`, return its column offset. +/// Returns `None` for any non-narrow / strided / out-of-range view. +fn qmatmul_last_dim_view_offset( + view: &Layout, + output_shape: &[usize], + matrix_cols: u32, +) -> Option { + if view.shape() != output_shape { + return None; + } + if view.strides().last().copied() != Some(1) { + return None; + } + let offset = u32::try_from(view.offset()).ok()?; + let output_cols = *output_shape.last()? as u32; + if offset.checked_add(output_cols)? > matrix_cols { + return None; + } + Some(offset) +} + +/// Replace every last-dimension-offset read of the qmatmul input with a +/// temporary accumulator slot, one per distinct column offset. +#[allow(clippy::too_many_arguments)] +fn replace_indexed_qmatmul_accumulators( + expr: &NaryExpr, + qmatmul_input_idx: usize, + output_rank: usize, + output_cols: u32, + matrix_cols: u32, + temp_input_base: usize, + accumulator_offsets: &mut Vec, + accumulator_map: &mut FxHashMap, +) -> Option { + compose::rewrite_loads(expr, &mut |input_idx, indices, mapped| { + if input_idx != qmatmul_input_idx { + return Some(NaryExpr::IndexedInput { + input_idx, + indices: mapped, + }); + } + let offset = extract_qmatmul_last_dim_offset(indices, output_rank)?; + if output_cols + .checked_add(offset) + .is_none_or(|cols| cols > matrix_cols) + { + return None; + } + let value_idx = *accumulator_map.entry(offset).or_insert_with(|| { + let value_idx = accumulator_offsets.len(); + accumulator_offsets.push(offset); + value_idx + }); + Some(NaryExpr::input(temp_input_base + value_idx, output_rank)) + }) +} + +fn extract_qmatmul_last_dim_offset(indices: &[NaryExpr], output_rank: usize) -> Option { + if indices.len() != output_rank { + return None; + } + for (dim, index) in indices[..output_rank - 1].iter().enumerate() { + if !matches!(index, NaryExpr::DimIndex(index_dim) if *index_dim == dim) { + return None; + } + } + extract_dim_plus_u32_offset(&indices[output_rank - 1], output_rank - 1) +} + +fn extract_dim_plus_u32_offset(expr: &NaryExpr, dim: usize) -> Option { + match expr { + NaryExpr::DimIndex(index_dim) if *index_dim == dim => Some(0), + NaryExpr::Op { children, function } + if function.op == NaryOp::Add && children.len() == 2 => + { + extract_dim_plus_u32_offset_pair(&children[0], &children[1], dim) + .or_else(|| extract_dim_plus_u32_offset_pair(&children[1], &children[0], dim)) + } + NaryExpr::Op { children, function } + if matches!(function.op, NaryOp::AddConst(NaryScalar::U32(_))) + && children.len() == 1 => + { + let NaryOp::AddConst(NaryScalar::U32(offset)) = function.op else { + unreachable!(); + }; + matches!(&children[0], NaryExpr::DimIndex(index_dim) if *index_dim == dim) + .then_some(offset) + } + _ => None, + } +} + +fn extract_dim_plus_u32_offset_pair( + dim_expr: &NaryExpr, + offset_expr: &NaryExpr, + dim: usize, +) -> Option { + let NaryExpr::DimIndex(index_dim) = dim_expr else { + return None; + }; + if *index_dim != dim { + return None; + } + let NaryExpr::Scalar(NaryScalar::U32(offset)) = offset_expr else { + return None; + }; + Some(*offset) +} + +/// Fold the temporary accumulator slots back onto the epilogue's value +/// inputs, which the qmatmul kernel binds first. +fn remap_temp_accumulator_inputs( + expr: &NaryExpr, + temp_input_base: usize, + accumulator_count: usize, +) -> NaryExpr { + compose::map_loads(expr, &mut |input_idx, _, indices| { + let input_idx = + if (temp_input_base..temp_input_base + accumulator_count).contains(&input_idx) { + input_idx - temp_input_base + } else { + input_idx + }; + NaryExpr::IndexedInput { input_idx, indices } + }) +} diff --git a/fusor-ml/core/src/compute_graph/resolve/egraph/structural_memo.rs b/fusor-ml/core/src/compute_graph/resolve/egraph/structural_memo.rs new file mode 100644 index 000000000..28cfde54f --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/egraph/structural_memo.rs @@ -0,0 +1,1227 @@ +//! Allocation-independent fusion-plan sharing. +//! +//! The value e-graph cannot equate two transformer layers: their activation +//! and weight buffers are different values. Fusion planning can still share +//! allocation-independent structural templates. A compact structural +//! interner canonicalizes those templates without creating a second e-graph; +//! the first occurrence records a rewrite that later occurrences instantiate +//! by rebinding dependency roles and the concrete QMatrix. + +use std::hash::Hash; + +use parking_lot::Mutex; +use rustc_hash::{FxHashMap, FxHashSet}; + +use super::super::ExecutionVariant; +use super::EGraphDriver; +use super::extract::ExtractState; +use super::interner::{ + PayloadKey, SpecId, TwoLane, local_hash, rebind_variant_dependencies, variant_dependencies, +}; +use super::lang::Prov; +use super::rules_fuse::FusionView; +use crate::compute_graph::NodeIndex; +use crate::quantized::QMatrix; +use crate::quantized::embedding::QEmbeddingOperation; +use crate::quantized::matmul::{ElementwiseEpilogue, QMatMulOperation}; +use crate::{DataTypeEnum, FusorConfig}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +struct PlanAtomId(u32); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(super) struct StructuralId(u32); + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct StructuralNode { + atom: PlanAtomId, + children: Box<[StructuralId]>, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct LayoutSpec { + datatype: DataTypeEnum, + offset: usize, + shape: Box<[usize]>, + strides: Box<[usize]>, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum PlanNodeKind { + Operator(SpecId), + Tensor, + Boundary, + Missing, + Frontier, +} + +/// Exact facts a generator is permitted to observe for one role in its +/// local window. Including role numbers preserves aliasing: `[x, x]` and +/// `[x, y]` never share a plan even when x and y have identical layouts. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct PlanAtom { + role: u32, + kind: PlanNodeKind, + /// Canonical local identity of an embedded QMatrix allocation. This + /// preserves the distinction between "two views of the same qmatmul" + /// and "two same-shaped qmatmuls with different weights" while still + /// allowing corresponding allocations in different layers to share. + matrix_alias: Option, + layout: Option, + reads: u32, + cached: bool, + externally_live: bool, + is_target: bool, +} + +pub(super) struct PlanInstance { + pub(super) root: StructuralId, + nodes: Vec, + role_of: FxHashMap, +} + +#[derive(Clone)] +pub(super) struct VariantTemplate { + variant: ExecutionVariant, + dependency_roles: Vec, + matrix_role: Option, + spec: Option, +} + +#[derive(Clone)] +pub(super) enum PlanDecision { + NoRewrite, + Rewrite(VariantTemplate), +} + +pub(super) enum PlanLookup { + Miss, + Hit(Option), +} + +#[derive(Debug, Clone, Copy, Default)] +pub(super) struct PlanSharingStats { + pub(super) windows: u64, + pub(super) unique_windows: u64, + pub(super) hits: u64, + pub(super) misses: u64, + pub(super) templates: u64, + pub(super) negative_templates: u64, + /// Rewrites reading past the window horizon, which plan fresh every + /// visit ([`VariantTemplate::capture`]). + pub(super) unshareable: u64, + /// Per-resolve misses answered by the device-scoped [`FusionPlanStore`]. + pub(super) store_hits: u64, + pub(super) store_misses: u64, +} + +/// Resolve-independent identity of one planning window: a two-lane +/// structural hash over the window's atoms and topology. Like +/// `FlushPlanKey`, the key is trusted without exact verification; +/// `FUSOR_VERIFY_PLAN_SHARING` regenerates and compares on every hit. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(super) struct WindowKey([u64; 2]); + +/// A stored template body with quantized-matrix identity erased. The store +/// outlives resolves on its device, so entries must hold no [`QMatrix`]: its +/// buffer Arc would pin dropped weights and its `Device` handle would cycle +/// the store back to `DeviceInner`. The template's `matrix_role` re-supplies +/// the concrete matrix at instantiation ([`VariantTemplate::capture`] +/// asserts one exists for every matrix-carrying variant). +#[derive(Clone)] +enum StoredBody { + /// Variant kinds that hold no buffer or device handles. + Plain(ExecutionVariant), + QMatMul(StoredQMatMul), + QEmbedding(StoredQEmbedding), +} + +#[derive(Clone)] +struct StoredQMatMul { + input_datatype: DataTypeEnum, + input: NodeIndex, + in_shape: Box<[usize]>, + out_shape: Box<[usize]>, + pre_element_wise_expr: Option, + post_element_wise_expr: Option, + post_accumulator_offsets: Box<[u32]>, +} + +#[derive(Clone)] +struct StoredQEmbedding { + indexes: NodeIndex, + out_shape: Box<[usize]>, + datatype: DataTypeEnum, +} + +impl StoredBody { + fn capture(variant: &ExecutionVariant) -> Self { + match variant { + ExecutionVariant::QMatMul(op) => { + let QMatMulOperation { + input_datatype, + input, + matrix: _, + in_shape, + out_shape, + pre_element_wise_expr, + post_element_wise_expr, + post_accumulator_offsets, + } = op.as_ref().clone(); + StoredBody::QMatMul(StoredQMatMul { + input_datatype, + input, + in_shape, + out_shape, + pre_element_wise_expr, + post_element_wise_expr, + post_accumulator_offsets, + }) + } + ExecutionVariant::QEmbedding(op) => { + let QEmbeddingOperation { + indexes, + matrix: _, + out_shape, + datatype, + } = op.clone(); + StoredBody::QEmbedding(StoredQEmbedding { + indexes, + out_shape, + datatype, + }) + } + ExecutionVariant::Elementwise(_) + | ExecutionVariant::Reduce(_) + | ExecutionVariant::Fold(_) + | ExecutionVariant::View(_) + | ExecutionVariant::MatMul(_) + | ExecutionVariant::RowProgram(_) + | ExecutionVariant::Attention(_) => StoredBody::Plain(variant.clone()), + ExecutionVariant::Tensor(_) + | ExecutionVariant::QMatrix(_) + | ExecutionVariant::Assign(_) + | ExecutionVariant::Region(_) => { + unreachable!("fusion templates never store this variant kind") + } + } + } + + fn rebuild(&self, matrix: Option) -> ExecutionVariant { + match self { + StoredBody::Plain(variant) => { + debug_assert!(matrix.is_none()); + variant.clone() + } + StoredBody::QMatMul(stored) => { + let StoredQMatMul { + input_datatype, + input, + in_shape, + out_shape, + pre_element_wise_expr, + post_element_wise_expr, + post_accumulator_offsets, + } = stored.clone(); + ExecutionVariant::QMatMul(Box::new(QMatMulOperation { + input_datatype, + input, + matrix: matrix.expect("stored qmatmul template requires a matrix role"), + in_shape, + out_shape, + pre_element_wise_expr, + post_element_wise_expr, + post_accumulator_offsets, + })) + } + StoredBody::QEmbedding(stored) => { + let StoredQEmbedding { + indexes, + out_shape, + datatype, + } = stored.clone(); + ExecutionVariant::QEmbedding(QEmbeddingOperation { + indexes, + matrix: matrix.expect("stored qembedding template requires a matrix role"), + out_shape, + datatype, + }) + } + } + } +} + +#[derive(Clone)] +struct StoredTemplate { + body: StoredBody, + dependency_roles: Vec, + matrix_role: Option, +} + +impl StoredTemplate { + fn from_template(template: &VariantTemplate) -> Self { + Self { + body: StoredBody::capture(&template.variant), + dependency_roles: template.dependency_roles.clone(), + matrix_role: template.matrix_role, + } + } + + fn instantiate(&self, instance: &PlanInstance, view: &FusionView<'_>) -> ExecutionVariant { + let matrix = self.matrix_role.map(|role| { + let inner = instance.nodes[role as usize]; + matrix_of( + view.variant_of(inner) + .expect("stored template matrix role must have a selected variant"), + ) + .expect("stored template matrix role must remain quantized") + .clone() + }); + let mut variant = self.body.rebuild(matrix); + let dependencies = self + .dependency_roles + .iter() + .map(|&role| instance.nodes[role as usize]) + .collect::>(); + rebind_variant_dependencies(&mut variant, &dependencies); + variant + } +} + +#[derive(Clone)] +enum StoredDecision { + NoRewrite, + Rewrite(StoredTemplate), +} + +/// Entries are individually cheap to regenerate; the cap only guards +/// runaway unique structure, so eviction is a wholesale reset. +const FUSION_PLAN_STORE_CAP: usize = 4096; + +/// Device-scoped fusion-plan decisions keyed by [`WindowKey`], shared across +/// resolves: the first resolve to plan a window pays generation, every later +/// isomorphic window on the device — next training step, next decode token — +/// instantiates the stored template. Templates are matrix-free (see +/// [`StoredBody`]) and every hit still re-validates kills and switch cost +/// against live state, exactly like intra-resolve sharing. +#[derive(Default)] +pub(crate) struct FusionPlanStore { + decisions: Mutex>, +} + +impl FusionPlanStore { + /// `None`: no stored decision. `Some(None)`: stored no-rewrite. + /// `Some(Some(variant))`: stored template instantiated for `instance`. + pub(super) fn instantiate( + &self, + key: WindowKey, + instance: &PlanInstance, + view: &FusionView<'_>, + ) -> Option> { + let decisions = self.decisions.lock(); + Some(match decisions.get(&key)? { + StoredDecision::NoRewrite => None, + StoredDecision::Rewrite(template) => Some(template.instantiate(instance, view)), + }) + } + + pub(super) fn record(&self, key: WindowKey, decision: &PlanDecision) { + let stored = match decision { + PlanDecision::NoRewrite => StoredDecision::NoRewrite, + PlanDecision::Rewrite(template) => { + StoredDecision::Rewrite(StoredTemplate::from_template(template)) + } + }; + let mut decisions = self.decisions.lock(); + if decisions.len() >= FUSION_PLAN_STORE_CAP && !decisions.contains_key(&key) { + tracing::debug!( + "fusion plan store reached {FUSION_PLAN_STORE_CAP} unique windows; resetting" + ); + decisions.clear(); + } + decisions.entry(key).or_insert(stored); + } +} + +/// Per-resolve memo. Liveness facts and read counts are part of every window +/// key, so repeated layers within this resolve share plans while windows +/// with different liveness never conflate. Per-resolve misses fall through +/// to the device-scoped [`FusionPlanStore`]. +pub(super) struct FusionPlanMemo { + atoms: Vec, + atom_ids: FxHashMap, + nodes: Vec, + node_ids: FxHashMap, + decisions: FxHashMap, + seen_windows: FxHashSet, + stats: PlanSharingStats, + stub_depth: u32, + /// Spike ledger: total time spent capturing windows, accumulated only + /// when `FUSOR_SPIKE_HOISTING` asked for it. + capture_time: Option, +} + +impl Default for FusionPlanMemo { + fn default() -> Self { + Self { + atoms: Vec::new(), + atom_ids: FxHashMap::default(), + nodes: Vec::new(), + node_ids: FxHashMap::default(), + decisions: FxHashMap::default(), + seen_windows: FxHashSet::default(), + stats: PlanSharingStats::default(), + stub_depth: WINDOW_STUB_DEPTH, + capture_time: None, + } + } +} + +impl FusionPlanMemo { + pub(super) fn for_config(config: &FusorConfig) -> Self { + Self { + stub_depth: config.spike_window_depth.unwrap_or(WINDOW_STUB_DEPTH), + capture_time: config.spike_hoisting.then(Default::default), + ..Default::default() + } + } + + pub(super) fn capture( + &mut self, + driver: &EGraphDriver, + state: &ExtractState, + view: &FusionView<'_>, + prov: Prov, + ) -> PlanInstance { + let start = self.capture_time.map(|_| std::time::Instant::now()); + self.stats.windows += 1; + let inner = driver.egraph.analysis.facts_of(prov).inner; + let stub_depth = self.stub_depth; + let mut builder = WindowBuilder { + memo: self, + driver, + state, + view, + stub_depth, + nodes: Vec::new(), + role_of: FxHashMap::default(), + local_ids: FxHashMap::default(), + matrix_aliases: FxHashMap::default(), + }; + let root = builder.add(inner, 0); + builder.memo.seen_windows.insert(root); + builder.memo.stats.unique_windows = builder.memo.seen_windows.len() as u64; + let instance = PlanInstance { + root, + nodes: builder.nodes, + role_of: builder.role_of, + }; + if let (Some(total), Some(start)) = (self.capture_time.as_mut(), start) { + *total += start.elapsed(); + } + instance + } + + pub(super) fn capture_time(&self) -> std::time::Duration { + self.capture_time.unwrap_or_default() + } + + pub(super) fn stub_depth(&self) -> u32 { + self.stub_depth + } + + pub(super) fn lookup(&mut self, instance: &PlanInstance, view: &FusionView<'_>) -> PlanLookup { + let Some(decision) = self.decisions.get(&instance.root) else { + self.stats.misses += 1; + return PlanLookup::Miss; + }; + self.stats.hits += 1; + match decision { + PlanDecision::NoRewrite => PlanLookup::Hit(None), + PlanDecision::Rewrite(template) => { + PlanLookup::Hit(Some(template.instantiate(instance, view))) + } + } + } + + pub(super) fn record( + &mut self, + instance: &PlanInstance, + view: &FusionView<'_>, + result: Option<&ExecutionVariant>, + ) -> Option<&PlanDecision> { + if !self.decisions.contains_key(&instance.root) { + let decision = match result { + None => { + self.stats.negative_templates += 1; + PlanDecision::NoRewrite + } + Some(variant) => { + let Some(template) = VariantTemplate::capture(variant, instance, view) else { + self.stats.unshareable += 1; + return None; + }; + self.stats.templates += 1; + PlanDecision::Rewrite(template) + } + }; + self.decisions.insert(instance.root, decision); + } + Some(&self.decisions[&instance.root]) + } + + pub(super) fn note_store_hit(&mut self) { + self.stats.store_hits += 1; + } + + pub(super) fn note_store_miss(&mut self) { + self.stats.store_misses += 1; + } + + /// Resolve-independent identity of one planning window. Per-resolve + /// `SpecId`s are replaced by their stable structural spec keys; roles, + /// facts, layouts and topology hash in exactly the interned content, so + /// equal keys reproduce equal role numbering on both sides. + pub(super) fn window_key(&self, root: StructuralId, driver: &EGraphDriver) -> WindowKey { + let mut lanes = TwoLane::new(); + let mut order: FxHashMap = FxHashMap::default(); + self.hash_window(root, driver, &mut lanes, &mut order); + WindowKey(lanes.finish().0) + } + + fn hash_window( + &self, + id: StructuralId, + driver: &EGraphDriver, + lanes: &mut TwoLane, + order: &mut FxHashMap, + ) { + if let Some(&back) = order.get(&id) { + // Shared subterm: a back-reference, disjoint from the kind tags. + lanes.write_u64(u64::MAX); + lanes.write_u64(back); + return; + } + order.insert(id, order.len() as u64); + let node = &self.nodes[id.0 as usize]; + let atom = &self.atoms[node.atom.0 as usize]; + lanes.write_u64(atom.role as u64); + match atom.kind { + PlanNodeKind::Operator(spec) => { + lanes.write_u64(0); + let PayloadKey(words) = driver.egraph.analysis.payloads.spec_key(spec); + lanes.write_u64(words[0]); + lanes.write_u64(words[1]); + } + PlanNodeKind::Tensor => lanes.write_u64(1), + PlanNodeKind::Boundary => lanes.write_u64(2), + PlanNodeKind::Missing => lanes.write_u64(3), + PlanNodeKind::Frontier => lanes.write_u64(4), + } + lanes.write_u64(local_hash(|hasher| { + atom.matrix_alias.hash(hasher); + atom.layout.hash(hasher); + atom.reads.hash(hasher); + atom.cached.hash(hasher); + atom.externally_live.hash(hasher); + atom.is_target.hash(hasher); + })); + lanes.write_u64(node.children.len() as u64); + for &child in node.children.iter() { + self.hash_window(child, driver, lanes, order); + } + } + + pub(super) fn stats(&self) -> PlanSharingStats { + self.stats + } + + pub(super) fn known_spec(&self, instance: &PlanInstance) -> Option { + let PlanDecision::Rewrite(template) = self.decisions.get(&instance.root)? else { + return None; + }; + template.spec + } + + pub(super) fn record_spec(&mut self, root: StructuralId, spec: SpecId) { + let Some(PlanDecision::Rewrite(template)) = self.decisions.get_mut(&root) else { + return; + }; + match template.spec { + Some(existing) => debug_assert_eq!(existing, spec), + None => template.spec = Some(spec), + } + } + + fn intern_atom(&mut self, atom: PlanAtom) -> PlanAtomId { + if let Some(&id) = self.atom_ids.get(&atom) { + return id; + } + let id = PlanAtomId(self.atoms.len() as u32); + self.atoms.push(atom.clone()); + self.atom_ids.insert(atom, id); + id + } + + fn intern_node(&mut self, node: StructuralNode) -> StructuralId { + if let Some(&id) = self.node_ids.get(&node) { + return id; + } + let id = StructuralId(self.nodes.len() as u32); + self.nodes.push(node.clone()); + self.node_ids.insert(node, id); + id + } +} + +struct WindowBuilder<'a> { + memo: &'a mut FusionPlanMemo, + driver: &'a EGraphDriver, + state: &'a ExtractState, + view: &'a FusionView<'a>, + stub_depth: u32, + nodes: Vec, + role_of: FxHashMap, + local_ids: FxHashMap, + matrix_aliases: FxHashMap, +} + +/// Fusion generators make one single-step decision per visit: they observe +/// the candidate node, its direct inputs' selected variants and facts, and +/// (through the variants they emit) the identities and layouts of the +/// inputs' inputs. Deeper structure is invisible to a generation step, and +/// every hit re-validates kills and switch cost against live state — so the +/// structural window cuts at that horizon. Cutting is what makes repeated +/// layers share: an unbounded walk would drag each window's whole upstream +/// cone in, making every layer's window unique and the walk quadratic. +/// `FUSOR_SPIKE_WINDOW_DEPTH` widens the horizon for measurement only — +/// widening is always sound (windows only get more specific), it just costs +/// capture time and sharing. +const WINDOW_STUB_DEPTH: u32 = 2; + +impl WindowBuilder<'_> { + fn add(&mut self, inner: NodeIndex, depth: u32) -> StructuralId { + if let Some(&id) = self.local_ids.get(&inner) { + return id; + } + let role = if let Some(&role) = self.role_of.get(&inner) { + role + } else { + let role = self.nodes.len() as u32; + self.nodes.push(inner); + self.role_of.insert(inner, role); + role + }; + + let prov = self.driver.prov_of.get(&inner).copied(); + let facts = prov.map(|prov| self.driver.egraph.analysis.facts_of(prov)); + let variant = self.view.variant_of(inner); + // Rules only look through these five pure operator families. Other + // operations are materialization frontiers: their output layout and + // liveness matter, but their private implementation/allocation does + // not. Embeddings and effectful/multi-output nodes stay frontiers; + // row programs have a structural equality and rebinding contract. + let opaque = variant.is_some_and(|variant| { + matches!( + variant, + ExecutionVariant::QMatrix(_) + | ExecutionVariant::Assign(_) + | ExecutionVariant::Region(_) + | ExecutionVariant::QEmbedding(_) + ) + }); + let matrix_alias = (!opaque) + .then_some(variant) + .flatten() + .and_then(matrix_of) + .map(|matrix| { + let allocation = std::sync::Arc::as_ptr(matrix.buffer()) as usize; + let next = self.matrix_aliases.len() as u32; + *self.matrix_aliases.entry(allocation).or_insert(next) + }); + let stub = depth >= self.stub_depth; + let (kind, dependencies) = if opaque || stub { + let kind = if opaque { + PlanNodeKind::Frontier + } else if let Some(variant) = variant { + match variant { + ExecutionVariant::Tensor(_) => PlanNodeKind::Tensor, + _ => { + let prov = prov.expect("selected execution variant has provenance"); + let payload = self + .state + .selected_enode(self.driver, prov) + .payload() + .expect("non-tensor execution variant has a payload"); + PlanNodeKind::Operator( + self.driver.egraph.analysis.payloads.spec_of(payload), + ) + } + } + } else if facts.is_some_and(|facts| facts.exec.is_none()) { + PlanNodeKind::Boundary + } else { + PlanNodeKind::Missing + }; + (kind, Vec::new()) + } else if let Some(variant) = variant { + let kind = match variant { + ExecutionVariant::Tensor(_) => PlanNodeKind::Tensor, + _ => { + let prov = prov.expect("selected execution variant has provenance"); + let payload = self + .state + .selected_enode(self.driver, prov) + .payload() + .expect("non-tensor execution variant has a payload"); + PlanNodeKind::Operator(self.driver.egraph.analysis.payloads.spec_of(payload)) + } + }; + (kind, variant_dependencies(variant)) + } else if facts.is_some_and(|facts| facts.exec.is_none()) { + (PlanNodeKind::Boundary, Vec::new()) + } else { + (PlanNodeKind::Missing, Vec::new()) + }; + + let layout = self.view.layout_of(inner).map(|info| LayoutSpec { + datatype: info.datatype(), + offset: info.layout().offset(), + shape: info.layout().shape().into(), + strides: info.layout().strides().into(), + }); + let atom = PlanAtom { + role, + kind, + matrix_alias, + layout, + reads: prov + .map(|prov| self.state.reads[prov.0 as usize]) + .unwrap_or(0), + cached: facts.is_some_and(|facts| facts.exec.is_none()), + externally_live: facts.is_some_and(|facts| facts.externally_live), + is_target: facts.is_some_and(|facts| facts.is_target), + }; + let atom = self.memo.intern_atom(atom); + let children: Vec = dependencies + .into_iter() + .map(|dependency| self.add(dependency, depth + 1)) + .collect(); + let id = self.memo.intern_node(StructuralNode { + atom, + children: children.into_boxed_slice(), + }); + self.local_ids.insert(inner, id); + id + } +} + +impl VariantTemplate { + /// `None` when the rewrite reads a node the window cannot name. Epilogue + /// generators fold a whole producer chain in one step and so reach the + /// chain's own operands, which sit past the horizon; those roots plan + /// fresh every visit rather than share a template that cannot say which + /// node a later instance should rebind to. + fn capture( + variant: &ExecutionVariant, + instance: &PlanInstance, + view: &FusionView<'_>, + ) -> Option { + assert!( + !matches!( + variant, + ExecutionVariant::Tensor(_) + | ExecutionVariant::QMatrix(_) + | ExecutionVariant::Assign(_) + | ExecutionVariant::Region(_) + ), + "fusion generators must produce a structurally rebindable variant" + ); + let dependency_roles = variant_dependencies(variant) + .into_iter() + .map(|dependency| instance.role_of.get(&dependency).copied()) + .collect::>>()?; + let matrix_role = matrix_of(variant).and_then(|matrix| { + instance + .nodes + .iter() + .position(|&inner| { + view.variant_of(inner) + .and_then(matrix_of) + .is_some_and(|candidate| same_matrix_allocation(matrix, candidate)) + }) + .map(|role| role as u32) + }); + assert!( + matrix_of(variant).is_none() || matrix_role.is_some(), + "quantized fusion template matrix must be part of its structural window" + ); + Some(Self { + variant: variant.clone(), + dependency_roles, + matrix_role, + spec: None, + }) + } + + fn instantiate(&self, instance: &PlanInstance, view: &FusionView<'_>) -> ExecutionVariant { + let dependencies = self + .dependency_roles + .iter() + .map(|&role| instance.nodes[role as usize]) + .collect::>(); + let mut variant = self.variant.clone(); + rebind_variant_dependencies(&mut variant, &dependencies); + if let Some(role) = self.matrix_role { + let inner = instance.nodes[role as usize]; + let matrix = matrix_of( + view.variant_of(inner) + .expect("template matrix role must have a selected variant"), + ) + .expect("template matrix role must remain quantized") + .clone(); + match &mut variant { + ExecutionVariant::QMatMul(operation) => operation.matrix = matrix, + ExecutionVariant::QEmbedding(operation) => operation.matrix = matrix, + _ => unreachable!("only quantized variants carry a matrix role"), + } + } + variant + } +} + +fn matrix_of(variant: &ExecutionVariant) -> Option<&crate::quantized::QMatrix> { + match variant { + ExecutionVariant::QMatrix(operation) => Some(&operation.matrix), + ExecutionVariant::QMatMul(operation) => Some(&operation.matrix), + ExecutionVariant::QEmbedding(operation) => Some(&operation.matrix), + _ => None, + } +} + +fn same_matrix_allocation(a: &crate::quantized::QMatrix, b: &crate::quantized::QMatrix) -> bool { + std::sync::Arc::ptr_eq(a.buffer(), b.buffer()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::compute_graph::resolve::Resolver; + use crate::compute_graph::resolve::egraph::rules_fuse::FusionCtx; + use crate::{Device, QMatrix, Tensor}; + + #[test] + fn repeated_windows_share_a_plan_but_rebind_distinct_values() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let left = Tensor::new::(&device, &[1.0, 2.0, 3.0, 4.0]); + let right = Tensor::new::(&device, &[5.0, 6.0, 7.0, 8.0]); + let left_out = (&left + 1.0) * 2.0; + let right_out = (&right + 1.0) * 2.0; + let targets = [left_out.data().key, right_out.data().key]; + + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.to_vec()); + for &target in &targets { + resolver.build_execution_graph(graph, target); + } + let driver = EGraphDriver::ingest(&resolver, graph); + let state = ExtractState::new(&driver); + let ctx = FusionCtx { + graph, + layouts: std::cell::RefCell::new(Default::default()), + }; + let view = FusionView::new(&driver, &state, &ctx); + let left_prov = driver.prov_of[&targets[0]]; + let right_prov = driver.prov_of[&targets[1]]; + + // Equal plans are not equal values: allocation identity keeps + // the executable e-classes separate. + assert_ne!(driver.class_for(targets[0]), driver.class_for(targets[1])); + + let mut memo = FusionPlanMemo::default(); + let left_instance = memo.capture(&driver, &state, &view, left_prov); + assert!(matches!( + memo.lookup(&left_instance, &view), + PlanLookup::Miss + )); + let left_variant = view + .generate_candidates(left_prov) + .into_iter() + .next() + .expect("left chain fuses"); + memo.record(&left_instance, &view, Some(&left_variant)); + + let right_instance = memo.capture(&driver, &state, &view, right_prov); + assert_eq!(left_instance.root, right_instance.root); + let PlanLookup::Hit(Some(rebound)) = memo.lookup(&right_instance, &view) else { + panic!("second isomorphic window should instantiate the first plan"); + }; + let fresh = view + .generate_candidates(right_prov) + .into_iter() + .next() + .expect("right chain fuses"); + match (rebound, fresh) { + ( + ExecutionVariant::Elementwise(rebound), + ExecutionVariant::Elementwise(fresh), + ) => { + assert_eq!(rebound, fresh); + assert!(rebound.inputs.iter().all(|input| { + right_instance.role_of.contains_key(input) + && !left_instance.role_of.contains_key(input) + })); + } + variants => panic!("unexpected generated variants: {variants:?}"), + } + let stats = memo.stats(); + assert_eq!(stats.templates, 1); + assert_eq!(stats.hits, 1); + }); + }); + } + + #[test] + fn plan_store_shares_templates_across_resolves() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + for step in 0..2u32 { + let input = Tensor::new::(&device, &[step as f32, 2.0, 3.0, 4.0]); + let out = (&input + 1.0) * 2.0; + let targets = [out.data().key]; + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.to_vec()); + for &target in &targets { + resolver.build_execution_graph(graph, target); + } + let driver = EGraphDriver::ingest(&resolver, graph); + let state = ExtractState::new(&driver); + let ctx = FusionCtx { + graph, + layouts: std::cell::RefCell::new(Default::default()), + }; + let view = FusionView::new(&driver, &state, &ctx); + let prov = driver.prov_of[&targets[0]]; + let mut memo = FusionPlanMemo::default(); + let instance = memo.capture(&driver, &state, &view, prov); + let key = memo.window_key(instance.root, &driver); + let store = device.fusion_plan_store(); + let fresh = view + .generate_candidates(prov) + .into_iter() + .next() + .expect("chain fuses"); + match store.instantiate(key, &instance, &view) { + None => { + assert_eq!(step, 0, "second resolve must hit the store"); + let decision = memo + .record(&instance, &view, Some(&fresh)) + .expect("chain rewrite stays inside its window"); + store.record(key, decision); + } + Some(Some(rebound)) => { + assert_eq!(step, 1, "first resolve cannot hit an empty store"); + let ( + ExecutionVariant::Elementwise(rebound), + ExecutionVariant::Elementwise(fresh), + ) = (rebound, fresh) + else { + panic!("unexpected variant kinds"); + }; + assert_eq!(rebound, fresh); + assert!( + rebound + .inputs + .iter() + .all(|input| { instance.role_of.contains_key(input) }) + ); + } + Some(None) => panic!("no-rewrite stored for a fusible window"), + } + }); + } + }); + } + + #[test] + fn plan_store_rebinds_weights_across_resolves() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let weights = [dense_qmatrix(&device, 0.25), dense_qmatrix(&device, 0.5)]; + for step in 0..2usize { + let activation = Tensor::new::( + &device, + &[vec![ + 1.0 + step as f32, + -2.0, + 3.0, + -4.0, + 5.0, + -6.0, + 7.0, + -8.0, + ]], + ); + let out = activation.q_mat_mul(&weights[step]); + let targets = [out.data().key]; + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.to_vec()); + for &target in &targets { + resolver.build_execution_graph(graph, target); + } + resolver.recognize_contractions(graph); + let driver = EGraphDriver::ingest(&resolver, graph); + let state = ExtractState::new(&driver); + let ctx = FusionCtx { + graph, + layouts: std::cell::RefCell::new(Default::default()), + }; + let view = FusionView::new(&driver, &state, &ctx); + let prov = driver.prov_of[&targets[0]]; + let mut memo = FusionPlanMemo::default(); + let instance = memo.capture(&driver, &state, &view, prov); + let key = memo.window_key(instance.root, &driver); + let store = device.fusion_plan_store(); + match store.instantiate(key, &instance, &view) { + None => { + assert_eq!(step, 0, "second resolve must hit the store"); + let variant = view.variant_of(targets[0]).unwrap().clone(); + let decision = memo + .record(&instance, &view, Some(&variant)) + .expect("qmatmul rewrite stays inside its window"); + store.record(key, decision); + } + Some(Some(ExecutionVariant::QMatMul(rebound))) => { + assert_eq!(step, 1, "first resolve cannot hit an empty store"); + assert!(std::sync::Arc::ptr_eq( + rebound.matrix.buffer(), + weights[1].buffer() + )); + assert!(!std::sync::Arc::ptr_eq( + rebound.matrix.buffer(), + weights[0].buffer() + )); + assert_eq!(rebound.input, activation.data().key); + } + other => panic!( + "unexpected stored decision (step {step}, some={})", + other.is_some() + ), + } + }); + } + }); + } + + #[test] + fn transformer_sized_repetition_plans_once() { + pollster::block_on(async { + const LAYERS: usize = 32; + let Ok(device) = Device::new().await else { + return; + }; + let inputs: Vec = (0..LAYERS) + .map(|layer| { + Tensor::new::( + &device, + &[ + layer as f32, + layer as f32 + 1.0, + layer as f32 + 2.0, + layer as f32 + 3.0, + ], + ) + }) + .collect(); + let outputs: Vec = inputs.iter().map(|input| (input + 1.0) * 2.0).collect(); + let targets: Vec = outputs.iter().map(|output| output.data().key).collect(); + + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.clone()); + for &target in &targets { + resolver.build_execution_graph(graph, target); + } + let driver = EGraphDriver::ingest(&resolver, graph); + let state = ExtractState::new(&driver); + let ctx = FusionCtx { + graph, + layouts: std::cell::RefCell::new(Default::default()), + }; + let view = FusionView::new(&driver, &state, &ctx); + let mut memo = FusionPlanMemo::default(); + + for (layer, &target) in targets.iter().enumerate() { + let prov = driver.prov_of[&target]; + let instance = memo.capture(&driver, &state, &view, prov); + match memo.lookup(&instance, &view) { + PlanLookup::Miss if layer == 0 => { + let variant = view + .generate_candidates(prov) + .into_iter() + .next() + .expect("repeated layer chain fuses"); + memo.record(&instance, &view, Some(&variant)); + } + PlanLookup::Hit(Some(_)) if layer > 0 => {} + _ => panic!("unexpected planning result for repeated layer {layer}"), + } + } + + let stats = memo.stats(); + assert_eq!(stats.windows, LAYERS as u64); + assert_eq!(stats.unique_windows, 1); + assert_eq!(stats.templates, 1); + assert_eq!(stats.misses, 1); + assert_eq!(stats.hits, (LAYERS - 1) as u64); + }); + }); + } + + fn dense_qmatrix(device: &Device, value: f32) -> QMatrix { + const N: usize = 4; + const K: usize = 8; + let bytes = std::iter::repeat_n(value, N * K) + .flat_map(f32::to_le_bytes) + .collect::>(); + QMatrix::from_parts(device, &bytes, Box::new([N, K]), fusor_gguf::GgmlType::F32).unwrap() + } + + #[test] + fn qmatmul_plan_rebinds_the_layer_weight_allocation() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let activations: Vec = (0..4) + .map(|offset| { + Tensor::new::( + &device, + &[vec![ + 1.0 + offset as f32, + -2.0, + 3.0, + -4.0, + 5.0, + -6.0, + 7.0, + -8.0, + ]], + ) + }) + .collect(); + let left_weight = dense_qmatrix(&device, 0.25); + let right_weight = dense_qmatrix(&device, 0.5); + let other_weight = dense_qmatrix(&device, 0.75); + let left = activations[0].q_mat_mul(&left_weight); + let right = activations[1].q_mat_mul(&right_weight); + let same_matrix_peer = activations[2].q_mat_mul(&left_weight); + let other_matrix_peer = activations[3].q_mat_mul(&other_weight); + // Separate operations over the exact same activation and weight + // become several observations of one semantic e-class. + let duplicate_a = activations[0].q_mat_mul(&left_weight); + let duplicate_b = activations[0].q_mat_mul(&left_weight); + let same_matrix_sum = &left + &same_matrix_peer; + let different_matrix_sum = &right + &other_matrix_peer; + let duplicate_sum = &duplicate_a + &duplicate_b; + let qmatmuls = [ + left.data().key, + right.data().key, + same_matrix_peer.data().key, + other_matrix_peer.data().key, + duplicate_a.data().key, + duplicate_b.data().key, + ]; + let targets = [ + same_matrix_sum.data().key, + different_matrix_sum.data().key, + duplicate_sum.data().key, + ]; + + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.to_vec()); + for &target in &targets { + resolver.build_execution_graph(graph, target); + } + resolver.recognize_contractions(graph); + resolver.recognize_embeddings(graph); + resolver.recognize_attention(graph); + assert_eq!( + resolver + .execution_graph + .node_indices() + .filter(|&node| matches!( + resolver.execution_graph[node].variant, + ExecutionVariant::QMatMul(_) + )) + .count(), + qmatmuls.len(), + "every observation of a shared e-class must be specialized" + ); + let driver = EGraphDriver::ingest(&resolver, graph); + let state = ExtractState::new(&driver); + let ctx = FusionCtx { + graph, + layouts: std::cell::RefCell::new(Default::default()), + }; + let view = FusionView::new(&driver, &state, &ctx); + assert_eq!( + variant_dependencies(view.variant_of(qmatmuls[0]).unwrap()), + vec![activations[0].data().key], + "a synchronized alternative must keep the first occurrence's concrete dependencies" + ); + assert_eq!( + variant_dependencies(view.variant_of(qmatmuls[2]).unwrap()), + vec![activations[2].data().key], + "a synchronized alternative must keep this occurrence's concrete dependencies" + ); + // Use the unique left-weight observation here; qmatmuls[0] + // deliberately shares a value e-class with the duplicate + // pair below and therefore has an aggregated read count. + let left_prov = driver.prov_of[&qmatmuls[2]]; + let right_prov = driver.prov_of[&qmatmuls[1]]; + assert_ne!( + driver.class_for(qmatmuls[2]), + driver.class_for(qmatmuls[1]), + "different weight buffers remain different values" + ); + + let mut memo = FusionPlanMemo::default(); + let left_instance = memo.capture(&driver, &state, &view, left_prov); + let left_variant = view.variant_of(qmatmuls[2]).unwrap().clone(); + memo.record(&left_instance, &view, Some(&left_variant)); + + let right_instance = memo.capture(&driver, &state, &view, right_prov); + assert_eq!(left_instance.root, right_instance.root); + let PlanLookup::Hit(Some(ExecutionVariant::QMatMul(rebound))) = + memo.lookup(&right_instance, &view) + else { + panic!("qmatmul planning template should be shared"); + }; + assert!(std::sync::Arc::ptr_eq( + rebound.matrix.buffer(), + right_weight.buffer() + )); + assert!(!std::sync::Arc::ptr_eq( + rebound.matrix.buffer(), + left_weight.buffer() + )); + + // Local matrix-alias identity is part of the plan key. Two + // distinct qmatmul nodes sharing one weight may take the + // same-base accumulator rewrite; same-shaped but different + // weights must not reuse that decision. + let same_sum = memo.capture(&driver, &state, &view, driver.prov_of[&targets[0]]); + let different_sum = + memo.capture(&driver, &state, &view, driver.prov_of[&targets[1]]); + assert_ne!(same_sum.root, different_sum.root); + }); + }); + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/execution.rs b/fusor-ml/core/src/compute_graph/resolve/execution.rs index 26c7cfde3..19cc8fd89 100644 --- a/fusor-ml/core/src/compute_graph/resolve/execution.rs +++ b/fusor-ml/core/src/compute_graph/resolve/execution.rs @@ -1,68 +1,91 @@ use super::*; -impl Resolver { - pub(super) fn release_dead_intermediates( - graph: &mut ComputeGraphInner, - produced_ops: &[&QueuedOperation], - remaining_consumers: &mut FxHashMap, - targets: &FxHashSet, - ) { - for op in produced_ops { - op.visit_dependencies(&mut |dep| { - if let Some(count) = remaining_consumers.get_mut(&dep) { - *count = count.saturating_sub(1); - if *count == 0 - && !targets.contains(&dep) - && !graph.has_live_lazy_descendant(dep) - { - // All consumers within this execution have been - // processed and no user-held lazy tensor still - // transitively depends on `dep` — free the cached - // buffer. The descendant check must include - // `live_descendant_count`, not just direct - // references: clearing `cached` on a node that still - // has an alive-uncached descendant flips it back to - // alive-uncached without propagating the transition, - // undercounting every ancestor's descendant counter. - if let Some(node) = graph.nodes.nodes.node_weight_mut(dep) { - node.cached = None; - } - } - } - }); - } +/// Wall-clock per optimizer phase, one resolve. +#[derive(Clone, Copy, Debug, Default)] +pub(super) struct OptimizePhases { + pub(super) recognition: Duration, + pub(super) extraction: Duration, + pub(super) physical: Duration, +} + +/// The remaining-consumer bookkeeping a release pass decrements. A live +/// resolve counts by inner node index; a replayed plan counts by plan slot. +pub(super) trait ConsumerCounts { + type Key: Copy; + + /// Drop one consumer of `key`, returning the node whose cached buffer is + /// now dead: `None` while consumers remain, while the key is untracked, + /// or when the node is an output of this execution. + fn consume(&mut self, key: Self::Key) -> Option; +} + +pub(super) struct NodeConsumers<'a> { + pub(super) counts: &'a mut FxHashMap, + pub(super) targets: &'a FxHashSet, +} + +impl ConsumerCounts for NodeConsumers<'_> { + type Key = NodeIndex; + + fn consume(&mut self, node: NodeIndex) -> Option { + let count = self.counts.get_mut(&node)?; + *count = count.saturating_sub(1); + (*count == 0 && !self.targets.contains(&node)).then_some(node) } +} - /// Like `release_dead_intermediates` but uses the compute graph's - /// `visit_dependencies` instead of an Operation's. Used for map-layout - /// and resize nodes that are resolved immediately without being lowered - /// to an Operation. - pub(super) fn release_dead_intermediates_from_graph( - graph: &mut ComputeGraphInner, - produced_nodes: &[NodeIndex], - remaining_consumers: &mut FxHashMap, - targets: &FxHashSet, - ) { - for &produced in produced_nodes { - let mut deps = Vec::new(); - graph.visit_dependencies(produced, &mut |dep| { - deps.push(dep); - }); - for dep in deps { - if let Some(count) = remaining_consumers.get_mut(&dep) { - *count = count.saturating_sub(1); - if *count == 0 - && !targets.contains(&dep) - && !graph.has_live_lazy_descendant(dep) - && let Some(node) = graph.nodes.nodes.node_weight_mut(dep) - { - node.cached = None; - } - } - } - } +pub(super) struct SlotConsumers<'a> { + pub(super) slots: &'a [NodeIndex], + pub(super) counts: &'a mut [u32], + pub(super) is_target: &'a [bool], +} + +impl ConsumerCounts for SlotConsumers<'_> { + type Key = u32; + + fn consume(&mut self, slot: u32) -> Option { + let slot = slot as usize; + let count = &mut self.counts[slot]; + *count = count.saturating_sub(1); + (*count == 0 && !self.is_target[slot]).then(|| self.slots[slot]) } +} +/// Free the cached buffers of the dependencies `visit` yields. A buffer is +/// released once all consumers within this execution have been processed and +/// no user-held lazy tensor still transitively depends on it. The descendant +/// check must include `live_descendant_count`, not just direct references: +/// clearing `cached` on a node that still has an alive-uncached descendant +/// flips it back to alive-uncached without propagating the transition, +/// undercounting every ancestor's descendant counter. Because +/// `has_live_lazy_descendant` is consulted here, reference-count drift +/// invisible to a replayed plan's structural fingerprint is handled exactly +/// as a full resolve handles it. +pub(super) fn release_consumed( + graph: &mut ComputeGraphInner, + counts: &mut C, + mut ledger: Option<&mut super::alloc_reuse::BufferLedger>, + visit: impl FnOnce(&mut dyn FnMut(C::Key)), +) { + visit(&mut |key| { + let Some(dep) = counts.consume(key) else { + return; + }; + if graph.has_live_lazy_descendant(dep) { + return; + } + if let Some(ledger) = ledger.as_deref_mut() + && let Some(cached) = graph.get_cached_result(dep) + { + ledger.note_released(dep, cached); + } + if let Some(node) = graph.nodes.nodes.node_weight_mut(dep) { + node.cached = None; + } + }); +} + +impl Resolver { pub(super) fn try_prepare_in_place_slice_assign_copy( graph: &ComputeGraphInner, operation: &crate::slice_assign::SliceAssignOperation, @@ -179,17 +202,36 @@ impl Resolver { ) -> Option { match &node.variant { ExecutionVariant::Elementwise(op) => { - Some(QueuedOperation::Generic(Arc::new(op.clone()))) + Some(QueuedOperation::Operation(Arc::new(op.clone()))) } - ExecutionVariant::MatMul(op) => Some(QueuedOperation::Generic(Arc::new(op.clone()))), - ExecutionVariant::Reduce(op) => Some(QueuedOperation::Generic(Arc::new(op.clone()))), - ExecutionVariant::GraphOp(op) => Some(QueuedOperation::Generic(op.clone())), - ExecutionVariant::View(op) => Some(QueuedOperation::Generic(Arc::new(op.clone()))), - ExecutionVariant::Assign(op) => Some(QueuedOperation::Generic(Arc::new(op.clone()))), + ExecutionVariant::MatMul(op) => Some(QueuedOperation::Operation(Arc::new(op.clone()))), + ExecutionVariant::Reduce(op) => Some(QueuedOperation::Operation(Arc::new(op.clone()))), + // A fold lowers through its reduce form. Multi-slot carriers have + // no kernel yet, and nothing constructs one, so reaching this with + // a general fold is a wiring bug rather than a missing feature. + ExecutionVariant::Fold(op) => { + let reduce = op.to_reduce().expect( + "a multi-slot fold reached lowering; only reduce-form folds are lowerable", + ); + Some(QueuedOperation::Operation(Arc::new(reduce))) + } + ExecutionVariant::RowProgram(op) => { + Some(QueuedOperation::Operation(Arc::new(op.clone()))) + } + ExecutionVariant::Attention(op) => { + Some(QueuedOperation::Operation(Arc::new(op.clone()))) + } + ExecutionVariant::View(op) => Some(QueuedOperation::Operation(Arc::new(op.clone()))), + ExecutionVariant::Assign(op) => Some(QueuedOperation::Operation(Arc::new(op.clone()))), ExecutionVariant::QEmbedding(op) => { - Some(QueuedOperation::Generic(Arc::new(op.clone()))) + Some(QueuedOperation::Operation(Arc::new(op.clone()))) + } + ExecutionVariant::Region(op) => Some(QueuedOperation::Merged( + merge_horizontal::MergedSegments::Region(vec![(node.inner_idx, op.clone())]), + )), + ExecutionVariant::QMatMul(op) => { + Some(QueuedOperation::Operation(Arc::new(op.as_ref().clone()))) } - ExecutionVariant::QMatMul(op) => Some(QueuedOperation::QMatMul(op.clone())), ExecutionVariant::QMatrix(op) => { // Skip materializing the dense tensor when every consumer // reads the block-quantized data directly (fused reduces and @@ -198,7 +240,7 @@ impl Resolver { if self.qmatrix_consumed_raw(exec_idx, node.inner_idx) { return None; } - Some(QueuedOperation::Generic(Arc::new(op.clone()))) + Some(QueuedOperation::Operation(Arc::new(op.clone()))) } ExecutionVariant::Tensor(_) => None, // Handled in execution loop } @@ -239,285 +281,27 @@ impl Resolver { // --- Rewrite Engine --- pub(super) fn optimize(&mut self, graph: &mut ComputeGraphInner) { - let profile_enabled = std::env::var_os("FUSOR_TRACE_OPTIMIZE").is_some(); - let mut profile = OptimizeProfile::default(); - // Rebuild composed contraction / normalization clusters into their - // specialized operations first, while they are still in the exact - // canonical form the API emitted (before view folding or fusion - // disturbs them). - self.recognize_contractions(graph); - self.recognize_embeddings(graph); - self.recognize_attention(graph); - self.fuse_row_programs(graph); - self.recognize_assign_chains(graph); - // The current rewrite rules can only start from Nary nodes (nary - // fusion, post-op reduce/matmul fusion) or MatMul nodes (pre-op - // unary fusion). Avoid scanning every QMatMul/attention node in - // decode graphs with hundreds of kernels. - let has_reduce = self.execution_graph.node_indices().any(|node| { - matches!( - self.execution_graph[node].variant, - ExecutionVariant::Reduce(_) - ) - }); - let has_matmul = self.execution_graph.node_indices().any(|node| { - matches!( - self.execution_graph[node].variant, - ExecutionVariant::MatMul(_) - ) - }); - let has_qmatmul = self.execution_graph.node_indices().any(|node| { - matches!( - self.execution_graph[node].variant, - ExecutionVariant::QMatMul(_) - ) - }); - let allow_qmatmul_elementwise_fusion = self.execution_graph.node_count() - <= DEFAULT_OPTIMIZE_NODE_LIMIT - || std::env::var_os("FUSOR_RESOLVE_QMATMUL_ELEMENTWISE_FUSION").is_some(); - let mut worklist: VecDeque = self - .execution_graph - .node_indices() - .filter(|&node| self.is_optimization_candidate(node)) - .collect(); - let mut in_worklist: FxHashSet = worklist.iter().copied().collect(); - - while let Some(node_idx) = worklist.pop_front() { - profile.iterations += 1; - in_worklist.remove(&node_idx); - - if !self.execution_graph.contains_node(node_idx) { - continue; - } - - // Edges are dependency -> consumer, and only downstream nodes can - // become newly fusible from these rewrites. - let consumers: Vec<_> = self - .execution_graph - .neighbors_directed(node_idx, petgraph::Direction::Outgoing) - .collect(); - - // 1. Fold view inputs into the nary body so fusion sees through - // layout changes - // 2. Fuse naries together (combine expression trees) - // 3. Try to fuse resulting nary into specialized ops (reduce, matmul, etc.) - let changed = self.try_fold_view_inputs(graph, node_idx); - - let start = profile_enabled.then(Instant::now); - let changed = changed | self.try_fuse_naries(graph, node_idx); - if let Some(start) = start { - profile.fuse_naries_count += 1; - profile.fuse_naries += start.elapsed(); - } - - let changed = if changed { - true - } else { - let start = profile_enabled.then(Instant::now); - let changed = has_reduce - && (self.try_fuse_into_reduce(graph, node_idx) - || self.try_fuse_producer_into_reduce(graph, node_idx)); - if let Some(start) = start { - profile.fuse_reduce_count += 1; - profile.fuse_reduce += start.elapsed(); - } - changed - }; - - let changed = if changed { - true - } else { - let start = profile_enabled.then(Instant::now); - let changed = (has_matmul || has_qmatmul) - && self.try_fuse_into_matmul(graph, node_idx, allow_qmatmul_elementwise_fusion); - if let Some(start) = start { - profile.fuse_matmul_count += 1; - profile.fuse_matmul += start.elapsed(); - } - changed - }; - - if changed { - profile.changed += 1; - // Re-add the current node to worklist if it still exists - if self.execution_graph.contains_node(node_idx) - && self.is_optimization_candidate(node_idx) - && in_worklist.insert(node_idx) - { - worklist.push_back(node_idx); - } - - // Re-add downstream fusion candidates that might now be fusible - // — both the consumers captured before this rewrite and any it - // created — descending through view nodes (e.g. the MapLayout - // broadcast `add_` inserts) that sit between a changed node and - // the next candidate. - self.enqueue_downstream_candidates( - consumers, - Self::is_optimization_candidate, - &mut worklist, - &mut in_worklist, - ); - if self.execution_graph.contains_node(node_idx) { - let new_consumers: Vec<_> = self - .execution_graph - .neighbors_directed(node_idx, petgraph::Direction::Outgoing) - .collect(); - self.enqueue_downstream_candidates( - new_consumers, - Self::is_optimization_candidate, - &mut worklist, - &mut in_worklist, - ); - } - } - } - if profile_enabled { - profile.print(); - } - } - - pub(super) fn optimize_large_graph(&mut self, graph: &mut ComputeGraphInner) { - self.recognize_contractions(graph); - self.recognize_embeddings(graph); - self.recognize_attention(graph); - self.fuse_row_programs(graph); - self.recognize_assign_chains(graph); - let has_qmatmul = self.execution_graph.node_indices().any(|node| { - matches!( - self.execution_graph[node].variant, - ExecutionVariant::QMatMul(_) - ) - }); - if !has_qmatmul { - return; - } - - let mut worklist = self - .execution_graph - .node_indices() - .filter(|&node| self.is_large_graph_nary_candidate(node)) - .collect::>(); - let mut in_worklist = worklist.iter().copied().collect::>(); - - while let Some(node_idx) = worklist.pop_front() { - in_worklist.remove(&node_idx); - if !self.execution_graph.contains_node(node_idx) { - continue; - } - - let consumers = self - .execution_graph - .neighbors_directed(node_idx, petgraph::Direction::Outgoing) - .collect::>(); - let mut changed = self.try_fold_view_inputs(graph, node_idx); - changed |= self.try_fuse_naries(graph, node_idx); - if !changed && self.execution_graph.contains_node(node_idx) { - changed = self.try_fuse_into_matmul(graph, node_idx, true); - } - - if changed { - if self.execution_graph.contains_node(node_idx) - && self.is_large_graph_nary_candidate(node_idx) - && in_worklist.insert(node_idx) - { - worklist.push_back(node_idx); - } - self.enqueue_downstream_candidates( - consumers, - Self::is_large_graph_nary_candidate, - &mut worklist, - &mut in_worklist, - ); - } + // Every graph uses the full optimizer. Recognition and fusion share + // one value e-graph; the structural plan memo reuses discovery across + // repeated layers while preserving allocation-distinct values. + self.optimize_operations(graph); + + // Region formation generalizes the sole-consumer nary gate: it fuses + // externally-live producers into multi-output regions. Codegen then + // selects tuned kernels from operation shape and device capabilities. + let phase_start = Instant::now(); + self.form_elementwise_regions(graph); + self.optimize_phases.physical += phase_start.elapsed(); + #[cfg(feature = "graphvis")] + if let Some(dir) = &graph.device().config().dump_stages { + super::visualize::dump_stage( + dir, + &self.execution_graph, + super::visualize::Stage::Regions, + ); } } - /// Re-enqueue downstream fusion candidates reachable from `seeds`, - /// descending through `MapLayout` view nodes. A rewrite (e.g. fusing an - /// `add` into a qmatmul epilogue) can make a candidate that sits *behind* a - /// broadcast/narrow view newly fusible; those views are not optimization - /// candidates themselves, so a plain direct-consumer scan would never reach - /// the candidate past them. - fn enqueue_downstream_candidates( - &self, - seeds: impl IntoIterator, - is_candidate: impl Fn(&Self, ExecutionNodeIndex) -> bool, - worklist: &mut VecDeque, - in_worklist: &mut FxHashSet, - ) { - let mut stack: Vec = seeds.into_iter().collect(); - let mut visited = FxHashSet::default(); - while let Some(node) = stack.pop() { - if !self.execution_graph.contains_node(node) || !visited.insert(node) { - continue; - } - if is_candidate(self, node) { - if in_worklist.insert(node) { - worklist.push_back(node); - } - } else if matches!( - self.execution_graph[node].variant, - ExecutionVariant::View(_) - ) { - stack.extend( - self.execution_graph - .neighbors_directed(node, petgraph::Direction::Outgoing), - ); - } - } - } - - pub(super) fn is_large_graph_nary_candidate(&self, node_idx: ExecutionNodeIndex) -> bool { - let ExecutionVariant::Elementwise(nary) = &self.execution_graph[node_idx].variant else { - return false; - }; - if nary.shape.last().copied().unwrap_or_default() >= LARGE_GRAPH_NARY_FUSION_MIN_LAST_DIM { - return true; - } - - nary.inputs.iter().any(|&input| { - let (base_inner, _) = self.walk_view_chain(input); - self.get_input_node_in_exec_graph(base_inner) - .is_some_and(|exec_idx| { - matches!( - self.execution_graph[exec_idx].variant, - ExecutionVariant::QMatMul(_) - ) - }) - }) - } - - pub(super) fn is_single_token_qmatmul_graph(&self) -> bool { - let mut qmatmul_count = 0usize; - let mut single_token_count = 0usize; - for node in self.execution_graph.node_indices() { - let ExecutionVariant::QMatMul(qmatmul) = &self.execution_graph[node].variant else { - continue; - }; - qmatmul_count += 1; - if qmatmul.in_shape.len() >= 2 - && qmatmul.in_shape[..qmatmul.in_shape.len() - 1] - .iter() - .product::() - == 1 - { - single_token_count += 1; - } - } - qmatmul_count >= 16 && single_token_count * 4 >= qmatmul_count * 3 - } - - pub(super) fn is_optimization_candidate(&self, node_idx: ExecutionNodeIndex) -> bool { - matches!( - self.execution_graph[node_idx].variant, - ExecutionVariant::Elementwise(_) - | ExecutionVariant::MatMul(_) - | ExecutionVariant::QMatMul(_) - | ExecutionVariant::Reduce(_) - ) - } - // Helpers pub(super) fn add_physical_dependencies( &self, @@ -528,6 +312,11 @@ impl Resolver { let inner_idx = self.execution_graph[node_idx].inner_idx; for &input in inputs { graph.add_dependency_edge(input, inner_idx); + if let Some(recorder) = &self.recorder { + // These edges are persistent inner-graph side effects of the + // optimizer; a replayed plan must re-add them, so record them. + recorder.borrow_mut().record_physical_edge(input, inner_idx); + } } } @@ -538,37 +327,16 @@ impl Resolver { self.node_mapping.get(&inner_input).copied() } - /// Walk through view nodes from `inner` down to the first non-view - /// node, composing each view's collapsed stage stack. Public tensor ops - /// collapse into single view nodes at construction, but composed - /// clusters (attention's attached GQA/transpose views) still layer view - /// nodes deliberately. Returns the base node and the composed layout - /// over the base's logical value space; the layout is `None` when - /// `inner` is not a view (identity). Views that don't collapse or - /// compose (or carry a fill region) act as chain breaks: the walk stops - /// without seeing through them. - pub(super) fn walk_view_chain(&self, mut inner: NodeIndex) -> (NodeIndex, Option) { - let mut composed: Option = None; - loop { - let Some(exec) = self.get_input_node_in_exec_graph(inner) else { - return (inner, composed); - }; + /// [`egraph::compose::walk_view_chain`] over the execution graph's + /// current forms. + pub(super) fn walk_view_chain(&self, inner: NodeIndex) -> (NodeIndex, Option) { + egraph::compose::walk_view_chain(inner, |inner| { + let exec = self.get_input_node_in_exec_graph(inner)?; let ExecutionVariant::View(view) = &self.execution_graph[exec].variant else { - return (inner, composed); - }; - let Some(collapsed) = view.composed_layout() else { - return (inner, composed); - }; - let next = match &composed { - None => collapsed, - Some(outer) => match crate::view::compose_layouts(outer, &collapsed) { - Some(layout) => layout, - None => return (inner, composed), - }, + return None; }; - composed = Some(next); - inner = view.input; - } + Some((view.composed_layout()?, view.input)) + }) } /// The layout a (possibly chained-view) node presents over its base @@ -581,53 +349,6 @@ impl Resolver { } } - pub(super) fn infer_layout_cached( - &mut self, - graph: &ComputeGraphInner, - inner_idx: NodeIndex, - ) -> Option { - self.layout_pass.visit(graph, inner_idx); - self.layout_pass.output_layout.get(&inner_idx).cloned() - } - - pub(super) fn try_normalize_qmatmul_post_extra( - &mut self, - graph: &ComputeGraphInner, - extra_inner: NodeIndex, - output_shape: &[usize], - ) -> Option { - let last_dim = *output_shape.last()?; - let extra_info = self.infer_layout_cached(graph, extra_inner)?; - if extra_info.datatype() != DataTypeEnum::F32 || extra_info.layout().shape() != output_shape - { - return None; - } - - let layout = extra_info.layout(); - let is_column_broadcast = layout.offset() == 0 - && layout.strides().last().copied() == Some(1) - && layout.shape().last().copied() == Some(last_dim) - && layout.strides()[..layout.strides().len().saturating_sub(1)] - .iter() - .all(|stride| *stride == 0); - if !is_column_broadcast { - return Some(extra_inner); - } - - let (base_inner, _) = self.walk_view_chain(extra_inner); - let base_info = self.infer_layout_cached(graph, base_inner)?; - let base_layout = base_info.layout(); - if base_info.datatype() == DataTypeEnum::F32 - && base_layout.shape() == [last_dim] - && base_layout.is_contiguous() - && base_layout.offset() == 0 - { - Some(base_inner) - } else { - Some(extra_inner) - } - } - pub(super) fn check_cached(&self, graph: &ComputeGraphInner, inner_idx: NodeIndex) -> bool { graph.get_cached_result(inner_idx).is_some() } @@ -636,6 +357,10 @@ impl Resolver { if !self.execution_graph.contains_node(node_idx) { return; } + let inner_idx = self.execution_graph[node_idx].inner_idx; + if self.targets.contains(&inner_idx) { + return; + } if self .execution_graph .neighbors_directed(node_idx, petgraph::Direction::Outgoing) @@ -648,6 +373,7 @@ impl Resolver { .neighbors_directed(node_idx, petgraph::Direction::Incoming) .collect(); self.execution_graph.remove_node(node_idx); + self.node_mapping.remove(&inner_idx); // Recursively check if dependencies are now dead for dep in incoming { self.remove_node_if_dead(dep); @@ -655,3 +381,42 @@ impl Resolver { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{Device, Tensor}; + + #[test] + fn dead_node_removal_clears_its_inner_mapping() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let input = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let intermediate = &input + 1.0; + let output = &intermediate * 2.0; + let target = output.data().key; + let intermediate_inner = intermediate.data().key; + + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, vec![target]); + resolver.build_execution_graph(graph, target); + let intermediate_exec = resolver.node_mapping[&intermediate_inner]; + let target_exec = resolver.node_mapping[&target]; + let edge = resolver + .execution_graph + .find_edge(intermediate_exec, target_exec) + .expect("intermediate feeds the target"); + resolver.execution_graph.remove_edge(edge); + resolver.remove_node_if_dead(intermediate_exec); + + assert!(!resolver.execution_graph.contains_node(intermediate_exec)); + assert!( + !resolver.node_mapping.contains_key(&intermediate_inner), + "removed execution nodes must not remain addressable" + ); + }); + }); + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/flush_replay.rs b/fusor-ml/core/src/compute_graph/resolve/flush_replay.rs new file mode 100644 index 000000000..658e440c7 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/flush_replay.rs @@ -0,0 +1,1309 @@ +//! Reusable materialization-plan replay for structurally repeated graphs. +//! +//! Training loops flush an isomorphic tape every step, while iterative dense +//! inference repeatedly resolves one isomorphic output graph. Both present +//! the same operations, shapes, and expressions over fresh nodes and buffers. +//! The full resolver pipeline (execution-graph build, recognition + fusion, +//! toposort, lowering, consumer counting, per-op input gathering, workgroup +//! solving, and kernel building) is fully deterministic given that structure, +//! so its outcome can be recorded once and replayed on later steps. +//! +//! A replayed materialization skips every deterministic pass and only re-runs +//! the intrinsically per-step work: output-buffer allocation, positional buffer +//! rebinding ([`DirectKernelTemplate::bind_buffers`]), bind-group creation, +//! command encoding, and liveness bookkeeping. +//! +//! Safety properties (see the module tests and `verify_integrity`): +//! - Plans are strictly bufferless and `NodeIndex`-free: templates hold +//! compile artifacts only, and all cross-step references are fingerprint +//! slot positions remapped through a fresh DFS each step. Stable-graph +//! index recycling can never alias, and no `Arc` is retained +//! across steps (the buffer pool's `strong_count == 1` recycling is never +//! starved). +//! - Replay mutates the inner graph exclusively through the blessed APIs: +//! `set_cached_result`, `add_dependency_edge` (re-adding the optimizer's +//! recorded physical edges), and the exact release predicate used by +//! `release_consumed` — evaluated live, never recorded, so +//! reference-count drift (e.g. a user cloning a mid-graph handle) is +//! handled identically to a full resolve. +//! - Tensor and quantized-matrix storage are separate binding roles. A +//! `QMatrix` node can therefore supply its raw block buffer to fused +//! consumers and later materialize a dense tensor at the same graph slot +//! without either binding being confused with the other. +//! - Buffer identity between slots is deliberately absent from the +//! fingerprint (buffers change every step), so the recorder must never +//! produce a plan whose bindings depend on two slots incidentally sharing +//! one buffer: `attribute_buffer` poisons the recording when distinct +//! slots resolve to the same pointer unless the plan itself re-creates the +//! sharing (view aliases, in-place outputs), and pins every tracked `Arc` +//! so pool recycling cannot alias pointers mid-recording. + +use std::hash::Hash; +use std::num::NonZeroUsize; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use fusor_tile_ir_runtime::{TwoLaneHasher, single_lane}; +use lru::LruCache; +use parking_lot::Mutex; +use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet}; +use web_time::Instant; + +use super::{ + CommandRecord, ComputeGraphInner, ComputeGraphNodeVariant, DispatchRecord, NodeIndex, + QueuedOperation, Resolver, +}; +use crate::mir::kernel_backend::{DirectKernel, DirectKernelTemplate}; +use crate::mir::operation::{Operation, hash_layout}; +use crate::tensor::TensorData; +use crate::{DataTypeEnum, Device, Layout}; + +/// Bump when anything about the recorded plan layout or the fingerprint +/// recipe changes, so stale entries can never be replayed. +/// v9: equivalent e-class observations replay as aliases of their physical +/// representative rather than pretending to be executable view nodes. +const REPLAY_RECIPE_VERSION: u64 = 9; + +/// Admission (see [`FlushPlanCache::admit_record`]) keeps one-shot +/// fingerprints out, so capacity only bounds host metadata for fingerprints +/// that provably recur. +const FLUSH_PLAN_CACHE_SIZE: usize = 64; + +/// Bound on the sighted-fingerprint admission set (16 bytes per key). On +/// overflow the set resets wholesale; recurring shapes then re-earn +/// admission with two fresh sightings, which is harmless next to unbounded +/// growth. +const SEEN_KEY_CAP: usize = 1 << 16; + +/// Structural fingerprint of one flush's pending subgraph. +#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] +pub(crate) struct FlushPlanKey([u64; 2]); + +/// The reusable-plan decision for one materialization boundary. Both batched +/// flushes and ordinary single-target resolves enter through this gate so an +/// isomorphic graph has one cache lifecycle: resolve once, record on the +/// second sighting, then replay. +pub(crate) enum ReplayAction { + Resolve, + Record { + key: FlushPlanKey, + fingerprint: FlushFingerprint, + }, + Replay { + plan: Arc, + fingerprint: FlushFingerprint, + }, +} + +/// Per-device LRU of flush plans. Lives on `DeviceInner` beside the +/// kernel cache so it is reachable under the compute-graph write lock. +pub(crate) struct FlushPlanCache { + plans: Mutex, FxBuildHasher>>, + /// Fingerprints sighted at least once; recording is admitted only on a + /// second sighting (see [`FlushPlanCache::admit_record`]). + seen: Mutex>, + replays: AtomicU64, + records: AtomicU64, +} + +impl Default for FlushPlanCache { + fn default() -> Self { + Self { + plans: Mutex::new(LruCache::with_hasher( + NonZeroUsize::new(FLUSH_PLAN_CACHE_SIZE).expect("cache size must be non-zero"), + FxBuildHasher, + )), + seen: Mutex::new(FxHashSet::default()), + replays: AtomicU64::new(0), + records: AtomicU64::new(0), + } + } +} + +impl FlushPlanCache { + pub(crate) fn get(&self, key: &FlushPlanKey) -> Option> { + self.plans.lock().get(key).cloned() + } + + /// Admission gate: record a plan only for a fingerprint sighted before. + /// Recording costs a fully instrumented resolve, and shape-drifting + /// workloads — decode's growing KV cache, varying batch sizes — mint a + /// fresh fingerprint every flush; one-shot keys pay one hash-set insert + /// here instead and never occupy a plan slot. + pub(crate) fn admit_record(&self, key: FlushPlanKey) -> bool { + let mut seen = self.seen.lock(); + if seen.len() >= SEEN_KEY_CAP && !seen.contains(&key) { + seen.clear(); + } + !seen.insert(key) + } + + pub(crate) fn insert(&self, key: FlushPlanKey, plan: Arc) { + self.records.fetch_add(1, Ordering::Relaxed); + self.plans.lock().put(key, plan); + } + + pub(crate) fn note_replay(&self) { + self.replays.fetch_add(1, Ordering::Relaxed); + } + + #[cfg(test)] + pub(crate) fn replay_count(&self) -> u64 { + self.replays.load(Ordering::Relaxed) + } + + #[cfg(test)] + pub(crate) fn record_count(&self) -> u64 { + self.records.load(Ordering::Relaxed) + } +} + +/// Choose whether `targets` need a full resolve, should record, or can replay. +/// +/// A failed replay validation is non-mutating and falls back to a normal +/// resolve. A cache miss records only for fingerprints sighted before, so +/// one-shot graphs never pay the instrumented recording resolve. +pub(crate) fn prepare_replay(graph: &mut ComputeGraphInner, targets: &[NodeIndex]) -> ReplayAction { + let Some(fingerprint) = fingerprint_pending(graph, targets) else { + return ReplayAction::Resolve; + }; + // A single-target resolve can run while another live output depends on + // part of the same graph. That outside descendant affects allocation + // claims but is intentionally absent from this target's fingerprint. + // Replay only closed materialization boundaries: every live uncached + // handle must belong to the fingerprinted dependency closure. + if graph.pending_sinks.iter().any(|(&node, _)| { + graph + .nodes + .nodes + .node_weight(node) + .is_some_and(|data| data.reference_count > 0 && data.cached.is_none()) + && !fingerprint.pos_of.contains_key(&node) + }) { + return ReplayAction::Resolve; + } + + let device = graph.device(); + let cache = device.flush_plan_cache(); + match cache.get(&fingerprint.key) { + Some(plan) => { + if validate_replay(graph, &device, &plan, &fingerprint) { + ReplayAction::Replay { plan, fingerprint } + } else { + ReplayAction::Resolve + } + } + None => { + if cache.admit_record(fingerprint.key) { + ReplayAction::Record { + key: fingerprint.key, + fingerprint, + } + } else { + ReplayAction::Resolve + } + } + } +} + +/// The canonical [`TwoLaneHasher`] producing the 128-bit plan key. Replay +/// correctness does not rest on this key alone: upfront validation re-checks +/// step kinds and boundary caching, and the recorder refuses plans whose +/// buffer provenance is ambiguous. +struct FingerprintHasher(TwoLaneHasher); + +impl FingerprintHasher { + fn new() -> Self { + Self(TwoLaneHasher::new()) + } + + fn write_u64(&mut self, value: u64) { + self.0.write_u64(value); + } + + fn finish(self) -> FlushPlanKey { + FlushPlanKey(self.0.finish()) + } +} + +/// The allocation-erased storage contract of one boundary or leaf: what the +/// recorded kernels were compiled against, with the buffer identity — pure +/// per-step data — left out. +fn storage_contract(data: &TensorData) -> u64 { + single_lane(|state| { + data.datatype().hash(state); + hash_layout(state, data.layout()); + data.buffer().size().hash(state); + }) +} + +/// The structural fingerprint of one flush: a 128-bit key plus the slot +/// assignment (DFS discovery order) mapping fingerprint positions to the +/// current step's `NodeIndex`es. Recomputed fresh every flush; positions — +/// never `NodeIndex` values — are what plans store. +pub(crate) struct FlushFingerprint { + pub(crate) key: FlushPlanKey, + slots: Vec, + pos_of: FxHashMap, + /// True for slots that were already cached when the fingerprint was + /// taken. Boundaries are hashed opaquely (dtype/layout/allocation size) + /// and never descended past — the cached-view-collapse rule. + boundary: Vec, +} + +struct FingerprintState { + hasher: FingerprintHasher, + slots: Vec, + pos_of: FxHashMap, + boundary: Vec, +} + +/// Fingerprint the pending subgraph, mirroring `build_execution_graph`'s +/// traversal exactly: DFS from the pending sinks in order, visiting +/// dependencies in `visit_dependencies` order, stopping at cached nodes. +/// Tensor and QMatrix allocation identity is deliberately absent: their +/// shape/storage contracts are hashed while current buffers are rebound by +/// fingerprint position during replay. +pub(crate) fn fingerprint_pending( + graph: &ComputeGraphInner, + pending: &[NodeIndex], +) -> Option { + let mut state = FingerprintState { + hasher: FingerprintHasher::new(), + slots: Vec::new(), + pos_of: FxHashMap::default(), + boundary: Vec::new(), + }; + + for &sink in pending { + let pos = fingerprint_visit(graph, sink, &mut state)?; + state.hasher.write_u64(0xE0); + state.hasher.write_u64(pos as u64); + } + state.hasher.write_u64(REPLAY_RECIPE_VERSION); + + Some(FlushFingerprint { + key: state.hasher.finish(), + slots: state.slots, + pos_of: state.pos_of, + boundary: state.boundary, + }) +} + +fn fingerprint_visit( + graph: &ComputeGraphInner, + node: NodeIndex, + state: &mut FingerprintState, +) -> Option { + if let Some(&pos) = state.pos_of.get(&node) { + return Some(pos); + } + let node_data = graph.nodes.nodes.node_weight(node)?; + let pos = u32::try_from(state.slots.len()).ok()?; + state.slots.push(node); + state.pos_of.insert(node, pos); + + if let Some(cached) = &node_data.cached { + // Cached boundary leaf: opaque. The buffer contents are per-step + // data; only the layout contract the recorded kernels were compiled + // against must match. The liveness bit is part of the structure too: + // allocation-reuse claims depend on whether the user still holds a + // handle to the boundary tensor, so "user kept a handle" must never + // collide with a plan whose claims assumed the buffer was dead. + state.boundary.push(true); + state.hasher.write_u64(0xB0); + state + .hasher + .write_u64((node_data.reference_count > 0) as u64); + state.hasher.write_u64(storage_contract(cached)); + return Some(pos); + } + state.boundary.push(false); + + // Recognition and fusion consult live references (`has_live_reference`), + // so the liveness bit is part of the structure. + let live_ref = node_data.reference_count > 0; + match &node_data.variant { + ComputeGraphNodeVariant::QMatrix(operation) => { + state.hasher.write_u64(0xA6); + state.hasher.write_u64(live_ref as u64); + state + .hasher + .write_u64(single_lane(|h| operation.hash_kernel_fields(h))); + } + ComputeGraphNodeVariant::Tensor(data) => { + state.hasher.write_u64(0xA1); + state.hasher.write_u64(live_ref as u64); + state.hasher.write_u64(storage_contract(data)); + } + variant => { + let (tag, op): (u64, &dyn Operation) = match variant { + ComputeGraphNodeVariant::Elementwise(op) => (0xA2, op), + ComputeGraphNodeVariant::Reduce(op) => (0xA3, op), + ComputeGraphNodeVariant::View(op) => (0xA4, op), + ComputeGraphNodeVariant::Assign(op) => (0xA5, op), + ComputeGraphNodeVariant::Tensor(_) | ComputeGraphNodeVariant::QMatrix(_) => { + unreachable!("handled above") + } + }; + state.hasher.write_u64(tag); + state.hasher.write_u64(live_ref as u64); + state + .hasher + .write_u64(single_lane(|h| op.hash_kernel_fields(h))); + } + } + + let mut dependencies = Vec::new(); + node_data.variant.visit_dependencies(&mut |dep| { + dependencies.push(dep); + }); + state.hasher.write_u64(dependencies.len() as u64); + for dep in dependencies { + let dep_pos = fingerprint_visit(graph, dep, state)?; + state.hasher.write_u64(dep_pos as u64); + } + Some(pos) +} + +/// The recorded outcome of one materialization. Strictly bufferless and +/// `NodeIndex`-free; every cross-step reference is a fingerprint slot +/// position. +pub(crate) struct FlushPlan { + node_count: u32, + /// Positions of the flush targets. Targets keep their cached results; + /// everything else may be released once its last consumer ran. + target_positions: Box<[u32]>, + /// Inner-graph dependency edges the optimizer added during the recording + /// resolve (`add_physical_dependencies`). Replay re-adds them so liveness + /// accounting sees post-fusion dependencies. + physical_edges: Box<[(u32, u32)]>, + steps: Box<[PlanStep]>, +} + +enum PlanStep { + /// Cache an input leaf's own tensor data (current step's contents). + TensorLeaf { pos: u32 }, + /// Zero-cost view alias re-derived via `try_map_tensor`. + ViewAlias { pos: u32, consumed: Box<[u32]> }, + /// A semantically equivalent observation aliases the representative that + /// performed the physical work. The original observation may be any + /// operation kind; replay copies the representative's tensor metadata. + SharedAlias { + pos: u32, + source: u32, + consumed: Box<[u32]>, + }, + /// In-place slice-assign buffer copies re-derived from the current graph. + CopyAssign { pos: u32, consumed: Box<[u32]> }, + /// One lowered operation's kernel dispatches. + Dispatch { + pos: u32, + kernels: Box<[PlanKernel]>, + scratch: Box<[BufferSpec]>, + output: OutputSpec, + /// Dependency positions (with multiplicity, in `visit_dependencies` + /// order) to feed the release accounting after this step. + consumed: Box<[u32]>, + }, + /// One horizontally merged dispatch caching several segments' outputs. + /// Every output is a freshly allocated buffer (never an alias), so slot + /// attribution stays 1 buffer <-> 1 slot. + MergedDispatch { + outputs: Box<[(u32, OutputSpec)]>, + kernels: Box<[PlanKernel]>, + scratch: Box<[BufferSpec]>, + consumed: Box<[u32]>, + }, +} + +struct PlanKernel { + template: DirectKernelTemplate, + /// Source per binding buffer, in `binding_buffers()` order. + bindings: Box<[BufferBinding]>, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum BufferBinding { + /// Tensor leaf, cached boundary, or materialized output at this slot. + Tensor(u32), + /// Raw block storage embedded in a QMatrix node at this slot. + QMatrix(u32), + /// Build-local buffer allocated once for this plan step. + Scratch(u32), +} + +#[derive(Clone, Copy)] +struct BufferSpec { + size: u64, + usage: wgpu::BufferUsages, +} + +struct OutputSpec { + layout: Layout, + datatype: DataTypeEnum, + source: OutputSource, +} + +enum OutputSource { + /// Freshly allocated output buffer (the common case: `Operation::inputs` + /// allocates the output and appends it as the last MIR value). + Fresh { + buffer_size: u64, + usage: wgpu::BufferUsages, + }, + /// Output aliases an existing slot's buffer (in-place operations). + Alias { slot: u32 }, +} + +/// Records one materialization into a [`FlushPlan`]. Every hook is a no-op +/// once poisoned. Poisoning is reserved for ambiguous alias provenance; raw +/// QMatrix storage and build-local scratch buffers are plan-expressible. +pub(crate) struct PlanRecorder { + pos_of: FxHashMap, + node_count: u32, + target_positions: Box<[u32]>, + /// Buffer provenance: raw buffer pointer -> the slot whose caching event + /// most recently produced that pointer. Seeded with boundary and tensor + /// leaf buffers; updated as steps record. Sound only because `pinned` + /// keeps every tracked `Arc` alive for the recording's duration (a + /// pointer can never be pool-recycled into a different buffer + /// mid-recording) and because incidental sharing poisons the recording + /// (see [`Self::attribute_buffer`]). + provenance: FxHashMap, + /// Strong clones of every buffer entered into `provenance`. Dropped with + /// the recorder at the end of the flush, so no buffer outlives the step. + pinned: Vec>, + physical_edges: Vec<(u32, u32)>, + steps: Vec, + poisoned: bool, +} + +impl PlanRecorder { + pub(crate) fn new( + graph: &ComputeGraphInner, + targets: &[NodeIndex], + fingerprint: FlushFingerprint, + ) -> Self { + let mut poisoned = false; + let target_positions = targets + .iter() + .map(|target| match fingerprint.pos_of.get(target) { + Some(&pos) => pos, + None => { + poisoned = true; + 0 + } + }) + .collect(); + let mut recorder = Self { + node_count: fingerprint.slots.len() as u32, + pos_of: fingerprint.pos_of, + target_positions, + provenance: FxHashMap::default(), + pinned: Vec::new(), + physical_edges: Vec::new(), + steps: Vec::new(), + poisoned, + }; + for (i, &node) in fingerprint.slots.iter().enumerate() { + let Some(node_data) = graph.nodes.nodes.node_weight(node) else { + continue; + }; + if let Some(cached) = &node_data.cached { + // Boundary sharing is incidental (e.g. a cached view aliasing + // its cached input): never reproduced by the plan, so it must + // poison rather than conflate the slots. + recorder.attribute_buffer(cached.buffer(), BufferBinding::Tensor(i as u32), false); + } else if let ComputeGraphNodeVariant::Tensor(data) = &node_data.variant { + recorder.attribute_buffer(data.buffer(), BufferBinding::Tensor(i as u32), false); + } else if let ComputeGraphNodeVariant::QMatrix(operation) = &node_data.variant { + recorder.attribute_buffer( + operation.matrix.buffer(), + BufferBinding::QMatrix(i as u32), + false, + ); + } + } + recorder + } + + /// Attribute `buffer` to slot `pos`, pinning the `Arc` so its raw pointer + /// stays unique for the rest of the recording. + /// + /// Kernel bindings are recorded by looking buffer pointers up in + /// `provenance`, so two distinct slots sharing one buffer are + /// indistinguishable at record time. That is only sound when the plan + /// itself re-creates the sharing every step (`structural`: view aliases, + /// in-place outputs — either slot resolves to the same buffer at replay). + /// Incidental sharing (two cached boundaries or tensor leaves that happen + /// to alias this step) is NOT part of the fingerprint — an isomorphic + /// later step may present distinct buffers, and a plan recorded here + /// would silently bind the wrong tensor. Poison instead. + fn attribute_buffer( + &mut self, + buffer: &Arc, + binding: BufferBinding, + structural: bool, + ) { + self.pinned.push(buffer.clone()); + if let Some(prev) = self + .provenance + .insert(Arc::as_ptr(buffer) as usize, binding) + && prev != binding + && !structural + { + self.poison(); + } + } + + /// Classify one output buffer as Fresh or a resolver-chosen Alias. + /// Returns `None` (and poisons) on any disagreement between pointer + /// provenance and the resolver's claim ledger. + /// Classify one output buffer as Fresh or a resolver-chosen Alias. + /// The provenance slot (the buffer's last structural writer) is the + /// replay-faithful alias target: several fingerprint slots can share one + /// buffer (a producer and its views), so the claim ledger's source node + /// attests that the reuse was deliberately chosen while provenance names + /// the slot whose replay buffer is this buffer. A provenance hit without + /// a chosen claim is accidental buffer sharing between distinct slots + /// (replaying it would bind one slot's data where another slot's is + /// expected), and a claim without a hit is an accounting inconsistency — + /// both poison the plan. + fn classify_output_source( + &mut self, + output_ptr: usize, + output: &TensorData, + claimed_from: Option, + ) -> Option { + match (self.provenance.get(&output_ptr).copied(), claimed_from) { + (Some(BufferBinding::Tensor(slot)), Some(_)) => Some(OutputSource::Alias { slot }), + (None, None) => Some(OutputSource::Fresh { + buffer_size: output.buffer().size(), + usage: output.buffer().usage(), + }), + _ => { + self.poison(); + None + } + } + } + + /// Raw pointers of every buffer this recorder holds a strong clone of. + pub(super) fn pinned_ptrs(&self) -> impl Iterator + '_ { + self.pinned + .iter() + .map(|buffer| Arc::as_ptr(buffer) as usize) + } + + fn pos(&mut self, node: NodeIndex) -> Option { + match self.pos_of.get(&node) { + Some(&pos) => Some(pos), + None => { + self.poison(); + None + } + } + } + + fn consumed_positions(&mut self, deps: &[NodeIndex]) -> Option> { + let mut consumed = Vec::with_capacity(deps.len()); + for &dep in deps { + consumed.push(self.pos(dep)?); + } + Some(consumed.into()) + } + + fn record_kernels( + &mut self, + kernels: &[DirectKernel], + ) -> (Box<[PlanKernel]>, Box<[BufferSpec]>) { + let mut scratch_by_ptr = FxHashMap::::default(); + let mut scratch = Vec::new(); + let mut plan_kernels = Vec::with_capacity(kernels.len()); + for kernel in kernels { + let buffers = kernel.binding_buffers(); + let mut bindings = Vec::with_capacity(buffers.len()); + for buffer in &buffers { + let ptr = Arc::as_ptr(buffer) as usize; + let binding = if let Some(&binding) = self.provenance.get(&ptr) { + binding + } else if let Some(&slot) = scratch_by_ptr.get(&ptr) { + BufferBinding::Scratch(slot) + } else { + let slot = scratch.len() as u32; + scratch.push(BufferSpec { + size: buffer.size(), + usage: buffer.usage(), + }); + scratch_by_ptr.insert(ptr, slot); + // Keep the allocation alive so another scratch buffer + // cannot recycle the same raw pointer while recording. + self.pinned.push(buffer.clone()); + BufferBinding::Scratch(slot) + }; + bindings.push(binding); + } + plan_kernels.push(PlanKernel { + template: kernel.to_template(), + bindings: bindings.into(), + }); + } + (plan_kernels.into(), scratch.into()) + } + + pub(super) fn poison(&mut self) { + self.poisoned = true; + } + + pub(super) fn record_physical_edge(&mut self, from: NodeIndex, to: NodeIndex) { + if self.poisoned { + return; + } + let Some(from) = self.pos(from) else { return }; + let Some(to) = self.pos(to) else { return }; + self.physical_edges.push((from, to)); + } + + pub(super) fn record_tensor_leaf(&mut self, node: NodeIndex, data: &TensorData) { + if self.poisoned { + return; + } + let Some(pos) = self.pos(node) else { return }; + // A leaf sharing a buffer with another slot (boundary or leaf) is + // incidental: poison rather than conflate. + self.attribute_buffer(data.buffer(), BufferBinding::Tensor(pos), false); + self.steps.push(PlanStep::TensorLeaf { pos }); + } + + pub(super) fn record_view_alias( + &mut self, + node: NodeIndex, + result: &TensorData, + deps: &[NodeIndex], + ) { + if self.poisoned { + return; + } + let Some(pos) = self.pos(node) else { return }; + let Some(consumed) = self.consumed_positions(deps) else { + return; + }; + // Structural: replay re-derives this alias from the input's current + // buffer via `try_map_tensor`, so the sharing holds every step. + self.attribute_buffer(result.buffer(), BufferBinding::Tensor(pos), true); + self.steps.push(PlanStep::ViewAlias { pos, consumed }); + } + + pub(super) fn record_shared_alias( + &mut self, + node: NodeIndex, + result: &TensorData, + source: NodeIndex, + ) { + if self.poisoned { + return; + } + let Some(pos) = self.pos(node) else { return }; + let Some(source) = self.pos(source) else { + return; + }; + self.attribute_buffer(result.buffer(), BufferBinding::Tensor(pos), true); + self.steps.push(PlanStep::SharedAlias { + pos, + source, + consumed: vec![source].into_boxed_slice(), + }); + } + + pub(super) fn record_copy_assign( + &mut self, + node: NodeIndex, + output: &TensorData, + op: &QueuedOperation, + ) { + if self.poisoned { + return; + } + let Some(pos) = self.pos(node) else { return }; + let mut deps = Vec::new(); + op.visit_dependencies(&mut |dep| deps.push(dep)); + let Some(consumed) = self.consumed_positions(&deps) else { + return; + }; + // Structural: replay re-derives the in-place output from the current + // graph via `try_prepare_in_place_slice_assign_copy`. + self.attribute_buffer(output.buffer(), BufferBinding::Tensor(pos), true); + self.steps.push(PlanStep::CopyAssign { pos, consumed }); + } + + pub(super) fn record_dispatch( + &mut self, + node: NodeIndex, + kernels: &[DirectKernel], + output: &TensorData, + op: &QueuedOperation, + claimed_from: Option, + ) { + if self.poisoned { + return; + } + let Some(pos) = self.pos(node) else { return }; + + // Alias-vs-Fresh classification requires BOTH the pointer-provenance + // hit AND the resolver's explicit claim to agree on the source slot: + // a provenance hit without a chosen claim is accidental buffer + // sharing between distinct slots (replaying it would bind one slot's + // data where another slot's is expected — poison), and a claim + // without a hit is an accounting inconsistency (poison). `pinned` keeps every + // tracked pointer alive, so hits are never pool-recycled coincidences. + let output_ptr = Arc::as_ptr(output.buffer()) as usize; + let Some(source) = self.classify_output_source(output_ptr, output, claimed_from) else { + return; + }; + let output_spec = OutputSpec { + layout: output.layout().clone(), + datatype: output.datatype(), + source, + }; + // Register the output before classifying bindings: the output buffer + // is itself a binding slot (appended by `Operation::inputs`). + // Structural: replay reproduces the alias through `OutputSource`. + self.attribute_buffer(output.buffer(), BufferBinding::Tensor(pos), true); + + let (plan_kernels, scratch) = self.record_kernels(kernels); + + let mut deps = Vec::new(); + op.visit_dependencies(&mut |dep| deps.push(dep)); + let Some(consumed) = self.consumed_positions(&deps) else { + return; + }; + self.steps.push(PlanStep::Dispatch { + pos, + kernels: plan_kernels, + scratch, + output: output_spec, + consumed, + }); + } + + /// Record one horizontally merged dispatch. Every segment output must be + /// a fresh buffer: a provenance hit means the output aliases an existing + /// slot, which the merged replay arm cannot re-create — poison. + pub(super) fn record_merged_dispatch( + &mut self, + segment_outputs: &[(NodeIndex, &TensorData, Option)], + kernels: &[DirectKernel], + op: &super::merge_horizontal::MergedSegments, + ) { + if self.poisoned { + return; + } + let mut outputs = Vec::with_capacity(segment_outputs.len()); + for (node, output, claimed_from) in segment_outputs { + let Some(pos) = self.pos(*node) else { return }; + let output_ptr = Arc::as_ptr(output.buffer()) as usize; + let Some(source) = self.classify_output_source(output_ptr, output, *claimed_from) + else { + return; + }; + outputs.push(( + pos, + OutputSpec { + layout: output.layout().clone(), + datatype: output.datatype(), + source, + }, + )); + // Structural: the merged replay arm re-creates this slot (fresh + // allocation or the recorded alias). + self.attribute_buffer(output.buffer(), BufferBinding::Tensor(pos), true); + } + let (plan_kernels, scratch) = self.record_kernels(kernels); + + let mut deps = Vec::new(); + op.visit_dependencies(&mut |dep| deps.push(dep)); + let Some(consumed) = self.consumed_positions(&deps) else { + return; + }; + self.steps.push(PlanStep::MergedDispatch { + outputs: outputs.into(), + kernels: plan_kernels, + scratch, + consumed, + }); + } + + pub(crate) fn finish(self) -> Option { + if self.poisoned { + return None; + } + Some(FlushPlan { + node_count: self.node_count, + target_positions: self.target_positions, + physical_edges: self.physical_edges.into(), + steps: self.steps.into(), + }) + } +} + +/// Replay a recorded plan against the current step's isomorphic graph. +/// Returns `false` (without having mutated anything) if upfront validation +/// fails; the caller then falls back to a full resolve. +fn validate_replay( + graph: &ComputeGraphInner, + device: &Device, + plan: &FlushPlan, + fingerprint: &FlushFingerprint, +) -> bool { + let slots = &fingerprint.slots; + if plan.node_count as usize != slots.len() { + return false; + } + // Upfront validation, before any mutation: every step's slot must hold + // the node kind the plan expects and must not be cached yet, and every + // boundary must be cached. After this point the plan cannot fail. + let check_slot = |pos: u32, kind: u8| -> bool { + let Some(node) = graph.nodes.nodes.node_weight(slots[pos as usize]) else { + return false; + }; + if node.cached.is_some() { + return false; + } + matches!( + (&node.variant, kind), + (ComputeGraphNodeVariant::Tensor(_), 0) + | (ComputeGraphNodeVariant::View(_), 1) + | (ComputeGraphNodeVariant::Assign(_), 2) + | (_, 3) + ) + }; + // An alias output overwrites its source slot's buffer. A source + // produced by an earlier plan step is exclusively owned by this replay. + // A boundary source (cached before the flush) is legal only when its + // liveness matches the recording: the fingerprint pins the node-level + // bits (reference count, descendants), and the strong count check below + // catches non-graph holders (an in-flight download, a raw clone) — any + // mismatch falls back to a full resolve, which re-decides the claims. + let mut produced = vec![false; slots.len()]; + for step in &plan.steps { + let alias_ok = |spec: &OutputSpec, produced: &[bool]| match spec.source { + OutputSource::Alias { slot } => { + if produced[slot as usize] { + return true; + } + let Some(node) = graph.nodes.nodes.node_weight(slots[slot as usize]) else { + return false; + }; + let Some(cached) = &node.cached else { + return false; + }; + let buffer = cached.buffer(); + let expected = 1 + u32::from(device.buffer_pool_is_tracked( + buffer.size(), + buffer.usage(), + buffer, + )); + Arc::strong_count(buffer) as u32 == expected + } + OutputSource::Fresh { .. } => true, + }; + let ok = match step { + PlanStep::TensorLeaf { pos } => { + produced[*pos as usize] = true; + check_slot(*pos, 0) + } + PlanStep::ViewAlias { pos, .. } => { + produced[*pos as usize] = true; + check_slot(*pos, 1) + } + PlanStep::SharedAlias { pos, source, .. } => { + let observation_is_uncached = graph + .nodes + .nodes + .node_weight(slots[*pos as usize]) + .is_some_and(|node| node.cached.is_none()); + let ok = produced[*source as usize] && observation_is_uncached; + produced[*pos as usize] = true; + ok + } + PlanStep::CopyAssign { pos, .. } => { + produced[*pos as usize] = true; + check_slot(*pos, 2) + } + PlanStep::Dispatch { pos, output, .. } => { + let ok = check_slot(*pos, 3) && alias_ok(output, &produced); + produced[*pos as usize] = true; + ok + } + PlanStep::MergedDispatch { outputs, .. } => { + let ok = outputs + .iter() + .all(|(pos, spec)| check_slot(*pos, 3) && alias_ok(spec, &produced)); + for (pos, _) in outputs.iter() { + produced[*pos as usize] = true; + } + ok + } + }; + if !ok { + return false; + } + } + for (i, &node) in slots.iter().enumerate() { + if fingerprint.boundary[i] && graph.get_cached_result(node).is_none() { + return false; + } + } + true +} + +/// Execute an already validated plan and append `tail` to its final command +/// encoder before submission. Validation and execution are split so callers +/// retain ownership of an `FnOnce` tail when a plan miss must fall back to the +/// full resolver. +pub(crate) fn execute_replay_with_tail( + graph: &mut ComputeGraphInner, + plan: &FlushPlan, + fingerprint: &FlushFingerprint, + tail: impl FnOnce(&TensorData, &mut wgpu::CommandEncoder) -> T, +) -> (super::ResolverResult, T) { + let slots = &fingerprint.slots; + let device = graph.device(); + let host_trace = device.config().trace_resolve_host; + let start = host_trace.then(Instant::now); + + // Re-add the optimizer's persistent physical dependency edges through the + // liveness-maintaining API. + for &(from, to) in plan.physical_edges.iter() { + graph.add_dependency_edge(slots[from as usize], slots[to as usize]); + } + + // Structural consumer counts, identical to the recording resolve's + // `remaining_consumers` map (keyed by slot position instead of node). + let mut counts = vec![0u32; slots.len()]; + for step in &plan.steps { + for &c in step_consumed(step) { + counts[c as usize] += 1; + } + } + let mut is_target = vec![false; slots.len()]; + for &t in plan.target_positions.iter() { + is_target[t as usize] = true; + } + let mut counts = super::execution::SlotConsumers { + slots, + counts: &mut counts, + is_target: &is_target, + }; + + // Buffers per slot, captured at each slot's caching event. Kept locally + // (not read back through `cached`) so mid-replay releases can't drop a + // buffer a later dispatch still binds — mirroring how the full resolve + // pins bound buffers in its command records. + let mut slot_buffers: Vec>> = vec![None; slots.len()]; + let mut qmatrix_buffers: Vec>> = vec![None; slots.len()]; + for (i, &node) in slots.iter().enumerate() { + if fingerprint.boundary[i] { + slot_buffers[i] = graph.get_cached_result(node).map(|d| d.buffer().clone()); + } else if let Some(node) = graph.nodes.nodes.node_weight(node) + && let ComputeGraphNodeVariant::QMatrix(operation) = &node.variant + { + qmatrix_buffers[i] = Some(operation.matrix.buffer().clone()); + } + } + + let kernel_cache = device.kernel_cache(); + let mut commands = Vec::::with_capacity(plan.steps.len()); + for step in &plan.steps { + match step { + PlanStep::TensorLeaf { pos } => { + let idx = slots[*pos as usize]; + let ComputeGraphNodeVariant::Tensor(data) = &graph + .nodes + .nodes + .node_weight(idx) + .expect("flush replay: validated slot disappeared") + .variant + else { + unreachable!("flush replay: validated tensor leaf changed kind"); + }; + let data = data.clone(); + slot_buffers[*pos as usize] = Some(data.buffer().clone()); + graph.set_cached_result(idx, data); + } + PlanStep::ViewAlias { pos, consumed } => { + let idx = slots[*pos as usize]; + let result = { + let node = graph + .nodes + .nodes + .node_weight(idx) + .expect("flush replay: validated slot disappeared"); + let ComputeGraphNodeVariant::View(view) = &node.variant else { + unreachable!("flush replay: validated view alias changed kind"); + }; + let input = graph + .get_cached_result(view.input) + .expect("flush replay: view alias input must be cached"); + view.try_map_tensor(input) + .expect("flush replay: recorded view alias must still map") + }; + slot_buffers[*pos as usize] = Some(result.buffer().clone()); + graph.set_cached_result(idx, result); + release_slots(graph, &mut counts, consumed); + } + PlanStep::SharedAlias { + pos, + source, + consumed, + } => { + let idx = slots[*pos as usize]; + let result = graph + .get_cached_result(slots[*source as usize]) + .expect("flush replay: shared representative must be cached"); + slot_buffers[*pos as usize] = Some(result.buffer().clone()); + graph.set_cached_result(idx, result.clone()); + release_slots(graph, &mut counts, consumed); + } + PlanStep::CopyAssign { pos, consumed } => { + let idx = slots[*pos as usize]; + let (output, copies) = { + let node = graph + .nodes + .nodes + .node_weight(idx) + .expect("flush replay: validated slot disappeared"); + let ComputeGraphNodeVariant::Assign(op) = &node.variant else { + unreachable!("flush replay: validated slice assign changed kind"); + }; + Resolver::try_prepare_in_place_slice_assign_copy(graph, op) + .expect("flush replay: recorded slice-assign copy must still apply") + }; + slot_buffers[*pos as usize] = Some(output.buffer().clone()); + graph.set_cached_result(idx, output); + commands.extend(copies.into_iter().map(CommandRecord::CopyBuffer)); + release_slots(graph, &mut counts, consumed); + } + PlanStep::Dispatch { + pos, + kernels, + scratch, + output, + consumed, + } => { + let idx = slots[*pos as usize]; + let output_data = match &output.source { + OutputSource::Fresh { buffer_size, usage } => { + let buffer = device.create_buffer(*buffer_size, *usage); + TensorData::new_from_parts( + &device, + buffer, + output.layout.clone(), + output.datatype, + ) + } + OutputSource::Alias { slot } => { + let buffer = slot_buffers[*slot as usize] + .clone() + .expect("flush replay: alias output slot has no buffer"); + TensorData::new_from_parts( + &device, + buffer, + output.layout.clone(), + output.datatype, + ) + } + }; + slot_buffers[*pos as usize] = Some(output_data.buffer().clone()); + let scratch_buffers = scratch + .iter() + .map(|spec| device.create_buffer(spec.size, spec.usage)) + .collect::>(); + for kernel in kernels.iter() { + let buffers = kernel + .bindings + .iter() + .map(|binding| match *binding { + BufferBinding::Tensor(slot) => slot_buffers[slot as usize] + .clone() + .expect("flush replay: tensor binding slot has no buffer"), + BufferBinding::QMatrix(slot) => qmatrix_buffers[slot as usize] + .clone() + .expect("flush replay: QMatrix binding slot has no buffer"), + BufferBinding::Scratch(slot) => scratch_buffers[slot as usize].clone(), + }) + .collect::>(); + let bound = kernel.template.bind_buffers(&buffers); + if let Some(dispatch) = bound.prepare_dispatch(kernel_cache) { + commands.push(CommandRecord::Dispatch(DispatchRecord { + dispatch, + name: bound.name().to_string(), + category: None, + })); + } + } + graph.set_cached_result(idx, output_data); + release_slots(graph, &mut counts, consumed); + } + PlanStep::MergedDispatch { + outputs, + kernels, + scratch, + consumed, + } => { + for (pos, output) in outputs.iter() { + let buffer = match &output.source { + OutputSource::Fresh { buffer_size, usage } => { + device.create_buffer(*buffer_size, *usage) + } + OutputSource::Alias { slot } => slot_buffers[*slot as usize] + .clone() + .expect("flush replay: alias output slot has no buffer"), + }; + let output_data = TensorData::new_from_parts( + &device, + buffer, + output.layout.clone(), + output.datatype, + ); + slot_buffers[*pos as usize] = Some(output_data.buffer().clone()); + graph.set_cached_result(slots[*pos as usize], output_data); + } + let scratch_buffers = scratch + .iter() + .map(|spec| device.create_buffer(spec.size, spec.usage)) + .collect::>(); + for kernel in kernels.iter() { + let buffers = kernel + .bindings + .iter() + .map(|binding| match *binding { + BufferBinding::Tensor(slot) => slot_buffers[slot as usize] + .clone() + .expect("flush replay: tensor binding slot has no buffer"), + BufferBinding::QMatrix(slot) => qmatrix_buffers[slot as usize] + .clone() + .expect("flush replay: QMatrix binding slot has no buffer"), + BufferBinding::Scratch(slot) => scratch_buffers[slot as usize].clone(), + }) + .collect::>(); + let bound = kernel.template.bind_buffers(&buffers); + if let Some(dispatch) = bound.prepare_dispatch(kernel_cache) { + commands.push(CommandRecord::Dispatch(DispatchRecord { + dispatch, + name: bound.name().to_string(), + category: None, + })); + } + } + release_slots(graph, &mut counts, consumed); + } + } + } + + let total_kernels = commands + .iter() + .filter(|command| matches!(command, CommandRecord::Dispatch(_))) + .count(); + let trace = device.config().trace_decode || device.config().trace_resolve; + if trace { + let trace_names = device.config().trace_decode_names; + let mut categories = FxHashMap::::default(); + let mut names = FxHashMap::::default(); + for command in &commands { + let CommandRecord::Dispatch(record) = command else { + continue; + }; + let category = record.name.split('_').take(2).collect::>().join("_"); + *categories.entry(category).or_default() += 1; + if trace_names { + *names.entry(record.name.clone()).or_default() += 1; + } + } + let mut categories = categories.into_iter().collect::>(); + categories.sort_by(|a, b| a.0.cmp(&b.0)); + tracing::info!("resolve_dispatch_categories {categories:?} replayed=true"); + if trace_names { + let mut names = names.into_iter().collect::>(); + names.sort_by(|a, b| a.0.cmp(&b.0)); + tracing::info!("resolve_dispatch_names {names:?} replayed=true"); + } + } + let mut command_encoder = encode_commands(&device, &commands, total_kernels); + let target = slots[plan.target_positions[0] as usize]; + let data = graph + .get_result(target) + .expect("flush replay: target result not cached"); + let tail_result = tail(&data, &mut command_encoder); + device.wgpu_queue().submit(Some(command_encoder.finish())); + device.reset_initialized_buffers(); + device.flush_plan_cache().note_replay(); + + if let Some(start) = start { + tracing::info!( + "resolve_host_profile queued_ops={} kernels={total_kernels} total={:?} replayed=true", + plan.steps.len(), + start.elapsed(), + ); + } + ( + super::ResolverResult { + data, + total_kernels, + }, + tail_result, + ) +} + +fn step_consumed(step: &PlanStep) -> &[u32] { + match step { + PlanStep::TensorLeaf { .. } => &[], + PlanStep::ViewAlias { consumed, .. } + | PlanStep::SharedAlias { consumed, .. } + | PlanStep::CopyAssign { consumed, .. } + | PlanStep::Dispatch { consumed, .. } + | PlanStep::MergedDispatch { consumed, .. } => consumed, + } +} + +/// Release the plan slots one step consumed, evaluating the live release +/// predicate against the current graph. +fn release_slots( + graph: &mut ComputeGraphInner, + counts: &mut super::execution::SlotConsumers<'_>, + consumed: &[u32], +) { + super::execution::release_consumed(graph, counts, None, |release| { + consumed.iter().copied().for_each(release) + }); +} + +/// Encode and submit the replayed command stream with the same pass/submit +/// chunking policy as the full resolver (single pass + single submit below +/// 1024 kernels; chunked with Metal waits above). +fn encode_commands( + device: &Device, + commands: &[CommandRecord], + total_kernels: usize, +) -> wgpu::CommandEncoder { + let command_encoder = + device + .wgpu_device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("Resolver Encoder"), + }); + super::queue_executor::encode_command_records( + device, + commands, + total_kernels, + None, + command_encoder, + |encoder, wait| { + device.wgpu_queue().submit(Some(encoder.finish())); + if wait { + device.poll_wait(); + } + }, + ) +} diff --git a/fusor-ml/core/src/compute_graph/resolve/fold_views.rs b/fusor-ml/core/src/compute_graph/resolve/fold_views.rs deleted file mode 100644 index 9de390b0c..000000000 --- a/fusor-ml/core/src/compute_graph/resolve/fold_views.rs +++ /dev/null @@ -1,182 +0,0 @@ -use crate::view::ViewOperation; - -use super::*; - -impl Resolver { - /// Fold view inputs of an n-ary node directly into its expression: each - /// `IndexedInput` through the view becomes a load of the view's base - /// node with the view's coordinate map applied (and a bounds-select - /// around partially-defined stages). This removes the view node from - /// between the producer and consumer, so the n-ary fusion passes see - /// through layout changes instead of stopping at them. - /// - /// Affine maps always fold — they rewrite to plain index arithmetic. - /// Maps that need delinearization (divmod address arithmetic, from a - /// reshape regrouping non-mergeable strides) re-derive coordinates on - /// every load, so they only fold where each element is loaded once; a - /// load re-read across unindexed dims (a contraction operand) keeps the - /// view node and materializes through the gather fallback instead. - pub(super) fn try_fold_view_inputs( - &mut self, - graph: &mut ComputeGraphInner, - node_idx: ExecutionNodeIndex, - ) -> bool { - let ExecutionVariant::Elementwise(nary) = self.execution_graph[node_idx].variant.clone() - else { - return false; - }; - - let mut expression = nary.expression.clone(); - let mut inputs = nary.inputs.clone(); - let mut folded = Vec::new(); - - for (slot, input_inner) in nary.inputs.iter().copied().enumerate() { - if self.check_cached(graph, input_inner) { - continue; - } - let Some(input_exec) = self.get_input_node_in_exec_graph(input_inner) else { - continue; - }; - if !self.execution_graph.contains_node(input_exec) { - continue; - } - let ExecutionVariant::View(view) = &self.execution_graph[input_exec].variant else { - continue; - }; - let needs_delinearize = view.stages.iter().any(|stage| { - crate::view::affine_dim_indices(&stage.layout, &stage.input_shape).is_none() - }); - if needs_delinearize && input_reread_factor(&expression, &nary.shape, slot) > 1 { - continue; - } - let view = view.clone(); - let Some(rewritten) = Self::rewrite_view_input(&expression, slot, &view) else { - continue; - }; - expression = rewritten; - inputs[slot] = view.input; - folded.push((input_exec, view.input)); - } - if folded.is_empty() { - return false; - } - - let (final_inputs, final_expression) = Self::deduplicate_inputs(inputs, expression); - let new_nary = ElementwiseOperation { - inputs: final_inputs, - expression: final_expression, - shape: nary.shape.clone(), - output_datatype: nary.output_datatype, - }; - self.execution_graph[node_idx].variant = ExecutionVariant::Elementwise(new_nary.clone()); - - for (view_exec, base_inner) in &folded { - if let Some(edge) = self.execution_graph.find_edge(*view_exec, node_idx) { - self.execution_graph.remove_edge(edge); - } - if let Some(base_exec) = self.get_input_node_in_exec_graph(*base_inner) - && self - .execution_graph - .find_edge(base_exec, node_idx) - .is_none() - { - self.execution_graph.add_edge(base_exec, node_idx, ()); - } - } - self.add_physical_dependencies(graph, node_idx, &new_nary.inputs); - for (view_exec, _) in folded { - self.remove_node_if_dead(view_exec); - } - true - } - - /// Rewrite every access to input `target_idx` through `view`'s - /// coordinate map: the original index expressions (the view's output - /// coordinates) walk down the stage stack to base coordinates, with - /// fill selects and in-bounds clamps where stages are partially defined - /// (both select branches evaluate). - fn rewrite_view_input( - expr: &NaryExpr, - target_idx: usize, - view: &ViewOperation, - ) -> Option { - Some(match expr { - NaryExpr::Op { children, function } => NaryExpr::Op { - children: children - .iter() - .map(|child| Self::rewrite_view_input(child, target_idx, view)) - .collect::>>()?, - function: function.clone(), - }, - NaryExpr::IndexedInput { input_idx, indices } => { - let indices: Vec = indices - .iter() - .map(|index| Self::rewrite_view_input(index, target_idx, view)) - .collect::>>()?; - if *input_idx != target_idx { - NaryExpr::IndexedInput { - input_idx: *input_idx, - indices, - } - } else { - view.value_expression(*input_idx, &indices)?.0 - } - } - NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), - NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), - }) - } -} - -/// The worst re-read factor across this slot's loads: the product of -/// index-space dims a load's coordinates never reference — each such dim -/// re-reads the same element once per step. -fn input_reread_factor(expr: &NaryExpr, shape: &[usize], slot: usize) -> usize { - fn collect_dims(expr: &NaryExpr, referenced: &mut [bool]) { - match expr { - NaryExpr::Op { children, .. } => { - for child in children { - collect_dims(child, referenced); - } - } - NaryExpr::IndexedInput { indices, .. } => { - for index in indices { - collect_dims(index, referenced); - } - } - NaryExpr::DimIndex(dim) => referenced[*dim] = true, - NaryExpr::Scalar(_) => {} - } - } - fn visit_loads(expr: &NaryExpr, shape: &[usize], slot: usize, worst: &mut usize) { - match expr { - NaryExpr::Op { children, .. } => { - for child in children { - visit_loads(child, shape, slot, worst); - } - } - NaryExpr::IndexedInput { input_idx, indices } => { - for index in indices { - visit_loads(index, shape, slot, worst); - } - if *input_idx == slot { - let mut referenced = vec![false; shape.len()]; - for index in indices { - collect_dims(index, &mut referenced); - } - let factor: usize = shape - .iter() - .zip(&referenced) - .filter(|(_, referenced)| !**referenced) - .map(|(size, _)| *size) - .product(); - *worst = (*worst).max(factor); - } - } - NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => {} - } - } - let mut worst = 1; - visit_loads(expr, shape, slot, &mut worst); - worst -} diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_basic.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_basic.rs deleted file mode 100644 index 89eca64d1..000000000 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_basic.rs +++ /dev/null @@ -1,574 +0,0 @@ -use super::*; - -impl Resolver { - pub(super) fn try_fuse_naries( - &mut self, - graph: &mut ComputeGraphInner, - node_idx: ExecutionNodeIndex, - ) -> bool { - let node_variant = self.execution_graph[node_idx].variant.clone(); - - let ExecutionVariant::Elementwise(nary) = node_variant else { - return false; - }; - - // Collect all fusible nary inputs - let mut expression = nary.expression.clone(); - let mut all_inputs = nary.inputs.clone(); - let mut fused_execs = Vec::new(); - - let max_fused_inputs = graph.device().nary_direct_input_binding_budget(); - - for &input_inner in nary.inputs.iter() { - if self.check_cached(graph, input_inner) { - continue; - } - let Some(input_exec) = self.get_input_node_in_exec_graph(input_inner) else { - continue; - }; - // Check if the node still exists (it may have been removed during optimization) - if !self.execution_graph.contains_node(input_exec) { - continue; - } - // Inlining duplicates the producer's work unless this node is its - // only consumer: the producer still materializes for everyone - // else (e.g. the residual stream feeds every later layer — fusing - // it forward would re-sum the whole prefix per layer). A user-held - // reference alone doesn't block fusion — only another consumer in - // this resolve does. - if self - .execution_graph - .neighbors_directed(input_exec, petgraph::Direction::Outgoing) - .count() - != 1 - { - continue; - } - let ExecutionVariant::Elementwise(input_nary) = - &self.execution_graph[input_exec].variant - else { - continue; - }; - - // Inline: offset input nary's indices to append after current inputs. - let offset = all_inputs.len(); - let inlined = Self::offset_input_indices(&input_nary.expression, offset); - // `input_inner` may appear in `all_inputs` at multiple slots — - // beyond the explicit `input_idx` slot from `nary.inputs`, earlier - // fusions in this same loop can have inlined chains that - // re-introduce `input_inner` at later slots. Substitute at every - // such slot so we don't leave dangling `IndexedInput` references - // pointing to a now-fused-away node. - let target_slots: Vec = all_inputs - .iter() - .enumerate() - .filter_map(|(slot, value)| (*value == input_inner).then_some(slot)) - .collect(); - let mut new_expression = expression.clone(); - let mut success = true; - for slot in target_slots { - let (next, s) = Self::substitute_input_in_expr(&new_expression, slot, &inlined); - new_expression = next; - success &= s; - } - - // Only fuse if substitution was successful - // If not, the expression still references the original input which must remain - if success { - // Count unique inputs after potential merge (duplicates share a binding). - let unique_inputs: FxHashSet<_> = all_inputs - .iter() - .chain(input_nary.inputs.iter()) - .copied() - .collect(); - - if unique_inputs.len() > max_fused_inputs { - // Skip fusion - would exceed GPU binding limit - continue; - } - - expression = new_expression; - all_inputs.extend(input_nary.inputs.iter().copied()); - fused_execs.push((input_exec, input_nary.inputs.clone())); - } - } - - if fused_execs.is_empty() { - return false; - } - - // Deduplicate and remove unused inputs - let (final_inputs, final_expression) = Self::deduplicate_inputs(all_inputs, expression); - - let new_nary = ElementwiseOperation { - inputs: final_inputs.clone(), - expression: final_expression, - shape: nary.shape.clone(), - output_datatype: nary.output_datatype, - }; - - self.execution_graph[node_idx].variant = ExecutionVariant::Elementwise(new_nary.clone()); - - // Update graph edges - for (input_exec, new_inputs) in fused_execs { - if let Some(edge) = self.execution_graph.find_edge(input_exec, node_idx) { - self.execution_graph.remove_edge(edge); - } - for &new_input in &new_inputs { - if let Some(exec) = self.get_input_node_in_exec_graph(new_input) - && self.execution_graph.find_edge(exec, node_idx).is_none() - { - self.execution_graph.add_edge(exec, node_idx, ()); - } - } - self.remove_node_if_dead(input_exec); - } - - self.add_physical_dependencies(graph, node_idx, &new_nary.inputs); - true - } - - /// Add offset to all input indices in an expression. - pub(super) fn offset_input_indices(expr: &NaryExpr, offset: usize) -> NaryExpr { - match expr { - NaryExpr::Op { children, function } => NaryExpr::Op { - children: children - .iter() - .map(|c| Self::offset_input_indices(c, offset)) - .collect(), - function: function.clone(), - }, - NaryExpr::IndexedInput { input_idx, indices } => NaryExpr::IndexedInput { - input_idx: input_idx + offset, - indices: indices - .iter() - .map(|c| Self::offset_input_indices(c, offset)) - .collect(), - }, - NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), - NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), - } - } - - /// Substitute IndexedInput(target_idx) with element-wise access with the replacement expression. - /// Returns (new_expression, success) where success is true if all references to target_idx - /// were successfully substituted. If false, the input should NOT be removed from the graph. - pub(super) fn substitute_input_in_expr( - expr: &NaryExpr, - target_idx: usize, - replacement: &NaryExpr, - ) -> (NaryExpr, bool) { - /// Helper to extract input_idx from an IndexedInput with element-wise access - fn get_elementwise_input_idx(expr: &NaryExpr) -> Option { - match expr { - NaryExpr::IndexedInput { input_idx, indices } - if NaryExpr::is_elementwise_indices(indices) => - { - Some(*input_idx) - } - _ => None, - } - } - - match expr { - NaryExpr::Op { children, function } => { - let mut all_success = true; - let new_children: Vec<_> = children - .iter() - .map(|c| { - let (new_c, success) = - Self::substitute_input_in_expr(c, target_idx, replacement); - all_success &= success; - new_c - }) - .collect(); - ( - NaryExpr::Op { - children: new_children, - function: function.clone(), - }, - all_success, - ) - } - NaryExpr::IndexedInput { input_idx, indices } => { - if *input_idx == target_idx { - // Check if this is element-wise access - if NaryExpr::is_elementwise_indices(indices) { - // Element-wise can be fully replaced with any expression - (replacement.clone(), true) - } else { - // Custom indexing can only substitute if replacement is also element-wise - if let Some(new_idx) = get_elementwise_input_idx(replacement) { - let mut all_success = true; - let new_indices: Vec<_> = indices - .iter() - .map(|c| { - let (new_c, success) = - Self::substitute_input_in_expr(c, target_idx, replacement); - all_success &= success; - new_c - }) - .collect(); - ( - NaryExpr::IndexedInput { - input_idx: new_idx, - indices: new_indices, - }, - all_success, - ) - } else { - // Cannot fuse complex expression into custom indexed input - let all_success = false; - let new_indices: Vec<_> = indices - .iter() - .map(|c| { - let (new_c, _) = - Self::substitute_input_in_expr(c, target_idx, replacement); - new_c - }) - .collect(); - ( - NaryExpr::IndexedInput { - input_idx: *input_idx, - indices: new_indices, - }, - all_success, - ) - } - } - } else { - // Recurse into the index expressions - let mut all_success = true; - let new_indices: Vec<_> = indices - .iter() - .map(|c| { - let (new_c, s) = - Self::substitute_input_in_expr(c, target_idx, replacement); - all_success &= s; - new_c - }) - .collect(); - ( - NaryExpr::IndexedInput { - input_idx: *input_idx, - indices: new_indices, - }, - all_success, - ) - } - } - NaryExpr::DimIndex(dim) => (NaryExpr::DimIndex(*dim), true), - NaryExpr::Scalar(value) => (NaryExpr::Scalar(*value), true), - } - } - - /// Remove unused inputs and deduplicate, returning new inputs and remapped expression. - pub(super) fn deduplicate_inputs( - inputs: Vec, - expr: NaryExpr, - ) -> (Vec, NaryExpr) { - // Collect which input indices are actually used - let mut seen_indices = FxHashSet::default(); - let mut used_indices = Vec::new(); - Self::collect_used_inputs(&expr, &mut seen_indices, &mut used_indices); - - // Build the input-index remap, collecting only used inputs. - let mut new_inputs = Vec::new(); - let mut old_to_new = FxHashMap::default(); - - for old_idx in used_indices { - let node = inputs[old_idx]; - // Check if this node already exists in new_inputs (deduplication) - let new_idx = if let Some(existing) = new_inputs.iter().position(|&n| n == node) { - existing - } else { - let idx = new_inputs.len(); - new_inputs.push(node); - idx - }; - old_to_new.insert(old_idx, new_idx); - } - - let new_expr = Self::remap_input_indices(&expr, &old_to_new); - (new_inputs, new_expr) - } - - pub(super) fn collect_used_inputs( - expr: &NaryExpr, - seen: &mut FxHashSet, - used: &mut Vec, - ) { - match expr { - NaryExpr::Op { children, .. } => { - for child in children { - Self::collect_used_inputs(child, seen, used); - } - } - NaryExpr::IndexedInput { input_idx, indices } => { - if seen.insert(*input_idx) { - used.push(*input_idx); - } - for c in indices { - Self::collect_used_inputs(c, seen, used); - } - } - NaryExpr::DimIndex(_) => {} - NaryExpr::Scalar(_) => {} - } - } - - pub(super) fn remap_input_indices( - expr: &NaryExpr, - mapping: &FxHashMap, - ) -> NaryExpr { - match expr { - NaryExpr::Op { children, function } => NaryExpr::Op { - children: children - .iter() - .map(|c| Self::remap_input_indices(c, mapping)) - .collect(), - function: function.clone(), - }, - NaryExpr::IndexedInput { input_idx, indices } => NaryExpr::IndexedInput { - input_idx: mapping[input_idx], - indices: indices - .iter() - .map(|c| Self::remap_input_indices(c, mapping)) - .collect(), - }, - NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), - NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), - } - } - - pub(super) fn replace_inputs_in_expr( - expr: &NaryExpr, - replacements: &[Option], - ) -> Option { - match expr { - NaryExpr::Op { children, function } => Some(NaryExpr::Op { - children: children - .iter() - .map(|child| Self::replace_inputs_in_expr(child, replacements)) - .collect::>>()?, - function: function.clone(), - }), - NaryExpr::IndexedInput { input_idx, indices } => { - if let Some(replacement) = replacements.get(*input_idx).and_then(|r| r.as_ref()) { - if NaryExpr::is_elementwise_indices(indices) { - Some(replacement.clone()) - } else { - None - } - } else { - Some(NaryExpr::IndexedInput { - input_idx: *input_idx, - indices: indices - .iter() - .map(|index| Self::replace_inputs_in_expr(index, replacements)) - .collect::>>()?, - }) - } - } - NaryExpr::DimIndex(dim) => Some(NaryExpr::DimIndex(*dim)), - NaryExpr::Scalar(value) => Some(NaryExpr::Scalar(*value)), - } - } - - pub(super) fn qmatmul_same_base(first: &QMatMulOperation, second: &QMatMulOperation) -> bool { - first.input_datatype == second.input_datatype - && first.input == second.input - && first.matrix == second.matrix - && first.in_shape == second.in_shape - && first.out_shape == second.out_shape - && first.pre_element_wise_expr == second.pre_element_wise_expr - && first.post_accumulator_offsets == second.post_accumulator_offsets - } - - pub(super) fn qmatmul_output_expr( - qmatmul: &QMatMulOperation, - extras: &mut Vec, - rank: usize, - ) -> Option { - if let Some(epilogue) = &qmatmul.post_element_wise_expr { - let value_arity = qmatmul.post_accumulator_offsets.len().max(1); - let mut mapping = Vec::with_capacity(value_arity + epilogue.extras.len()); - mapping.extend(0..value_arity); - mapping.extend((0..epilogue.extras.len()).map(|i| extras.len() + value_arity + i)); - extras.extend(epilogue.extras.iter().copied()); - Some(epilogue.expression.remap_inputs(&mapping)) - } else { - Some(NaryExpr::input(0, rank)) - } - } - - /// Try to extract a unary function chain from a node variant. - /// Only Nary ops with a single input and element-wise access can be converted. - pub(super) fn try_get_unary_chain(variant: &ExecutionVariant) -> Option { - match variant { - ExecutionVariant::Elementwise(nary) => nary.try_extract_unary_chain(), - _ => None, - } - } - - pub(super) fn try_fuse_into_reduce( - &mut self, - graph: &mut ComputeGraphInner, - node_idx: ExecutionNodeIndex, - ) -> bool { - let node_variant = self.execution_graph[node_idx].variant.clone(); - - let Some(el_op) = Self::try_get_unary_chain(&node_variant) else { - return false; - }; - - let input_inner = el_op.value; - if self.check_cached(graph, input_inner) { - return false; - } - - let Some(input_exec_idx) = self.get_input_node_in_exec_graph(input_inner) else { - return false; - }; - - let input_variant = self.execution_graph[input_exec_idx].variant.clone(); - let ExecutionVariant::Reduce(reduce_op) = input_variant else { - return false; - }; - - let mut new_reduce = reduce_op.clone(); - let mut existing_post = new_reduce.post_element_wise.functions.clone(); - existing_post.extend(el_op.functions.functions.iter().cloned()); - new_reduce.post_element_wise = - UnaryFunctionChain::new(existing_post, reduce_op.post_element_wise.input_datatype()); - - self.execution_graph[node_idx].variant = ExecutionVariant::Reduce(new_reduce.clone()); - - for &reduce_input_inner in &reduce_op.inputs { - if let Some(reduce_input_exec) = self.get_input_node_in_exec_graph(reduce_input_inner) { - self.execution_graph - .add_edge(reduce_input_exec, node_idx, ()); - } - } - - if let Some(edge) = self.execution_graph.find_edge(input_exec_idx, node_idx) { - self.execution_graph.remove_edge(edge); - } - self.add_physical_dependencies(graph, node_idx, &reduce_op.inputs); - self.remove_node_if_dead(input_exec_idx); - true - } - - /// Inline elementwise producers into a reduce's fused expression: the - /// reduce evaluates the producer at every index-space coordinate, so a - /// producer consumed only by this reduce never needs to materialize. - /// Composed contractions that recognition did not claim collapse to a - /// single map-reduce kernel here, where the tiled lowering can stage - /// their reused inputs through workgroup memory. - pub(super) fn try_fuse_producer_into_reduce( - &mut self, - graph: &mut ComputeGraphInner, - node_idx: ExecutionNodeIndex, - ) -> bool { - let ExecutionVariant::Reduce(reduce) = self.execution_graph[node_idx].variant.clone() - else { - return false; - }; - - let mut expression = reduce.expression.clone(); - let mut all_inputs = reduce.inputs.clone(); - let mut fused_execs = Vec::new(); - let max_fused_inputs = graph.device().nary_direct_input_binding_budget(); - - for &input_inner in reduce.inputs.iter() { - if self.check_cached(graph, input_inner) { - continue; - } - let Some(input_exec) = self.get_input_node_in_exec_graph(input_inner) else { - continue; - }; - if !self.execution_graph.contains_node(input_exec) { - continue; - } - // Same sole-consumer rule as nary fusion: inlining a shared - // producer would re-evaluate it for every other consumer. - if self - .execution_graph - .neighbors_directed(input_exec, petgraph::Direction::Outgoing) - .count() - != 1 - { - continue; - } - let ExecutionVariant::Elementwise(input_nary) = - &self.execution_graph[input_exec].variant - else { - continue; - }; - // The reduce evaluates this input across the full index space; - // a producer with any other shape reads out of range. - if input_nary.shape != reduce.shape { - continue; - } - - let offset = all_inputs.len(); - let inlined = Self::offset_input_indices(&input_nary.expression, offset); - let target_slots: Vec = all_inputs - .iter() - .enumerate() - .filter_map(|(slot, value)| (*value == input_inner).then_some(slot)) - .collect(); - let mut new_expression = expression.clone(); - let mut success = true; - for slot in target_slots { - let (next, s) = Self::substitute_input_in_expr(&new_expression, slot, &inlined); - new_expression = next; - success &= s; - } - - if success { - let unique_inputs: FxHashSet<_> = all_inputs - .iter() - .chain(input_nary.inputs.iter()) - .copied() - .collect(); - if unique_inputs.len() > max_fused_inputs { - continue; - } - - expression = new_expression; - all_inputs.extend(input_nary.inputs.iter().copied()); - fused_execs.push((input_exec, input_nary.inputs.clone())); - } - } - - if fused_execs.is_empty() { - return false; - } - - let (final_inputs, final_expression) = Self::deduplicate_inputs(all_inputs, expression); - - let mut new_reduce = reduce.clone(); - new_reduce.inputs = final_inputs; - new_reduce.expression = final_expression; - let new_inputs = new_reduce.inputs.clone(); - self.execution_graph[node_idx].variant = ExecutionVariant::Reduce(new_reduce); - - for (input_exec, producer_inputs) in fused_execs { - if let Some(edge) = self.execution_graph.find_edge(input_exec, node_idx) { - self.execution_graph.remove_edge(edge); - } - for &new_input in &producer_inputs { - if let Some(exec) = self.get_input_node_in_exec_graph(new_input) - && self.execution_graph.find_edge(exec, node_idx).is_none() - { - self.execution_graph.add_edge(exec, node_idx, ()); - } - } - self.remove_node_if_dead(input_exec); - } - - self.add_physical_dependencies(graph, node_idx, &new_inputs); - true - } -} diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs deleted file mode 100644 index d1436ffa5..000000000 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_matmul.rs +++ /dev/null @@ -1,956 +0,0 @@ -use super::*; - -impl Resolver { - pub(super) fn try_fuse_into_matmul( - &mut self, - graph: &mut ComputeGraphInner, - node_idx: ExecutionNodeIndex, - allow_qmatmul_elementwise_fusion: bool, - ) -> bool { - let node_variant = self.execution_graph[node_idx].variant.clone(); - - // Post-op: fuse elementwise after matmul (dense or quantized). - if let Some(el_op) = Self::try_get_unary_chain(&node_variant) { - let input_inner = el_op.value; - if !self.check_cached(graph, input_inner) - && let Some(input_exec_idx) = self.get_input_node_in_exec_graph(input_inner) - { - let input_variant = self.execution_graph[input_exec_idx].variant.clone(); - // An un-flattened operand was chosen for the coop kernel, - // which hosts no element-wise chains: fusing one here would - // demote the matmul to the generic divmod-per-load reduce. - if let ExecutionVariant::MatMul(matmul_op) = input_variant - && matmul_op.a.is_plain() - && matmul_op.b.is_plain() - { - let mut new_matmul = matmul_op.clone(); - let mut existing_post = new_matmul.post_element_wise.functions.clone(); - existing_post.extend(el_op.functions.functions.iter().cloned()); - new_matmul.post_element_wise = UnaryFunctionChain::new( - existing_post, - matmul_op.post_element_wise.input_datatype(), - ); - - self.execution_graph[node_idx].variant = - ExecutionVariant::MatMul(new_matmul.clone()); - - let (first_inner, second_inner) = (matmul_op.first, matmul_op.second); - if let Some(idx) = self.get_input_node_in_exec_graph(first_inner) { - self.execution_graph.add_edge(idx, node_idx, ()); - } - if let Some(idx) = self.get_input_node_in_exec_graph(second_inner) { - self.execution_graph.add_edge(idx, node_idx, ()); - } - if let Some(edge) = self.execution_graph.find_edge(input_exec_idx, node_idx) { - self.execution_graph.remove_edge(edge); - } - self.add_physical_dependencies(graph, node_idx, &[first_inner, second_inner]); - self.remove_node_if_dead(input_exec_idx); - return true; - } - } - } - - // Post-op (QMatMul): fuse a general element-wise expression after - // qmatmul. This handles composite expressions like GELU and ordered - // extra inputs whose layouts match the output visitation shape. - if allow_qmatmul_elementwise_fusion - && let ExecutionVariant::Elementwise(nary) = &node_variant - { - // Split/gate expressions built from `narrow` views of a qmatmul - // output (e.g. SwiGLU's gate/up halves) reach the qmatmul through - // MapLayout chains with distinct last-dimension column offsets. - // Absorb them into the accumulator-offset post epilogue before the - // per-input scan below. - if self.try_fuse_qmatmul_narrow_accumulators(graph, node_idx, nary) { - return true; - } - for (candidate_input_idx, &input_inner) in nary.inputs.iter().enumerate() { - if self.get_input_node_in_exec_graph(input_inner).is_none() { - continue; - } - let (qmatmul_inner, map_chain) = self.walk_view_chain(input_inner); - let Some(qmatmul_exec_idx) = self.get_input_node_in_exec_graph(qmatmul_inner) - else { - continue; - }; - let ExecutionVariant::QMatMul(qmatmul_op) = - self.execution_graph[qmatmul_exec_idx].variant.clone() - else { - continue; - }; - if map_chain.is_none() - && !self.check_cached(graph, input_inner) - && qmatmul_op.post_element_wise_expr.is_none() - && qmatmul_op.in_shape[..qmatmul_op.in_shape.len() - 1] - .iter() - .product::() - == 1 - && let Some((expression, accumulator_offsets, extras)) = self - .try_extract_indexed_qmatmul_post_expr( - graph, - nary, - candidate_input_idx, - &qmatmul_op.out_shape, - ) - { - let Some(input_datatype) = nary - .expression - .elementwise_input_datatype(candidate_input_idx) - else { - continue; - }; - if input_datatype != crate::DataTypeEnum::F32 - || nary.output_datatype != crate::DataTypeEnum::F32 - { - continue; - } - if !qmatmul_op.supports_indexed_post_accumulator_offsets( - &graph.device(), - &nary.shape, - &accumulator_offsets, - ) { - continue; - } - - let post_element_wise_expr = ElementwiseEpilogue { - expression, - extras: extras.clone(), - input_datatype, - output_datatype: nary.output_datatype, - }; - - let mut new_q = qmatmul_op.clone(); - new_q.out_shape = nary.shape.clone(); - new_q.post_element_wise_expr = Some(post_element_wise_expr); - new_q.post_accumulator_offsets = accumulator_offsets.into_boxed_slice(); - - if !new_q.fits_binding_budget(&graph.device()) { - continue; - } - - self.commit_qmatmul_post_fusion(graph, node_idx, &nary.inputs, new_q); - return true; - } - let Some(mapped_layout) = - Self::apply_view_chain(&Layout::contiguous(&qmatmul_op.out_shape), &map_chain) - else { - continue; - }; - if mapped_layout != Layout::contiguous(&nary.shape) { - continue; - } - if !nary.expression.uses_input(candidate_input_idx) - || nary - .expression - .uses_custom_indexing_for_input(candidate_input_idx) - { - continue; - }; - let Some(input_datatype) = nary - .expression - .elementwise_input_datatype(candidate_input_idx) - else { - continue; - }; - let mut extras = Vec::new(); - let mut replacements = vec![None; nary.inputs.len()]; - let mut valid_expression = true; - for (input_idx, &nary_input) in nary.inputs.iter().enumerate() { - let (base_inner, chain) = self.walk_view_chain(nary_input); - let base_qmatmul = - self.get_input_node_in_exec_graph(base_inner) - .and_then(|exec| match &self.execution_graph[exec].variant { - ExecutionVariant::QMatMul(op) => Some(op.clone()), - _ => None, - }); - if let Some(base_qmatmul) = base_qmatmul - && Self::qmatmul_same_base(&qmatmul_op, &base_qmatmul) - { - let alias_layout = Self::apply_view_chain( - &Layout::contiguous(&base_qmatmul.out_shape), - &chain, - ); - if alias_layout == Some(Layout::contiguous(&nary.shape)) - && !nary.expression.uses_custom_indexing_for_input(input_idx) - { - replacements[input_idx] = Self::qmatmul_output_expr( - &base_qmatmul, - &mut extras, - nary.shape.len(), - ); - continue; - } - valid_expression = false; - break; - } - - let Some(extra) = - self.try_normalize_qmatmul_post_extra(graph, nary_input, &nary.shape) - else { - valid_expression = false; - break; - }; - replacements[input_idx] = - Some(NaryExpr::input(extras.len() + 1, nary.shape.len())); - extras.push(extra); - } - if !valid_expression { - continue; - } - let Some(expression) = - Self::replace_inputs_in_expr(&nary.expression, &replacements) - else { - continue; - }; - if self.check_cached(graph, input_inner) - || input_datatype != crate::DataTypeEnum::F32 - || nary.output_datatype != crate::DataTypeEnum::F32 - || !qmatmul_op.supports_elementwise_epilogue_fusion(&graph.device()) - { - continue; - } - - let post_element_wise_expr = ElementwiseEpilogue { - expression, - extras: extras.clone(), - input_datatype: qmatmul_op - .post_element_wise_expr - .as_ref() - .map(|existing| existing.input_datatype) - .unwrap_or(input_datatype), - output_datatype: nary.output_datatype, - }; - - let mut new_q = qmatmul_op.clone(); - new_q.post_element_wise_expr = Some(post_element_wise_expr); - - if !new_q.fits_binding_budget(&graph.device()) { - continue; - } - - self.commit_qmatmul_post_fusion(graph, node_idx, &nary.inputs, new_q); - return true; - } - } - - // Pre-op (QMatMul): fuse a general element-wise expression upstream - // of a single-row qmatmul input. For batched/tiled qmatmul, the - // transformed activation tile is reloaded for each output-column - // tile, so expensive expressions like GELU would be recomputed many - // times. Keep those chains materialized once instead. - if allow_qmatmul_elementwise_fusion - && let ExecutionVariant::QMatMul(qmatmul_op) = &node_variant - && qmatmul_op.in_shape[..qmatmul_op.in_shape.len() - 1] - .iter() - .product::() - == 1 - && qmatmul_op.supports_elementwise_epilogue_fusion(&graph.device()) - && !self.check_cached(graph, qmatmul_op.input) - && let Some(input_exec) = self.get_input_node_in_exec_graph(qmatmul_op.input) - { - let (nary_inner, nary_map_chain) = self.walk_view_chain(qmatmul_op.input); - let Some(nary_exec) = self.get_input_node_in_exec_graph(nary_inner) else { - return false; - }; - let ExecutionVariant::Elementwise(nary) = - self.execution_graph[nary_exec].variant.clone() - else { - return false; - }; - let mapped_layout = - Self::apply_view_chain(&Layout::contiguous(&nary.shape), &nary_map_chain); - if mapped_layout != Some(Layout::contiguous(&qmatmul_op.in_shape)) { - return false; - } - - for (candidate_input_idx, &primary_input) in nary.inputs.iter().enumerate() { - if !nary.expression.uses_input(candidate_input_idx) - || nary - .expression - .uses_custom_indexing_for_input(candidate_input_idx) - { - continue; - } - let Some(input_datatype) = nary - .expression - .elementwise_input_datatype(candidate_input_idx) - else { - continue; - }; - if input_datatype != crate::DataTypeEnum::F32 - || nary.output_datatype != crate::DataTypeEnum::F32 - { - continue; - } - - let (primary_inner, primary_chain) = self.walk_view_chain(primary_input); - let Some(primary_info) = self.infer_layout_cached(graph, primary_inner) else { - continue; - }; - let Some(primary_layout) = - Self::apply_view_chain(primary_info.layout(), &primary_chain) - else { - continue; - }; - if primary_layout != Layout::contiguous(&nary.shape) { - continue; - } - - let mut mapping = vec![usize::MAX; nary.inputs.len()]; - let mut extras = Vec::new(); - let mut valid_expression = true; - for (input_idx, &nary_input) in nary.inputs.iter().enumerate() { - let (base_inner, chain) = self.walk_view_chain(nary_input); - if base_inner == primary_inner { - let alias_layout = Self::apply_view_chain(primary_info.layout(), &chain); - if alias_layout == Some(Layout::contiguous(&nary.shape)) - && !nary.expression.uses_custom_indexing_for_input(input_idx) - { - mapping[input_idx] = 0; - continue; - } - valid_expression = false; - break; - } - - let Some(extra) = - self.try_normalize_qmatmul_post_extra(graph, nary_input, &nary.shape) - else { - valid_expression = false; - break; - }; - mapping[input_idx] = extras.len() + 1; - extras.push(extra); - } - if !valid_expression { - continue; - } - let expression = nary.expression.remap_inputs(&mapping); - - let pre_element_wise_expr = - if let Some(existing) = &qmatmul_op.pre_element_wise_expr { - if existing.input_datatype != nary.output_datatype { - continue; - } - let mut mapping = Vec::with_capacity(1 + existing.extras.len()); - mapping.push(0); - mapping.extend((0..existing.extras.len()).map(|i| i + 1 + extras.len())); - let shifted_existing = existing.expression.remap_inputs(&mapping); - let (expression, success) = - Self::substitute_input_in_expr(&shifted_existing, 0, &expression); - if !success { - continue; - } - let mut combined_extras = extras.clone(); - combined_extras.extend(existing.extras.clone()); - ElementwiseEpilogue { - expression, - extras: combined_extras, - input_datatype, - output_datatype: existing.output_datatype, - } - } else { - ElementwiseEpilogue { - expression, - extras: extras.clone(), - input_datatype, - output_datatype: nary.output_datatype, - } - }; - - let mut new_q = qmatmul_op.clone(); - let deps_extras = pre_element_wise_expr.extras.clone(); - new_q.input = primary_inner; - new_q.pre_element_wise_expr = Some(pre_element_wise_expr); - - if !new_q.fits_binding_budget(&graph.device()) { - continue; - } - - if let Some(edge) = self.execution_graph.find_edge(input_exec, node_idx) { - self.execution_graph.remove_edge(edge); - } - if let Some(new) = self.get_input_node_in_exec_graph(new_q.input) { - self.execution_graph.add_edge(new, node_idx, ()); - } - for extra in &deps_extras { - if let Some(idx) = self.get_input_node_in_exec_graph(*extra) - && self.execution_graph.find_edge(idx, node_idx).is_none() - { - self.execution_graph.add_edge(idx, node_idx, ()); - } - } - self.execution_graph[node_idx].variant = ExecutionVariant::QMatMul(new_q.clone()); - self.remove_node_if_dead(input_exec); - let mut deps = vec![new_q.input]; - deps.extend(deps_extras); - self.add_physical_dependencies(graph, node_idx, &deps); - return true; - } - } - - // Pre-op: fuse elementwise before matmul inputs. Skipped for - // un-flattened operands: pre chains would demote the matmul off the - // coop kernel they were chosen for. - if let ExecutionVariant::MatMul(matmul_op) = &node_variant - && matmul_op.a.is_plain() - && matmul_op.b.is_plain() - { - let mut new_matmul = matmul_op.clone(); - let mut changed = false; - - // Check first input - if !self.check_cached(graph, matmul_op.first) - && let Some(first_exec) = self.get_input_node_in_exec_graph(matmul_op.first) - && let Some(el_op) = - Self::try_get_unary_chain(&self.execution_graph[first_exec].variant) - { - new_matmul.first = el_op.value; - let mut functions = el_op.functions.functions.clone(); - functions.extend(new_matmul.pre_element_wise[0].functions.iter().cloned()); - new_matmul.pre_element_wise[0] = - UnaryFunctionChain::new(functions, el_op.functions.input_datatype()); - changed = true; - } - - // Check second input - if !self.check_cached(graph, matmul_op.second) - && let Some(second_exec) = self.get_input_node_in_exec_graph(matmul_op.second) - && let Some(el_op) = - Self::try_get_unary_chain(&self.execution_graph[second_exec].variant) - { - new_matmul.second = el_op.value; - let mut functions = el_op.functions.functions.clone(); - functions.extend(new_matmul.pre_element_wise[1].functions.iter().cloned()); - new_matmul.pre_element_wise[1] = - UnaryFunctionChain::new(functions, el_op.functions.input_datatype()); - changed = true; - } - - if changed { - self.execution_graph[node_idx].variant = - ExecutionVariant::MatMul(new_matmul.clone()); - - if new_matmul.first != matmul_op.first { - let old = self.get_input_node_in_exec_graph(matmul_op.first).unwrap(); - if let Some(edge) = self.execution_graph.find_edge(old, node_idx) { - self.execution_graph.remove_edge(edge); - } - if let Some(new) = self.get_input_node_in_exec_graph(new_matmul.first) { - self.execution_graph.add_edge(new, node_idx, ()); - } - self.remove_node_if_dead(old); - } - if new_matmul.second != matmul_op.second { - let old = self.get_input_node_in_exec_graph(matmul_op.second).unwrap(); - if let Some(edge) = self.execution_graph.find_edge(old, node_idx) { - self.execution_graph.remove_edge(edge); - } - if let Some(new) = self.get_input_node_in_exec_graph(new_matmul.second) { - self.execution_graph.add_edge(new, node_idx, ()); - } - self.remove_node_if_dead(old); - } - self.add_physical_dependencies( - graph, - node_idx, - &[new_matmul.first, new_matmul.second], - ); - return true; - } - } - - false - } - - fn qmatmul_dependencies(qmatmul: &QMatMulOperation) -> Vec { - let mut deps = vec![qmatmul.input]; - if let Some(pre) = &qmatmul.pre_element_wise_expr { - deps.extend(pre.extras.iter().copied()); - } - if let Some(post) = &qmatmul.post_element_wise_expr { - deps.extend(post.extras.iter().copied()); - } - deps - } - - /// Replace the n-ary node at `node_idx` with a fused qmatmul `new_q`, - /// rewiring the execution-graph edges: drop the edges from the original - /// n-ary inputs that the fused operation no longer reads, add edges from - /// every dependency (activation input + epilogue extras), and prune any - /// inputs that became dead. Shared by every qmatmul post-epilogue path. - fn commit_qmatmul_post_fusion( - &mut self, - graph: &mut ComputeGraphInner, - node_idx: ExecutionNodeIndex, - nary_inputs: &[NodeIndex], - new_q: Box, - ) { - let deps = Self::qmatmul_dependencies(&new_q); - self.execution_graph[node_idx].variant = ExecutionVariant::QMatMul(new_q); - - for input in nary_inputs { - if deps.contains(input) { - continue; - } - if let Some(input_exec) = self.get_input_node_in_exec_graph(*input) - && let Some(edge) = self.execution_graph.find_edge(input_exec, node_idx) - { - self.execution_graph.remove_edge(edge); - } - } - for dep in &deps { - if let Some(idx) = self.get_input_node_in_exec_graph(*dep) - && self.execution_graph.find_edge(idx, node_idx).is_none() - { - self.execution_graph.add_edge(idx, node_idx, ()); - } - } - self.add_physical_dependencies(graph, node_idx, &deps); - for input in nary_inputs { - if deps.contains(input) { - continue; - } - if let Some(input_exec) = self.get_input_node_in_exec_graph(*input) { - self.remove_node_if_dead(input_exec); - } - } - } - - /// Absorb a split/gate n-ary whose inputs are `narrow` (MapLayout) views of - /// a single-row qmatmul output into that qmatmul's accumulator-offset post - /// epilogue. Each distinct last-dimension column offset (e.g. the gate half - /// at 0 and the up half at `pair_len`) becomes one accumulator value, so a - /// SwiGLU-style `silu(gate) * up` resolves to a single dynamic qmatmul - /// kernel where the backend supports it. Returns `false` (leaving the nodes - /// untouched) when the pattern, dtype, layout, accumulator offsets, or - /// binding budget are unsupported. - fn try_fuse_qmatmul_narrow_accumulators( - &mut self, - graph: &mut ComputeGraphInner, - node_idx: ExecutionNodeIndex, - nary: &ElementwiseOperation, - ) -> bool { - if nary.output_datatype != crate::DataTypeEnum::F32 { - return false; - } - - // Find the qmatmul reached through a narrow MapLayout view. A direct - // (chain-less) reference is the indexed-input form handled below. - let mut base = None; - for &input in &nary.inputs { - let (base_inner, chain) = self.walk_view_chain(input); - if chain.is_none() { - continue; - } - let Some(exec) = self.get_input_node_in_exec_graph(base_inner) else { - continue; - }; - if let ExecutionVariant::QMatMul(op) = &self.execution_graph[exec].variant { - // A qmatmul that already carries a post epilogue isn't a clean - // accumulator-offset base; leave it to the general scan. - if op.post_element_wise_expr.is_some() { - continue; - } - base = Some((base_inner, op.clone())); - break; - } - } - let Some((qmatmul_inner, qmatmul_op)) = base else { - return false; - }; - if self.check_cached(graph, qmatmul_inner) { - return false; - } - - let Some((expression, accumulator_offsets, extras)) = self - .try_extract_mapped_qmatmul_post_expr( - graph, - nary, - qmatmul_inner, - &qmatmul_op.out_shape, - ) - else { - return false; - }; - - if !qmatmul_op.supports_indexed_post_accumulator_offsets( - &graph.device(), - &nary.shape, - &accumulator_offsets, - ) { - return false; - } - - let post_element_wise_expr = ElementwiseEpilogue { - expression, - extras, - input_datatype: crate::DataTypeEnum::F32, - output_datatype: nary.output_datatype, - }; - - let mut new_q = qmatmul_op; - new_q.out_shape = nary.shape.clone(); - new_q.post_element_wise_expr = Some(post_element_wise_expr); - new_q.post_accumulator_offsets = accumulator_offsets.into_boxed_slice(); - - if !new_q.fits_binding_budget(&graph.device()) { - return false; - } - - self.commit_qmatmul_post_fusion(graph, node_idx, &nary.inputs, new_q); - true - } - - /// Build the post epilogue expression, accumulator column offsets, and - /// extra-tensor dependencies for an n-ary whose inputs are last-dimension - /// `narrow` views of `qmatmul_inner`. Inputs that view the qmatmul become - /// accumulator values (indices `0..offsets.len()`, deduplicated by column - /// offset); every other input becomes a normalized extra tensor (indices - /// after the accumulators). Returns `None` when an input isn't a clean - /// last-dimension narrow, uses custom indexing, or can't be normalized. - fn try_extract_mapped_qmatmul_post_expr( - &mut self, - graph: &mut ComputeGraphInner, - nary: &ElementwiseOperation, - qmatmul_inner: NodeIndex, - qmatmul_out_shape: &[usize], - ) -> Option<(NaryExpr, Vec, Vec)> { - if nary.shape.len() != qmatmul_out_shape.len() { - return None; - } - // The accumulator-offset epilogue is only lowered by the single-row - // qgemv path, so every leading dimension must collapse to one row. - if qmatmul_out_shape[..qmatmul_out_shape.len() - 1] - .iter() - .product::() - != 1 - { - return None; - } - let output_cols = nary.shape.last().copied()? as u32; - let matrix_cols = qmatmul_out_shape.last().copied()? as u32; - // A full-width (or wider) output isn't a split; the general scan owns - // that case. - if output_cols >= matrix_cols { - return None; - } - - let qmatmul_out_layout = Layout::contiguous(qmatmul_out_shape); - let rank = nary.shape.len(); - - enum MappedInput { - Accumulator(usize), - Extra(usize), - } - - let mut accumulator_offsets = Vec::new(); - let mut accumulator_map = FxHashMap::default(); - let mut extras = Vec::new(); - let mut mapped = Vec::with_capacity(nary.inputs.len()); - for (input_idx, &nary_input) in nary.inputs.iter().enumerate() { - if !nary.expression.uses_input(input_idx) { - mapped.push(None); - continue; - } - if nary.expression.uses_custom_indexing_for_input(input_idx) { - return None; - } - let (base_inner, chain) = self.walk_view_chain(nary_input); - if base_inner == qmatmul_inner { - let view = Self::apply_view_chain(&qmatmul_out_layout, &chain)?; - let offset = Self::qmatmul_last_dim_view_offset(&view, &nary.shape, matrix_cols)?; - let value_idx = *accumulator_map.entry(offset).or_insert_with(|| { - let idx = accumulator_offsets.len(); - accumulator_offsets.push(offset); - idx - }); - mapped.push(Some(MappedInput::Accumulator(value_idx))); - } else { - let extra = - self.try_normalize_qmatmul_post_extra(graph, nary_input, &nary.shape)?; - let pos = extras.len(); - extras.push(extra); - mapped.push(Some(MappedInput::Extra(pos))); - } - } - - // Two distinct column offsets are the smallest split worth folding into - // the accumulator-offset path; a single offset is either the default - // full-width store or a partial column the qgemv path can't cover. - if accumulator_offsets.len() < 2 { - return None; - } - - let accumulator_count = accumulator_offsets.len(); - let mut replacements = vec![None; nary.inputs.len()]; - for (input_idx, kind) in mapped.into_iter().enumerate() { - match kind { - Some(MappedInput::Accumulator(value_idx)) => { - replacements[input_idx] = Some(NaryExpr::input(value_idx, rank)); - } - Some(MappedInput::Extra(pos)) => { - replacements[input_idx] = Some(NaryExpr::input(accumulator_count + pos, rank)); - } - None => {} - } - } - - let expression = Self::replace_inputs_in_expr(&nary.expression, &replacements)?; - Some((expression, accumulator_offsets, extras)) - } - - /// If `view` is a contiguous last-dimension narrow of a single-row qmatmul - /// output whose shape matches `output_shape`, return its column offset. - /// Returns `None` for any non-narrow / strided / out-of-range view. - fn qmatmul_last_dim_view_offset( - view: &Layout, - output_shape: &[usize], - matrix_cols: u32, - ) -> Option { - if view.shape() != output_shape { - return None; - } - if view.strides().last().copied() != Some(1) { - return None; - } - let offset = u32::try_from(view.offset()).ok()?; - let output_cols = *output_shape.last()? as u32; - if offset.checked_add(output_cols)? > matrix_cols { - return None; - } - Some(offset) - } - - fn try_extract_indexed_qmatmul_post_expr( - &mut self, - graph: &mut ComputeGraphInner, - nary: &ElementwiseOperation, - qmatmul_input_idx: usize, - qmatmul_out_shape: &[usize], - ) -> Option<(NaryExpr, Vec, Vec)> { - if nary.output_datatype != crate::DataTypeEnum::F32 - || nary.shape.len() != qmatmul_out_shape.len() - || nary.shape.as_ref() == qmatmul_out_shape - { - return None; - } - let output_cols = nary.shape.last().copied()? as u32; - let matrix_cols = qmatmul_out_shape.last().copied()? as u32; - if output_cols >= matrix_cols { - return None; - } - - let temp_input_base = nary.inputs.len(); - let mut accumulator_offsets = Vec::new(); - let mut accumulator_map = FxHashMap::default(); - let expression = Self::replace_indexed_qmatmul_accumulators( - &nary.expression, - qmatmul_input_idx, - nary.shape.len(), - output_cols, - matrix_cols, - temp_input_base, - &mut accumulator_offsets, - &mut accumulator_map, - )?; - if accumulator_offsets.len() < 2 { - return None; - } - - let mut replacements = vec![None; nary.inputs.len()]; - let mut extras = Vec::new(); - for (input_idx, &input) in nary.inputs.iter().enumerate() { - if input_idx == qmatmul_input_idx || !nary.expression.uses_input(input_idx) { - continue; - } - if nary.expression.uses_custom_indexing_for_input(input_idx) { - return None; - } - let extra = self.try_normalize_qmatmul_post_extra(graph, input, &nary.shape)?; - replacements[input_idx] = Some(NaryExpr::input( - accumulator_offsets.len() + extras.len(), - nary.shape.len(), - )); - extras.push(extra); - } - - let expression = Self::replace_inputs_in_expr(&expression, &replacements)?; - let expression = Self::remap_temp_accumulator_inputs( - &expression, - temp_input_base, - accumulator_offsets.len(), - ); - Some((expression, accumulator_offsets, extras)) - } - - #[allow(clippy::too_many_arguments)] - fn replace_indexed_qmatmul_accumulators( - expr: &NaryExpr, - qmatmul_input_idx: usize, - output_rank: usize, - output_cols: u32, - matrix_cols: u32, - temp_input_base: usize, - accumulator_offsets: &mut Vec, - accumulator_map: &mut FxHashMap, - ) -> Option { - match expr { - NaryExpr::Op { children, function } => Some(NaryExpr::Op { - children: children - .iter() - .map(|child| { - Self::replace_indexed_qmatmul_accumulators( - child, - qmatmul_input_idx, - output_rank, - output_cols, - matrix_cols, - temp_input_base, - accumulator_offsets, - accumulator_map, - ) - }) - .collect::>>()?, - function: function.clone(), - }), - NaryExpr::IndexedInput { input_idx, indices } if *input_idx == qmatmul_input_idx => { - let offset = Self::extract_qmatmul_last_dim_offset(indices, output_rank)?; - if output_cols - .checked_add(offset) - .is_none_or(|cols| cols > matrix_cols) - { - return None; - } - let value_idx = if let Some(value_idx) = accumulator_map.get(&offset) { - *value_idx - } else { - let value_idx = accumulator_offsets.len(); - accumulator_offsets.push(offset); - accumulator_map.insert(offset, value_idx); - value_idx - }; - Some(NaryExpr::input(temp_input_base + value_idx, output_rank)) - } - NaryExpr::IndexedInput { input_idx, indices } => Some(NaryExpr::IndexedInput { - input_idx: *input_idx, - indices: indices - .iter() - .map(|index| { - Self::replace_indexed_qmatmul_accumulators( - index, - qmatmul_input_idx, - output_rank, - output_cols, - matrix_cols, - temp_input_base, - accumulator_offsets, - accumulator_map, - ) - }) - .collect::>>()?, - }), - NaryExpr::DimIndex(dim) => Some(NaryExpr::DimIndex(*dim)), - NaryExpr::Scalar(value) => Some(NaryExpr::Scalar(*value)), - } - } - - fn extract_qmatmul_last_dim_offset(indices: &[NaryExpr], output_rank: usize) -> Option { - if indices.len() != output_rank { - return None; - } - for (dim, index) in indices[..output_rank - 1].iter().enumerate() { - if !matches!(index, NaryExpr::DimIndex(index_dim) if *index_dim == dim) { - return None; - } - } - Self::extract_dim_plus_u32_offset(&indices[output_rank - 1], output_rank - 1) - } - - fn extract_dim_plus_u32_offset(expr: &NaryExpr, dim: usize) -> Option { - match expr { - NaryExpr::DimIndex(index_dim) if *index_dim == dim => Some(0), - NaryExpr::Op { children, function } - if function.op == NaryOp::Add && children.len() == 2 => - { - Self::extract_dim_plus_u32_offset_pair(&children[0], &children[1], dim).or_else( - || Self::extract_dim_plus_u32_offset_pair(&children[1], &children[0], dim), - ) - } - NaryExpr::Op { children, function } - if matches!(function.op, NaryOp::AddConst(NaryScalar::U32(_))) - && children.len() == 1 => - { - let NaryOp::AddConst(NaryScalar::U32(offset)) = function.op else { - unreachable!(); - }; - matches!(&children[0], NaryExpr::DimIndex(index_dim) if *index_dim == dim) - .then_some(offset) - } - _ => None, - } - } - - fn extract_dim_plus_u32_offset_pair( - dim_expr: &NaryExpr, - offset_expr: &NaryExpr, - dim: usize, - ) -> Option { - let NaryExpr::DimIndex(index_dim) = dim_expr else { - return None; - }; - if *index_dim != dim { - return None; - } - let NaryExpr::Scalar(NaryScalar::U32(offset)) = offset_expr else { - return None; - }; - Some(*offset) - } - - fn remap_temp_accumulator_inputs( - expr: &NaryExpr, - temp_input_base: usize, - accumulator_count: usize, - ) -> NaryExpr { - match expr { - NaryExpr::Op { children, function } => NaryExpr::Op { - children: children - .iter() - .map(|child| { - Self::remap_temp_accumulator_inputs( - child, - temp_input_base, - accumulator_count, - ) - }) - .collect(), - function: function.clone(), - }, - NaryExpr::IndexedInput { input_idx, indices } => { - let input_idx = - if (temp_input_base..temp_input_base + accumulator_count).contains(input_idx) { - input_idx - temp_input_base - } else { - *input_idx - }; - NaryExpr::IndexedInput { - input_idx, - indices: indices - .iter() - .map(|index| { - Self::remap_temp_accumulator_inputs( - index, - temp_input_base, - accumulator_count, - ) - }) - .collect(), - } - } - NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), - NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), - } - } -} diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_region.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_region.rs new file mode 100644 index 000000000..23c548fdb --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_region.rs @@ -0,0 +1,545 @@ +//! Multi-output region formation. +//! +//! Runs after e-graph fusion for every resolve. Grows regions from each +//! unclaimed elementwise sink by absorbing elementwise producers whose +//! *every* consumer already sits in the region — the generalization of the +//! sole-consumer rule: an externally live value (flush target / user-held +//! node) no longer blocks fusion, it is emitted as one of the region's +//! outputs while later statements read it from a register. +//! +//! Contracting a region is acyclic by construction: an absorbed producer has +//! no out-edges leaving the member set, so redirecting its remaining +//! (incoming) edges to the region node cannot create a cycle. + +use super::*; +use crate::region::{ElementwiseRegionOperation, RegionStatement}; + +/// Where a member's input slot points after region rewriting. +enum RegionSlot { + /// Deduplicated external input slot. + External(usize), + /// Another member's register value: the final `IndexedInput` slot + /// (`inputs.len() + statement position`). + Register(usize), +} + +/// Why producers were rejected at region-growth fixpoints, accumulated per +/// resolve and reported under `FUSOR_TRACE_RESOLVE`. Only the final +/// (fixpoint-reaching) probe of each region contributes, so the counts +/// describe the frontier that actually blocked growth, not transient +/// re-scans. +#[derive(Debug, Default, Clone, Copy)] +struct RegionRejects { + non_elementwise: u32, + /// Breakdown of `non_elementwise` by the producer's variant, in order: + /// tensor/view inputs, reduces, row programs, matmul-family, other. + non_elementwise_kinds: [u32; 5], + shape_mismatch: u32, + producer_cached: u32, + outside_consumer: u32, + custom_indexed_read: u32, + binding_budget: u32, +} + +impl RegionRejects { + fn add(&mut self, other: &Self) { + self.non_elementwise += other.non_elementwise; + for (slot, value) in self + .non_elementwise_kinds + .iter_mut() + .zip(other.non_elementwise_kinds) + { + *slot += value; + } + self.shape_mismatch += other.shape_mismatch; + self.producer_cached += other.producer_cached; + self.outside_consumer += other.outside_consumer; + self.custom_indexed_read += other.custom_indexed_read; + self.binding_budget += other.binding_budget; + } +} + +impl Resolver { + pub(super) fn form_elementwise_regions(&mut self, graph: &mut ComputeGraphInner) { + let budget = graph.device().nary_direct_input_binding_budget(); + let Ok(order) = toposort(&self.execution_graph, None) else { + return; + }; + let position: FxHashMap = order + .iter() + .enumerate() + .map(|(pos, &node)| (node, pos)) + .collect(); + let mut claimed: FxHashSet = FxHashSet::default(); + let trace = graph.device().config().trace_resolve; + let mut rejects = RegionRejects::default(); + let mut regions_formed = 0usize; + let mut statements_fused = 0usize; + + for &sink in order.iter().rev() { + if claimed.contains(&sink) || !self.execution_graph.contains_node(sink) { + continue; + } + let ExecutionVariant::Elementwise(sink_op) = &self.execution_graph[sink].variant else { + continue; + }; + if self.check_cached(graph, self.execution_graph[sink].inner_idx) { + continue; + } + let shape = sink_op.shape.clone(); + + // Fixpoint absorption. Outside consumers admitted through the + // precise reachability fallback (rather than the topological + // fast path) are remembered: later candidates must not sit + // downstream of them, or contraction would close a cycle. + let mut member_set: FxHashSet = FxHashSet::default(); + let mut risky_outside: FxHashSet = FxHashSet::default(); + member_set.insert(sink); + loop { + let mut probe = RegionRejects::default(); + let candidate = self.find_absorbable_producer( + graph, + &member_set, + &claimed, + &shape, + budget, + &position, + position[&sink], + &risky_outside, + &mut probe, + ); + match candidate { + Some((producer, borderline)) => { + member_set.insert(producer); + risky_outside.extend(borderline); + } + None => { + // The fixpoint probe: these producers are what + // actually blocked further growth. + rejects.add(&probe); + break; + } + } + } + if member_set.len() == 1 { + continue; + } + + regions_formed += 1; + statements_fused += member_set.len(); + // Topological member order = statement order. + let mut members: Vec = member_set.iter().copied().collect(); + members.sort_by_key(|node| position[node]); + self.finalize_region(graph, sink, &members, &member_set, shape); + claimed.extend(member_set); + } + + if trace { + tracing::info!( + "region_fusion regions={regions_formed} statements={statements_fused} rejects={rejects:?}" + ); + } + } + + /// One producer of the current member set that passes every absorption + /// gate (returned with any outside consumers admitted through the + /// reachability fallback), or `None` at fixpoint. + #[allow(clippy::too_many_arguments)] + fn find_absorbable_producer( + &self, + graph: &ComputeGraphInner, + member_set: &FxHashSet, + claimed: &FxHashSet, + shape: &[usize], + budget: usize, + position: &FxHashMap, + sink_position: usize, + risky_outside: &FxHashSet, + rejects: &mut RegionRejects, + ) -> Option<(ExecutionNodeIndex, Vec)> { + for &member in member_set { + for producer in self + .execution_graph + .neighbors_directed(member, petgraph::Direction::Incoming) + { + if member_set.contains(&producer) || claimed.contains(&producer) { + continue; + } + let producer_inner = self.execution_graph[producer].inner_idx; + let ExecutionVariant::Elementwise(producer_op) = + &self.execution_graph[producer].variant + else { + rejects.non_elementwise += 1; + let kind = match &self.execution_graph[producer].variant { + ExecutionVariant::Tensor(_) | ExecutionVariant::View(_) => 0, + ExecutionVariant::Reduce(_) => 1, + ExecutionVariant::RowProgram(_) => 2, + ExecutionVariant::MatMul(_) + | ExecutionVariant::QMatMul(_) + | ExecutionVariant::Attention(_) => 3, + _ => 4, + }; + rejects.non_elementwise_kinds[kind] += 1; + continue; + }; + if producer_op.shape.as_ref() != shape { + rejects.shape_mismatch += 1; + continue; + } + if self.check_cached(graph, producer_inner) { + rejects.producer_cached += 1; + continue; + } + // Outside consumers no longer block absorption: the + // producer's value becomes a region output (the binding + // budget below already accounts for it). The hazard is a + // cycle: a consumer that transitively feeds the region. + // Fast path: every member feeds the sink, so the sink + // carries the region's maximum topological position, and + // consumers strictly after it are provably downstream. A + // consumer at or before the sink in the linearization is + // only *possibly* cyclic — one arbitrary topological order + // proves nothing about independence — so those fall through + // to a bounded reachability probe: reject only when the + // consumer actually reaches a member. Admitted borderline + // consumers are returned so the fixpoint can keep later + // candidates from sitting downstream of them. + let mut borderline: Vec = Vec::new(); + for consumer in self + .execution_graph + .neighbors_directed(producer, petgraph::Direction::Outgoing) + { + if member_set.contains(&consumer) { + continue; + } + if position + .get(&consumer) + .is_some_and(|&pos| pos > sink_position) + { + continue; + } + borderline.push(consumer); + } + if borderline + .iter() + .any(|&consumer| self.region_probe_reaches(consumer, member_set)) + { + rejects.outside_consumer += 1; + continue; + } + // A candidate downstream of a previously-admitted borderline + // consumer would order the region after that consumer while + // the consumer waits on the region: a cycle. + if !risky_outside.is_empty() + && self.region_probe_reached_from(producer, risky_outside) + { + rejects.outside_consumer += 1; + continue; + } + // Every member reads the producer elementwise: register + // values have no coordinates to custom-index with. + let read_elementwise = member_set.iter().all(|&reader| { + let ExecutionVariant::Elementwise(reader_op) = + &self.execution_graph[reader].variant + else { + return true; + }; + reader_op + .inputs + .iter() + .enumerate() + .filter(|(_, input)| **input == producer_inner) + .all(|(slot, _)| !reader_op.expression.uses_custom_indexing_for_input(slot)) + }); + if !read_elementwise { + rejects.custom_indexed_read += 1; + continue; + } + // Binding budget: distinct external inputs + live outputs of + // the grown region must fit one dispatch. + if self.region_binding_count(graph, member_set, Some(producer)) > budget { + rejects.binding_budget += 1; + continue; + } + return Some((producer, borderline)); + } + } + None + } + + /// Bounded DFS along outgoing edges: does `from` reach any member? + /// Exceeding the probe cap answers `true` (conservative reject). + fn region_probe_reaches( + &self, + from: ExecutionNodeIndex, + targets: &FxHashSet, + ) -> bool { + const CAP: usize = 512; + let mut stack = vec![from]; + let mut visited: FxHashSet = FxHashSet::default(); + while let Some(node) = stack.pop() { + if !visited.insert(node) { + continue; + } + if visited.len() > CAP { + return true; + } + if targets.contains(&node) { + return true; + } + stack.extend( + self.execution_graph + .neighbors_directed(node, petgraph::Direction::Outgoing), + ); + } + false + } + + /// Bounded DFS along incoming edges: is `node` downstream of any source? + /// Exceeding the probe cap answers `true` (conservative reject). + fn region_probe_reached_from( + &self, + node: ExecutionNodeIndex, + sources: &FxHashSet, + ) -> bool { + const CAP: usize = 512; + let mut stack = vec![node]; + let mut visited: FxHashSet = FxHashSet::default(); + while let Some(current) = stack.pop() { + if !visited.insert(current) { + continue; + } + if visited.len() > CAP { + return true; + } + if current != node && sources.contains(¤t) { + return true; + } + stack.extend( + self.execution_graph + .neighbors_directed(current, petgraph::Direction::Incoming), + ); + } + false + } + + /// Distinct external inputs + emitted outputs for `member_set` + /// (optionally grown by `extra`). + fn region_binding_count( + &self, + graph: &ComputeGraphInner, + member_set: &FxHashSet, + extra: Option, + ) -> usize { + let mut inner_members: FxHashSet = FxHashSet::default(); + let all = member_set.iter().copied().chain(extra); + for member in all.clone() { + inner_members.insert(self.execution_graph[member].inner_idx); + } + let mut external: FxHashSet = FxHashSet::default(); + let mut outputs = 0usize; + for member in all { + let node = &self.execution_graph[member]; + let ExecutionVariant::Elementwise(op) = &node.variant else { + continue; + }; + for input in &op.inputs { + if !inner_members.contains(input) { + external.insert(*input); + } + } + if self.region_member_is_live(graph, member, member_set, extra) { + outputs += 1; + } + } + external.len() + outputs + } + + /// Whether a member's value must be written out: it is externally live + /// (user-held / pending sink) or it is the region sink (consumers + /// outside the region read it). + fn region_member_is_live( + &self, + graph: &ComputeGraphInner, + member: ExecutionNodeIndex, + member_set: &FxHashSet, + extra: Option, + ) -> bool { + let inner = self.execution_graph[member].inner_idx; + if graph + .nodes + .nodes + .node_weight(inner) + .is_some_and(|node| node.reference_count > 0) + { + return true; + } + self.execution_graph + .neighbors_directed(member, petgraph::Direction::Outgoing) + .any(|consumer| !member_set.contains(&consumer) && Some(consumer) != extra) + } + + fn finalize_region( + &mut self, + graph: &mut ComputeGraphInner, + sink: ExecutionNodeIndex, + members: &[ExecutionNodeIndex], + member_set: &FxHashSet, + shape: Box<[usize]>, + ) { + // Statement position per member (topological order). + let statement_pos: FxHashMap = members + .iter() + .enumerate() + .map(|(pos, &member)| (member, pos)) + .collect(); + let mut inner_of: FxHashMap = members + .iter() + .map(|&member| (self.execution_graph[member].inner_idx, member)) + .collect(); + // A member's input can name a coalesced *observation* of another + // member's value rather than that member's own index: ingestion puts + // semantically identical nodes in one e-class, keeps one execution + // node, and records the rest in `shared_outputs`. Those aliases must + // map to the producing member's register too — leaving one as an + // external input makes the region read a value it computes itself, + // which nothing in the queue produces. + for (owner, aliases) in &self.shared_outputs { + if let Some(&member) = inner_of.get(owner) { + for &alias in aliases { + inner_of.entry(alias).or_insert(member); + } + } + } + + // Deduplicated external inputs; slots assigned in first-use order. + let mut inputs: Vec = Vec::new(); + let mut input_slot: FxHashMap = FxHashMap::default(); + let mut statements = Vec::with_capacity(members.len()); + for &member in members { + let ExecutionVariant::Elementwise(op) = self.execution_graph[member].variant.clone() + else { + unreachable!("region members are elementwise"); + }; + // Map each member input slot to its region slot. Register slots + // are provisional (`inputs.len()` is not final until all members + // are processed), so store statement positions and fix up below. + let slot_map: Vec = op + .inputs + .iter() + .map(|input| match inner_of.get(input) { + Some(&producer) => RegionSlot::Register(statement_pos[&producer]), + None => { + let next = inputs.len(); + let slot = *input_slot.entry(*input).or_insert_with(|| { + inputs.push(*input); + next + }); + RegionSlot::External(slot) + } + }) + .collect(); + let member_inner = self.execution_graph[member].inner_idx; + let live = member == sink + || graph + .nodes + .nodes + .node_weight(member_inner) + .is_some_and(|node| node.reference_count > 0) + || self + .execution_graph + .neighbors_directed(member, petgraph::Direction::Outgoing) + .any(|consumer| !member_set.contains(&consumer)); + statements.push(( + op.expression, + op.output_datatype, + slot_map, + member_inner, + live, + )); + } + let input_count = inputs.len(); + let statements: Vec = statements + .into_iter() + .map( + |(expression, datatype, slot_map, member_inner, live)| RegionStatement { + expression: Self::remap_region_expr(&expression, &slot_map, input_count), + datatype, + output: live.then_some(member_inner), + }, + ) + .collect(); + + let op = ElementwiseRegionOperation { + inputs, + statements, + shape, + }; + + // Rewrite the graph: the sink node becomes the region; absorbed + // members redirect their incoming external edges to the sink and + // disappear. `node_mapping` keeps absorbed inner nodes reachable + // (they resolve to the region node, which caches their outputs). + let external_inputs = op.inputs.clone(); + for &member in members { + if member == sink { + continue; + } + let incoming: Vec = self + .execution_graph + .neighbors_directed(member, petgraph::Direction::Incoming) + .filter(|producer| !member_set.contains(producer)) + .collect(); + for producer in incoming { + if !self + .execution_graph + .neighbors_directed(sink, petgraph::Direction::Incoming) + .any(|existing| existing == producer) + { + self.execution_graph.add_edge(producer, sink, ()); + } + } + // A member absorbed as a region output keeps its downstream + // ordering: its consumers now read the region's cached result. + let outgoing: Vec = self + .execution_graph + .neighbors_directed(member, petgraph::Direction::Outgoing) + .filter(|consumer| !member_set.contains(consumer)) + .collect(); + for consumer in outgoing { + if !self + .execution_graph + .neighbors_directed(sink, petgraph::Direction::Outgoing) + .any(|existing| existing == consumer) + { + self.execution_graph.add_edge(sink, consumer, ()); + } + } + let member_inner = self.execution_graph[member].inner_idx; + self.execution_graph.remove_node(member); + self.node_mapping.insert(member_inner, sink); + } + self.execution_graph[sink].variant = ExecutionVariant::Region(op); + self.add_physical_dependencies(graph, sink, &external_inputs); + } + + /// Rewrite a member expression into region slots: external reads keep + /// their (recursively remapped) index expressions; register reads drop + /// their identity indices (guaranteed by the absorption gate) and point + /// past the input slots into the statement `extras`. + fn remap_region_expr(expr: &NaryExpr, slot_map: &[RegionSlot], input_count: usize) -> NaryExpr { + egraph::compose::map_loads( + expr, + &mut |input_idx, _, indices| match &slot_map[input_idx] { + RegionSlot::Register(statement) => NaryExpr::IndexedInput { + input_idx: input_count + statement, + indices: Vec::new(), + }, + RegionSlot::External(slot) => NaryExpr::IndexedInput { + input_idx: *slot, + indices, + }, + }, + ) + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs b/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs index 8d0d58ecd..e6a704ea8 100644 --- a/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs +++ b/fusor-ml/core/src/compute_graph/resolve/fusion_row.rs @@ -10,8 +10,6 @@ //! collapse to one kernel this way — and so does any other normalization- //! shaped cluster the tensor API emits. -use std::sync::Arc; - use petgraph::algo::toposort; use crate::{ @@ -74,6 +72,20 @@ impl ClusterBuilder<'_> { self.externals.len() - 1 } } + + /// Member insertion is idempotent: one node can be reached through + /// several operand walks, but is absorbed and removed exactly once. + fn add_member(&mut self, node: NodeIndex) { + if !self.members.contains(&node) { + self.members.push(node); + } + } + + fn add_members(&mut self, nodes: impl IntoIterator) { + for node in nodes { + self.add_member(node); + } + } } /// Rewrite the slot references of one absorbed node's expression into the @@ -115,31 +127,14 @@ fn rewrite_slots(expr: &NaryExpr, rewrites: &[SlotRewrite]) -> Option /// Remap construction-time scalar slots to their final indices. fn finalize_slots(expr: &NaryExpr, external_count: usize) -> NaryExpr { - match expr { - NaryExpr::Op { children, function } => NaryExpr::Op { - children: children - .iter() - .map(|child| finalize_slots(child, external_count)) - .collect(), - function: function.clone(), - }, - NaryExpr::IndexedInput { input_idx, indices } => { - let input_idx = if *input_idx >= SCALAR_SLOT_BASE { - external_count + (*input_idx - SCALAR_SLOT_BASE) - } else { - *input_idx - }; - NaryExpr::IndexedInput { - input_idx, - indices: indices - .iter() - .map(|index| finalize_slots(index, external_count)) - .collect(), - } - } - NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), - NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), - } + egraph::compose::map_loads(expr, &mut |input_idx, _, indices| { + let input_idx = if input_idx >= SCALAR_SLOT_BASE { + external_count + (input_idx - SCALAR_SLOT_BASE) + } else { + input_idx + }; + NaryExpr::IndexedInput { input_idx, indices } + }) } impl Resolver { @@ -175,7 +170,7 @@ impl Resolver { rootless += 1; } } - if std::env::var_os("FUSOR_TRACE_ROW_FUSION").is_some() { + if graph.device().config().trace_row_fusion { eprintln!( "row_fusion: {reduce_count} reduces, {} roots, {rootless} rootless", roots.len() @@ -327,6 +322,7 @@ impl Resolver { out } + #[allow(clippy::too_many_arguments)] fn collect_operand( &self, graph: &ComputeGraphInner, @@ -374,7 +370,7 @@ impl Resolver { nodes.push(node); node = view.input; } - let Some(layout) = layout else { + let Some(mut layout) = layout else { break 'scalar; }; let reduce = loop { @@ -396,6 +392,22 @@ impl Resolver { let (_, input) = unary_elementwise(nary).unwrap(); node = input; } + // A pure view *between* unary chain links (the + // `sum_keepdim` unsqueeze under `div_scalar`/`sqrt` in + // layer norm) composes into the running layout; unaries + // are pointwise, so their position relative to pure + // layout stages cannot change per-row values. + ExecutionVariant::View(view) => { + let Some(collapsed) = view.composed_layout() else { + break 'scalar; + }; + layout = match crate::view::compose_layouts(&layout, &collapsed) { + Some(layout) => layout, + None => break 'scalar, + }; + nodes.push(node); + node = view.input; + } _ => break 'scalar, } }; @@ -462,7 +474,7 @@ impl Resolver { return false; }; - let trace = std::env::var_os("FUSOR_TRACE_ROW_FUSION").is_some(); + let trace = graph.device().config().trace_row_fusion; // Fusing pays for itself only when at least one reduction folds in. if builder.phases.is_empty() { return false; @@ -524,7 +536,7 @@ impl Resolver { graph, root_idx, &externals, - ExecutionVariant::GraphOp(Arc::new(operation)), + ExecutionVariant::RowProgram(operation), ); true } @@ -588,7 +600,7 @@ impl Resolver { rewrites.push(self.absorb_operand(graph, builder, input)?); } let expr = rewrite_slots(&nary.expression, &rewrites)?; - builder.members.push(inner); + builder.add_member(inner); builder.full_exprs.insert(inner, expr.clone()); Some(expr) } @@ -624,7 +636,13 @@ impl Resolver { // An optional unary chain on the reduced value (mean scaling, eps, // rsqrt...) folds into the phase's post chain, innermost first. + // Pure views interleaved with the chain (the `sum_keepdim` + // unsqueeze) compose into the layout; they are tracked separately + // from `chain_nodes` so phase deduplication stays keyed on real + // chain nodes. + let mut layout = layout; let mut chain_nodes = Vec::new(); + let mut sandwich_views = Vec::new(); let mut chain: Vec = Vec::new(); let reduce = loop { if self.check_cached(graph, node) { @@ -639,6 +657,12 @@ impl Resolver { chain_nodes.push(node); node = input; } + ExecutionVariant::View(view) => { + let collapsed = view.composed_layout()?; + layout = crate::view::compose_layouts(&layout, &collapsed)?; + sandwich_views.push(node); + node = view.input; + } _ => return None, } }; @@ -669,7 +693,8 @@ impl Resolver { .iter() .position(|(base, _)| *base == chain_nodes.first().copied().unwrap_or(node)) { - builder.members.extend(views); + builder.add_members(views); + builder.add_members(sandwich_views); return Some(scalar_ref(phase, builder.shape.len())); } @@ -685,15 +710,16 @@ impl Resolver { ); let phase_key = chain_nodes.first().copied().unwrap_or(node); - builder.members.extend(views); - builder.members.extend(chain_nodes); - builder.members.push(node); + builder.add_members(views); + builder.add_members(sandwich_views); + builder.add_members(chain_nodes); + builder.add_member(node); let phase_index = builder.phases.len(); builder.phases.push(( phase_key, RowReduce { expression, - function: reduce.function.clone(), + combine: crate::row_program::RowCombine::BuiltIn(reduce.function.clone()), post_chain, }, )); diff --git a/fusor-ml/core/src/compute_graph/resolve/key_goldens.rs b/fusor-ml/core/src/compute_graph/resolve/key_goldens.rs new file mode 100644 index 000000000..dfda1a277 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/key_goldens.rs @@ -0,0 +1,413 @@ +//! Byte goldens for the resolver's structural cache keys. +//! +//! Every key here selects a cache entry that is trusted without re-deriving +//! what it stands for: kernel plans replayed by positional rebind, merged +//! dispatch plans shared across processes through the persistent store, and +//! flush plans replayed instead of a resolve. A recipe change that keeps +//! correctness intact is therefore still a regression — it invalidates the +//! warm stores of every machine that ever ran an older build — and no +//! correctness test can see it. These goldens pin the exact key words for a +//! spread of real graphs so a refactor of the folds has to be byte-preserving. +//! +//! Re-capture procedure after an INTENTIONAL recipe change: run the test and +//! paste the measured lines from the failure message into the tables below, +//! bumping the affected recipe version (`REPLAY_RECIPE_VERSION` for flush +//! plans, the leading literal in `kernel_cache_key_with_dispatch` for kernel +//! plans) in the same change. + +use fusor_gguf::GgmlType; + +use super::merge_horizontal::MergedSegments; +use super::*; +use crate::reduce::{ReduceFunction, ReduceOp, ReduceOperation}; +use crate::region::{ElementwiseRegionOperation, RegionStatement}; +use crate::row_program::RowProgramOperation; +use crate::{Device, QMatrix, Tensor}; + +/// `structural_kernel_key` over each single-operation lowering arm, keyed by +/// the graph that produced it. +const KERNEL_KEY_GOLDENS: &[(&str, &str)] = &[ + ( + "step_elementwise", + "KernelCacheKey([11452713101039835292, 469725187975047675])", + ), + ( + "step_reduce", + "KernelCacheKey([14261650954393790832, 3278663041261894351])", + ), + ( + "step_view", + "KernelCacheKey([13199296625049102793, 2216308711984315176])", + ), + ( + "step_assign", + "KernelCacheKey([3412231204244641173, 10875987364889405171])", + ), + ( + "decode_elementwise", + "KernelCacheKey([4670963369128312198, 12134719529705967332])", + ), + ( + "decode_reduce", + "KernelCacheKey([10432744781750632184, 17896500942395396182])", + ), + ( + "quantized_elementwise", + "KernelCacheKey([8460052255698420948, 15923808416343184946])", + ), +]; + +/// `merged_plan_cache_key` over both fold arms: the region arm hashes kernel +/// fields and MIR values inline, every other segment kind folds its own +/// structural kernel key. +const MERGED_KEY_GOLDENS: &[(&str, &str)] = &[ + ( + "region_x1", + "KernelCacheKey([16903751727878911052, 1182021694381039177])", + ), + ( + "region_x2", + "KernelCacheKey([10258795302309731177, 1156476658944166214])", + ), + ( + "row_x1", + "KernelCacheKey([5374707752665978227, 8083608762445340972])", + ), + ( + "row_x2", + "KernelCacheKey([11580284453418073024, 16462936542216973224])", + ), +]; + +/// `FlushPlanKey` over whole pending subgraphs: a transformer step, a decode +/// token, and a quantized decode token. +const FLUSH_KEY_GOLDENS: &[(&str, &str)] = &[ + ( + "transformer_step", + "FlushPlanKey([13924756313179939098, 12746842292723456953])", + ), + ( + "decode_token", + "FlushPlanKey([2248803744916415424, 13788545245063220203])", + ), + ( + "quantized_decode", + "FlushPlanKey([9842396758258840322, 12416291958004416485])", + ), +]; + +fn assert_goldens(what: &str, goldens: &[(&str, &str)], measured: &[(&str, String)]) { + let expected: Vec<(&str, String)> = goldens + .iter() + .map(|(name, key)| (*name, (*key).to_string())) + .collect(); + if expected != measured { + let lines = measured + .iter() + .map(|(name, key)| format!(" (\"{name}\", \"{key}\"),")) + .collect::>() + .join("\n"); + panic!("{what} key bytes changed. measured values:\n{lines}"); + } +} + +/// Cache everything under `node` so `Operation::inputs` can gather it, +/// exactly as the resolver's queue does before it keys a kernel plan: leaves +/// contribute their own storage, intermediates a freshly allocated output. +fn cache_dependencies(graph: &mut ComputeGraphInner, node: NodeIndex) { + let mut deps = Vec::new(); + graph.visit_dependencies(node, &mut |dep| deps.push(dep)); + for dep in deps { + if graph.get_cached_result(dep).is_some() { + continue; + } + match &graph + .nodes + .nodes + .node_weight(dep) + .expect("live node") + .variant + { + // Raw block storage is gathered straight off the node. + ComputeGraphNodeVariant::QMatrix(_) => continue, + ComputeGraphNodeVariant::Tensor(data) => { + let data = data.clone(); + graph.set_cached_result(dep, data); + continue; + } + _ => {} + } + cache_dependencies(graph, dep); + let output = { + let variant = &graph + .nodes + .nodes + .node_weight(dep) + .expect("live node") + .variant; + let operation = operation_of(variant); + let inputs = operation.inputs(graph); + operation.output(graph, &inputs) + }; + let MirValue::Tensor(output) = output else { + panic!("golden dependency output is not a tensor"); + }; + graph.set_cached_result(dep, output); + } +} + +fn operation_of(variant: &ComputeGraphNodeVariant) -> &dyn Operation { + match variant { + ComputeGraphNodeVariant::Elementwise(op) => op, + ComputeGraphNodeVariant::Reduce(op) => op, + ComputeGraphNodeVariant::View(op) => op, + ComputeGraphNodeVariant::Assign(op) => op, + ComputeGraphNodeVariant::Tensor(_) | ComputeGraphNodeVariant::QMatrix(_) => { + panic!("golden target is not a lowered operation") + } + } +} + +/// The structural kernel key of `tensor`'s own operation, gathered and solved +/// the way the queue executor does. +fn kernel_key(device: &Device, tensor: &Tensor) -> String { + let node = tensor.data().key; + device.compute_graph().with_mut(|graph| { + cache_dependencies(graph, node); + let variant = &graph + .nodes + .nodes + .node_weight(node) + .expect("live node") + .variant; + let operation = operation_of(variant); + let inputs = operation.inputs(graph); + let workgroup = operation + .workgroup_shape_constraints(device) + .solve(device.max_subgroup_size(), &device.limits()) + .expect("golden operation has a solvable workgroup shape"); + format!( + "{:?}", + structural_kernel_key(operation, &inputs, &workgroup) + ) + }) +} + +fn flush_key(device: &Device, targets: &[NodeIndex]) -> String { + device.compute_graph().with_mut(|graph| { + let fingerprint = + flush_replay::fingerprint_pending(graph, targets).expect("golden graph fingerprints"); + format!("{:?}", fingerprint.key) + }) +} + +fn tensor_value(device: &Device, shape: &[usize]) -> MirValue { + MirValue::Tensor(TensorData::new_for_shape( + device, + shape, + crate::DataTypeEnum::F32, + )) +} + +fn region_segment(inputs: usize, outputs: usize, shape: &[usize]) -> ElementwiseRegionOperation { + ElementwiseRegionOperation { + inputs: (0..inputs).map(NodeIndex::new).collect(), + statements: (0..outputs) + .map(|statement| RegionStatement { + expression: crate::nary_wise::NaryExpr::input(statement % inputs, inputs), + datatype: crate::DataTypeEnum::F32, + output: Some(NodeIndex::new(100 + statement)), + }) + .collect(), + shape: shape.to_vec().into_boxed_slice(), + } +} + +fn row_segment(axis_len: usize) -> RowProgramOperation { + RowProgramOperation::from_reduce(&ReduceOperation { + inputs: vec![NodeIndex::new(0)], + expression: crate::nary_wise::NaryExpr::indexed_input( + 0, + vec![ + crate::nary_wise::NaryExpr::DimIndex(0), + crate::nary_wise::NaryExpr::DimIndex(1), + ], + ), + shape: vec![1, axis_len].into_boxed_slice(), + function: ReduceFunction { + name: Some("sum".to_string()), + op: ReduceOp::Sum, + initial_value: crate::nary_wise::NaryScalar::F32(0.0), + datatype: crate::DataTypeEnum::F32, + }, + post_element_wise: crate::nary_wise::UnaryFunctionChain::empty(crate::DataTypeEnum::F32), + axis: 1, + }) +} + +pub(super) fn q4k_weight(device: &Device, rows: usize, cols: usize) -> QMatrix { + let blocks = rows * cols / 256; + let bytes: Vec = (0..blocks * 144) + .map(|index| ((index * 31 + 7) % 251) as u8) + .collect(); + QMatrix::from_parts( + device, + &bytes, + vec![rows, cols].into_boxed_slice(), + GgmlType::Q4K, + ) + .expect("q4k weight") +} + +/// A transformer-step-shaped graph: a projection composed as broadcast +/// multiply plus reduction, a residual add, a reshape, and a slice assign. +fn transformer_step(device: &Device) -> Tensor { + let patch = Tensor::splat::(device, 0.0, &[1, 4, 64]); + step_view(device).slice_assign([0..1, 0..4, 0..64], &patch) +} + +fn step_view(device: &Device) -> Tensor { + step_elementwise(device).reshape([2, 4, 64]) +} + +fn step_elementwise(device: &Device) -> Tensor { + let x = Tensor::splat::(device, 0.0, &[8, 64]); + let w = Tensor::splat::(device, 0.0, &[64, 64]); + &x.mat_mul(&w) + &x +} + +fn step_reduce(device: &Device) -> Tensor { + let x = Tensor::splat::(device, 0.0, &[8, 64]); + let w = Tensor::splat::(device, 0.0, &[64, 64]); + x.mat_mul(&w) +} + +/// A decode-token-shaped graph: one row against a resident weight. +fn decode_token(device: &Device) -> Tensor { + let projected = decode_reduce(device); + &projected * &projected +} + +fn decode_reduce(device: &Device) -> Tensor { + let x = Tensor::splat::(device, 0.0, &[1, 512]); + let w = Tensor::splat::(device, 0.0, &[512, 512]); + x.mat_mul(&w) +} + +/// A quantized-decode-shaped graph: one row against a Q4K weight. +fn quantized_decode(device: &Device) -> Tensor { + let x = Tensor::splat::(device, 0.0, &[1, 512]); + let weight = q4k_weight(device, 512, 512); + let projected = x.q_mat_mul(&weight); + &projected + &projected +} + +#[test] +fn structural_kernel_key_bytes_are_stable() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let cases: &[(&str, fn(&Device) -> Tensor)] = &[ + ("step_elementwise", step_elementwise), + ("step_reduce", step_reduce), + ("step_view", step_view), + ("step_assign", transformer_step), + ("decode_elementwise", decode_token), + ("decode_reduce", decode_reduce), + ("quantized_elementwise", quantized_decode), + ]; + let measured: Vec<(&str, String)> = cases + .iter() + .map(|(name, build)| (*name, kernel_key(&device, &build(&device)))) + .collect(); + assert_goldens("structural_kernel_key", KERNEL_KEY_GOLDENS, &measured); + }); +} + +#[test] +fn merged_plan_cache_key_bytes_are_stable() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let region_inputs = |count: usize| -> Vec> { + (0..count) + .map(|_| vec![tensor_value(&device, &[8, 64]); 2]) + .collect() + }; + let row_inputs = |count: usize| -> Vec> { + (0..count) + .map(|_| { + vec![ + tensor_value(&device, &[1, 512]), + tensor_value(&device, &[1, 1]), + ] + }) + .collect() + }; + let cases: Vec<(&str, MergedSegments, Vec>)> = vec![ + ( + "region_x1", + MergedSegments::Region(vec![(NodeIndex::new(1), region_segment(1, 1, &[8, 64]))]), + region_inputs(1), + ), + ( + "region_x2", + MergedSegments::Region(vec![ + (NodeIndex::new(1), region_segment(1, 1, &[8, 64])), + (NodeIndex::new(2), region_segment(1, 2, &[8, 64])), + ]), + region_inputs(2), + ), + ( + "row_x1", + MergedSegments::Row(vec![(NodeIndex::new(1), row_segment(512))]), + row_inputs(1), + ), + ( + "row_x2", + MergedSegments::Row(vec![ + (NodeIndex::new(1), row_segment(512)), + (NodeIndex::new(2), row_segment(256)), + ]), + row_inputs(2), + ), + ]; + let measured: Vec<(&str, String)> = cases + .iter() + .map(|(name, merged, inputs)| { + ( + *name, + format!( + "{:?}", + super::queue_executor::merged_plan_cache_key(merged, inputs) + ), + ) + }) + .collect(); + assert_goldens("merged_plan_cache_key", MERGED_KEY_GOLDENS, &measured); + }); +} + +#[test] +fn flush_plan_key_bytes_are_stable() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let cases: &[(&str, fn(&Device) -> Tensor)] = &[ + ("transformer_step", transformer_step), + ("decode_token", decode_token), + ("quantized_decode", quantized_decode), + ]; + let measured: Vec<(&str, String)> = cases + .iter() + .map(|(name, build)| { + let target = build(&device); + (*name, flush_key(&device, &[target.data().key])) + }) + .collect(); + assert_goldens("FlushPlanKey", FLUSH_KEY_GOLDENS, &measured); + }); +} diff --git a/fusor-ml/core/src/compute_graph/resolve/merge_horizontal.rs b/fusor-ml/core/src/compute_graph/resolve/merge_horizontal.rs new file mode 100644 index 000000000..20528355c --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/merge_horizontal.rs @@ -0,0 +1,792 @@ +//! Horizontal merging of independent same-category operations into one +//! dispatch. +//! +//! Dense training tapes are dominated by launch overhead: dozens of tiny +//! elementwise kernels, per-parameter Adam updates, and same-shape row +//! reduces, each ~5-10us of fixed launch cost. This pass walks the queue in +//! canonical dependency-depth order and groups independent operations of the +//! same category (Adam updates / elementwise naries / chunked-map row +//! programs) into one kernel whose body is a build-time-unrolled chain of +//! per-segment regions, each guarded by a uniform linear-workgroup-id range +//! compare. +//! +//! Safety and gating: +//! - Every graph exposes compatible matmul, row and elementwise operations. +//! Per-operation shape and binding limits decide eligibility. +//! - Grouping is dependency-sound by a wave discipline: an operation joins +//! the open wave of its category only if it does not (transitively) depend +//! on any open-wave member of any category; a dependency on an open wave +//! flushes that wave (emitting the merged dispatch) first. The emitted +//! queue is therefore always a valid topological order. +//! - Merged outputs are one fresh buffer per segment (never a concatenated +//! buffer), so flush-replay slot attribution stays 1 buffer <-> 1 slot. + +use super::*; +use crate::row_program::RowProgramOperation; + +const CAT_ROW: usize = 2; +/// Dense matmuls bound for the single-pass cooperative-matrix kernel. +const CAT_MATMUL: usize = 3; +/// Dense matmuls bound for the split-K route (weight-gradient shapes). +/// A separate wave from [`CAT_MATMUL`]: the training tape's split-K matmuls +/// are gradient sinks nothing but the optimizer consumes, so their wave +/// survives the entire backward walk and merges across layers, while the +/// forward/dX matmuls flush at every chain dependency. +const CAT_MATMUL_SPLITK: usize = 4; +/// Elementwise regions: multi-output regions from `fusion_region` and +/// single elementwise operations hosted as one-statement regions. +const CAT_REGION: usize = 0; +const CATEGORY_COUNT: usize = 6; + +/// Return a deterministic, fusion-friendly topological order. +/// +/// `petgraph::toposort` returns an arbitrary valid linear extension. Feeding +/// that order directly to the greedy horizontal merger can interleave one +/// branch's consumer before another branch's independent producer, creating +/// false category-level dependency cycles and singleton flushes. Grouping by +/// dependency depth preserves topological validity while presenting every +/// ready antichain to the merger before any newly-ready consumers. +pub(super) fn fusion_toposort(graph: &ExecutionGraph) -> Vec { + let mut order = + toposort(graph, None).unwrap_or_else(|_| panic!("Cycle detected in execution graph")); + let mut depth = FxHashMap::::default(); + for &node in &order { + let node_depth = graph + .neighbors_directed(node, petgraph::Direction::Incoming) + .filter_map(|parent| depth.get(&parent).copied()) + .max() + .map_or(0, |parent_depth| parent_depth + 1); + depth.insert(node, node_depth); + } + order.sort_by_key(|&node| (depth[&node], graph[node].inner_idx.index())); + order +} + +/// Segments of one merged dispatch, in queue order. All members are +/// mutually independent. +#[derive(Debug)] +pub(super) enum MergedSegments { + Row(Vec<(NodeIndex, RowProgramOperation)>), + MatMul(Vec<(NodeIndex, crate::matmul::MatMulOperation)>), + Region(Vec<(NodeIndex, crate::region::ElementwiseRegionOperation)>), +} + +impl MergedSegments { + pub(super) fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { + match self { + Self::Row(segments) => { + for (_, op) in segments { + op.visit_dependencies(f); + } + } + Self::MatMul(segments) => { + for (_, op) in segments { + op.visit_dependencies(f); + } + } + Self::Region(segments) => { + for (_, op) in segments { + op.visit_dependencies(f); + } + } + } + } + + /// Segment (node, op) views in queue order. + pub(super) fn segment_ops(&self) -> Vec<(NodeIndex, &dyn Operation)> { + match self { + Self::Row(segments) => segments + .iter() + .map(|(node, op)| (*node, op as &dyn Operation)) + .collect(), + Self::MatMul(segments) => segments + .iter() + .map(|(node, op)| (*node, op as &dyn Operation)) + .collect(), + Self::Region(_) => { + unreachable!("region segments are gathered without the Operation trait") + } + } + } + + fn representative(&self) -> NodeIndex { + match self { + Self::Row(segments) => segments[0].0, + Self::MatMul(segments) => segments[0].0, + Self::Region(segments) => segments[0].0, + } + } +} + +/// One mergeable operation, categorized. +enum SegOp { + Row(RowProgramOperation), + Region(crate::region::ElementwiseRegionOperation), + /// Carries the profile key so the flush can partition the wave into + /// same-profile dispatches. + MatMul( + Box, + crate::matmul::MatmulMergeKey, + ), +} + +impl SegOp { + fn category(&self) -> usize { + match self { + Self::Row(_) => CAT_ROW, + Self::Region(_) => CAT_REGION, + Self::MatMul(_, key) => { + if key.splits().is_some() { + CAT_MATMUL_SPLITK + } else { + CAT_MATMUL + } + } + } + } + + /// Storage bindings this segment will declare in a merged kernel. + fn bindings(&self) -> usize { + match self { + Self::Row(op) => op.inputs.len() + 1, + Self::Region(op) => op.binding_count(), + Self::MatMul(..) => MATMUL_SEGMENT_BINDINGS, + } + } +} + +fn output_observations( + outputs: impl IntoIterator, + shared_outputs: &FxHashMap>, +) -> Vec { + let mut observations = Vec::new(); + let mut seen = FxHashSet::default(); + let mut pending: Vec = Vec::new(); + for output in outputs { + if seen.insert(output) { + observations.push(output); + } + pending.push(output); + } + // Observations chain: follow them to the end, or a merged dispatch will + // leave the far end of a chain unobserved. + while let Some(current) = pending.pop() { + let Some(aliases) = shared_outputs.get(¤t) else { + continue; + }; + for &alias in aliases { + if seen.insert(alias) { + observations.push(alias); + pending.push(alias); + } + } + } + observations +} + +/// Bindings per merged-matmul segment: `a`, `b`, and the output (whose +/// allocation carries any split-K scratch). +const MATMUL_SEGMENT_BINDINGS: usize = 3; + +pub(super) struct HorizontalMerger { + device: crate::Device, + /// Max total storage bindings per merged dispatch. + budget: usize, + /// Current open-wave generation per category (starts at 1; 0 = "none"). + open_gen: [u32; CATEGORY_COUNT], + /// For every processed node: the latest wave generation (per category) + /// it transitively depends on. + dep_gen: FxHashMap, + /// Wave membership (including trailing views): node -> (category, + /// generation it was added under). + member: FxHashMap, + waves: [Vec<(NodeIndex, SegOp)>; CATEGORY_COUNT], + /// Zero-cost view aliases of wave members, deferred with their wave and + /// emitted right after its merged dispatch. + trailing: [Vec<(NodeIndex, QueuedOperation)>; CATEGORY_COUNT], + wave_bindings: [usize; CATEGORY_COUNT], + /// Distinct input nodes of the open region wave: shared inputs (a + /// learning-rate tensor read by every optimizer segment) bind once, so + /// the budget only charges new nodes. + region_wave_inputs: FxHashSet, + /// `before[k][c]`: open wave `k` must flush before open wave `c` + /// (members of `c` depend on members of `k`). Kept acyclic. + before: [[bool; CATEGORY_COUNT]; CATEGORY_COUNT], +} + +impl HorizontalMerger { + pub(super) fn new(device: &crate::Device) -> Self { + Self { + device: device.clone(), + // Total storage bindings per merged dispatch. The nary budget is + // inputs-only (it assumes one extra output binding); merged + // kernels bind everything explicitly, and Metal rejects pipeline + // layouts at the full 31-buffer limit, so stay at the budget. + budget: device.nary_direct_input_binding_budget(), + open_gen: [1; CATEGORY_COUNT], + dep_gen: FxHashMap::default(), + member: FxHashMap::default(), + waves: Default::default(), + trailing: Default::default(), + wave_bindings: [0; CATEGORY_COUNT], + region_wave_inputs: FxHashSet::default(), + before: [[false; CATEGORY_COUNT]; CATEGORY_COUNT], + } + } + + /// Whether wave `from` must (transitively) flush before wave `to`. + fn reaches(&self, from: usize, to: usize) -> bool { + if self.before[from][to] { + return true; + } + (0..CATEGORY_COUNT).any(|middle| { + middle != from && middle != to && self.before[from][middle] && self.before[middle][to] + }) + } + + fn categorize(&self, node: &ExecutionNode) -> Option { + match &node.variant { + // Regions already have a valid single-region queue lowering. The + // merger may combine independent regions; there is no element cap + // because multi-output regions have no tiled fallback. + ExecutionVariant::Region(op) => Some(SegOp::Region(op.clone())), + ExecutionVariant::Elementwise(op) => { + // Ops at or above the register-reuse tiled path's engagement + // element count stay out of merges so that plan never + // regresses; the bound is derived from the same policy the + // tiled planner reads, so the two cannot drift apart. + let elements: usize = op.shape.iter().product(); + let bound = self.device.dispatch_policy().merge_elements_bound(); + (elements < bound && op.inputs.len() + 1 < self.budget).then(|| { + SegOp::Region(crate::region::ElementwiseRegionOperation::from_nary( + op.clone(), + node.inner_idx, + )) + }) + } + ExecutionVariant::Reduce(op) => { + self.row_candidate(RowProgramOperation::from_reduce(op)) + } + ExecutionVariant::RowProgram(op) => self.row_candidate(op.clone()), + ExecutionVariant::MatMul(op) => { + // Only matmuls bound for the cooperative-matrix kernel + // produce a profile; generic-path contractions and + // epilogue-fused matmuls lower standalone. + if self.device.config().trace_matmul_merge { + eprintln!( + "matmul_merge_candidate name={} profile={:?}", + op.name(), + op.merge_profile(&self.device) + ); + } + let key = op.merge_profile(&self.device)?; + Some(SegOp::MatMul(Box::new(op.clone()), key)) + } + _ => None, + } + } + + fn row_candidate(&self, row: RowProgramOperation) -> Option { + (row.dynamic_axis.is_none() + && row.mergeable_chunked_map() + && row.inputs.len() + 1 < self.budget) + .then_some(SegOp::Row(row)) + } + + /// Feed one toposorted node. `lowered` is the node's normal lowering (for + /// the non-merged path). Emits queue entries into `out`. + pub(super) fn push( + &mut self, + node: &ExecutionNode, + lowered: Option, + shared_outputs: &FxHashMap>, + out: &mut Vec<(NodeIndex, QueuedOperation)>, + ) { + // Latest wave generation (per category) this node depends on, + // through direct wave members and transitively via `dep_gen`. + let mut dep = [0u32; CATEGORY_COUNT]; + let mut visit = |input: NodeIndex| { + if let Some(&(cat, generation)) = self.member.get(&input) { + dep[cat] = dep[cat].max(generation); + } + if let Some(gens) = self.dep_gen.get(&input) { + for (slot, generation) in dep.iter_mut().zip(gens) { + *slot = (*slot).max(*generation); + } + } + }; + match &node.variant { + ExecutionVariant::Tensor(_) => {} + ExecutionVariant::QMatrix(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::Elementwise(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::Reduce(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::Fold(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::View(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::Assign(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::Region(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::MatMul(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::QMatMul(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::QEmbedding(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::RowProgram(op) => op.visit_dependencies(&mut visit), + ExecutionVariant::Attention(op) => op.visit_dependencies(&mut visit), + } + + match self.categorize(node) { + Some(seg) => { + let cat = seg.category(); + // Multi-output regions are observed through every emitted + // inner node, not only the sink that represents the region + // in the execution graph. A later segment may depend on any + // of those observations and must therefore flush this wave. + let outputs: Vec = match &seg { + SegOp::Region(op) => op + .statements + .iter() + .filter_map(|statement| statement.output) + .collect(), + _ => vec![node.inner_idx], + }; + let observations = output_observations(outputs, shared_outputs); + debug_assert!(observations.contains(&node.inner_idx)); + // Flush the own-category wave first if this op depends on it + // (this also flushes the wave's ordered predecessors). + if dep[cat] == self.open_gen[cat] { + self.flush(cat, out); + } + // Region segments only pay for input nodes the open wave + // does not already bind: the builder deduplicates shared + // read-only inputs (a learning-rate tensor read by every + // optimizer segment) into one binding. + let fresh_bindings = |merger: &Self, seg: &SegOp| match seg { + SegOp::Region(op) => { + op.inputs + .iter() + .filter(|input| !merger.region_wave_inputs.contains(*input)) + .collect::>() + .len() + + op.output_count() + } + other => other.bindings(), + }; + if self.wave_bindings[cat] + fresh_bindings(self, &seg) > self.budget { + self.flush(cat, out); + } + // Recompute after a possible flush: a cleared wave shares + // nothing yet. + let seg_bindings = fresh_bindings(self, &seg); + if let SegOp::Region(op) = &seg { + self.region_wave_inputs.extend(op.inputs.iter().copied()); + } + // Dependencies on other open waves become flush-order + // constraints (that wave's merged dispatch is emitted before + // ours) unless that would create a cycle. + for (other, &generation) in dep.iter().enumerate() { + if other != cat && generation == self.open_gen[other] { + if self.reaches(cat, other) { + self.flush(other, out); + } else { + self.before[other][cat] = true; + } + } + } + self.wave_bindings[cat] += seg_bindings; + self.waves[cat].push((node.inner_idx, seg)); + for observation in observations { + self.member.insert(observation, (cat, self.open_gen[cat])); + self.dep_gen.insert(observation, dep); + } + } + None => { + let observations = + output_observations(std::iter::once(node.inner_idx), shared_outputs); + for &observation in &observations { + self.dep_gen.insert(observation, dep); + } + let open: Vec = (0..CATEGORY_COUNT) + .filter(|&cat| dep[cat] == self.open_gen[cat]) + .collect(); + // A zero-cost view of an open-wave value defers with that + // wave instead of flushing it: it is emitted right after the + // wave's merged dispatch, so its input is cached in time. + if matches!(node.variant, ExecutionVariant::View(_)) + && !open.is_empty() + && let Some(op) = lowered + { + // Attach to the wave every other open dependency is + // ordered before; bail to flushing if none dominates. + if let Some(&last) = open + .iter() + .find(|&&cat| open.iter().all(|&o| o == cat || self.reaches(o, cat))) + { + for observation in observations { + self.member.insert(observation, (last, self.open_gen[last])); + } + self.trailing[last].push((node.inner_idx, op)); + return; + } + for cat in open { + self.flush(cat, out); + } + out.push((node.inner_idx, op)); + return; + } + for cat in open { + self.flush(cat, out); + } + if let Some(op) = lowered { + out.push((node.inner_idx, op)); + } + } + } + } + + pub(super) fn finish(&mut self, out: &mut Vec<(NodeIndex, QueuedOperation)>) { + for cat in 0..CATEGORY_COUNT { + self.flush(cat, out); + } + } + + fn flush(&mut self, cat: usize, out: &mut Vec<(NodeIndex, QueuedOperation)>) { + // Ordered predecessors first. Clear this wave's constraint edges + // before recursing so a predecessor's own flush can never loop back. + let predecessors: Vec = (0..CATEGORY_COUNT) + .filter(|&other| other != cat && self.before[other][cat]) + .collect(); + for other in 0..CATEGORY_COUNT { + self.before[cat][other] = false; + self.before[other][cat] = false; + } + for other in predecessors { + self.flush(other, out); + } + let wave = std::mem::take(&mut self.waves[cat]); + let trailing = std::mem::take(&mut self.trailing[cat]); + if cat == CAT_REGION { + self.region_wave_inputs.clear(); + } + self.wave_bindings[cat] = 0; + self.open_gen[cat] += 1; + if wave.is_empty() { + debug_assert!(trailing.is_empty()); + return; + } + if wave.len() == 1 { + let (node, seg) = wave.into_iter().next().expect("length checked"); + let op: QueuedOperation = match seg { + SegOp::Row(op) => QueuedOperation::Operation(Arc::new(op)), + SegOp::MatMul(op, _) => QueuedOperation::Operation(Arc::new(*op)), + // A lone one-statement region lowers through the standalone + // elementwise path (which keeps the register-reuse tiled + // plan); genuine multi-output regions have no standalone + // lowering and stay merged. + SegOp::Region(op) => { + if op.statements.len() == 1 { + let nary = op + .into_nary() + .expect("single-statement regions always emit their output"); + QueuedOperation::Operation(Arc::new(nary)) + } else { + QueuedOperation::Merged(MergedSegments::Region(vec![(node, op)])) + } + } + }; + out.push((node, op)); + out.extend(trailing); + return; + } + if cat == CAT_MATMUL || cat == CAT_MATMUL_SPLITK { + // Partition into same-profile groups (first-occurrence order): + // only identical profiles share a guarded dispatch, so the + // segment bodies differ solely in their storage bindings. All + // wave members are mutually independent, so any group order is + // a valid topological order. + type MatmulGroup = Vec<(NodeIndex, crate::matmul::MatMulOperation)>; + let mut groups: Vec<(crate::matmul::MatmulMergeKey, MatmulGroup)> = Vec::new(); + for (node, seg) in wave { + let SegOp::MatMul(op, key) = seg else { + unreachable!("wave category mismatch"); + }; + match groups.iter_mut().find(|(existing, _)| *existing == key) { + Some((_, group)) => group.push((node, *op)), + None => groups.push((key, vec![(node, *op)])), + } + } + for (key, group) in groups { + if self.device.config().trace_matmul_merge { + eprintln!( + "matmul_merge_flush cat={cat} size={} key={key:?}", + group.len() + ); + } + if group.len() == 1 { + let (node, op) = group.into_iter().next().expect("length checked"); + out.push((node, QueuedOperation::Operation(Arc::new(op)))); + } else { + let merged = MergedSegments::MatMul(group); + out.push((merged.representative(), QueuedOperation::Merged(merged))); + } + } + out.extend(trailing); + return; + } + let merged = match cat { + CAT_REGION => MergedSegments::Region( + wave.into_iter() + .map(|(node, seg)| match seg { + SegOp::Region(op) => (node, op), + _ => unreachable!("wave category mismatch"), + }) + .collect(), + ), + _ => MergedSegments::Row( + wave.into_iter() + .map(|(node, seg)| match seg { + SegOp::Row(op) => (node, op), + _ => unreachable!("wave category mismatch"), + }) + .collect(), + ), + }; + out.push((merged.representative(), QueuedOperation::Merged(merged))); + out.extend(trailing); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::kernel_selection::CooperativeMatrixKind; + use crate::nary_wise::{NaryExpr, NaryScalar}; + use crate::reduce::{ReduceFunction, ReduceOp, ReduceOperation}; + use crate::region::{ElementwiseRegionOperation, RegionStatement}; + use crate::{Device, Tensor}; + + #[test] + fn fusion_toposort_is_invariant_to_valid_topological_interleaving() { + let region_node = |output: usize, input: usize| ExecutionNode { + inner_idx: NodeIndex::new(output), + variant: ExecutionVariant::Region(ElementwiseRegionOperation { + inputs: vec![NodeIndex::new(input)], + statements: vec![RegionStatement { + expression: NaryExpr::input(0, 1), + datatype: crate::DataTypeEnum::F32, + output: Some(NodeIndex::new(output)), + }], + shape: vec![64].into_boxed_slice(), + }), + }; + let row_node = |output: usize, input: usize| ExecutionNode { + inner_idx: NodeIndex::new(output), + variant: ExecutionVariant::RowProgram(RowProgramOperation::from_reduce( + &ReduceOperation { + inputs: vec![NodeIndex::new(input)], + expression: NaryExpr::indexed_input( + 0, + vec![NaryExpr::DimIndex(0), NaryExpr::DimIndex(1)], + ), + shape: vec![1, 64].into_boxed_slice(), + function: ReduceFunction { + name: Some("sum".to_string()), + op: ReduceOp::Sum, + initial_value: NaryScalar::F32(0.0), + datatype: crate::DataTypeEnum::F32, + }, + post_element_wise: crate::nary_wise::UnaryFunctionChain::empty( + crate::DataTypeEnum::F32, + ), + axis: 1, + }, + )), + }; + + let schedule = |nodes: Vec| { + let mut graph = ExecutionGraph::default(); + let mut by_inner = FxHashMap::default(); + for node in nodes { + let inner = node.inner_idx; + let exec = graph.add_node(node); + by_inner.insert(inner, exec); + } + let execs = graph.node_indices().collect::>(); + for exec in execs { + let mut dependencies = Vec::new(); + match &graph[exec].variant { + ExecutionVariant::Region(op) => { + op.visit_dependencies(&mut |dep| dependencies.push(dep)) + } + ExecutionVariant::RowProgram(op) => { + op.visit_dependencies(&mut |dep| dependencies.push(dep)) + } + _ => unreachable!("test graph only contains regions and row programs"), + } + for dependency in dependencies { + if let Some(&parent) = by_inner.get(&dependency) { + graph.add_edge(parent, exec, ()); + } + } + } + + fusion_toposort(&graph) + .into_iter() + .map(|exec| graph[exec].inner_idx.index()) + .collect::>() + }; + + // Both insertion orders encode the same two independent + // region -> row -> region branches. + let layered = schedule(vec![ + region_node(10, 0), + region_node(20, 1), + row_node(11, 10), + row_node(21, 20), + region_node(12, 11), + region_node(22, 21), + ]); + let interleaved = schedule(vec![ + region_node(10, 0), + row_node(11, 10), + region_node(12, 11), + region_node(20, 1), + row_node(21, 20), + region_node(22, 21), + ]); + + assert_eq!(interleaved, layered); + assert_eq!(layered, [10, 20, 11, 21, 12, 22]); + } + + #[test] + fn region_output_alias_orders_a_dependent_region_after_its_wave() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let region = |inputs: Vec, outputs: &[usize]| ElementwiseRegionOperation { + inputs, + statements: outputs + .iter() + .map(|&output| RegionStatement { + expression: NaryExpr::Scalar(NaryScalar::F32(0.0)), + datatype: crate::DataTypeEnum::F32, + output: Some(NodeIndex::new(output)), + }) + .collect(), + shape: vec![1].into_boxed_slice(), + }; + let first = ExecutionNode { + inner_idx: NodeIndex::new(10), + variant: ExecutionVariant::Region(region(vec![NodeIndex::new(0)], &[10, 11])), + }; + let second = ExecutionNode { + inner_idx: NodeIndex::new(20), + variant: ExecutionVariant::Region(region(vec![NodeIndex::new(11)], &[20])), + }; + + let mut merger = HorizontalMerger::new(&device); + let mut output = Vec::new(); + let shared_outputs = FxHashMap::default(); + merger.push(&first, None, &shared_outputs, &mut output); + merger.push(&second, None, &shared_outputs, &mut output); + merger.finish(&mut output); + + assert_eq!( + output.len(), + 2, + "a consumer of any emitted region output must run in a later dispatch" + ); + }); + } + + #[test] + fn shared_output_alias_orders_its_consumer_after_the_producer_wave() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let region = |inputs: Vec, output: usize| ElementwiseRegionOperation { + inputs, + statements: vec![RegionStatement { + expression: NaryExpr::Scalar(NaryScalar::F32(0.0)), + datatype: crate::DataTypeEnum::F32, + output: Some(NodeIndex::new(output)), + }], + shape: vec![1].into_boxed_slice(), + }; + let first = ExecutionNode { + inner_idx: NodeIndex::new(10), + variant: ExecutionVariant::Region(region(vec![NodeIndex::new(0)], 10)), + }; + let second = ExecutionNode { + inner_idx: NodeIndex::new(20), + variant: ExecutionVariant::Region(region(vec![NodeIndex::new(11)], 20)), + }; + let shared_outputs = + FxHashMap::from_iter([(NodeIndex::new(10), vec![NodeIndex::new(11)])]); + + let mut merger = HorizontalMerger::new(&device); + let mut output = Vec::new(); + merger.push(&first, None, &shared_outputs, &mut output); + merger.push(&second, None, &shared_outputs, &mut output); + merger.finish(&mut output); + + assert_eq!( + output.len(), + 2, + "a consumer of a shared observation must run after its representative" + ); + }); + } + + #[test] + fn standard_graph_merges_independent_qkv_matmuls() { + pollster::block_on(async { + const D: usize = 64; + + let Ok(device) = Device::new().await else { + return; + }; + let coop_viable = device + .coop_token(CooperativeMatrixKind::F32F32M8N8K8) + .is_some() + && device + .subgroup_config() + .is_some_and(|config| config.is_fixed()); + if !coop_viable { + return; + } + + let input_values = (0..D * D) + .map(|index| (index % 17) as f32 * 0.125 - 1.0) + .collect::>(); + let diagonal = |scale: f32| { + (0..D * D) + .map(|index| { + let row = index / D; + let column = index % D; + if row == column { scale } else { 0.0 } + }) + .collect::>() + }; + let input = Tensor::from_slice(&device, [D, D], &input_values); + let q = input.mat_mul(&Tensor::from_slice(&device, [D, D], &diagonal(1.0))); + let k = input.mat_mul(&Tensor::from_slice(&device, [D, D], &diagonal(2.0))); + let v = input.mat_mul(&Tensor::from_slice(&device, [D, D], &diagonal(-0.5))); + let output = &(&q + &k) + &v; + + assert!( + output.resolves_in::<2>(), + "Q/K/V should share one matmul dispatch followed by one nary dispatch", + ); + let values = output.as_slice::<2, f32>().await.unwrap(); + for row in 0..D { + for column in 0..D { + let expected = input_values[row * D + column] * 2.5; + assert!( + (values[[row, column]] - expected).abs() < 1e-4, + "mismatch at [{row}, {column}]", + ); + } + } + }); + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/mod.rs b/fusor-ml/core/src/compute_graph/resolve/mod.rs index 89c7213be..76ec04bb8 100644 --- a/fusor-ml/core/src/compute_graph/resolve/mod.rs +++ b/fusor-ml/core/src/compute_graph/resolve/mod.rs @@ -1,42 +1,57 @@ -use std::{collections::VecDeque, str::FromStr, sync::Arc}; +//! One materialization pipeline for the lazy compute graph. +//! +//! A resolve builds the temporary execution graph, recognizes specialized +//! operations and applies policy-driven fusion through the +//! equality-saturation optimizer (see [`egraph`]), lowers nodes into an +//! operation queue, builds complete kernel plans, and encodes the resulting +//! command records. Flush replay skips deterministic planning but rejoins the +//! same command-record encoder. + +use std::sync::Arc; use web_time::{Duration, Instant}; use crate::{ DataTypeEnum, Layout, - compute_graph::layout_pass::LayoutPass, mir::{inputs::MirValue, kernel_backend::PreparedDirectDispatch, operation::Operation}, - nary_wise::{ - ElementwiseOperation, ExtractedUnaryChain, NaryExpr, NaryOp, NaryScalar, UnaryFunctionChain, - }, - quantized::matmul::{ElementwiseEpilogue, QMatMulOperation}, + nary_wise::{ElementwiseOperation, NaryExpr, NaryOp, NaryScalar}, + quantized::matmul::QMatMulOperation, tensor::TensorData, }; use petgraph::algo::toposort; use petgraph::stable_graph::StableGraph; use rustc_hash::{FxHashMap, FxHashSet}; -use super::{ - ComputeGraphInner, ComputeGraphNode, ComputeGraphNodeVariant, GraphOperation, NodeIndex, -}; +use super::{ComputeGraphInner, ComputeGraphNode, ComputeGraphNodeVariant, NodeIndex}; use crate::{ MatMulOperation, ReduceOperation, dequantize::DequantizeOperation, quantized::embedding::QEmbeddingOperation, slice_assign::SliceAssignOperation, view::ViewOperation, }; +mod alloc_reuse; mod cluster_match; +mod egraph; mod execution; -mod fold_views; -mod fusion_basic; -mod fusion_matmul; +pub(crate) mod flush_replay; +mod fusion_region; mod fusion_row; -mod plan_cache; +#[cfg(test)] +mod key_goldens; +pub(crate) mod merge_horizontal; +pub(crate) mod plan_cache; +mod queue_executor; mod recognize; mod recognize_attention; mod recognize_cat; +#[cfg(test)] +mod recognize_gates; +mod sink_views; mod run; +#[cfg(feature = "graphvis")] +mod visualize; +pub(crate) use egraph::FusionPlanStore; pub(crate) use plan_cache::structural_kernel_key; pub(crate) struct ResolverResult { @@ -71,22 +86,17 @@ enum CommandRecord { } enum QueuedOperation { - Generic(Arc), - QMatMul(Box), + Operation(Arc), + /// Independent compatible operations merged into one dispatch (see + /// `merge_horizontal`). + Merged(merge_horizontal::MergedSegments), } impl QueuedOperation { fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { match self { - Self::Generic(operation) => operation.visit_dependencies(f), - Self::QMatMul(operation) => operation.visit_dependencies(f), - } - } - - fn inputs(&self, graph: &ComputeGraphInner) -> Vec { - match self { - Self::Generic(operation) => operation.inputs(graph), - Self::QMatMul(operation) => operation.inputs(graph), + Self::Operation(operation) => operation.visit_dependencies(f), + Self::Merged(merged) => merged.visit_dependencies(f), } } } @@ -129,57 +139,6 @@ struct ResolveHostProfile { profile_readback: Duration, } -#[derive(Default)] -struct ResolveHostCategoryProfile { - count: usize, - inputs: Duration, - output: Duration, - workgroup: Duration, - build_kernel: Duration, - prepare_dispatch: Duration, -} - -#[derive(Default)] -struct OptimizeProfile { - iterations: usize, - changed: usize, - fuse_naries_count: usize, - fuse_naries: Duration, - fuse_reduce_count: usize, - fuse_reduce: Duration, - fuse_matmul_count: usize, - fuse_matmul: Duration, -} - -impl OptimizeProfile { - fn print(&self) { - tracing::info!( - "resolve_optimize_profile iterations={} changed={} \ -fuse_naries_count={} fuse_naries={:?} \ -fuse_reduce_count={} fuse_reduce={:?} \ -fuse_matmul_count={} fuse_matmul={:?}", - self.iterations, - self.changed, - self.fuse_naries_count, - self.fuse_naries, - self.fuse_reduce_count, - self.fuse_reduce, - self.fuse_matmul_count, - self.fuse_matmul, - ); - } -} - -const DEFAULT_OPTIMIZE_NODE_LIMIT: usize = 512; -const LARGE_GRAPH_NARY_FUSION_MIN_LAST_DIM: usize = 512; - -fn optimize_node_limit() -> usize { - std::env::var("FUSOR_RESOLVE_OPTIMIZE_MAX_NODES") - .ok() - .and_then(|value| usize::from_str(&value).ok()) - .unwrap_or(DEFAULT_OPTIMIZE_NODE_LIMIT) -} - impl ResolveHostProfile { fn print(&self, total: Duration, queued_ops: usize, kernels: usize) { tracing::info!( @@ -209,52 +168,7 @@ timestamp_setup={:?} encode={:?} submit={:?} profile_readback={:?}", } } -fn print_host_category_profile(profile: FxHashMap<&'static str, ResolveHostCategoryProfile>) { - let mut profile = profile - .into_iter() - .map(|(category, profile)| { - ( - category, - profile.count, - profile.inputs, - profile.output, - profile.workgroup, - profile.build_kernel, - profile.prepare_dispatch, - ) - }) - .collect::>(); - profile.sort_by_key(|entry| std::cmp::Reverse(entry.5)); - tracing::info!("resolve_host_category_profile {profile:?}"); -} - -fn node_category_inner(variant: &ComputeGraphNodeVariant) -> &'static str { - match variant { - ComputeGraphNodeVariant::Tensor(_) => "tensor", - ComputeGraphNodeVariant::QMatrix(_) => "q_matrix", - ComputeGraphNodeVariant::Elementwise(_) => "elementwise", - ComputeGraphNodeVariant::Reduce(_) => "reduce", - ComputeGraphNodeVariant::View(_) => "view", - ComputeGraphNodeVariant::Assign(_) => "assign", - } -} - #[allow(dead_code, reason = "execution-side category labeling for profiling")] -fn node_category(variant: &ExecutionVariant) -> &'static str { - match variant { - ExecutionVariant::Elementwise(_) => "nary", - ExecutionVariant::Assign(_) => "slice_assign", - ExecutionVariant::View(_) => "view", - ExecutionVariant::QMatrix(_) => "dequantize", - ExecutionVariant::QEmbedding(_) => "q_embedding", - ExecutionVariant::MatMul(_) => "matmul", - ExecutionVariant::QMatMul(_) => "q_matmul", - ExecutionVariant::Tensor(_) => "tensor", - ExecutionVariant::Reduce(_) => "reduce", - ExecutionVariant::GraphOp(op) => op.category(), - } -} - /// What an execution-graph node lowers to. The graph vocabulary (the first /// six variants) enters verbatim; the region variants exist only here — /// recognition rebuilds them from composed clusters, and fusion enriches @@ -265,13 +179,63 @@ pub(crate) enum ExecutionVariant { QMatrix(DequantizeOperation), Elementwise(ElementwiseOperation), Reduce(ReduceOperation), + /// A fold with a named carrier and an explicit combine. `Reduce` is the + /// single-slot, built-in-combine special case; a fold that has one lowers + /// through it, and multi-slot folds are the new capability. + Fold(crate::fold::FoldOperation), View(ViewOperation), Assign(SliceAssignOperation), + /// Multi-output elementwise region formed by `fusion_region` on the + /// dense branch; never present in the inner graph. + Region(crate::region::ElementwiseRegionOperation), // Recognized regions. MatMul(MatMulOperation), QMatMul(Box), QEmbedding(QEmbeddingOperation), - GraphOp(Arc), + RowProgram(crate::row_program::RowProgramOperation), + Attention(crate::flash_attention::FlashAttentionOperation), +} + +impl ExecutionVariant { + /// Dependencies in dependency-slot order: the order the e-graph mirrors + /// as e-node children. Tensor leaves have none. + pub(super) fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { + match self { + Self::Tensor(_) => {} + Self::QMatrix(op) => op.visit_dependencies(f), + Self::Elementwise(op) => op.visit_dependencies(f), + Self::Reduce(op) => op.visit_dependencies(f), + Self::Fold(op) => op.visit_dependencies(f), + Self::View(op) => op.visit_dependencies(f), + Self::Assign(op) => op.visit_dependencies(f), + Self::Region(op) => op.visit_dependencies(f), + Self::MatMul(op) => op.visit_dependencies(f), + Self::QMatMul(op) => op.visit_dependencies(f), + Self::QEmbedding(op) => op.visit_dependencies(f), + Self::RowProgram(op) => op.visit_dependencies(f), + Self::Attention(op) => op.visit_dependencies(f), + } + } + + /// The same slots as [`Self::visit_dependencies`], in the same order, as + /// rebindable references. + pub(super) fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + match self { + Self::Tensor(_) => {} + Self::QMatrix(op) => op.visit_dependencies_mut(f), + Self::Elementwise(op) => op.visit_dependencies_mut(f), + Self::Reduce(op) => op.visit_dependencies_mut(f), + Self::Fold(op) => op.visit_dependencies_mut(f), + Self::View(op) => op.visit_dependencies_mut(f), + Self::Assign(op) => op.visit_dependencies_mut(f), + Self::Region(op) => op.visit_dependencies_mut(f), + Self::MatMul(op) => op.visit_dependencies_mut(f), + Self::QMatMul(op) => op.visit_dependencies_mut(f), + Self::QEmbedding(op) => op.visit_dependencies_mut(f), + Self::RowProgram(op) => op.visit_dependencies_mut(f), + Self::Attention(op) => op.visit_dependencies_mut(f), + } + } } impl From for ExecutionVariant { @@ -306,20 +270,40 @@ fn padded_query_buffer_size(size: u64) -> u64 { } #[cfg(not(target_arch = "wasm32"))] -fn print_gpu_kernel_profile( +fn collect_gpu_kernel_profile( records: &[DispatchMetadata], timestamps: &[u64], timestamp_period_ns: f64, - timestamp_mode: &str, -) { + timestamp_mode: &'static str, +) -> crate::KernelProfile { let mut category_profile = FxHashMap::::default(); let mut name_profile = FxHashMap::::default(); let mut accounted_ns = 0.0; + let mut unmeasured = 0usize; + let mut span_begin = u64::MAX; + let mut span_end = 0u64; for (index, record) in records.iter().enumerate() { - let begin = timestamps.get(index * 2).copied().unwrap_or_default(); - let end = timestamps.get(index * 2 + 1).copied().unwrap_or(begin); - let ns = end.saturating_sub(begin) as f64 * timestamp_period_ns; + // A slot the GPU never sampled resolves as zero, and an invalid sample as + // `MTLCounterErrorValue`. Both are indistinguishable from a free dispatch if + // summed, so they are counted as unmeasured instead of folded in as 0 ns. + let sample = timestamps + .get(index * 2) + .zip(timestamps.get(index * 2 + 1)) + .filter(|(begin, end)| { + **begin != 0 + && **end != 0 + && **begin != u64::MAX + && **end != u64::MAX + && end >= begin + }); + let Some((&begin, &end)) = sample else { + unmeasured += 1; + continue; + }; + span_begin = span_begin.min(begin); + span_end = span_end.max(end); + let ns = (end - begin) as f64 * timestamp_period_ns; accounted_ns += ns; if let Some(category) = &record.category { category_profile @@ -332,62 +316,153 @@ fn print_gpu_kernel_profile( } } - let span_ns = match (timestamps.first(), timestamps.last()) { - (Some(first), Some(last)) => last.saturating_sub(*first) as f64 * timestamp_period_ns, - _ => 0.0, - }; + // Only the sampled dispatches bound the span; a partially measured resolve has no + // honest wall span, so report it as absent rather than as a subset. + let span_ms = (unmeasured == 0 && span_begin <= span_end) + .then(|| (span_end - span_begin) as f64 * timestamp_period_ns / 1_000_000.0); - let mut categories = category_profile - .into_iter() - .map(|(name, aggregate)| { - ( - name, - aggregate.count, - aggregate.total_ns / 1_000_000.0, - aggregate.total_ns / aggregate.count as f64 / 1_000.0, - aggregate.max_ns / 1_000.0, - ) - }) - .collect::>(); - categories.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal)); - - let mut names = name_profile - .into_iter() - .map(|(name, aggregate)| { - ( + let rows = |profile: FxHashMap| { + let mut rows = profile + .into_iter() + .map(|(name, aggregate)| crate::KernelProfileRow { name, - aggregate.count, - aggregate.total_ns / 1_000_000.0, - aggregate.total_ns / aggregate.count as f64 / 1_000.0, - aggregate.max_ns / 1_000.0, - ) - }) - .collect::>(); - names.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal)); - names.truncate(32); + count: aggregate.count, + total_ms: aggregate.total_ns / 1_000_000.0, + average_us: aggregate.total_ns / aggregate.count as f64 / 1_000.0, + max_us: aggregate.max_ns / 1_000.0, + }) + .collect::>(); + rows.sort_by(|a, b| { + b.total_ms + .partial_cmp(&a.total_ms) + .unwrap_or(std::cmp::Ordering::Equal) + }); + rows + }; + let categories = rows(category_profile); + let mut top_names = rows(name_profile); + top_names.truncate(32); - tracing::info!( - "resolve_gpu_kernel_profile mode={} kernels={} accounted_ms={:.3} span_ms={:.3} timestamp_period_ns={:.3}", + let profile = crate::KernelProfile { timestamp_mode, - records.len(), - accounted_ns / 1_000_000.0, - span_ns / 1_000_000.0, - timestamp_period_ns + kernels: records.len(), + unmeasured_kernels: unmeasured, + accounted_ms: accounted_ns / 1_000_000.0, + span_ms, + timestamp_period_ns, + categories, + top_names, + }; + log_gpu_kernel_profile(&profile); + profile +} + +#[cfg(not(target_arch = "wasm32"))] +fn log_gpu_kernel_profile(profile: &crate::KernelProfile) { + let tuples = |rows: &[crate::KernelProfileRow]| { + rows.iter() + .map(|row| { + ( + row.name.clone(), + row.count, + row.total_ms, + row.average_us, + row.max_us, + ) + }) + .collect::>() + }; + tracing::info!( + "resolve_gpu_kernel_profile mode={} kernels={} unmeasured={} accounted_ms={:.3} span_ms={} timestamp_period_ns={:.3}", + profile.timestamp_mode, + profile.kernels, + profile.unmeasured_kernels, + profile.accounted_ms, + profile + .span_ms + .map_or_else(|| "absent".to_string(), |span| format!("{span:.3}")), + profile.timestamp_period_ns ); + let categories = tuples(&profile.categories); tracing::info!("resolve_gpu_kernel_categories {categories:?}"); + let names = tuples(&profile.top_names); tracing::info!("resolve_gpu_kernel_top_names {names:?}"); } +#[cfg(all(test, not(target_arch = "wasm32")))] +mod profile_tests { + use super::{DispatchMetadata, collect_gpu_kernel_profile}; + + fn records(count: usize) -> Vec { + (0..count) + .map(|i| DispatchMetadata { + name: Some(format!("kernel_{i}")), + category: Some("matmul_x".to_string()), + }) + .collect() + } + + #[test] + fn every_sampled_dispatch_is_accounted() { + let profile = collect_gpu_kernel_profile( + &records(3), + &[100, 400, 500, 900, 1000, 1600], + 1.0, + "pass_boundary", + ); + assert_eq!(profile.kernels, 3); + assert_eq!(profile.unmeasured_kernels, 0); + assert_eq!(profile.accounted_ms, (300 + 400 + 600) as f64 / 1_000_000.0); + assert_eq!(profile.span_ms, Some(1500.0 / 1_000_000.0)); + assert_eq!(profile.categories[0].count, 3); + } + + #[test] + fn unsampled_dispatches_are_reported_not_summed_as_zero() { + // Slots the GPU never wrote resolve as zero; the middle dispatch must not + // land in the aggregates as a free kernel, and the span no longer covers + // the whole resolve. + let profile = collect_gpu_kernel_profile( + &records(3), + &[100, 400, 0, 0, 1000, 1600], + 1.0, + "pass_boundary", + ); + assert_eq!(profile.unmeasured_kernels, 1); + assert_eq!(profile.accounted_ms, (300 + 600) as f64 / 1_000_000.0); + assert_eq!(profile.span_ms, None); + assert_eq!(profile.categories[0].count, 2); + } + + #[test] + fn truncated_readback_reports_the_missing_tail() { + let profile = collect_gpu_kernel_profile(&records(2), &[100, 400], 1.0, "inside_pass"); + assert_eq!(profile.kernels, 2); + assert_eq!(profile.unmeasured_kernels, 1); + assert_eq!(profile.span_ms, None); + } +} + pub(crate) struct Resolver { execution_graph: ExecutionGraph, node_mapping: FxHashMap, - // Persistent memoized layout inference: the inner graph's node variants - // are immutable during optimization (rewrites only touch the execution - // graph and dependency edges), so layouts computed once stay valid for - // the whole resolve. - layout_pass: LayoutPass, targets: Vec, resolved_set: FxHashSet, + // Materialization-plan recorder, armed on the first occurrence of a + // structurally cacheable target set. Dense and quantized graphs share it. + // `RefCell` because some hook sites (`add_physical_dependencies`) only + // hold `&self`. + recorder: Option>, + // Compatible independent operations may merge into one dispatch. + /// One semantic e-class may satisfy several lazy graph observations. + /// Keys are the execution nodes that materialize; values receive the + /// same allocation without another dispatch. + shared_outputs: FxHashMap>, + /// Unary chains moved into a matmul epilogue this resolve. + sunk_chains: usize, + /// Wall-clock spent in each optimizer sub-phase of this resolve, for the + /// host-cost ledger printed under `FUSOR_TRACE_RESOLVE_HOST`. + optimize_phases: execution::OptimizePhases, } impl Resolver { @@ -413,8 +488,31 @@ impl Resolver { targets, execution_graph: Default::default(), node_mapping: Default::default(), - layout_pass: Default::default(), resolved_set, + recorder: None, + shared_outputs: Default::default(), + sunk_chains: 0, + optimize_phases: Default::default(), } } + + /// A batch resolver that additionally records a replayable + /// [`flush_replay::FlushPlan`] of everything it resolves. + pub(crate) fn new_batch_with_recording( + graph: &mut ComputeGraphInner, + targets: Vec, + fingerprint: flush_replay::FlushFingerprint, + ) -> Self { + let recorder = flush_replay::PlanRecorder::new(graph, &targets, fingerprint); + let mut resolver = Self::new_batch(graph, targets); + resolver.recorder = Some(std::cell::RefCell::new(recorder)); + resolver + } + + /// The recorded plan, if recording was armed and never poisoned. + pub(crate) fn take_recorded_plan(&mut self) -> Option { + self.recorder + .take() + .and_then(|recorder| recorder.into_inner().finish()) + } } diff --git a/fusor-ml/core/src/compute_graph/resolve/plan_cache.rs b/fusor-ml/core/src/compute_graph/resolve/plan_cache.rs index 9f710e383..9f71f5eb3 100644 --- a/fusor-ml/core/src/compute_graph/resolve/plan_cache.rs +++ b/fusor-ml/core/src/compute_graph/resolve/plan_cache.rs @@ -1,24 +1,136 @@ -//! Structural plan-cache keys for resolved operations. +//! The canonical structural folds behind the resolver's cache keys. +//! +//! The single-operation plan key and the horizontally merged plan key fold +//! the same item — an operation's type and kernel fields plus the MIR values +//! it binds — and differ only in how much dispatch identity survives, which +//! [`Identity`] names: a merged segment shares one grid the merged builder +//! derives from the whole wave, so per-segment geometry cannot key it. +//! +//! These keys select cached kernel plans that are replayed by positional +//! rebind and shared across processes through the persistent plan store, so +//! the exact bytes are load-bearing; `key_goldens` pins them. +use std::hash::Hash; + +use rustc_hash::FxHasher; + +use super::merge_horizontal::MergedSegments; use crate::mir::inputs::MirValue; use crate::mir::kernel_backend::{KernelCacheKey, KernelVariantKey}; -use crate::mir::operation::Operation; +use crate::mir::operation::{Operation, hash_mir_value}; use crate::mir::workgroup_shape::WorkgroupShape; -struct DirectPlanCacheKernelVariant; +/// A variant marker's `TypeId` is hashed into every key it stamps, so its +/// declaration site — module path included — is part of the recipe. +struct KernelPlanCacheVariant; + +/// How much of the dispatch an operation's structural key keeps. +enum Identity<'a> { + /// The exact dispatch this operation will run: its solved workgroup + /// shape and grid are baked into the generated kernel. + Dispatch(&'a WorkgroupShape), + /// Dispatch geometry erased, for segments whose grid is decided by the + /// merged builder rather than by the segment itself. + Erased, +} + +/// One work item's structural key at the requested dispatch identity. +fn item_key( + operation: &dyn Operation, + identity: Identity<'_>, + inputs: &[MirValue], + variant: KernelVariantKey, +) -> KernelCacheKey { + let (workgroup, dispatch_size) = match identity { + Identity::Dispatch(workgroup) => { + (Some(workgroup), operation.dispatch_size(workgroup, inputs)) + } + Identity::Erased => (None, [0; 3]), + }; + operation.kernel_cache_key_with_dispatch(variant, workgroup, dispatch_size, inputs) +} pub(crate) fn structural_kernel_key( operation: &dyn Operation, inputs: &[MirValue], workgroup: &WorkgroupShape, ) -> KernelCacheKey { - let dispatch_size = operation.dispatch_size(workgroup, inputs); - operation.kernel_cache_key_with_dispatch( - KernelVariantKey::of::(), - Some(workgroup), - dispatch_size, + let operation_key = item_key( + operation, + Identity::Dispatch(workgroup), inputs, - ) + KernelVariantKey::of::(), + ); + KernelCacheKey::from_hash_inputs(|state| { + operation_key.hash(state); + }) +} + +/// A structural plan-cache key for one horizontally merged dispatch: the +/// wave discriminant plus every segment's own structural key, so isomorphic +/// waves across resolves and processes share one plan. Region segments merge +/// without the `Operation` trait and fold their kernel fields inline. +pub(super) fn merged_segments_key( + variant: KernelVariantKey, + merged: &MergedSegments, + segment_inputs: &[Vec], +) -> KernelCacheKey { + KernelCacheKey::from_hash_inputs(|state| { + variant.hash(state); + std::mem::discriminant(merged).hash(state); + match merged { + MergedSegments::Region(segments) => { + hash_merged_segments(state, segments.iter().map(|(_, op)| op), segment_inputs) + } + _ => { + segment_inputs.len().hash(state); + for ((_, op), inputs) in merged.segment_ops().iter().zip(segment_inputs) { + item_key(*op, Identity::Erased, inputs, variant).hash(state); + } + } + } + }) +} + +/// The kernel-field surface merged segments key on. Region segments merge +/// without the `Operation` trait, so the fold names this surface directly. +pub(crate) trait SegmentFields { + fn hash_kernel_fields(&self, state: &mut FxHasher); +} + +impl SegmentFields for crate::matmul::MatMulOperation { + fn hash_kernel_fields(&self, state: &mut FxHasher) { + Operation::hash_kernel_fields(self, state); + } +} + +impl SegmentFields for crate::row_program::RowProgramOperation { + fn hash_kernel_fields(&self, state: &mut FxHasher) { + Operation::hash_kernel_fields(self, state); + } +} + +impl SegmentFields for crate::region::ElementwiseRegionOperation { + fn hash_kernel_fields(&self, state: &mut FxHasher) { + crate::region::ElementwiseRegionOperation::hash_kernel_fields(self, state); + } +} + +/// Hash one merged dispatch's cache-key material: every segment's kernel +/// fields plus every MIR input value layout. +pub(crate) fn hash_merged_segments<'a, S: SegmentFields + 'a>( + state: &mut FxHasher, + segments: impl ExactSizeIterator, + segment_inputs: &[Vec], +) { + segments.len().hash(state); + for (op, inputs) in segments.zip(segment_inputs) { + op.hash_kernel_fields(state); + inputs.len().hash(state); + for input in inputs { + hash_mir_value(state, input); + } + } } #[cfg(test)] @@ -136,7 +248,7 @@ mod tests { // reference. Covers both build arms, including the fused-epilogue qmatmul the // old fast cache could not key. #[test] - fn direct_plan_cache_rebind_matches_fresh_build() { + fn kernel_plan_cache_rebind_matches_fresh_build() { pollster::block_on(async { let Ok(device) = Device::new().await else { return; @@ -186,7 +298,7 @@ mod tests { } #[test] - fn direct_plan_cache_distinguishes_same_shape_generic_ops() { + fn kernel_plan_cache_distinguishes_same_shape_generic_ops() { pollster::block_on(async { let Ok(device) = Device::new().await else { return; @@ -205,7 +317,7 @@ mod tests { } #[test] - fn direct_plan_cache_rebinds_repeated_binding_slots_positionally() { + fn kernel_plan_cache_rebinds_repeated_binding_slots_positionally() { pollster::block_on(async { let Ok(device) = Device::new().await else { return; @@ -231,7 +343,7 @@ mod tests { } #[test] - fn direct_plan_cache_distinguishes_same_shape_slice_assign_ranges() { + fn kernel_plan_cache_distinguishes_same_shape_slice_assign_ranges() { pollster::block_on(async { let Ok(device) = Device::new().await else { return; diff --git a/fusor-ml/core/src/compute_graph/resolve/queue_executor.rs b/fusor-ml/core/src/compute_graph/resolve/queue_executor.rs new file mode 100644 index 000000000..09512f34e --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/queue_executor.rs @@ -0,0 +1,968 @@ +//! Execution of the resolver's lowered operation queue. +//! +//! The executor owns serial input gathering and output allocation, parallel +//! kernel-plan building, then ordered recording and command preparation. It +//! is used by every resolved graph; horizontal merging only changes individual +//! queue entries. + +use super::merge_horizontal::MergedSegments; +use super::*; +use crate::mir::kernel_backend::DirectKernel; + +/// Every node that observes `node`'s value, transitively. +/// +/// Observations chain: one pass can record `a` as an observation of `b` and a +/// later one record `b` as an observation of `c`, and then caching `c` has to +/// reach `a` too. Following only the direct entries leaves the far end of such +/// a chain uncached, and its readers resolve to nothing. +fn observations_of( + node: NodeIndex, + shared_outputs: &FxHashMap>, +) -> Vec { + let mut observations = Vec::new(); + let mut pending = vec![node]; + let mut seen = FxHashSet::default(); + while let Some(current) = pending.pop() { + let Some(direct) = shared_outputs.get(¤t) else { + continue; + }; + for &observation in direct { + if seen.insert(observation) { + observations.push(observation); + pending.push(observation); + } + } + } + observations +} + +fn cache_output( + graph: &mut ComputeGraphInner, + node: NodeIndex, + result: &TensorData, + shared_outputs: &FxHashMap>, +) { + graph.set_cached_result(node, result.clone()); + for observation in observations_of(node, shared_outputs) { + graph.set_cached_result(observation, result.clone()); + } +} + +fn record_shared_outputs( + recorder: &std::cell::RefCell, + node: NodeIndex, + result: &TensorData, + shared_outputs: &FxHashMap>, +) { + for observation in observations_of(node, shared_outputs) { + recorder + .borrow_mut() + .record_shared_alias(observation, result, node); + } +} + +/// One entry of the three-phase queue, preserving queue order. +enum QueueStep { + View { + node: NodeIndex, + result: TensorData, + deps: Vec, + }, + CopyAssign { + node: NodeIndex, + copies: Vec, + op: QueuedOperation, + }, + Work(usize), +} + +enum QueueWorkKind { + Operation { + inputs: Vec, + workgroup_shape: crate::mir::workgroup_shape::WorkgroupShape, + resolved: TensorData, + /// Node whose dead buffer this output claimed, if any. + claimed_from: Option, + }, + Merged { + segment_inputs: Vec>, + /// One entry per segment output, in segment/statement order (regions + /// contribute several outputs per segment), with the node whose dead + /// buffer the output claimed, if any. + outputs: Vec<(NodeIndex, TensorData, Option)>, + }, +} + +struct QueueWork { + node: NodeIndex, + op: QueuedOperation, + kind: QueueWorkKind, + built: std::sync::Mutex>, +} + +struct BuiltWork { + kernels: Vec, + prepared: Vec>, +} + +#[cfg(not(target_arch = "wasm32"))] +const MIN_PARALLEL_BUILD_QUEUE: usize = 16; +#[cfg(not(target_arch = "wasm32"))] +const MIN_PARALLEL_BUILD_REMAINDER: usize = 4; +#[cfg(not(target_arch = "wasm32"))] +const COLD_BUILD_THRESHOLD: std::time::Duration = std::time::Duration::from_millis(1); + +#[cfg(not(target_arch = "wasm32"))] +fn should_parallelize_build_remainder( + worker_count: usize, + remaining: usize, + probe_elapsed: std::time::Duration, +) -> bool { + worker_count > 1 + && remaining >= MIN_PARALLEL_BUILD_REMAINDER + && probe_elapsed >= COLD_BUILD_THRESHOLD +} + +pub(super) fn merged_plan_cache_key( + merged: &MergedSegments, + segment_inputs: &[Vec], +) -> crate::mir::kernel_backend::KernelCacheKey { + // Declared here because its `TypeId` — declaration site and all — stamps + // every merged plan key already in the persistent store. + struct MergedPlanKernelVariant; + super::plan_cache::merged_segments_key( + crate::mir::kernel_backend::KernelVariantKey::of::(), + merged, + segment_inputs, + ) +} + +fn build_queue_work( + work: &QueueWork, + graph: &ComputeGraphInner, + device: &crate::Device, +) -> BuiltWork { + let build_timer = std::time::Instant::now(); + let kernels = match (&work.op, &work.kind) { + ( + QueuedOperation::Operation(operation), + QueueWorkKind::Operation { + inputs, + workgroup_shape, + .. + }, + ) => { + let build_kernels = || { + operation + .build_direct_kernel_plan(graph, workgroup_shape, inputs) + .unwrap_or_else(|error| panic!("{error}")) + .into_kernels() + }; + let kernel_key = structural_kernel_key(operation.as_ref(), inputs, workgroup_shape); + let kernels = super::run::resolve_cached_kernel_plan( + device.kernel_cache(), + kernel_key, + super::run::kernel_plan_binding_buffers(inputs), + build_kernels, + ); + kernels + } + (QueuedOperation::Merged(merged), QueueWorkKind::Merged { segment_inputs, .. }) => { + // Merged kernels go through the same plan cache as single ops: + // buffers are presented flattened in segment order, and the + // insert path verifies that order matches the kernel's true + // binding order (folded or deduplicated plans silently skip). + let expected: Vec> = segment_inputs + .iter() + .flatten() + .filter_map(|value| match value { + MirValue::Tensor(tensor) => Some(tensor.buffer().clone()), + MirValue::QMatrix(matrix) => Some(matrix.buffer().clone()), + MirValue::Integer(_) | MirValue::Float(_) => None, + }) + .collect(); + let plan_key = merged_plan_cache_key(merged, segment_inputs); + if let Some(kernels) = device.kernel_cache().kernel_plan_cache().get_many( + device.kernel_cache(), + plan_key, + &[&expected], + ) { + return finish_queue_build(build_timer, kernels, device); + } + let built = match merged { + MergedSegments::Row(segments) => { + crate::row_program::build_merged_row_program_kernel( + graph, + &segments + .iter() + .map(|(_, op)| op.clone()) + .collect::>(), + segment_inputs, + ) + } + MergedSegments::MatMul(segments) => crate::matmul::build_merged_matmul_kernel( + graph, + &segments + .iter() + .map(|(_, op)| op.clone()) + .collect::>(), + segment_inputs, + ), + MergedSegments::Region(segments) => crate::nary_direct::build_merged_region_kernel( + graph, + &segments + .iter() + .map(|(_, op)| op.clone()) + .collect::>(), + segment_inputs, + ), + }; + match built { + Some(kernel) => { + device.kernel_cache().kernel_plan_cache().insert_many( + plan_key, + std::slice::from_ref(&kernel), + &[&expected], + ); + vec![kernel] + } + None if matches!(merged, MergedSegments::Region(_)) => { + // Region fallback: one standalone region kernel per + // segment. Replay records the resulting kernel batch. + let MergedSegments::Region(segments) = merged else { + unreachable!("matched above"); + }; + let kernels = segments + .iter() + .zip(segment_inputs) + .map(|((_, op), inputs)| { + crate::nary_direct::build_merged_region_kernel( + graph, + std::slice::from_ref(op), + std::slice::from_ref(inputs), + ) + .unwrap_or_else(|| { + panic!("region fallback did not provide a kernel: {}", op.name()) + }) + }) + .collect(); + kernels + } + None => { + // Fallback: per-segment kernels. Replay records the + // resulting kernel batch with the same output slots. + let max_subgroup_size = device.max_subgroup_size(); + let kernels = merged + .segment_ops() + .into_iter() + .zip(segment_inputs) + .flat_map(|((_, op), inputs)| { + let constraints = op.workgroup_shape_constraints(device); + let workgroup_shape = constraints + .solve(max_subgroup_size, &device.limits()) + .unwrap_or_else(|| { + panic!("failed to solve workgroup shape for merged fallback") + }); + op.build_direct_kernel_plan(graph, &workgroup_shape, inputs) + .unwrap_or_else(|error| panic!("{error}")) + .into_kernels() + }) + .collect(); + kernels + } + } + } + _ => unreachable!("queue work kind matches its queued operation"), + }; + finish_queue_build(build_timer, kernels, device) +} + +/// Prepare dispatches (which also compiles shaders and pipelines, here on +/// the parallel build workers) and assemble the phase-2 result. +fn finish_queue_build( + build_timer: std::time::Instant, + kernels: Vec, + device: &crate::Device, +) -> BuiltWork { + let prepared = kernels + .iter() + .map(|kernel| { + kernel + .prepare_dispatch(device.kernel_cache()) + .map(|dispatch| (dispatch, kernel.name().to_string())) + }) + .collect(); + if device.config().trace_build_times { + let total = build_timer.elapsed(); + if total.as_millis() >= 2 { + eprintln!( + "build_time total={total:?} first={}", + kernels.first().map(|k| k.name()).unwrap_or("") + ); + } + } + BuiltWork { kernels, prepared } +} + +impl Resolver { + /// Three-phase queue execution for all resolved graphs: serial input + /// gathering and output caching (queue order), parallel kernel building + /// and dispatch preparation, then serial recording, encoding, and + /// release accounting in exactly the original queue order. + #[allow(clippy::too_many_arguments)] + pub(super) fn execute_queue( + recorder: Option<&std::cell::RefCell>, + graph: &mut ComputeGraphInner, + device: &crate::Device, + max_subgroup_size: u32, + queued_operations: Vec<(NodeIndex, QueuedOperation)>, + remaining_consumers: &mut FxHashMap, + target_set: &FxHashSet, + shared_outputs: &FxHashMap>, + ledger: &mut super::alloc_reuse::BufferLedger, + commands: &mut Vec, + host_profile: &mut ResolveHostProfile, + host_trace: bool, + on_dispatch_name: &mut dyn FnMut(&str) -> Option, + ) { + // Phase 1: gather inputs, allocate outputs, cache results. + let gather_start = host_trace.then(Instant::now); + let mut steps = Vec::with_capacity(queued_operations.len()); + let mut work: Vec = Vec::new(); + for (node, queued_operation) in queued_operations { + let view_result = if let Some(node_data) = graph.nodes.nodes.node_weight(node) { + match &node_data.variant { + ComputeGraphNodeVariant::View(view) => graph + .get_cached_result(view.input) + .and_then(|input| view.try_map_tensor(input)), + _ => None, + } + } else { + None + }; + if let Some(result) = view_result { + let mut deps = Vec::new(); + graph.visit_dependencies(node, &mut |dep| deps.push(dep)); + cache_output(graph, node, &result, shared_outputs); + ledger.note_transient(result.buffer()); + ledger.consume(graph, &deps, target_set); + steps.push(QueueStep::View { node, result, deps }); + continue; + } + let slice_copy = graph.nodes.nodes.node_weight(node).and_then(|node_data| { + let ComputeGraphNodeVariant::Assign(slice_assign) = &node_data.variant else { + return None; + }; + Self::try_prepare_in_place_slice_assign_copy(graph, slice_assign) + }); + if let Some((output, copies)) = slice_copy { + cache_output(graph, node, &output, shared_outputs); + ledger.note_transient(output.buffer()); + for copy in &copies { + ledger.note_transient(©.source); + ledger.note_transient(©.destination); + } + let mut deps = Vec::new(); + queued_operation.visit_dependencies(&mut |dep| deps.push(dep)); + ledger.consume(graph, &deps, target_set); + steps.push(QueueStep::CopyAssign { + node, + copies, + op: queued_operation, + }); + continue; + } + match &queued_operation { + QueuedOperation::Operation(operation) => { + let mut inputs = operation.inputs(graph); + let output_value = operation.output(graph, &inputs); + let MirValue::Tensor(mut resolved) = output_value else { + panic!("Kernel input value is not a tensor"); + }; + // Cache the output before the death accounting: a + // source is only releasable once every alive-uncached + // descendant (this very operation) is cached. + cache_output(graph, node, &resolved, shared_outputs); + let mut deps = Vec::new(); + queued_operation.visit_dependencies(&mut |dep| deps.push(dep)); + ledger.consume(graph, &deps, target_set); + let mut claimed_from = None; + if ledger.enabled() { + let out_ptr = Arc::as_ptr(resolved.buffer()) as usize; + let forbidden: FxHashSet = inputs + .iter() + .filter_map(|value| match value { + MirValue::Tensor(tensor) => { + let ptr = Arc::as_ptr(tensor.buffer()) as usize; + (ptr != out_ptr).then_some(ptr) + } + _ => None, + }) + .collect(); + if let Some(swapped) = ledger.try_claim(node, &resolved, &forbidden) { + for value in inputs.iter_mut() { + if let MirValue::Tensor(tensor) = value + && Arc::as_ptr(tensor.buffer()) as usize == out_ptr + { + *value = swapped.clone().into(); + } + } + resolved = swapped; + claimed_from = ledger.chosen_source(node); + cache_output(graph, node, &resolved, shared_outputs); + } + } + ledger.note_alloc(&resolved); + for value in &inputs { + if let MirValue::Tensor(tensor) = value { + ledger.note_transient(tensor.buffer()); + } + } + ledger.note_transient(resolved.buffer()); + let constraints = operation.workgroup_shape_constraints(device); + let workgroup_shape = constraints + .solve(max_subgroup_size, &device.limits()) + .unwrap_or_else(|| { + panic!( + "Failed to find a valid workgroup shape for constraints {constraints:?}" + ) + }); + steps.push(QueueStep::Work(work.len())); + work.push(QueueWork { + node, + op: queued_operation, + kind: QueueWorkKind::Operation { + inputs, + workgroup_shape, + resolved, + claimed_from, + }, + built: std::sync::Mutex::new(None), + }); + } + QueuedOperation::Merged(merged) => { + let mut segment_inputs: Vec> = Vec::new(); + let mut outputs: Vec<(NodeIndex, TensorData, Option)> = Vec::new(); + if let MergedSegments::Region(segments) = merged { + let device = graph.device(); + // A segment may write an output over one of its own + // input buffers only when no other segment of this + // dispatch binds that buffer (concurrent workgroups) + // — count cached-buffer pointers across the whole + // dispatch and require the source to be unique. + let mut dispatch_ptr_uses: FxHashMap = FxHashMap::default(); + for (_, op) in segments { + for idx in &op.inputs { + if let Some(cached) = graph.get_cached_result(*idx) { + *dispatch_ptr_uses + .entry(Arc::as_ptr(cached.buffer()) as usize) + .or_insert(0) += 1; + } + } + } + // Segments share one unsynchronized dispatch, so a + // scratch claim must avoid every segment's reads, not + // just the claiming segment's own. + let dispatch_reads: FxHashSet = + dispatch_ptr_uses.keys().copied().collect(); + for (_, op) in segments { + let values: Vec = op + .inputs + .iter() + .map(|idx| { + graph + .get_result(*idx) + .expect("region inputs resolve before the region") + .into() + }) + .collect(); + // Register the gathered input clones before any + // claim so the reference accounting that guards + // in-place claims sees them. + for value in &values { + if let MirValue::Tensor(tensor) = value { + ledger.note_transient(tensor.buffer()); + } + } + let mut values = values; + let reads = op.input_read_summary(); + let mut slot_claimed = vec![false; op.inputs.len()]; + // Cache every output before the death accounting: + // sources are only releasable once this region + // (their last alive-uncached descendant) counts + // as cached. + let mut fresh_outputs = Vec::new(); + for statement in &op.statements { + let Some(out_node) = statement.output else { + continue; + }; + let output = TensorData::new_for_shape( + &device, + &op.shape, + statement.datatype, + ); + cache_output(graph, out_node, &output, shared_outputs); + fresh_outputs.push(output); + } + { + let mut deps = Vec::new(); + op.visit_dependencies(&mut |dep| deps.push(dep)); + ledger.consume(graph, &deps, target_set); + } + let mut fresh_outputs = fresh_outputs.into_iter(); + for (position, statement) in op.statements.iter().enumerate() { + let Some(out_node) = statement.output else { + continue; + }; + let mut output = fresh_outputs + .next() + .expect("one fresh output per statement"); + let mut claimed_from = None; + // Write in place over an input this statement + // is the last reader of: per-thread the load + // precedes the store and threads own disjoint + // elements, so identity reads stay exact. + for (slot, source) in op.inputs.iter().enumerate() { + if slot_claimed[slot] + || !reads[slot].identity_only + || reads[slot].last_reader != Some(position) + { + continue; + } + let unique = graph + .get_cached_result(*source) + .map(|cached| Arc::as_ptr(cached.buffer()) as usize) + .and_then(|ptr| dispatch_ptr_uses.get(&ptr)) + == Some(&1); + if !unique { + continue; + } + if let Some(swapped) = ledger.try_claim_in_place( + out_node, &output, *source, graph, target_set, + ) { + output = swapped; + claimed_from = Some(*source); + slot_claimed[slot] = true; + break; + } + } + if claimed_from.is_none() + && let Some(swapped) = + ledger.try_claim(out_node, &output, &dispatch_reads) + { + output = swapped; + claimed_from = ledger.chosen_source(out_node); + } + if claimed_from.is_some() { + cache_output(graph, out_node, &output, shared_outputs); + } + ledger.note_alloc(&output); + ledger.note_transient(output.buffer()); + values.push(output.clone().into()); + outputs.push((out_node, output, claimed_from)); + } + segment_inputs.push(values); + } + } else { + for (seg_node, op) in merged.segment_ops() { + let inputs = op.inputs(graph); + let MirValue::Tensor(output) = op.output(graph, &inputs) else { + panic!("merged segment output is not a tensor"); + }; + cache_output(graph, seg_node, &output, shared_outputs); + ledger.note_alloc(&output); + for value in &inputs { + if let MirValue::Tensor(tensor) = value { + ledger.note_transient(tensor.buffer()); + } + } + ledger.note_transient(output.buffer()); + outputs.push((seg_node, output, None)); + segment_inputs.push(inputs); + } + let mut deps = Vec::new(); + queued_operation.visit_dependencies(&mut |dep| deps.push(dep)); + ledger.consume(graph, &deps, target_set); + } + steps.push(QueueStep::Work(work.len())); + work.push(QueueWork { + node, + op: queued_operation, + kind: QueueWorkKind::Merged { + segment_inputs, + outputs, + }, + built: std::sync::Mutex::new(None), + }); + } + } + } + // Allocation is complete: releases past this point free buffers no + // claim can use anymore. + ledger.freeze(); + if let Some(start) = gather_start { + host_profile.inputs += start.elapsed(); + } + + // Phase 2: build kernels and prepare dispatches. Builds are pure + // functions of (operation, layouts, buffers); the shared kernel + // caches are internally synchronized. + let build_start = host_trace.then(Instant::now); + #[cfg(target_arch = "wasm32")] + for item in &work { + *item.built.lock().unwrap() = Some(build_queue_work(item, graph, device)); + } + #[cfg(not(target_arch = "wasm32"))] + { + let workers = std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1) + .min(work.len()) + .max(1); + // Tiny queues build serially. Larger queues probe on the caller + // thread until a build is measurably cold. A warm plan-cache hit + // costs much less than creating a fresh worker cohort, while the + // first shader/pipeline miss still moves the remaining cold work + // onto parallel workers. + if workers <= 1 || work.len() < MIN_PARALLEL_BUILD_QUEUE { + for item in &work { + *item.built.lock().unwrap() = Some(build_queue_work(item, graph, device)); + } + } else { + let mut next_index = 0; + while let Some(item) = work.get(next_index) { + let probe_start = std::time::Instant::now(); + *item.built.lock().unwrap() = Some(build_queue_work(item, graph, device)); + next_index += 1; + if should_parallelize_build_remainder( + workers, + work.len() - next_index, + probe_start.elapsed(), + ) { + break; + } + } + + if next_index < work.len() { + let remaining_workers = workers.min(work.len() - next_index); + let next = std::sync::atomic::AtomicUsize::new(next_index); + let graph_ref: &ComputeGraphInner = graph; + std::thread::scope(|scope| { + for _ in 0..remaining_workers { + scope.spawn(|| { + loop { + let index = + next.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let Some(item) = work.get(index) else { break }; + let built = build_queue_work(item, graph_ref, device); + *item.built.lock().unwrap() = Some(built); + } + }); + } + }); + } + } + } + if let Some(start) = build_start { + host_profile.build_kernel += start.elapsed(); + } + + // Phase 3: record, encode, and release in queue order. + let encode_start = host_trace.then(Instant::now); + let mut consumers = super::execution::NodeConsumers { + counts: remaining_consumers, + targets: target_set, + }; + for step in steps { + match step { + QueueStep::View { node, result, deps } => { + if let Some(recorder) = recorder { + recorder + .borrow_mut() + .record_view_alias(node, &result, &deps); + record_shared_outputs(recorder, node, &result, shared_outputs); + } + super::execution::release_consumed( + graph, + &mut consumers, + Some(ledger), + |release| deps.into_iter().for_each(release), + ); + } + QueueStep::CopyAssign { node, copies, op } => { + if let Some(recorder) = recorder { + let output = graph + .get_cached_result(node) + .expect("copy-assign output cached in phase 1") + .clone(); + recorder.borrow_mut().record_copy_assign(node, &output, &op); + } + commands.extend(copies.into_iter().map(CommandRecord::CopyBuffer)); + super::execution::release_consumed( + graph, + &mut consumers, + Some(ledger), + |release| op.visit_dependencies(release), + ); + } + QueueStep::Work(index) => { + let item = &work[index]; + let built = item + .built + .lock() + .unwrap() + .take() + .expect("queue work built in phase 2"); + if let Some(recorder) = recorder { + match (&item.op, &item.kind) { + ( + QueuedOperation::Operation(_), + QueueWorkKind::Operation { + resolved, + claimed_from, + .. + }, + ) => { + recorder.borrow_mut().record_dispatch( + item.node, + &built.kernels, + resolved, + &item.op, + *claimed_from, + ); + record_shared_outputs( + recorder, + item.node, + resolved, + shared_outputs, + ); + } + ( + QueuedOperation::Merged(merged), + QueueWorkKind::Merged { outputs, .. }, + ) => { + let node_outputs: Vec<(NodeIndex, &TensorData, Option)> = + outputs + .iter() + .map(|(node, output, claimed)| (*node, output, *claimed)) + .collect(); + recorder.borrow_mut().record_merged_dispatch( + &node_outputs, + &built.kernels, + merged, + ); + for (node, output, _) in outputs { + record_shared_outputs(recorder, *node, output, shared_outputs); + } + } + _ => unreachable!("queue work kind matches its queued operation"), + } + } + for (dispatch, name) in built.prepared.into_iter().flatten() { + let category = on_dispatch_name(&name); + commands.push(CommandRecord::Dispatch(DispatchRecord { + dispatch, + name, + category, + })); + } + super::execution::release_consumed( + graph, + &mut consumers, + Some(ledger), + |release| item.op.visit_dependencies(release), + ); + } + } + } + if let Some(start) = encode_start { + host_profile.prepare_dispatch += start.elapsed(); + } + } +} + +/// Where a profiled resolve writes its per-dispatch timestamps. +pub(super) struct TimestampPlan<'a> { + pub(super) query_set: &'a wgpu::QuerySet, + /// Timestamps ride inside the shared passes; without the feature every + /// dispatch takes its own pass and the writes land on its boundaries. + pub(super) inside_pass: bool, +} + +/// Encode a command stream with the resolver's pass and submit chunking. +/// +/// Intermediate chunks are handed to `submit_chunk`; the final encoder is +/// returned so the caller can append tail work before its synchronization +/// boundary and final submit. +pub(super) fn encode_command_records( + device: &crate::Device, + commands: &[CommandRecord], + total_kernels: usize, + timestamps: Option>, + mut command_encoder: wgpu::CommandEncoder, + mut submit_chunk: impl FnMut(wgpu::CommandEncoder, bool), +) -> wgpu::CommandEncoder { + let dispatches_per_pass = super::run::dispatches_per_pass(device, total_kernels); + let dispatches_per_submit = super::run::dispatches_per_submit(device, total_kernels); + let wait_after_chunk_submit = device.backend() == wgpu::Backend::Metal; + let mut command_index = 0usize; + let mut dispatch_index = 0usize; + let mut dispatches_in_submit = 0usize; + let mut encoder_has_commands = false; + let mut pass_segments = 0usize; + let mut copy_records = 0usize; + + while command_index < commands.len() { + if encoder_has_commands && dispatches_in_submit >= dispatches_per_submit { + let next_encoder = + device + .wgpu_device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("Resolver Encoder"), + }); + let ready_encoder = std::mem::replace(&mut command_encoder, next_encoder); + submit_chunk(ready_encoder, wait_after_chunk_submit); + encoder_has_commands = false; + dispatches_in_submit = 0; + } + + match &commands[command_index] { + CommandRecord::CopyBuffer(copy) => { + command_encoder.copy_buffer_to_buffer( + ©.source, + copy.source_offset, + ©.destination, + copy.destination_offset, + copy.size, + ); + copy_records += 1; + encoder_has_commands = true; + command_index += 1; + } + CommandRecord::Dispatch(record) => { + if let Some(plan) = ×tamps + && !plan.inside_pass + { + let mut pass = + command_encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some(record.name.as_str()), + timestamp_writes: Some(wgpu::ComputePassTimestampWrites { + query_set: plan.query_set, + beginning_of_pass_write_index: Some((dispatch_index * 2) as u32), + end_of_pass_write_index: Some((dispatch_index * 2 + 1) as u32), + }), + }); + record.dispatch.run(&mut pass); + drop(pass); + pass_segments += 1; + dispatch_index += 1; + dispatches_in_submit += 1; + encoder_has_commands = true; + command_index += 1; + continue; + } + + let mut pass = command_encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("Resolver Direct Kernels"), + timestamp_writes: None, + }); + pass_segments += 1; + let mut pass_dispatches = 0usize; + while command_index < commands.len() + && pass_dispatches < dispatches_per_pass + && dispatches_in_submit < dispatches_per_submit + { + let CommandRecord::Dispatch(record) = &commands[command_index] else { + break; + }; + if let Some(plan) = ×tamps { + pass.write_timestamp(plan.query_set, (dispatch_index * 2) as u32); + } + pass.push_debug_group(&record.name); + record.dispatch.run(&mut pass); + pass.pop_debug_group(); + if let Some(plan) = ×tamps { + pass.write_timestamp(plan.query_set, (dispatch_index * 2 + 1) as u32); + } + dispatch_index += 1; + dispatches_in_submit += 1; + command_index += 1; + pass_dispatches += 1; + encoder_has_commands = true; + } + } + } + } + if cfg!(target_arch = "wasm32") || device.config().trace_resolve_host { + tracing::info!( + "resolve_pass_layout kernels={total_kernels} passes={pass_segments} copies={copy_records}" + ); + } + + command_encoder +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::{Device, Tensor}; + + #[test] + fn warm_build_probe_keeps_queue_serial() { + assert!(!should_parallelize_build_remainder( + 8, + 32, + COLD_BUILD_THRESHOLD - std::time::Duration::from_nanos(1), + )); + } + + #[test] + fn cold_build_probe_parallelizes_useful_remainder() { + assert!(should_parallelize_build_remainder( + 8, + MIN_PARALLEL_BUILD_REMAINDER, + COLD_BUILD_THRESHOLD, + )); + } + + #[test] + fn cold_build_probe_does_not_spawn_for_tiny_remainder() { + assert!(!should_parallelize_build_remainder( + 8, + MIN_PARALLEL_BUILD_REMAINDER - 1, + COLD_BUILD_THRESHOLD, + )); + assert!(!should_parallelize_build_remainder( + 1, + 32, + COLD_BUILD_THRESHOLD, + )); + } + + #[test] + fn shared_eclass_dispatches_once_and_caches_every_observation() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let input = Tensor::new(&device, &[1.0f32, 2.0, 3.0, 4.0]); + let left = &input * 2.0; + let right = &input * 2.0; + let targets = vec![left.data().key, right.data().key]; + let (kernels, same_buffer) = device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.clone()); + let mut removed = Vec::new(); + let result = resolver.run(graph, &mut removed); + let left = graph.get_cached_result(targets[0]).unwrap(); + let right = graph.get_cached_result(targets[1]).unwrap(); + ( + result.total_kernels, + std::sync::Arc::ptr_eq(left.buffer(), right.buffer()), + ) + }); + device.poll_wait(); + assert_eq!(kernels, 1); + assert!(same_buffer); + }); + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize.rs b/fusor-ml/core/src/compute_graph/resolve/recognize.rs index f1083b32d..b2778b990 100644 --- a/fusor-ml/core/src/compute_graph/resolve/recognize.rs +++ b/fusor-ml/core/src/compute_graph/resolve/recognize.rs @@ -1,16 +1,17 @@ -//! Pattern recognition over the composed 3-op graph. +//! Contraction matchers over the composed 3-op graph. //! //! The tensor API expresses contractions as `Elementwise(Mul) + Reduce(Sum)` //! over a shared index space (see `Tensor::mat_mul` / `Tensor::q_mat_mul`). -//! This pass runs before view folding and n-ary fusion, while the composed -//! cluster is still in the exact canonical form the API emitted, and rebuilds -//! the specialized operation so the existing kernel paths (and their epilogue -//! fusion) take over. Anything it does not recognize lowers through the -//! generic elementwise + reduce kernels — slower, but correct. +//! The matchers and builders are consumed by a linear recognition sweep over +//! the identity forms emitted by the tensor interface. Keeping recognition +//! outside the equality-saturation pass avoids rebuilding the same large +//! decode e-graph for every generated token; `egraph/HOISTING_SPIKE.md` +//! measures what that would cost. +//! Anything they do not recognize lowers through the generic elementwise + +//! reduce kernels — slower, but correct. use crate::{ - MatMulOperation, ReduceOperation, dequantize::DequantizeOperation, - quantized::matmul::QMatMulOperation, reduce::ReduceOp, + MatMulOperation, ReduceOperation, quantized::matmul::QMatMulOperation, reduce::ReduceOp, }; use super::*; @@ -172,44 +173,12 @@ impl ComputeGraphInner { _ => None, } } - - /// Recognize a composed contraction rooted at `key` directly on the inner - /// graph (the single-target fast path, no resolver involved). The multiply - /// must feed only this reduce. - pub(crate) fn match_direct_qmatmul(&self, key: NodeIndex) -> Option { - let ComputeGraphNodeVariant::Reduce(reduce) = &self.nodes.nodes.node_weight(key)?.variant - else { - return None; - }; - let value = reduce.plain_input()?; - if self.get_cached_result(value).is_some() || self.has_live_reference(value) { - return None; - } - let ComputeGraphNodeVariant::Elementwise(nary) = - &self.nodes.nodes.node_weight(value)?.variant - else { - return None; - }; - if self - .nodes - .nodes - .neighbors_directed(value, petgraph::Direction::Outgoing) - .count() - != 1 - { - return None; - } - let contraction = match_contraction(reduce, nary)?; - let (operation, _) = contraction.to_q_mat_mul(|key| self.dequantize_variant(key))?; - Some(operation) - } } impl Resolver { - /// Recognize composed contraction clusters in the execution graph and - /// rebuild them as `MatMul` / `QMatMul` nodes. Runs as a linear sweep - /// before any other rewrite, while the clusters are still in the exact - /// form the API emitted. + /// Recognize canonical contraction clusters with one linear graph sweep. + /// This is equivalent to the former native-egg recognition rule, but the + /// API emits an exact canonical form so saturation adds no value here. pub(super) fn recognize_contractions(&mut self, graph: &mut ComputeGraphInner) { let reduces: Vec = self .execution_graph @@ -231,36 +200,62 @@ impl Resolver { graph: &mut ComputeGraphInner, node_idx: ExecutionNodeIndex, ) -> bool { + let trace = std::env::var_os("FUSOR_TRACE_RECOGNIZE").is_some(); let ExecutionVariant::Reduce(reduce) = &self.execution_graph[node_idx].variant else { return false; }; let Some(value) = reduce.plain_input() else { + if trace { + eprintln!("recognize-matmul: no plain input"); + } return false; }; if self.check_cached(graph, value) || graph.has_live_reference(value) { + if trace { + eprintln!( + "recognize-matmul: cached={} live={}", + self.check_cached(graph, value), + graph.has_live_reference(value) + ); + } return false; } let Some(nary_exec) = self.get_input_node_in_exec_graph(value) else { + if trace { + eprintln!("recognize-matmul: producer not in exec graph"); + } return false; }; let ExecutionVariant::Elementwise(nary) = &self.execution_graph[nary_exec].variant else { + if trace { + eprintln!("recognize-matmul: producer not elementwise"); + } return false; }; - // The multiply must exist solely for this reduce: consumed elsewhere - // it has to materialize anyway. if self .execution_graph .neighbors_directed(nary_exec, petgraph::Direction::Outgoing) .count() != 1 { + if trace { + eprintln!( + "recognize-matmul: multi-consumer producer shape={:?}", + nary.shape + ); + } return false; } let Some(contraction) = match_contraction(reduce, nary) else { + if trace { + eprintln!( + "recognize-matmul: contraction mismatch shape={:?}", + nary.shape + ); + } return false; }; - let _ = nary_exec; if let Some((operation, activation)) = contraction.to_q_mat_mul(|key| graph.dequantize_variant(key)) { @@ -273,8 +268,26 @@ impl Resolver { return true; } if let Some((mut operation, _)) = contraction.to_mat_mul(&graph.device()) { - self.try_unflatten_matmul_input(graph, &mut operation); + let device = graph.device(); + try_unflatten_matmul_input_with( + &mut operation, + &device, + |node| self.check_cached(graph, node), + |node| graph.has_live_reference(node), + |node| match &graph.nodes.nodes.node_weight(node)?.variant { + ComputeGraphNodeVariant::View(view) => Some(view.clone()), + _ => None, + }, + ); let inputs = [operation.first, operation.second]; + if trace { + eprintln!( + "recognize-matmul: COMMIT m={} k={} n={}", + operation.a.rows(), + operation.a.cols(), + operation.b.cols() + ); + } self.commit_recognized( graph, node_idx, @@ -286,102 +299,7 @@ impl Resolver { false } - /// Read a recognized matmul's A operand through its un-flattened - /// producer. Conv's im2col flatten regroups a windowed view's dims - /// across overlapping strides, which no single strided layout can - /// express: the view keeps a stage boundary and would materialize - /// through the gather fallback. When the flat `[M, K]` operand is - /// exactly such a reinterpret over an affine stage, point the matmul at - /// the producer, carry the affine stage as the operand's base map, and - /// let the kernels divmod the flat coordinates back apart per load — an - /// implicit GEMM with no gather dispatch. - fn try_unflatten_matmul_input( - &mut self, - graph: &ComputeGraphInner, - operation: &mut crate::MatMulOperation, - ) { - if !operation.a.is_plain() || operation.a.batch_dims != 0 { - return; - } - // Only the cooperative-matrix kernel reads an un-flattened operand - // faster than gather-then-matmul: its tile staging amortizes the - // per-load coordinate decomposition. The generic reduce re-derives - // coordinates for every load and measures slower than the gather at - // every meaningful size, so anything bound for it keeps the - // materialized matrix. - if !operation.hardware_matmul_statically_viable(&graph.device()) { - return; - } - let (m, k) = (operation.a.rows(), operation.a.cols()); - // An already-materialized (or externally held) operand is cheaper to - // read flat than to re-derive coordinates for. - if self.check_cached(graph, operation.first) || graph.has_live_reference(operation.first) { - return; - } - let Some(node) = graph.nodes.nodes.node_weight(operation.first) else { - return; - }; - let ComputeGraphNodeVariant::View(view) = &node.variant else { - return; - }; - // The stack must be an affine relayout under a flat [M, K] - // reinterpret, both pure relayouts (no fill regions). - let [windowed, flat] = view.stages.as_slice() else { - return; - }; - if !windowed.is_fully_defined() - || !flat.is_fully_defined() - || !flat.layout.is_contiguous() - || flat.layout.offset() != 0 - || flat.layout.shape() != [m, k] - { - return; - } - // The kernels substitute the windowed map as affine per-dim index - // arithmetic; validate it here so the lowering can rely on it. - if crate::view::affine_dim_indices(&windowed.layout, &windowed.input_shape).is_none() { - return; - } - let operand_shape = windowed.shape(); - // The producer's dims must split cleanly into an `M` prefix and a - // `K` suffix for the per-side flat-coordinate decomposition. - let mut product = 1usize; - let mut k_start = operand_shape.len(); - while k_start > 0 && product < k { - k_start -= 1; - let Some(next) = product.checked_mul(operand_shape[k_start]) else { - return; - }; - product = next; - } - if product != k - || k_start == 0 - || k_start == operand_shape.len() - || operand_shape[..k_start].iter().product::() != m - { - return; - } - // The flat row/column coordinates decompose with u32 arithmetic. - let probe = NaryExpr::DimIndex(0); - if crate::view::row_major_indices_from_flat(probe.clone(), &operand_shape[..k_start]) - .is_none() - || crate::view::row_major_indices_from_flat(probe, &operand_shape[k_start..]).is_none() - { - return; - } - operation.first = view.input; - operation.a = crate::matmul::MatrixOperand { - shape: operand_shape.into(), - batch_dims: 0, - row_dims: k_start, - base_map: Some(crate::matmul::OperandBaseMap { - layout: windowed.layout.clone(), - base_shape: windowed.input_shape.clone(), - }), - }; - } - - /// Sweep elementwise nodes for quantized embedding gathers. + /// Recognize quantized row gathers with one linear graph sweep. pub(super) fn recognize_embeddings(&mut self, graph: &mut ComputeGraphInner) { let candidates: Vec = self .execution_graph @@ -394,18 +312,13 @@ impl Resolver { }) .collect(); for node in candidates { - if !self.execution_graph.contains_node(node) { - continue; + if self.execution_graph.contains_node(node) { + self.try_recognize_q_embedding(graph, node); } - self.try_recognize_q_embedding(graph, node); } } - /// Quantized row gather: `Elementwise([table, idx], table[idx[i], j])` - /// over a `[count, hidden]` space (see `QMatrix::index_select_rows_to`). - /// Rebuilds the block-amortized embedding kernel; dense-storage tables - /// stay on the generic elementwise path, which reads them directly. - pub(super) fn try_recognize_q_embedding( + fn try_recognize_q_embedding( &mut self, graph: &mut ComputeGraphInner, node_idx: ExecutionNodeIndex, @@ -416,7 +329,6 @@ impl Resolver { if nary.inputs.len() != 2 || nary.shape.len() != 2 { return false; } - // Peel the optional cast from the load type to the requested type. let gather = match &nary.expression { NaryExpr::Op { children, function } if function.op == crate::nary_wise::NaryOp::Cast && children.len() == 1 => @@ -449,15 +361,15 @@ impl Resolver { let Some(dequantize) = graph.dequantize_variant(nary.inputs[0]) else { return false; }; - if crate::quantized::dequantize::quant_format(&dequantize.matrix).is_none() { - return false; - } - if dequantize.matrix.shape().len() != 2 || dequantize.matrix.shape()[1] != nary.shape[1] { + if crate::quantized::dequantize::quant_format(&dequantize.matrix).is_none() + || dequantize.matrix.shape().len() != 2 + || dequantize.matrix.shape()[1] != nary.shape[1] + { return false; } let indexes = nary.inputs[1]; - let operation = crate::quantized::embedding::QEmbeddingOperation::new( + let operation = QEmbeddingOperation::new( indexes, nary.shape[0], dequantize.matrix.clone(), @@ -471,39 +383,91 @@ impl Resolver { ); true } +} - /// Replace a recognized cluster's root with the rebuilt operation: drop - /// every edge from the cluster's intermediates, wire the operation's - /// dependencies directly, and let the now-unconsumed intermediates fall - /// out of the execution graph. - pub(super) fn commit_recognized( - &mut self, - graph: &mut ComputeGraphInner, - node_idx: ExecutionNodeIndex, - dependencies: &[NodeIndex], - variant: ExecutionVariant, - ) { - self.execution_graph[node_idx].variant = variant; - - let previous: Vec = self - .execution_graph - .neighbors_directed(node_idx, petgraph::Direction::Incoming) - .collect(); - for &prev in &previous { - if let Some(edge) = self.execution_graph.find_edge(prev, node_idx) { - self.execution_graph.remove_edge(edge); - } - } - for &dependency in dependencies { - if let Some(exec) = self.get_input_node_in_exec_graph(dependency) - && self.execution_graph.find_edge(exec, node_idx).is_none() - { - self.execution_graph.add_edge(exec, node_idx, ()); - } - } - self.add_physical_dependencies(graph, node_idx, dependencies); - for prev in previous { - self.remove_node_if_dead(prev); - } +/// Read a recognized matmul's A operand through its un-flattened producer. +/// Native egg appliers own their context and cannot borrow the live compute +/// graph, so every graph observation is supplied explicitly. +pub(super) fn try_unflatten_matmul_input_with( + operation: &mut crate::MatMulOperation, + device: &crate::Device, + check_cached: impl Fn(NodeIndex) -> bool, + has_live_reference: impl Fn(NodeIndex) -> bool, + view_for: impl Fn(NodeIndex) -> Option, +) { + if !operation.a.is_plain() || operation.a.batch_dims != 0 { + return; } + // Only the cooperative-matrix kernel reads an un-flattened operand + // faster than gather-then-matmul: its tile staging amortizes the + // per-load coordinate decomposition. The generic reduce re-derives + // coordinates for every load and measures slower than the gather at + // every meaningful size, so anything bound for it keeps the + // materialized matrix. + if !operation.hardware_matmul_statically_viable(device) { + return; + } + let (m, k) = (operation.a.rows(), operation.a.cols()); + // An already-materialized (or externally held) operand is cheaper to + // read flat than to re-derive coordinates for. + if check_cached(operation.first) || has_live_reference(operation.first) { + return; + } + let Some(view) = view_for(operation.first) else { + return; + }; + // The stack must be an affine relayout under a flat [M, K] + // reinterpret, both pure relayouts (no fill regions). + let [windowed, flat] = view.stages.as_slice() else { + return; + }; + if !windowed.is_fully_defined() + || !flat.is_fully_defined() + || !flat.layout.is_contiguous() + || flat.layout.offset() != 0 + || flat.layout.shape() != [m, k] + { + return; + } + // The kernels substitute the windowed map as affine per-dim index + // arithmetic; validate it here so the lowering can rely on it. + if crate::view::affine_dim_indices(&windowed.layout, &windowed.input_shape).is_none() { + return; + } + let operand_shape = windowed.shape(); + // The producer's dims must split cleanly into an `M` prefix and a + // `K` suffix for the per-side flat-coordinate decomposition. + let mut product = 1usize; + let mut k_start = operand_shape.len(); + while k_start > 0 && product < k { + k_start -= 1; + let Some(next) = product.checked_mul(operand_shape[k_start]) else { + return; + }; + product = next; + } + if product != k + || k_start == 0 + || k_start == operand_shape.len() + || operand_shape[..k_start].iter().product::() != m + { + return; + } + // The flat row/column coordinates decompose with u32 arithmetic. + let probe = NaryExpr::DimIndex(0); + if crate::view::row_major_indices_from_flat(probe.clone(), &operand_shape[..k_start]).is_none() + || crate::view::row_major_indices_from_flat(probe, &operand_shape[k_start..]).is_none() + { + return; + } + operation.first = view.input; + operation.a = crate::matmul::MatrixOperand { + shape: operand_shape.into(), + batch_dims: 0, + row_dims: k_start, + base_map: Some(crate::matmul::OperandBaseMap { + layout: windowed.layout.clone(), + base_shape: windowed.input_shape.clone(), + }), + }; } diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs b/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs index 5f096f78b..d8c6e0f72 100644 --- a/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs +++ b/fusor-ml/core/src/compute_graph/resolve/recognize_attention.rs @@ -1,7 +1,7 @@ //! Recognition of composed attention clusters. //! //! Runs third, after contractions and normalizations: by then the canonical -//! cluster from `Tensor::flash_attention` has collapsed to +//! cluster from `Tensor::attention` has collapsed to //! `MatMul(Softmax(scale·MatMul(q, kᵀ) [+ mask]), v)` with the GQA-expand / //! transpose / mask broadcast views still attached to the original //! q/k/v/mask nodes. Recognition rebuilds the attention row program when its @@ -14,7 +14,9 @@ use crate::{ view::ViewOperation, }; -use super::cluster_match::{binary_elementwise, layout_matches, unary_elementwise}; +use super::cluster_match::{ + binary_elementwise, keepdim_broadcast_layout, layout_matches, unary_elementwise, +}; use super::*; struct MatchedAttention { @@ -65,14 +67,16 @@ fn match_causal_select(nary: &ElementwiseOperation) -> Option { if indices.len() != nary.shape.len() || !NaryExpr::is_elementwise_indices(indices) { return None; } - let negative_infinity = match on_false { - NaryExpr::Scalar(crate::nary_wise::NaryScalar::F32(value)) => *value == f32::NEG_INFINITY, + let masked_score = match on_false { + NaryExpr::Scalar(crate::nary_wise::NaryScalar::F32(value)) => { + *value == crate::composite::attention::MASKED_SCORE_F32 + } NaryExpr::Scalar(crate::nary_wise::NaryScalar::F16(value)) => { - *value == half::f16::NEG_INFINITY + *value == crate::composite::attention::MASKED_SCORE_F16 } _ => false, }; - negative_infinity.then_some(nary.inputs[0]) + masked_score.then_some(nary.inputs[0]) } impl Resolver { @@ -87,11 +91,50 @@ impl Resolver { ) }) .collect(); + // The paired KV-side contraction pattern roots at a slice-assign + // chain (composed as elementwise region-selects); claim it before + // the single-contraction scans so the halves are not recognized + // separately. + let assign_candidates: Vec = self + .execution_graph + .node_indices() + .filter(|&node| { + matches!( + self.execution_graph[node].variant, + ExecutionVariant::Elementwise(_) + ) + }) + .collect(); + for node in assign_candidates { + if !self.execution_graph.contains_node(node) { + continue; + } + self.try_recognize_attention_grad_pair(graph, node); + } for node in candidates { if !self.execution_graph.contains_node(node) { continue; } - self.try_recognize_attention(graph, node); + if self.try_recognize_attention(graph, node) { + continue; + } + self.try_recognize_attention_grad(graph, node); + } + let lse_candidates: Vec = self + .execution_graph + .node_indices() + .filter(|&node| { + matches!( + self.execution_graph[node].variant, + ExecutionVariant::Elementwise(_) + ) + }) + .collect(); + for node in lse_candidates { + if !self.execution_graph.contains_node(node) { + continue; + } + self.try_recognize_score_lse(graph, node); } } @@ -123,32 +166,46 @@ impl Resolver { if let Some(mask) = matched.mask { dependencies.push(mask); } - // The recognized cluster lowers through the generic attention row - // program; shapes it cannot host stay composed (matmul + softmax - // row program + matmul). - let Some(operation) = crate::row_program::attention_row_program( - &graph.device(), - crate::row_program::AttentionInputs { - q: matched.q, - k: matched.k, - v: matched.v, - mask: matched.mask, - q_shape: &matched.q_shape, - k_shape: &matched.k_shape, - v_shape: &matched.v_shape, - mask_shape: matched.mask_shape.as_deref(), - scale: matched.scale, - input_dtype: matched.datatype, - causal: matched.causal, - }, - ) else { + let inputs = crate::row_program::AttentionInputs { + q: matched.q, + k: matched.k, + v: matched.v, + mask: matched.mask, + q_shape: &matched.q_shape, + k_shape: &matched.k_shape, + v_shape: &matched.v_shape, + mask_shape: matched.mask_shape.as_deref(), + scale: matched.scale, + input_dtype: matched.datatype, + causal: matched.causal, + }; + // Multi-row shapes need cross-row K/V reuse: the fused flash kernel + // claims them whenever the device and shape qualify. Everything else + // — decode's single row, ragged or oversized extents, non-f32 — + // lowers through the generic attention row program as before. + if matched.q_shape[2] > 1 + && let Some(operation) = crate::flash_attention::FlashAttentionOperation::try_new_output( + &graph.device(), + &inputs, + ) + { + self.commit_recognized( + graph, + node_idx, + &dependencies, + ExecutionVariant::Attention(operation), + ); + return true; + } + let Some(operation) = crate::row_program::attention_row_program(&graph.device(), inputs) + else { return false; }; self.commit_recognized( graph, node_idx, &dependencies, - ExecutionVariant::GraphOp(Arc::new(operation)), + ExecutionVariant::RowProgram(operation), ); true } @@ -190,7 +247,7 @@ impl Resolver { let expected_consumers = |node: NodeIndex| if node == softmax.input { 2 } else { 1 }; // Causal masking: select(kv_pos <= q_pos, scaled, -inf) — pure index - // arithmetic emitted by `flash_attention_causal`. + // arithmetic emitted by `attention_causal`. let mut causal = false; if let Some(scaled) = self.inner_nary(scores_inner).and_then(match_causal_select) { if !self.exclusively_consumed(graph, scores_inner, expected_consumers(scores_inner)) { @@ -383,3 +440,637 @@ impl Resolver { Some((broadcast_view.input, vec![b, num_kv_heads, s, d])) } } + +/// One matched scaled-masked score cluster +/// (`scale·q·kᵀ [+ mask | causal-select]`). +struct MatchedScores { + q: NodeIndex, + k: NodeIndex, + mask: Option, + causal: bool, + scale: f32, + /// `[batch, heads, q_len, kv_len]`. + shape: [usize; 4], + head_dim: usize, + kv_heads: usize, + datatype: DataTypeEnum, +} + +impl MatchedScores { + fn dims(&self) -> [usize; 6] { + let [batch, heads, q_len, kv_len] = self.shape; + [batch, heads, self.kv_heads, q_len, kv_len, self.head_dim] + } +} + +/// A probability-jacobian factor: `p ∘ (grad_o·vᵀ − bcast(dsum)) · scale`. +struct MatchedDs { + scores: MatchedScores, + lse: NodeIndex, + dsum: NodeIndex, + grad_o: NodeIndex, + v: NodeIndex, +} + +impl Resolver { + /// Match the canonical score cluster rooted at `root_inner`. No + /// consumption constraints: patterns share these intermediates by + /// construction, and whichever nodes end up unconsumed after their + /// roots are rewritten cascade away. + fn match_score_cluster( + &self, + graph: &ComputeGraphInner, + root_inner: NodeIndex, + ) -> Option { + let root = self.inner_nary(root_inner)?; + let shape: [usize; 4] = root.shape.as_ref().try_into().ok()?; + let [batch, num_heads, q_seq_len, kv_seq_len] = shape; + let datatype = root.output_datatype; + let mut causal = false; + let mut mask = None; + let mut scores_inner = root_inner; + if let Some(scaled) = self.inner_nary(scores_inner).and_then(match_causal_select) { + causal = true; + scores_inner = scaled; + } else if let Some((NaryOp::Add, lhs, rhs)) = + self.inner_nary(scores_inner).and_then(binary_elementwise) + { + let mut matched = None; + for (scaled_side, mask_side) in [(lhs, rhs), (rhs, lhs)] { + let Some(view) = self.inner_view(mask_side) else { + continue; + }; + let Some(stage) = view.plain().filter(|stage| stage.is_fully_defined()) else { + continue; + }; + let expected = + Layout::from_parts(0, shape.to_vec().into(), [0, 0, kv_seq_len, 1].into()); + if !layout_matches(Some(&stage.layout), &expected) + || stage.input_shape.as_ref() != [q_seq_len, kv_seq_len] + { + continue; + } + matched = Some((scaled_side, view.input)); + break; + } + let (scaled_side, mask_node) = matched?; + mask = Some(mask_node); + scores_inner = scaled_side; + } + let scale = { + let nary = self.inner_nary(scores_inner)?; + let (function, _) = unary_elementwise(nary)?; + match function.op { + NaryOp::MulConst(NaryScalar::F32(scale)) => scale, + NaryOp::MulConst(NaryScalar::F16(scale)) => scale.to_f32(), + _ => return None, + } + }; + let qk_inner = self.match_unary(scores_inner, |function| { + matches!(function.op, NaryOp::MulConst(_)) + })?; + let qk = self.inner_matmul(qk_inner)?; + if !qk.pre_element_wise[0].functions.is_empty() + || !qk.pre_element_wise[1].functions.is_empty() + || !qk.post_element_wise.functions.is_empty() + || !qk.a.is_plain() + || !qk.b.is_plain() + { + return None; + } + let q = qk.first; + let q_shape = qk.a.shape.to_vec(); + if q_shape.len() != 4 || q_shape[..3] != [batch, num_heads, q_seq_len] { + return None; + } + let head_dim = q_shape[3]; + let expanded_shape = [batch, num_heads, kv_seq_len, head_dim]; + let kt = self.inner_view(qk.second)?; + let kt_stage = kt.plain().filter(|stage| stage.is_fully_defined())?; + let expected_kt = Layout::contiguous(&expanded_shape).transpose(2, 3); + if !layout_matches(Some(&kt_stage.layout), &expected_kt) + || kt_stage.input_shape.as_ref() != expanded_shape + { + return None; + } + let (k, k_shape) = self.peel_gqa_expand(graph, kt.input, &expanded_shape)?; + Some(MatchedScores { + q, + k, + mask, + causal, + scale, + shape, + head_dim, + kv_heads: k_shape[1], + datatype, + }) + } + + /// Peel a keepdim row statistic broadcast back over `shape`, returning + /// the rank-3 base node. + fn peel_row_broadcast(&self, view_inner: NodeIndex, shape: &[usize; 4]) -> Option { + let (base, layout) = self.walk_view_chain(view_inner); + if base == view_inner { + return None; + } + layout_matches(layout.as_ref(), &keepdim_broadcast_layout(shape, 3)).then_some(base) + } + + /// Probabilities recomputed from a row statistic: + /// `exp(scores − bcast(rowstat))`. Returns the score cluster and the + /// statistic node. + fn match_prob_cluster( + &self, + graph: &ComputeGraphInner, + prob_inner: NodeIndex, + ) -> Option<(MatchedScores, NodeIndex)> { + let shifted_inner = self.match_unary(prob_inner, |function| function.op == NaryOp::Exp)?; + let (sub_op, scores_root, stat_view) = self + .inner_nary(shifted_inner) + .and_then(binary_elementwise)?; + if sub_op != NaryOp::Sub { + return None; + } + let scores = self.match_score_cluster(graph, scores_root)?; + let stat = self.peel_row_broadcast(stat_view, &scores.shape)?; + Some((scores, stat)) + } + + /// Match `p ∘ (grad_o·vᵀ − bcast(dsum)) · scale` rooted at `ds_inner`. + fn match_ds_cluster( + &self, + graph: &ComputeGraphInner, + ds_inner: NodeIndex, + ) -> Option { + let ds_scale = { + let nary = self.inner_nary(ds_inner)?; + let (function, _) = unary_elementwise(nary)?; + match function.op { + NaryOp::MulConst(NaryScalar::F32(scale)) => scale, + NaryOp::MulConst(NaryScalar::F16(scale)) => scale.to_f32(), + _ => return None, + } + }; + let mul_inner = self.match_unary(ds_inner, |function| { + matches!(function.op, NaryOp::MulConst(_)) + })?; + let (mul_op, lhs, rhs) = self.inner_nary(mul_inner).and_then(binary_elementwise)?; + if mul_op != NaryOp::Mul { + return None; + } + for (p_side, sub_side) in [(lhs, rhs), (rhs, lhs)] { + let Some((scores, lse)) = self.match_prob_cluster(graph, p_side) else { + continue; + }; + if scores.scale != ds_scale { + continue; + } + let Some((NaryOp::Sub, dp_root, dsum_view)) = + self.inner_nary(sub_side).and_then(binary_elementwise) + else { + continue; + }; + let Some(dsum) = self.peel_row_broadcast(dsum_view, &scores.shape) else { + continue; + }; + let Some(dp) = self.inner_matmul(dp_root) else { + continue; + }; + if !dp.pre_element_wise[0].functions.is_empty() + || !dp.pre_element_wise[1].functions.is_empty() + || !dp.post_element_wise.functions.is_empty() + || !dp.a.is_plain() + || !dp.b.is_plain() + { + continue; + } + let [batch, heads, _q_len, kv_len] = scores.shape; + let expanded = [batch, heads, kv_len, scores.head_dim]; + let Some(vt) = self.inner_view(dp.second) else { + continue; + }; + let Some(vt_stage) = vt.plain().filter(|stage| stage.is_fully_defined()) else { + continue; + }; + let expected = Layout::contiguous(&expanded).transpose(2, 3); + if !layout_matches(Some(&vt_stage.layout), &expected) + || vt_stage.input_shape.as_ref() != expanded + { + continue; + } + return Some(MatchedDs { + scores, + lse, + dsum, + grad_o: dp.first, + v: vt.input, + }); + } + None + } + + /// A transpose view over the score space `[b, h, q, kv] → [b, h, kv, q]`, + /// returning the viewed node. + fn peel_score_transpose(&self, inner: NodeIndex, shape: &[usize; 4]) -> Option { + let view = self.inner_view(inner)?; + let stage = view.plain().filter(|stage| stage.is_fully_defined())?; + let expected = Layout::contiguous(shape).transpose(2, 3); + (layout_matches(Some(&stage.layout), &expected) && stage.input_shape.as_ref() == shape) + .then_some(view.input) + } + + /// Recognize the probability-contraction patterns rooted at a matmul: + /// `ds·k`, `dsᵀ·q`, and `pᵀ·x` all stream tile recomputation from the + /// row statistics instead of materializing the probability matrices. + fn try_recognize_attention_grad( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> bool { + use crate::flash_attention::{AttentionKernel, AttentionPatternNodes}; + let ExecutionVariant::MatMul(out) = &self.execution_graph[node_idx].variant else { + return false; + }; + if !out.pre_element_wise[0].functions.is_empty() + || !out.pre_element_wise[1].functions.is_empty() + || !out.post_element_wise.functions.is_empty() + || !out.a.is_plain() + || !out.b.is_plain() + { + return false; + } + let (a_inner, b_inner) = (out.first, out.second); + let b_shape = out.b.shape.to_vec(); + let datatype = out.datatype; + + // dq-shaped: MatMul(ds, k). + if let Some(ds) = self.match_ds_cluster(graph, a_inner) + && b_inner == ds.scores.k + && datatype == ds.scores.datatype + { + let nodes = AttentionPatternNodes { + q: ds.scores.q, + k: ds.scores.k, + v: Some(ds.v), + grad_o: Some(ds.grad_o), + lse: Some(ds.lse), + dsum: Some(ds.dsum), + mask: ds.scores.mask, + }; + if let Some(operation) = crate::flash_attention::FlashAttentionOperation::try_new( + &graph.device(), + AttentionKernel::GradQ, + nodes, + ds.scores.dims(), + ds.scores.scale, + ds.scores.causal, + datatype, + ) { + let dependencies = grad_dependencies(&operation); + self.commit_recognized( + graph, + node_idx, + &dependencies, + ExecutionVariant::Attention(operation), + ); + return true; + } + } + + // Transposed-operand shapes: MatMul(dsᵀ, q) and MatMul(pᵀ, x). + let Some(ds_shape) = self.transposed_operand_shape(a_inner) else { + return false; + }; + let Some(src) = self.peel_score_transpose(a_inner, &ds_shape) else { + return false; + }; + if let Some(ds) = self.match_ds_cluster(graph, src) + && ds.scores.shape == ds_shape + && b_inner == ds.scores.q + && datatype == ds.scores.datatype + { + let nodes = AttentionPatternNodes { + q: ds.scores.q, + k: ds.scores.k, + v: Some(ds.v), + grad_o: Some(ds.grad_o), + lse: Some(ds.lse), + dsum: Some(ds.dsum), + mask: ds.scores.mask, + }; + if let Some(operation) = crate::flash_attention::FlashAttentionOperation::try_new( + &graph.device(), + AttentionKernel::GradK, + nodes, + ds.scores.dims(), + ds.scores.scale, + ds.scores.causal, + datatype, + ) { + let dependencies = grad_dependencies(&operation); + self.commit_recognized( + graph, + node_idx, + &dependencies, + ExecutionVariant::Attention(operation), + ); + return true; + } + } + if let Some((scores, lse)) = self.match_prob_cluster(graph, src) + && scores.shape == ds_shape + && datatype == scores.datatype + && b_shape.len() == 4 + && b_shape[3] == scores.head_dim + { + let nodes = AttentionPatternNodes { + q: scores.q, + k: scores.k, + v: None, + grad_o: Some(b_inner), + lse: Some(lse), + dsum: None, + mask: scores.mask, + }; + if let Some(operation) = crate::flash_attention::FlashAttentionOperation::try_new( + &graph.device(), + AttentionKernel::GradV, + nodes, + scores.dims(), + scores.scale, + scores.causal, + datatype, + ) { + let dependencies = grad_dependencies(&operation); + self.commit_recognized( + graph, + node_idx, + &dependencies, + ExecutionVariant::Attention(operation), + ); + return true; + } + } + false + } + + /// The `[b, h, q, kv]` score space a transposed matmul operand views, + /// derived from the view node's stage. + fn transposed_operand_shape(&self, inner: NodeIndex) -> Option<[usize; 4]> { + let view = self.inner_view(inner)?; + let stage = view.plain().filter(|stage| stage.is_fully_defined())?; + stage.input_shape.as_ref().try_into().ok() + } + + /// Recognize the row log-sum-exp pattern rooted at an elementwise add: + /// `max(s, 3) + ln Σ exp(s − bcast(max))` streams the scores without + /// materializing them. + fn try_recognize_score_lse( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> bool { + use crate::flash_attention::{AttentionKernel, AttentionPatternNodes}; + let ExecutionVariant::Elementwise(add) = &self.execution_graph[node_idx].variant else { + return false; + }; + let Some((NaryOp::Add, lhs, rhs)) = binary_elementwise(add) else { + return false; + }; + let datatype = add.output_datatype; + for (m_side, log_side) in [(lhs, rhs), (rhs, lhs)] { + let Some((max_axis, max_value)) = + self.match_reduce(m_side, crate::reduce::ReduceOp::Max) + else { + continue; + }; + if max_axis != 3 { + continue; + } + let Some(sum_inner) = self.match_unary(log_side, |function| function.op == NaryOp::Log) + else { + continue; + }; + let Some((sum_axis, exp_inner)) = + self.match_reduce(sum_inner, crate::reduce::ReduceOp::Sum) + else { + continue; + }; + if sum_axis != 3 { + continue; + } + let Some(shifted) = self.match_unary(exp_inner, |function| function.op == NaryOp::Exp) + else { + continue; + }; + let Some((NaryOp::Sub, scores_root, max_view)) = + self.inner_nary(shifted).and_then(binary_elementwise) + else { + continue; + }; + if scores_root != max_value { + continue; + } + let Some(scores) = self.match_score_cluster(graph, scores_root) else { + continue; + }; + let Some(m_base) = self.peel_row_broadcast(max_view, &scores.shape) else { + continue; + }; + if m_base != m_side || datatype != scores.datatype { + continue; + } + let nodes = AttentionPatternNodes { + q: scores.q, + k: scores.k, + v: None, + grad_o: None, + lse: None, + dsum: None, + mask: scores.mask, + }; + if let Some(operation) = crate::flash_attention::FlashAttentionOperation::try_new( + &graph.device(), + AttentionKernel::LogSumExp, + nodes, + scores.dims(), + scores.scale, + scores.causal, + datatype, + ) { + let dependencies = grad_dependencies(&operation); + self.commit_recognized( + graph, + node_idx, + &dependencies, + ExecutionVariant::Attention(operation), + ); + return true; + } + } + false + } +} + +/// The operation's dependencies in `visit_dependencies` order. +fn grad_dependencies( + operation: &crate::flash_attention::FlashAttentionOperation, +) -> Vec { + use crate::mir::operation::Operation; + let mut dependencies = Vec::new(); + operation.visit_dependencies(&mut |node| dependencies.push(node)); + dependencies +} + +impl Resolver { + /// Match the dk-shaped contraction `dsᵀ · q` at a matmul node. + fn match_grad_k_root(&self, graph: &ComputeGraphInner, inner: NodeIndex) -> Option { + let matmul = self.inner_matmul(inner)?; + if !matmul.pre_element_wise[0].functions.is_empty() + || !matmul.pre_element_wise[1].functions.is_empty() + || !matmul.post_element_wise.functions.is_empty() + || !matmul.a.is_plain() + || !matmul.b.is_plain() + { + return None; + } + let (a_inner, b_inner, datatype) = (matmul.first, matmul.second, matmul.datatype); + let shape = self.transposed_operand_shape(a_inner)?; + let src = self.peel_score_transpose(a_inner, &shape)?; + let ds = self.match_ds_cluster(graph, src)?; + (ds.scores.shape == shape && b_inner == ds.scores.q && datatype == ds.scores.datatype) + .then_some(ds) + } + + /// Match the dv-shaped contraction `pᵀ · x` at a matmul node, returning + /// the score cluster, the row statistic, and the free operand. + fn match_grad_v_root( + &self, + graph: &ComputeGraphInner, + inner: NodeIndex, + ) -> Option<(MatchedScores, NodeIndex, NodeIndex)> { + let matmul = self.inner_matmul(inner)?; + if !matmul.pre_element_wise[0].functions.is_empty() + || !matmul.pre_element_wise[1].functions.is_empty() + || !matmul.post_element_wise.functions.is_empty() + || !matmul.a.is_plain() + || !matmul.b.is_plain() + { + return None; + } + let (a_inner, b_inner, datatype) = (matmul.first, matmul.second, matmul.datatype); + let b_shape = matmul.b.shape.to_vec(); + let shape = self.transposed_operand_shape(a_inner)?; + let src = self.peel_score_transpose(a_inner, &shape)?; + let (scores, lse) = self.match_prob_cluster(graph, src)?; + (scores.shape == shape + && datatype == scores.datatype + && b_shape.len() == 4 + && b_shape[3] == scores.head_dim) + .then_some((scores, lse, b_inner)) + } + + /// Recognize both KV-side contractions landing in one combined tensor + /// (`dk` rows then `dv` rows along the sequence axis via a slice-assign + /// chain over a zero base): the paired streaming kernel computes both in + /// one dispatch, sharing the probability recomputation. Composed + /// slice-assigns are elementwise region-selects over + /// `[destination, value]` (see `recognize_cat`). + fn try_recognize_attention_grad_pair( + &mut self, + graph: &mut ComputeGraphInner, + node_idx: ExecutionNodeIndex, + ) -> bool { + use crate::flash_attention::{AttentionKernel, AttentionPatternNodes}; + struct Link { + destination: NodeIndex, + value: NodeIndex, + slices: Box<[std::ops::Range]>, + } + let link_of = |resolver: &Self, inner: NodeIndex| -> Option { + let nary = resolver.inner_nary(inner)?; + let slices = super::recognize_cat::match_slice_assign(nary)?; + Some(Link { + destination: nary.inputs[0], + value: nary.inputs[1], + slices, + }) + }; + let outer_inner = match &self.execution_graph[node_idx].variant { + ExecutionVariant::Elementwise(_) => self.execution_graph[node_idx].inner_idx, + _ => return false, + }; + let Some(outer) = link_of(self, outer_inner) else { + return false; + }; + let Some(inner) = link_of(self, outer.destination) else { + return false; + }; + let ds = match self.match_grad_k_root(graph, inner.value) { + Some(ds) => ds, + None => return false, + }; + let (scores_v, lse_v, x) = match self.match_grad_v_root(graph, outer.value) { + Some(matched) => matched, + None => return false, + }; + // Both halves must come from the same cluster identity. + if scores_v.q != ds.scores.q + || scores_v.k != ds.scores.k + || scores_v.mask != ds.scores.mask + || scores_v.causal != ds.scores.causal + || scores_v.scale != ds.scores.scale + || lse_v != ds.lse + || x != ds.grad_o + { + return false; + } + // The chain must be exactly `zeros -> assign dk half -> assign dv + // half` with the halves in kernel order (dk rows first). + let [batch, heads, _q_len, kv_len] = ds.scores.shape; + let head_dim = ds.scores.head_dim; + let expect_inner: Box<[std::ops::Range]> = + [0..batch, 0..heads, 0..kv_len, 0..head_dim].into(); + let expect_outer: Box<[std::ops::Range]> = + [0..batch, 0..heads, kv_len..2 * kv_len, 0..head_dim].into(); + if inner.slices != expect_inner || outer.slices != expect_outer { + return false; + } + // Intermediates die once the root is rewired: the inner assign must + // feed only the outer, and the zero base only the inner. + if !self.exclusively_consumed(graph, outer.destination, 1) + || !self.exclusively_consumed(graph, inner.destination, 1) + { + return false; + } + let nodes = AttentionPatternNodes { + q: ds.scores.q, + k: ds.scores.k, + v: Some(ds.v), + grad_o: Some(ds.grad_o), + lse: Some(ds.lse), + dsum: Some(ds.dsum), + mask: ds.scores.mask, + }; + if let Some(operation) = crate::flash_attention::FlashAttentionOperation::try_new( + &graph.device(), + AttentionKernel::GradKV, + nodes, + ds.scores.dims(), + ds.scores.scale, + ds.scores.causal, + ds.scores.datatype, + ) { + let dependencies = grad_dependencies(&operation); + self.commit_recognized( + graph, + node_idx, + &dependencies, + ExecutionVariant::Attention(operation), + ); + return true; + } + false + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs b/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs index cc5efae09..c9f333ea0 100644 --- a/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs +++ b/fusor-ml/core/src/compute_graph/resolve/recognize_cat.rs @@ -19,14 +19,6 @@ use crate::view::{AffineIndex, affine_dim_indices}; use super::*; -/// Every destination element evaluates all lifted arms (both sides of a -/// select execute), so a K-way lift multiplies branch arithmetic by K. -/// Chains longer than this keep their materialized branches and fuse through -/// the regular n-ary path instead. -const MAX_LIFTED_BRANCHES: usize = 4; -/// Per-arm expression size cap for the non-collapsed (select chain) form. -const MAX_ARM_OPS: usize = 32; - /// One matched slice-assign link: an `Elementwise` node whose expression is /// exactly `slice_assign_expression(slices)` over `[destination, value]`. struct AssignLink { @@ -73,31 +65,20 @@ enum Rep { } fn apply_reps(expr: &NaryExpr, reps: &[Rep]) -> NaryExpr { - match expr { - NaryExpr::Op { children, function } => NaryExpr::Op { - children: children.iter().map(|c| apply_reps(c, reps)).collect(), - function: function.clone(), + egraph::compose::map_loads(expr, &mut |input_idx, _, indices| match &reps[input_idx] { + // Inline replacements are only built for slots accessed + // element-wise, where the whole load is the producer's value. + Rep::Inline(expr) => expr.clone(), + Rep::Slot(slot) => NaryExpr::IndexedInput { + input_idx: *slot, + indices, }, - NaryExpr::IndexedInput { input_idx, indices } => { - let indices: Vec = indices.iter().map(|c| apply_reps(c, reps)).collect(); - match &reps[*input_idx] { - // Inline replacements are only built for slots accessed - // element-wise, where the whole load is the producer's value. - Rep::Inline(expr) => expr.clone(), - Rep::Slot(slot) => NaryExpr::IndexedInput { - input_idx: *slot, - indices, - }, - } - } - NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), - NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), - } + }) } /// Recover the slice ranges from a candidate slice-assign expression, then /// verify the match by regenerating the canonical expression and comparing. -fn match_slice_assign(nary: &ElementwiseOperation) -> Option]>> { +pub(super) fn match_slice_assign(nary: &ElementwiseOperation) -> Option]>> { if nary.inputs.len() != 2 { return None; } @@ -194,21 +175,8 @@ fn slices_tile(out_shape: &[usize], chain: &[&AssignLink]) -> bool { } } -fn expr_op_count(expr: &NaryExpr) -> usize { - match expr { - NaryExpr::Op { children, .. } => 1 + children.iter().map(expr_op_count).sum::(), - NaryExpr::IndexedInput { indices, .. } => { - 1 + indices.iter().map(expr_op_count).sum::() - } - NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => 0, - } -} - impl Resolver { pub(super) fn recognize_assign_chains(&mut self, graph: &mut ComputeGraphInner) { - if std::env::var_os("FUSOR_RESOLVE_SKIP_ASSIGN_CHAINS").is_some() { - return; - } let mut links: FxHashMap = FxHashMap::default(); for exec in self.execution_graph.node_indices() { let node = &self.execution_graph[exec]; @@ -300,9 +268,6 @@ impl Resolver { chain.push(prev); } chain.reverse(); - if chain.len() > MAX_LIFTED_BRANCHES { - return; - } let base_inner = chain[0].destination; let out_shape = tail.shape.clone(); @@ -324,9 +289,6 @@ impl Resolver { let expression = if collapsed { arms.swap_remove(0).1 } else { - if arms.iter().any(|(_, arm)| expr_op_count(arm) > MAX_ARM_OPS) { - return; - } let mut expression = NaryExpr::input(base_slot, rank); for (condition, arm) in arms { expression = @@ -335,7 +297,8 @@ impl Resolver { expression }; - let (final_inputs, final_expression) = Self::deduplicate_inputs(state.inputs, expression); + let (final_inputs, final_expression) = + egraph::compose::deduplicate_inputs(state.inputs, expression); if final_inputs.len() > graph.device().nary_direct_input_binding_budget() { return; } diff --git a/fusor-ml/core/src/compute_graph/resolve/recognize_gates.rs b/fusor-ml/core/src/compute_graph/resolve/recognize_gates.rs new file mode 100644 index 000000000..79e4793d0 --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/recognize_gates.rs @@ -0,0 +1,531 @@ +//! Emitter ↔ recognizer binding gates. +//! +//! Every composite the tensor API emits — attention forward and backward, +//! RMS norm, softmax, RoPE, cat/slice-assign chains, the contractions +//! themselves — is a plain cluster of the 3-op vocabulary that some resolver +//! matcher has to claim. A miss is invisible to every other test in the +//! suite: the composed form computes bit-comparable values through the +//! generic elementwise + reduce kernels, several times slower. These gates +//! emit each composite through the public API, run the recognition phase +//! over it, and pin exactly what survives — an emitter and its matcher can +//! only drift apart by turning one of them red. +//! +//! Generalizes the regenerate-and-compare check `match_slice_assign` already +//! performs against `slice_assign_expression`. + +use std::ops::Range; + +use super::key_goldens::q4k_weight; +use super::recognize_cat::match_slice_assign; +use crate::composite::attention::{MASKED_SCORE_F16, MASKED_SCORE_F32}; +use crate::flash_attention::AttentionKernel; +use crate::{Device, Tensor}; + +use super::*; + +/// One surviving execution node: what it lowers to, and the inner nodes it +/// reads in dependency order. +struct Recognized { + label: &'static str, + inputs: Vec, +} + +fn attention_label(kind: AttentionKernel) -> &'static str { + match kind { + AttentionKernel::Output => "attention:output", + AttentionKernel::LogSumExp => "attention:log_sum_exp", + AttentionKernel::GradQ => "attention:grad_q", + AttentionKernel::GradK => "attention:grad_k", + AttentionKernel::GradV => "attention:grad_v", + AttentionKernel::GradKV => "attention:grad_kv", + } +} + +fn label(variant: &ExecutionVariant) -> &'static str { + match variant { + ExecutionVariant::Tensor(_) => "tensor", + ExecutionVariant::QMatrix(_) => "qmatrix", + ExecutionVariant::Elementwise(_) => "elementwise", + ExecutionVariant::Reduce(_) => "reduce", + ExecutionVariant::Fold(_) => "fold", + ExecutionVariant::View(_) => "view", + ExecutionVariant::Assign(_) => "assign", + ExecutionVariant::Region(_) => "region", + ExecutionVariant::MatMul(_) => "matmul", + ExecutionVariant::QMatMul(_) => "qmatmul", + ExecutionVariant::QEmbedding(_) => "qembedding", + ExecutionVariant::RowProgram(_) => "row_program", + ExecutionVariant::Attention(operation) => attention_label(operation.kind), + } +} + +/// Build the execution graph for `targets` and run the recognition phase +/// over it — the same passes, in the same order, that `optimize_operations` +/// runs before it hands the graph to extraction. +fn recognize(device: &Device, targets: &[&Tensor]) -> Vec { + let targets: Vec = targets.iter().map(|tensor| tensor.data().key).collect(); + device.compute_graph().with_mut(|graph| { + let mut resolver = Resolver::new_batch(graph, targets.clone()); + for &target in &targets { + resolver.build_execution_graph(graph, target); + } + resolver.recognize_all(graph); + let mut recognized: Vec = resolver + .execution_graph + .node_weights() + .map(|node| { + let mut inputs = Vec::new(); + node.variant + .visit_dependencies(&mut |input| inputs.push(input)); + Recognized { + label: label(&node.variant), + inputs, + } + }) + .collect(); + recognized.sort_by_key(|node| node.label); + recognized + }) +} + +fn labels(recognized: &[Recognized]) -> Vec<&'static str> { + recognized.iter().map(|node| node.label).collect() +} + +/// The dependency list of the one node carrying `label`. +fn inputs_of<'a>(recognized: &'a [Recognized], label: &str) -> &'a [NodeIndex] { + let mut matches = recognized.iter().filter(|node| node.label == label); + let node = matches + .next() + .unwrap_or_else(|| panic!("no {label} node in {:?}", labels(recognized))); + assert!( + matches.next().is_none(), + "more than one {label} node in {:?}", + labels(recognized) + ); + &node.inputs +} + +fn splat(device: &Device, value: f32, shape: &[usize]) -> Tensor { + Tensor::splat::(device, value, shape) +} + +fn splat_f16(device: &Device, value: f32, shape: &[usize]) -> Tensor { + Tensor::splat::(device, half::f16::from_f32(value), shape) +} + +#[test] +fn contraction_composites_recognize() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let a = splat(&device, 0.5, &[8, 16]); + let b = splat(&device, 0.25, &[16, 8]); + let recognized = recognize(&device, &[&a.mat_mul(&b)]); + assert_eq!(labels(&recognized), ["matmul", "tensor", "tensor"]); + assert_eq!(inputs_of(&recognized, "matmul"), [a.key(), b.key()]); + + let a = splat(&device, 0.5, &[2, 8, 16]); + let b = splat(&device, 0.25, &[2, 16, 8]); + let recognized = recognize(&device, &[&a.mat_mul(&b)]); + assert_eq!(labels(&recognized), ["matmul", "tensor", "tensor"]); + + let x = splat(&device, 0.5, &[1, 512]); + let weight = q4k_weight(&device, 512, 512); + let recognized = recognize(&device, &[&x.q_mat_mul(&weight)]); + assert_eq!(labels(&recognized), ["qmatmul", "tensor"]); + assert_eq!(inputs_of(&recognized, "qmatmul"), [x.key()]); + }); +} + +#[test] +fn quantized_row_gather_composite_recognizes() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let indexes = Tensor::new::(&device, &[0u32, 3, 5]); + let table = q4k_weight(&device, 512, 256); + let recognized = recognize(&device, &[&table.index_select_rows(&indexes)]); + assert_eq!(labels(&recognized), ["qembedding", "tensor"]); + assert_eq!(inputs_of(&recognized, "qembedding"), [indexes.key()]); + }); +} + +#[test] +fn softmax_composite_recognizes() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let x = splat(&device, 0.5, &[8, 64]); + let recognized = recognize(&device, &[&x.softmax(1)]); + assert_eq!(labels(&recognized), ["row_program", "tensor"]); + assert_eq!(inputs_of(&recognized, "row_program"), [x.key()]); + + let x = splat(&device, 0.5, &[2, 4, 64]); + let recognized = recognize(&device, &[&x.softmax_last_dim()]); + assert_eq!(labels(&recognized), ["row_program", "tensor"]); + }); +} + +#[test] +fn rms_norm_composite_recognizes() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let x = splat(&device, 0.5, &[8, 64]); + let weight = splat(&device, 1.5, &[64]); + let recognized = recognize(&device, &[&x.rms_norm_fused_no_bias(&weight, 1e-5)]); + assert_eq!( + labels(&recognized), + ["row_program", "tensor", "tensor", "view"] + ); + + let bias = splat(&device, 0.125, &[64]); + let recognized = recognize(&device, &[&x.rms_norm_fused(&weight, Some(&bias), 1e-5)]); + assert_eq!( + labels(&recognized), + ["row_program", "tensor", "tensor", "tensor", "view", "view"] + ); + }); +} + +#[test] +fn attention_forward_composite_recognizes() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let q = splat(&device, 0.5, &[2, 8, 32, 64]); + let k = splat(&device, 0.25, &[2, 8, 32, 64]); + let v = splat(&device, 0.125, &[2, 8, 32, 64]); + let recognized = recognize(&device, &[&q.attention(&k, &v, 0.125, None)]); + assert_eq!( + labels(&recognized), + ["attention:output", "tensor", "tensor", "tensor"] + ); + assert_eq!( + inputs_of(&recognized, "attention:output"), + [q.key(), k.key(), v.key()] + ); + + // One query row is decode's shape: no cross-row K/V reuse to buy, so + // the cluster lands on the attention row program instead. + let q = splat(&device, 0.5, &[2, 8, 1, 64]); + let recognized = recognize(&device, &[&q.attention(&k, &v, 0.125, None)]); + assert_eq!( + labels(&recognized), + ["row_program", "tensor", "tensor", "tensor"] + ); + assert_eq!( + inputs_of(&recognized, "row_program"), + [q.key(), k.key(), v.key()] + ); + }); +} + +#[test] +fn attention_mask_composites_recognize() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let q = splat(&device, 0.5, &[2, 8, 32, 64]); + let k = splat(&device, 0.25, &[2, 8, 32, 64]); + let v = splat(&device, 0.125, &[2, 8, 32, 64]); + let recognized = recognize(&device, &[&q.attention_causal(&k, &v, 0.125)]); + assert_eq!( + labels(&recognized), + ["attention:output", "tensor", "tensor", "tensor"] + ); + + let mask = splat(&device, -1.0, &[32, 32]); + let recognized = recognize(&device, &[&q.attention(&k, &v, 0.125, Some(&mask))]); + assert_eq!( + labels(&recognized), + ["attention:output", "tensor", "tensor", "tensor", "tensor"] + ); + assert_eq!( + inputs_of(&recognized, "attention:output"), + [q.key(), k.key(), v.key(), mask.key()] + ); + + if !device.f16_supported() { + return; + } + let q = splat_f16(&device, 0.5, &[2, 8, 32, 64]); + let k = splat_f16(&device, 0.25, &[2, 8, 32, 64]); + let v = splat_f16(&device, 0.125, &[2, 8, 32, 64]); + let recognized = recognize(&device, &[&q.attention_causal(&k, &v, 0.125)]); + assert_eq!( + labels(&recognized), + ["attention:output", "tensor", "tensor", "tensor"] + ); + }); +} + +/// The composed causal select writes a finite stand-in for `-inf` that the +/// flash kernels duplicate as their own private `MASKED_SCORE` literal +/// (`tile-ir-kernels/src/kernels/attention.rs`). The recognizer compares +/// against the emitter's constant, so only the kernel copy can drift — pin +/// the exact bytes on this side. +#[test] +fn masked_score_constants_are_the_kernel_literals() { + assert_eq!(MASKED_SCORE_F32.to_bits(), (-3.0e38f32).to_bits()); + assert_eq!(MASKED_SCORE_F16.to_bits(), half::f16::MIN.to_bits()); + assert_eq!(MASKED_SCORE_F16.to_f32(), -65504.0); +} + +#[test] +fn attention_gqa_expand_peels_back_to_the_kv_tensors() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let q = splat(&device, 0.5, &[2, 8, 32, 64]); + let k = splat(&device, 0.25, &[2, 2, 32, 64]); + let v = splat(&device, 0.125, &[2, 2, 32, 64]); + let recognized = recognize(&device, &[&q.attention(&k, &v, 0.125, None)]); + // The stride-0 group broadcast and its flat reinterpret are peeled: + // no view survives, and the kernel reads the unexpanded K/V nodes. + assert_eq!( + labels(&recognized), + ["attention:output", "tensor", "tensor", "tensor"] + ); + assert_eq!( + inputs_of(&recognized, "attention:output"), + [q.key(), k.key(), v.key()] + ); + + let q = splat(&device, 0.5, &[2, 8, 1, 64]); + let recognized = recognize(&device, &[&q.attention(&k, &v, 0.125, None)]); + assert_eq!( + labels(&recognized), + ["row_program", "tensor", "tensor", "tensor"] + ); + assert_eq!( + inputs_of(&recognized, "row_program"), + [q.key(), k.key(), v.key()] + ); + }); +} + +#[test] +fn attention_lse_composite_recognizes() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let q = splat(&device, 0.5, &[2, 8, 32, 64]); + let k = splat(&device, 0.25, &[2, 8, 32, 64]); + let recognized = recognize(&device, &[&q.attention_lse(&k, 0.125, None, true)]); + assert_eq!( + labels(&recognized), + ["attention:log_sum_exp", "tensor", "tensor"] + ); + assert_eq!( + inputs_of(&recognized, "attention:log_sum_exp"), + [q.key(), k.key()] + ); + }); +} + +#[test] +fn attention_backward_composites_recognize() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let shape = [2usize, 8, 32, 64]; + let q = splat(&device, 0.5, &shape); + let k = splat(&device, 0.25, &shape); + let v = splat(&device, 0.125, &shape); + let o = splat(&device, 0.375, &shape); + let grad_o = splat(&device, 0.0625, &shape); + let lse = splat(&device, 1.0, &[2, 8, 32]); + let (dq, dk, dv) = q.attention_grads(&k, &v, &o, &grad_o, &lse, 0.125, None, true); + let recognized = recognize(&device, &[&dq, &dk, &dv]); + // `dsum = sum(grad_o ∘ o)` is a genuine input to both kernels, so an + // elementwise + reduce pair survives beside them; the dk/dv halves + // read back out of the paired kernel's output as views. + assert_eq!( + labels(&recognized), + [ + "attention:grad_kv", + "attention:grad_q", + "elementwise", + "reduce", + "tensor", + "tensor", + "tensor", + "tensor", + "tensor", + "tensor", + "view", + "view", + ] + ); + let dsum = *inputs_of(&recognized, "attention:grad_q").last().unwrap(); + assert_eq!( + inputs_of(&recognized, "attention:grad_q"), + [q.key(), k.key(), v.key(), grad_o.key(), lse.key(), dsum] + ); + assert_eq!( + inputs_of(&recognized, "attention:grad_kv"), + [q.key(), k.key(), v.key(), grad_o.key(), lse.key(), dsum] + ); + + // The same cluster differentiated against the composed log-sum-exp + // keeps every pattern: the statistic itself streams too. + let lse = q.attention_lse(&k, 0.125, None, true); + let (dq, dk, dv) = q.attention_grads(&k, &v, &o, &grad_o, &lse, 0.125, None, true); + let recognized = recognize(&device, &[&dq, &dk, &dv]); + assert_eq!( + labels(&recognized), + [ + "attention:grad_kv", + "attention:grad_q", + "attention:log_sum_exp", + "elementwise", + "reduce", + "tensor", + "tensor", + "tensor", + "tensor", + "tensor", + "view", + "view", + ] + ); + }); +} + +/// The regenerate-and-compare binding in its original form: every shape the +/// emitter can produce must round-trip back to the exact ranges it was given. +#[test] +fn slice_assign_emitter_round_trips_through_its_matcher() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let cases: &[(&[usize], &[Range])] = &[ + (&[4, 8], &[0..2, 0..8]), + (&[4, 8], &[2..4, 0..8]), + (&[4, 8], &[1..3, 2..5]), + (&[2, 3, 4, 5], &[0..2, 1..3, 0..4, 2..5]), + (&[6], &[2..5]), + ]; + for (shape, slices) in cases { + let base = splat(&device, 0.5, shape); + let value_shape: Vec = slices.iter().map(|slice| slice.len()).collect(); + let value = splat(&device, 1.5, &value_shape); + let assigned = base.slice_assign(slices.to_vec(), &value); + let nary = device.compute_graph().with_mut(|graph| { + match &graph + .nodes + .nodes + .node_weight(assigned.key()) + .expect("live node") + .variant + { + ComputeGraphNodeVariant::Elementwise(nary) => nary.clone(), + other => panic!("slice_assign emitted {other:?}"), + } + }); + assert_eq!( + match_slice_assign(&nary).as_deref(), + Some(*slices), + "slice_assign {slices:?} over {shape:?}" + ); + } + }); +} + +#[test] +fn slice_assign_chain_recognizes() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // `Tensor::cat` in its composed form: a zero base, then one + // slice-assign per chunk. Recognition lifts every branch into one + // kernel over the destination index space. + let base = splat(&device, 0.0, &[4, 8]); + let left = splat(&device, 0.5, &[2, 8]); + let right = splat(&device, 1.5, &[2, 8]); + let cat = base + .slice_assign([0..2, 0..8], &left.exp()) + .slice_assign([2..4, 0..8], &right.exp()); + let recognized = recognize(&device, &[&cat]); + assert_eq!( + labels(&recognized), + ["elementwise", "tensor", "tensor", "tensor"] + ); + // Slots follow first appearance in the folded select chain, so the + // last chunk's operand leads and the base is the final fallthrough. + assert_eq!( + inputs_of(&recognized, "elementwise"), + [right.key(), left.key(), base.key()] + ); + }); +} + +/// RoPE has no matcher: the emitter writes the fused kernel's expression +/// directly, and a paired call must stay one kernel writing both halves of a +/// single allocation. Splitting it back into per-tensor ops is the same +/// silent slowdown a recognition miss is. +#[test] +fn rope_composites_stay_single_kernels() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let q = splat(&device, 0.5, &[1, 8, 4, 64]); + let cos = splat(&device, 0.25, &[4, 32]); + let sin = splat(&device, 0.125, &[4, 32]); + let recognized = recognize(&device, &[&q.rope_fused(&cos, &sin)]); + assert_eq!( + labels(&recognized), + ["elementwise", "tensor", "tensor", "tensor"] + ); + assert_eq!( + inputs_of(&recognized, "elementwise"), + [q.key(), cos.key(), sin.key()] + ); + + let k = splat(&device, 0.75, &[1, 2, 4, 64]); + let (rope_q, rope_k) = q.rope_normal_pair_fused(&k, &cos, &sin); + let recognized = recognize(&device, &[&rope_q, &rope_k]); + assert_eq!( + labels(&recognized), + [ + "elementwise", + "tensor", + "tensor", + "tensor", + "tensor", + "view", + "view" + ] + ); + assert_eq!( + inputs_of(&recognized, "elementwise"), + [q.key(), k.key(), cos.key(), sin.key()] + ); + }); +} diff --git a/fusor-ml/core/src/compute_graph/resolve/run.rs b/fusor-ml/core/src/compute_graph/resolve/run.rs index 12cc100c3..4d0b2e6bf 100644 --- a/fusor-ml/core/src/compute_graph/resolve/run.rs +++ b/fusor-ml/core/src/compute_graph/resolve/run.rs @@ -16,14 +16,10 @@ impl Resolver { _removed: &mut Vec, tail: impl FnOnce(&TensorData, &mut wgpu::CommandEncoder) -> T, ) -> (ResolverResult, T) { - let host_trace = - cfg!(target_arch = "wasm32") || std::env::var_os("FUSOR_TRACE_RESOLVE_HOST").is_some(); - let host_category_trace = std::env::var_os("FUSOR_TRACE_RESOLVE_HOST_CATEGORIES").is_some(); + let device = graph.device(); + let host_trace = cfg!(target_arch = "wasm32") || device.config().trace_resolve_host; let host_total_start = host_trace.then(Instant::now); let mut host_profile = ResolveHostProfile::default(); - let mut host_category_profile = - FxHashMap::<&'static str, ResolveHostCategoryProfile>::default(); - let device = graph.device(); let max_subgroup_size = device.max_subgroup_size(); // Pass 1: Build execution graph for all targets @@ -36,40 +32,39 @@ impl Resolver { if let Some(start) = start { host_profile.build_execution_graph += start.elapsed(); } + #[cfg(feature = "graphvis")] + if let Some(dir) = &device.config().dump_stages { + super::visualize::dump_stage( + dir, + &self.execution_graph, + super::visualize::Stage::Built, + ); + } } // Pass 2: Apply Rewrite Rules { let start = host_trace.then(Instant::now); - let optimize_limit = optimize_node_limit(); - let skip_large_graph_optimize = - optimize_limit != 0 && self.execution_graph.node_count() > optimize_limit; - let skip_decode_optimize = skip_large_graph_optimize - && self.is_single_token_qmatmul_graph() - && std::env::var_os("FUSOR_RESOLVE_OPTIMIZE_DECODE_GRAPHS").is_none(); - if std::env::var_os("FUSOR_RESOLVE_SKIP_OPTIMIZE").is_none() { - if skip_large_graph_optimize { - self.optimize_large_graph(graph); - } else { - self.optimize(graph); - } - } + let node_count = self.execution_graph.node_count(); + self.optimize(graph); if let Some(start) = start { host_profile.optimize += start.elapsed(); } - if host_trace && skip_large_graph_optimize { + if host_trace { + let phases = self.optimize_phases; tracing::info!( - "resolve_host_profile optimize_large_graph node_count={} limit={optimize_limit} skipped_decode={skip_decode_optimize}", - self.execution_graph.node_count(), + "resolve_optimize_phases node_count={node_count} recognition={:?} extraction={:?} physical={:?}", + phases.recognition, + phases.extraction, + phases.physical, ); } } - // Pass 3: Topological Sort + // Pass 3: Fusion-friendly topological sort. let sorted_nodes = { let start = host_trace.then(Instant::now); - let sorted_nodes = toposort(&self.execution_graph, None) - .unwrap_or_else(|_| panic!("Cycle detected in execution graph")); + let sorted_nodes = merge_horizontal::fusion_toposort(&self.execution_graph); if let Some(start) = start { host_profile.toposort += start.elapsed(); } @@ -78,23 +73,30 @@ impl Resolver { // Pass 4: Execution // Extract operations in order. - let target_set: FxHashSet = self.targets.iter().copied().collect(); + let mut target_set: FxHashSet = self.targets.iter().copied().collect(); + target_set.extend(self.shared_outputs.keys().copied()); let mut queued_operations = Vec::with_capacity(sorted_nodes.len()); { let start = host_trace.then(Instant::now); + let mut merger = merge_horizontal::HorizontalMerger::new(&device); for idx in sorted_nodes { let node = &self.execution_graph[idx]; // Handle Tensor caching explicitly here if let ExecutionVariant::Tensor(data) = &node.variant { + if let Some(recorder) = &self.recorder { + recorder + .borrow_mut() + .record_tensor_leaf(node.inner_idx, data); + } graph.set_cached_result(node.inner_idx, data.clone()); continue; } - if let Some(op) = self.lower_node(idx, node) { - queued_operations.push((node.inner_idx, op)); - } + let lowered = self.lower_node(idx, node); + merger.push(node, lowered, &self.shared_outputs, &mut queued_operations); } + merger.finish(&mut queued_operations); if let Some(start) = start { host_profile.queue_lowering += start.elapsed(); } @@ -135,325 +137,47 @@ impl Resolver { command_encoder }; - let trace = std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_RESOLVE").is_some(); - let trace_names = std::env::var_os("FUSOR_TRACE_DECODE_NAMES").is_some(); - let profile_gpu_kernels = std::env::var_os("FUSOR_TRACE_GPU_KERNELS").is_some(); + let trace = device.config().trace_decode || device.config().trace_resolve; + let trace_names = device.config().trace_decode_names; + let profile_gpu_kernels = device.config().trace_gpu_kernels; let collect_dispatch_metadata = trace || profile_gpu_kernels; let mut commands = Vec::::with_capacity(queued_operations.len()); let mut dispatch_categories = FxHashMap::::default(); let mut dispatch_names = FxHashMap::::default(); - let plan_cache_enabled = device.kernel_cache().direct_plan_cache().enabled(); - for (node, queued_operation) in queued_operations { - let operation_category = host_category_trace - .then(|| { - graph - .nodes - .nodes - .node_weight(node) - .map(|node| node_category_inner(&node.variant)) - }) - .flatten(); - // A view that composes with its input's buffer layout isn't a - // kernel. Resolve it immediately as a zero-cost buffer view; - // anything else (fill regions, non-composable reshapes) falls - // through to the gather kernel below. - let view_result = if let Some(node_data) = graph.nodes.nodes.node_weight(node) { - match &node_data.variant { - ComputeGraphNodeVariant::View(view) => graph - .get_cached_result(view.input) - .and_then(|input| view.try_map_tensor(input)), - _ => None, - } - } else { - None - }; - if let Some(result) = view_result { - let start = host_trace.then(Instant::now); - // Cache the result - graph.set_cached_result(node, result); - // Map-layout nodes are resolved immediately — release any - // input buffers that are no longer needed. - // Use graph.visit_dependencies for map_layout since they - // are not lowered to Operations. - Self::release_dead_intermediates_from_graph( - graph, - &[node], - &mut remaining_consumers, - &target_set, - ); - if let Some(start) = start { - host_profile.map_layout += start.elapsed(); - } - } else { - let slice_copy = graph.nodes.nodes.node_weight(node).and_then(|node_data| { - let ComputeGraphNodeVariant::Assign(slice_assign) = &node_data.variant else { - return None; - }; - Self::try_prepare_in_place_slice_assign_copy(graph, slice_assign) - }); - if let Some((output, copies)) = slice_copy { - graph.set_cached_result(node, output); - commands.extend(copies.into_iter().map(CommandRecord::CopyBuffer)); - let start = host_trace.then(Instant::now); - Self::release_dead_intermediates( - graph, - &[&queued_operation], - &mut remaining_consumers, - &target_set, - ); - if let Some(start) = start { - host_profile.release += start.elapsed(); - } - continue; - } - - let start = host_trace.then(Instant::now); - let new_inputs = queued_operation.inputs(graph); - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.inputs += elapsed; - if let Some(category) = operation_category { - host_category_profile.entry(category).or_default().inputs += elapsed; - } - } - if let QueuedOperation::QMatMul(qmatmul) = &queued_operation { - let start = host_trace.then(Instant::now); - let result = qmatmul.output(graph, &new_inputs); - let MirValue::Tensor(resolved) = result else { - panic!("QMatMul output value is not a tensor"); - }; - graph.set_cached_result(node, resolved.clone()); - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.output += elapsed; - if let Some(category) = operation_category { - host_category_profile.entry(category).or_default().output += elapsed; - } - } - - let start = host_trace.then(Instant::now); - let constraints = qmatmul.workgroup_shape_constraints(&device); - let workgroup_shape = constraints - .solve(max_subgroup_size, &device.limits()) - .unwrap_or_else(|| { - panic!( - "Failed to find a valid qmatmul workgroup shape for constraints {constraints:?}" - ) - }); - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.workgroup += elapsed; - if let Some(category) = operation_category { - host_category_profile.entry(category).or_default().workgroup += elapsed; - } - } - - let start = host_trace.then(Instant::now); - let build_kernels = || { - qmatmul - .build_direct_kernels(graph, &workgroup_shape, &new_inputs) - .unwrap_or_else(|error| panic!("{error}")) - .into_kernels() - }; - let kernels = if plan_cache_enabled { - let kernel_key = - structural_kernel_key(qmatmul.as_ref(), &new_inputs, &workgroup_shape); - resolve_cached_direct_plan( - device.kernel_cache().direct_plan_cache(), - kernel_key, - direct_plan_binding_buffers(&new_inputs), - build_kernels, - ) - } else { - build_kernels() - }; - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.build_kernel += elapsed; - if let Some(category) = operation_category { - let profile = host_category_profile.entry(category).or_default(); - profile.count += kernels.len(); - profile.build_kernel += elapsed; - } - } - - for direct_kernel in kernels { - let start = host_trace.then(Instant::now); - if let Some(dispatch) = - direct_kernel.prepare_dispatch(device.kernel_cache()) - { - let name = direct_kernel.name().to_string(); - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.prepare_dispatch += elapsed; - if let Some(category) = operation_category { - host_category_profile - .entry(category) - .or_default() - .prepare_dispatch += elapsed; - } - } - let category = collect_dispatch_metadata.then(|| { - let category = dispatch_category(&name); - if trace { - *dispatch_categories.entry(category.clone()).or_default() += 1; - if trace_names { - *dispatch_names.entry(name.clone()).or_default() += 1; - } - } - category - }); - commands.push(CommandRecord::Dispatch(DispatchRecord { - dispatch, - name, - category, - })); - } else if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.prepare_dispatch += elapsed; - if let Some(category) = operation_category { - host_category_profile - .entry(category) - .or_default() - .prepare_dispatch += elapsed; - } - } - } - - let start = host_trace.then(Instant::now); - Self::release_dead_intermediates( - graph, - &[&queued_operation], - &mut remaining_consumers, - &target_set, - ); - if let Some(start) = start { - host_profile.release += start.elapsed(); - } - continue; - } - - let QueuedOperation::Generic(operation) = &queued_operation else { - unreachable!("qmatmul resolver arm returned above"); - }; - let start = host_trace.then(Instant::now); - let result = operation.output(graph, &new_inputs); - let MirValue::Tensor(resolved) = result else { - panic!("Kernel input value is not a tensor"); - }; - graph.set_cached_result(node, resolved.clone()); - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.output += elapsed; - if let Some(category) = operation_category { - host_category_profile.entry(category).or_default().output += elapsed; - } - } - - let start = host_trace.then(Instant::now); - let constraints = operation.workgroup_shape_constraints(&device); - let workgroup_shape = constraints - .solve(max_subgroup_size, &device.limits()) - .unwrap_or_else(|| { - panic!( - "Failed to find a valid workgroup shape for constraints {constraints:?}" - ) - }); - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.workgroup += elapsed; - if let Some(category) = operation_category { - host_category_profile.entry(category).or_default().workgroup += elapsed; - } - } - let start = host_trace.then(Instant::now); - let build_kernels = || { - vec![ - operation - .build_direct_kernel(graph, &workgroup_shape, &new_inputs) - .unwrap_or_else(|| { - panic!( - "operation did not provide a direct kernel: {}", - operation.name() - ) - }), - ] - }; - let kernels = if plan_cache_enabled { - let kernel_key = - structural_kernel_key(operation.as_ref(), &new_inputs, &workgroup_shape); - resolve_cached_direct_plan( - device.kernel_cache().direct_plan_cache(), - kernel_key, - direct_plan_binding_buffers(&new_inputs), - build_kernels, - ) - } else { - build_kernels() - }; - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.build_kernel += elapsed; - if let Some(category) = operation_category { - let profile = host_category_profile.entry(category).or_default(); - profile.count += kernels.len(); - profile.build_kernel += elapsed; - } - } - for direct_kernel in kernels { - let start = host_trace.then(Instant::now); - if let Some(dispatch) = direct_kernel.prepare_dispatch(device.kernel_cache()) { - let name = direct_kernel.name().to_string(); - if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.prepare_dispatch += elapsed; - if let Some(category) = operation_category { - host_category_profile - .entry(category) - .or_default() - .prepare_dispatch += elapsed; - } - } - let category = collect_dispatch_metadata.then(|| { - let category = dispatch_category(&name); - if trace { - *dispatch_categories.entry(category.clone()).or_default() += 1; - if trace_names { - *dispatch_names.entry(name.clone()).or_default() += 1; - } - } - category - }); - commands.push(CommandRecord::Dispatch(DispatchRecord { - dispatch, - name, - category, - })); - } else if let Some(start) = start { - let elapsed = start.elapsed(); - host_profile.prepare_dispatch += elapsed; - if let Some(category) = operation_category { - host_category_profile - .entry(category) - .or_default() - .prepare_dispatch += elapsed; + // Every graph takes the three-phase queue runner: serial input + // gathering, parallel kernel building, serial recording/encoding. + // Compatible independent operations may merge across the queue; each + // operation kind retains its own shape, dependency and binding gates. + let mut ledger = super::alloc_reuse::BufferLedger::new(&device, Some(&remaining_consumers)); + if let Some(recorder) = &self.recorder { + ledger.register_recorder_pins(recorder.borrow().pinned_ptrs()); + } + Self::execute_queue( + self.recorder.as_ref(), + graph, + &device, + max_subgroup_size, + queued_operations, + &mut remaining_consumers, + &target_set, + &self.shared_outputs, + &mut ledger, + &mut commands, + &mut host_profile, + host_trace, + &mut |name: &str| { + collect_dispatch_metadata.then(|| { + let category = dispatch_category(name); + if trace { + *dispatch_categories.entry(category.clone()).or_default() += 1; + if trace_names { + *dispatch_names.entry(name.to_string()).or_default() += 1; } } - } - let start = host_trace.then(Instant::now); - Self::release_dead_intermediates( - graph, - &[&queued_operation], - &mut remaining_consumers, - &target_set, - ); - if let Some(start) = start { - host_profile.release += start.elapsed(); - } - }; - } - + category + }) + }, + ); let total_kernels = commands .iter() .filter(|command| matches!(command, CommandRecord::Dispatch(_))) @@ -531,89 +255,22 @@ impl Resolver { if !commands.is_empty() { let encode_start = host_trace.then(Instant::now); - let mut dispatch_index = 0usize; - let mut command_index = 0usize; - let dispatches_per_pass = dispatches_per_pass(total_kernels); - while command_index < commands.len() { - match &commands[command_index] { - CommandRecord::CopyBuffer(copy) => { - command_encoder.copy_buffer_to_buffer( - ©.source, - copy.source_offset, - ©.destination, - copy.destination_offset, - copy.size, - ); - command_index += 1; - } - CommandRecord::Dispatch(_) => { - if let Some((query_set, _, _, _)) = &query_resources - && !profile_inside_pass_timestamps - { - if let CommandRecord::Dispatch(record) = &commands[command_index] { - let mut pass = command_encoder.begin_compute_pass( - &wgpu::ComputePassDescriptor { - label: Some(record.name.as_str()), - timestamp_writes: Some(wgpu::ComputePassTimestampWrites { - query_set, - beginning_of_pass_write_index: Some( - (dispatch_index * 2) as u32, - ), - end_of_pass_write_index: Some( - (dispatch_index * 2 + 1) as u32, - ), - }), - }, - ); - record.dispatch.run(&mut pass); - } - dispatch_index += 1; - command_index += 1; - continue; - } - - let mut pass = - command_encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { - label: Some("Resolver Direct Kernels"), - timestamp_writes: None, - }); - let mut pass_dispatches = 0usize; - while command_index < commands.len() { - if pass_dispatches >= dispatches_per_pass { - break; - } - let CommandRecord::Dispatch(record) = &commands[command_index] else { - break; - }; - if let Some((query_set, _, _, _)) = &query_resources { - pass.write_timestamp(query_set, (dispatch_index * 2) as u32); - } - pass.push_debug_group(&record.name); - record.dispatch.run(&mut pass); - pass.pop_debug_group(); - if let Some((query_set, _, _, _)) = &query_resources { - pass.write_timestamp(query_set, (dispatch_index * 2 + 1) as u32); - } - dispatch_index += 1; - command_index += 1; - pass_dispatches += 1; - } + command_encoder = super::queue_executor::encode_command_records( + &device, + &commands, + total_kernels, + query_resources.as_ref().map(|(query_set, _, _, _)| { + super::queue_executor::TimestampPlan { + query_set, + inside_pass: profile_inside_pass_timestamps, } - } - } + }), + command_encoder, + |encoder, wait| { + submit_resolver_encoder(&device, encoder, wait, host_trace, &mut host_profile); + }, + ); - if let Some((query_set, query_buffer, readback_buffer, raw_query_size)) = - &query_resources - { - command_encoder.resolve_query_set(query_set, 0..query_count, query_buffer, 0); - command_encoder.copy_buffer_to_buffer( - query_buffer, - 0, - readback_buffer, - 0, - *raw_query_size, - ); - } if let Some(start) = encode_start { host_profile.encode += start.elapsed(); } @@ -625,15 +282,40 @@ impl Resolver { let tail_result = tail(&data, &mut command_encoder); // Submit any remaining commands. - let submit_start = host_trace.then(Instant::now); - device.wgpu_queue().submit(Some(command_encoder.finish())); - if let Some(start) = submit_start { - host_profile.submit += start.elapsed(); - } + submit_resolver_encoder( + &device, + command_encoder, + false, + host_trace, + &mut host_profile, + ); #[cfg(not(target_arch = "wasm32"))] { - if let Some((_, _, readback_buffer, raw_query_size)) = &query_resources { + if let Some((query_set, query_buffer, readback_buffer, raw_query_size)) = + &query_resources + { let profile_readback_start = host_trace.then(Instant::now); + // A command buffer's counter samples are not guaranteed visible to + // `resolve_query_set` until it has completed: Metal's writeback of the + // final compute encoder's stage-boundary samples races a resolve encoded + // behind it, leaving those slots zero. Resolving from a command buffer + // submitted after the sampling work retired removes the race. + device.poll_wait(); + let mut resolve_encoder = + device + .wgpu_device() + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("Resolver Timestamp Resolve"), + }); + resolve_encoder.resolve_query_set(query_set, 0..query_count, query_buffer, 0); + resolve_encoder.copy_buffer_to_buffer( + query_buffer, + 0, + readback_buffer, + 0, + *raw_query_size, + ); + device.wgpu_queue().submit(Some(resolve_encoder.finish())); let slice = readback_buffer.slice(..*raw_query_size); let (sender, receiver) = std::sync::mpsc::channel(); slice.map_async(wgpu::MapMode::Read, move |result| { @@ -644,7 +326,7 @@ impl Resolver { Ok(Ok(())) => { let view = slice.get_mapped_range(); let timestamps = bytemuck::cast_slice::(&view); - print_gpu_kernel_profile( + let profile = collect_gpu_kernel_profile( &dispatch_metadata, timestamps, device.wgpu_queue().get_timestamp_period() as f64, @@ -654,6 +336,7 @@ impl Resolver { "pass_boundary" }, ); + device.record_kernel_profile(profile); drop(view); readback_buffer.unmap(); } @@ -673,9 +356,6 @@ impl Resolver { if let Some(start) = host_total_start { host_profile.print(start.elapsed(), queued_operation_count, total_kernels); - if host_category_trace { - print_host_category_profile(host_category_profile); - } } ( ResolverResult { @@ -687,8 +367,8 @@ impl Resolver { } } -fn resolve_cached_direct_plan( - plan_cache: &fusor_tile_ir_runtime::DirectPlanCache, +pub(super) fn resolve_cached_kernel_plan( + kernel_cache: &fusor_tile_ir_runtime::KernelCache, cache_key: crate::mir::kernel_backend::KernelCacheKey, binding_buffers: Vec>>, build: impl FnOnce() -> Vec, @@ -697,14 +377,17 @@ fn resolve_cached_direct_plan( .iter() .map(Vec::as_slice) .collect::>(); - plan_cache - .try_get_or_insert_many(cache_key, &binding_slices, || { + kernel_cache + .kernel_plan_cache() + .try_get_or_insert_many(kernel_cache, cache_key, &binding_slices, || { Ok::<_, std::convert::Infallible>(build()) }) - .expect("infallible direct plan cache build failed") + .expect("infallible kernel plan cache build failed") } -fn direct_plan_binding_buffers(inputs: &[MirValue]) -> Vec>> { +pub(super) fn kernel_plan_binding_buffers( + inputs: &[MirValue], +) -> Vec>> { let buffers = inputs .iter() .filter_map(|input| match input { @@ -716,13 +399,41 @@ fn direct_plan_binding_buffers(inputs: &[MirValue]) -> Vec usize { - if let Ok(value) = std::env::var("FUSOR_RESOLVE_DISPATCHES_PER_PASS") - && let Ok(parsed) = value.parse::() - && parsed > 0 - { - return parsed; +pub(super) fn dispatches_per_pass(device: &crate::Device, total_kernels: usize) -> usize { + if let Some(value) = device.config().resolve_dispatches_per_pass { + return value; } if total_kernels >= 1024 { 1 } else { usize::MAX } } + +pub(super) fn dispatches_per_submit(device: &crate::Device, total_kernels: usize) -> usize { + if let Some(value) = device.config().resolve_dispatches_per_submit { + return value; + } + + // Chunked submits exist to bound in-flight memory on giant training + // graphs; small inference graphs must stay a single submit. + if device.backend() == wgpu::Backend::Metal && total_kernels >= 1024 { + 256 + } else { + usize::MAX + } +} + +fn submit_resolver_encoder( + device: &crate::Device, + command_encoder: wgpu::CommandEncoder, + wait: bool, + host_trace: bool, + host_profile: &mut ResolveHostProfile, +) { + let submit_start = host_trace.then(Instant::now); + device.wgpu_queue().submit(Some(command_encoder.finish())); + if wait { + device.poll_wait(); + } + if let Some(start) = submit_start { + host_profile.submit += start.elapsed(); + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/sink_views.rs b/fusor-ml/core/src/compute_graph/resolve/sink_views.rs new file mode 100644 index 000000000..b2c02182f --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/sink_views.rs @@ -0,0 +1,175 @@ +//! Sinking unary chains that sit behind pure views into their producer's +//! epilogue. +//! +//! A matmul's post epilogue runs over the workgroup's output tile before the +//! store, so an activation fused there costs nothing: no second kernel, no +//! second trip through memory for an activation-sized tensor. The e-graph +//! generator already fuses `unary(matmul)`, but only when the unary chain +//! reads the matmul *directly* — and it rarely does. Every convolution +//! reassembles its `(rows, out_channels)` matmul into +//! `(batch, out_channels, ...spatial)` first, so the activation the caller +//! writes lands behind a reshape and a permute and the fusion never fires. +//! +//! Pure views only relabel coordinates, and a unary function commutes with +//! relabelling: `f(view(x)) == view(f(x))`. So the chain can move to the +//! matmul and the views stay exactly where they are. This runs as a +//! recognizer rather than a generator because it rewrites *two* nodes — the +//! matmul gains the epilogue and the elementwise node collapses into an +//! observation of its own input — and a generator may only return a new +//! variant for the node it was asked about. + +use super::{ExecutionNodeIndex, ExecutionVariant, Resolver}; +use crate::compute_graph::{ComputeGraphInner, NodeIndex}; +use crate::nary_wise::UnaryFunctionChain; + +impl Resolver { + /// Move unary chains that read a matmul through pure views into that + /// matmul's post epilogue. + pub(super) fn sink_unary_chains_into_matmuls(&mut self, graph: &mut ComputeGraphInner) { + let before = self.sunk_chains; + let candidates: Vec = self + .execution_graph + .node_indices() + .filter(|&node| { + matches!( + self.execution_graph[node].variant, + ExecutionVariant::Elementwise(_) + ) + }) + .collect(); + for node in candidates { + self.try_sink_unary_chain(graph, node); + } + if graph.device().config().trace_resolve && self.sunk_chains != before { + tracing::info!("sink_views chains={}", self.sunk_chains - before); + } + } + + fn try_sink_unary_chain(&mut self, graph: &mut ComputeGraphInner, node: ExecutionNodeIndex) { + if !self.execution_graph.contains_node(node) { + return; + } + let ExecutionVariant::Elementwise(nary) = &self.execution_graph[node].variant else { + return; + }; + let output_datatype = nary.output_datatype; + let Some(chain) = nary.try_extract_unary_chain() else { + return; + }; + if chain.functions.functions.is_empty() { + return; + } + // A chain whose own value is observed elsewhere still collapses: the + // views below it produce the same buffer, and the node stays readable + // as an observation of its input. What must not be observed is the + // matmul's *un-activated* output, so every hop from the matmul up to + // this node has to feed only the next hop. + let Some(producer) = self.private_view_chain_to_matmul(graph, chain.value, node) else { + return; + }; + if self.check_cached(graph, self.execution_graph[producer].inner_idx) { + return; + } + let ExecutionVariant::MatMul(matmul) = &self.execution_graph[producer].variant else { + return; + }; + // Only dtype-preserving chains belong after the cooperative store. + if matmul.datatype != output_datatype { + return; + } + let mut fused = matmul.clone(); + let mut functions = fused.post_element_wise.functions.clone(); + functions.extend(chain.functions.functions.iter().cloned()); + fused.post_element_wise = + UnaryFunctionChain::new(functions, fused.post_element_wise.input_datatype()); + self.execution_graph[producer].variant = ExecutionVariant::MatMul(fused); + self.alias_to_input(graph, node, chain.value); + self.sunk_chains += 1; + } + + /// The matmul at the base of a chain of pure views reaching `input`, when + /// every node on that chain — the matmul included — is read by exactly one + /// consumer, ending at `consumer`. `None` when anything else observes an + /// intermediate value, when a hop is not a layout-only view, or when the + /// base is not a matmul. + fn private_view_chain_to_matmul( + &self, + graph: &ComputeGraphInner, + input: NodeIndex, + consumer: ExecutionNodeIndex, + ) -> Option { + let mut current = self.get_input_node_in_exec_graph(input)?; + let mut reader = consumer; + loop { + if self + .execution_graph + .neighbors_directed(current, petgraph::Direction::Outgoing) + .any(|other| other != reader) + { + return None; + } + // Every node up to (not including) the chain itself takes the + // activation into its value, so none of them may be observed from + // outside this resolve: a handle the caller still holds, or a + // value this resolve was asked to produce. + let inner = self.execution_graph[current].inner_idx; + if self.targets.contains(&inner) + || graph + .nodes + .nodes + .node_weight(inner) + .is_some_and(|node| node.reference_count > 0) + { + return None; + } + match &self.execution_graph[current].variant { + ExecutionVariant::MatMul(_) => return Some(current), + ExecutionVariant::View(view) => { + // Layout-only: a stage stack that composes to one layout + // relabels coordinates and nothing more. + view.composed_layout()?; + let next = self.get_input_node_in_exec_graph(view.input)?; + reader = current; + current = next; + } + _ => return None, + } + } + } + + /// Collapse `node` into an observation of `input`: its value is now + /// exactly what `input` produces. Consumers read `input`'s execution node, + /// and `node`'s own index stays resolvable through `shared_outputs`. + fn alias_to_input( + &mut self, + graph: &mut ComputeGraphInner, + node: ExecutionNodeIndex, + input: NodeIndex, + ) { + let Some(representative) = self.get_input_node_in_exec_graph(input) else { + return; + }; + let inner = self.execution_graph[node].inner_idx; + let consumers: Vec = self + .execution_graph + .neighbors_directed(node, petgraph::Direction::Outgoing) + .collect(); + for consumer in consumers { + if consumer != representative + && self + .execution_graph + .find_edge(representative, consumer) + .is_none() + { + self.execution_graph.add_edge(representative, consumer, ()); + } + } + self.execution_graph.remove_node(node); + self.node_mapping.remove(&inner); + self.shared_outputs.entry(input).or_default().push(inner); + graph.add_dependency_edge(input, inner); + if let Some(recorder) = &self.recorder { + recorder.borrow_mut().record_physical_edge(input, inner); + } + } +} diff --git a/fusor-ml/core/src/compute_graph/resolve/visualize.rs b/fusor-ml/core/src/compute_graph/resolve/visualize.rs new file mode 100644 index 000000000..c8e4ea6de --- /dev/null +++ b/fusor-ml/core/src/compute_graph/resolve/visualize.rs @@ -0,0 +1,157 @@ +//! Graphviz dumps of the execution graph between resolver stages. +//! +//! The lazy-graph `graphvis` renders what the user built; this renders what +//! the resolver did to it. One digraph per stage boundary makes each stage's +//! job legible: recognition collapsing a broadcast-mul + sum cluster into one +//! `matmul`, extraction folding producers into epilogues, region formation +//! and horizontal merge grouping survivors into dispatches. + +use std::fmt::Write as _; +use std::path::Path; + +use petgraph::visit::{EdgeRef, IntoEdgeReferences}; + +use super::{ExecutionGraph, ExecutionVariant}; + +/// Stage boundaries, in the order [`super::run`] reaches them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum Stage { + /// Pass 1 output: one node per lazy-graph node, nothing fused. + Built, + /// After the pre-ingest recognizers (matmul, embedding, attention, row + /// programs, assign chains). + Recognized, + /// After e-graph extraction and its deltas are applied: fusion committed, + /// killed producers gone. + Extracted, + /// After multi-output elementwise regions form. + Regions, +} + +impl Stage { + fn slug(self) -> &'static str { + match self { + Stage::Built => "1-built", + Stage::Recognized => "2-recognized", + Stage::Extracted => "3-extracted", + Stage::Regions => "4-regions", + } + } +} + +/// Node fill per variant family: leaves grey, recognized regions saturated, +/// the three-op core in one hue so fusion progress reads at a glance. +fn style(variant: &ExecutionVariant) -> (&'static str, &'static str) { + match variant { + ExecutionVariant::Tensor(_) => ("box", "#e8e8e8"), + ExecutionVariant::QMatrix(_) => ("box", "#d8d0e8"), + ExecutionVariant::Elementwise(_) => ("ellipse", "#cfe4f7"), + ExecutionVariant::Reduce(_) => ("ellipse", "#a8cbe8"), + // Folds share the reduce hue; a multi-slot carrier is the distinction + // worth seeing, so the label carries it. + ExecutionVariant::Fold(_) => ("ellipse", "#8fbcd9"), + ExecutionVariant::View(_) => ("ellipse", "#eef3f7"), + ExecutionVariant::Assign(_) => ("ellipse", "#f7e4cf"), + ExecutionVariant::Region(_) => ("octagon", "#bfe8cf"), + ExecutionVariant::MatMul(_) => ("doubleoctagon", "#f7cfcf"), + ExecutionVariant::QMatMul(_) => ("doubleoctagon", "#e8bfd8"), + ExecutionVariant::QEmbedding(_) => ("doubleoctagon", "#e8dcbf"), + ExecutionVariant::RowProgram(_) => ("doubleoctagon", "#cfd8f7"), + ExecutionVariant::Attention(_) => ("doubleoctagon", "#f7bfbf"), + } +} + +fn label(variant: &ExecutionVariant) -> String { + match variant { + ExecutionVariant::Tensor(data) => format!("tensor\\n{:?}", data.layout().shape()), + ExecutionVariant::QMatrix(_) => "qmatrix".to_string(), + ExecutionVariant::Elementwise(op) => { + format!("elementwise x{}\\n{:?}", op.inputs.len(), op.shape) + } + ExecutionVariant::Reduce(op) => format!("reduce {}", op.function.name()), + ExecutionVariant::Fold(op) => format!( + "fold x{}{}\\n{} outputs", + op.carrier.len(), + op.block.map_or(String::new(), |block| format!(" /{block}")), + op.outputs.len(), + ), + ExecutionVariant::View(_) => "view".to_string(), + ExecutionVariant::Assign(_) => "slice_assign".to_string(), + ExecutionVariant::Region(op) => format!("region\\n{} statements", op.statements.len()), + ExecutionVariant::MatMul(_) => "matmul".to_string(), + ExecutionVariant::QMatMul(op) => format!( + "qmatmul\\n{:?} -> {:?}{}{}", + op.in_shape, + op.out_shape, + op.pre_element_wise_expr + .as_ref() + .map_or("", |_| "\\n+pre epilogue"), + op.post_element_wise_expr + .as_ref() + .map_or("", |_| "\\n+post epilogue"), + ), + ExecutionVariant::QEmbedding(_) => "qembedding".to_string(), + ExecutionVariant::RowProgram(op) => format!("row_program\\n{} steps", op.steps.len()), + ExecutionVariant::Attention(op) => format!("attention\\n{:?}", op.kind), + } +} + +/// One stage's execution graph as a Graphviz digraph. +pub(crate) fn execution_graph_dot(graph: &ExecutionGraph, stage: Stage) -> String { + let mut out = String::new(); + let _ = writeln!(out, "digraph \"{}\" {{", stage.slug()); + let _ = writeln!(out, " rankdir=BT;"); + let _ = writeln!( + out, + " label=\"{} — {} nodes\";\n labelloc=t;", + stage.slug(), + graph.node_count() + ); + let _ = writeln!( + out, + " node [style=filled, fontname=\"Helvetica\", fontsize=10];" + ); + for idx in graph.node_indices() { + let node = &graph[idx]; + let (shape, fill) = style(&node.variant); + let _ = writeln!( + out, + " n{} [label=\"{}\\n#{}\", shape={shape}, fillcolor=\"{fill}\"];", + idx.index(), + label(&node.variant), + node.inner_idx.index(), + ); + } + for edge in graph.edge_references() { + let _ = writeln!( + out, + " n{} -> n{};", + edge.source().index(), + edge.target().index() + ); + } + let _ = writeln!(out, "}}"); + out +} + +/// Resolves are numbered so a decode trace's per-token graphs stay apart. +static RESOLVE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +/// Write one stage dump. Failures are traced, never fatal: this is a +/// debugging aid and must not change whether a resolve succeeds. +pub(crate) fn dump_stage(dir: &Path, graph: &ExecutionGraph, stage: Stage) { + use std::sync::atomic::Ordering; + let resolve = if stage == Stage::Built { + RESOLVE.fetch_add(1, Ordering::Relaxed) + } else { + RESOLVE.load(Ordering::Relaxed).saturating_sub(1) + }; + if let Err(error) = std::fs::create_dir_all(dir) { + tracing::warn!("dump_stages: {dir:?}: {error}"); + return; + } + let path = dir.join(format!("resolve{resolve:04}-{}.dot", stage.slug())); + if let Err(error) = std::fs::write(&path, execution_graph_dot(graph, stage)) { + tracing::warn!("dump_stages: {path:?}: {error}"); + } +} diff --git a/fusor-ml/core/src/compute_graph/tests.rs b/fusor-ml/core/src/compute_graph/tests.rs index 73bf7f53e..acf760d1e 100644 --- a/fusor-ml/core/src/compute_graph/tests.rs +++ b/fusor-ml/core/src/compute_graph/tests.rs @@ -1,4 +1,5 @@ use crate::{Device, StrideSpec, Tensor}; +use fusor_gguf::{BlockQ4K, GgmlType}; // Build a small intermediate that requires a real kernel (not a Tensor input). // `x` materializes via `(input * 2.0) + 1.0`, which fuses to a single nary. @@ -202,6 +203,93 @@ fn auto_flush_resolves_pending_siblings() { }); } +#[test] +fn qmatmul_triggers_auto_flush() { + pollster::block_on(async { + const N: usize = 4; + const K: usize = 8; + + let Ok(device) = Device::new().await else { + return; + }; + + let siblings: Vec<_> = (0..3).map(|_| build_intermediate(&device).sin()).collect(); + let weight_bytes: Vec = (0..N * K) + .map(|i| 0.1 + i as f32 * 0.05) + .flat_map(f32::to_le_bytes) + .collect(); + let weight = crate::QMatrix::from_parts( + &device, + &weight_bytes, + vec![N, K].into_boxed_slice(), + GgmlType::F32, + ) + .unwrap(); + let output = + Tensor::new::(&device, &[[1.0f32, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0]]) + .q_mat_mul(&weight); + + device.compute_graph().set_flush_threshold(1); + let (_, kernels) = output.data.materialize(); + assert!(kernels > 0, "QMatMul dispatched no kernels"); + + for sibling in siblings { + assert!( + device.compute_graph().is_cached_for_test(sibling.key()), + "QMatMul should auto-flush pending sibling outputs", + ); + } + }); +} + +#[test] +fn qmatmul_releases_a_dead_cached_activation() { + pollster::block_on(async { + const N: usize = 64; + const K: usize = 4; + + let Ok(device) = Device::new().await else { + return; + }; + let activation = build_intermediate(&device); + let activation_key = activation.key(); + let _ = activation.data.materialize(); + assert!(device.compute_graph().is_cached_for_test(activation_key)); + + let weight_bytes = (0..N * K) + .map(|i| 0.1 + i as f32 * 0.05) + .flat_map(f32::to_le_bytes) + .collect::>(); + let weight = crate::QMatrix::from_parts( + &device, + &weight_bytes, + vec![N, K].into_boxed_slice(), + GgmlType::F32, + ) + .unwrap(); + let output = activation.q_mat_mul(&weight); + drop(activation); + + let (_, kernels) = output.data.materialize(); + assert!( + (1..=2).contains(&kernels), + "QMatMul should use a singular or M-padded two-kernel plan", + ); + assert!( + !device.compute_graph().is_cached_for_test(activation_key), + "the direct queue should release its dead cached activation", + ); + + let values = output.as_slice::<2, f32>().await.unwrap(); + let expected = [3.0f32, 5.0, 7.0, 9.0] + .iter() + .enumerate() + .map(|(k, value)| value * (0.1 + k as f32 * 0.05)) + .sum::(); + assert!((values[[0, 0]] - expected).abs() < 1e-4); + }); +} + // --- split + op + cat lowering (resolve/recognize_cat.rs) --- /// Narrow a 2D tensor along a dimension as a view. @@ -219,6 +307,23 @@ fn narrow2(tensor: &Tensor, dim: usize, start: usize, length: usize) -> Tensor { tensor.restride(specs) } +#[test] +fn dense_offset_view_downloads_the_compact_range() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let input_values = (0..32).map(|value| value as f32).collect::>(); + let input = Tensor::from_slice(&device, [4, 8], &input_values); + let row = narrow2(&input, 0, 2, 1); + + let downloaded = row.as_slice::<2, f32>().await.unwrap(); + for column in 0..8 { + assert_eq!(downloaded[[0, column]], (16 + column) as f32); + } + }); +} + fn cat_test_input(device: &Device) -> (Tensor, Vec>) { let rows: Vec> = (0..4) .map(|r| (0..8).map(|c| (r * 8 + c) as f32 * 0.1).collect()) @@ -385,6 +490,79 @@ fn three_way_chunk_cat_collapses() { }); } +#[test] +fn wide_deep_chunk_cat_has_no_branch_or_expression_cap() { + pollster::block_on(async { + const CHUNKS: usize = 8; + const ROWS_PER_CHUNK: usize = 2; + const COLS: usize = 4; + const OPS_PER_BRANCH: usize = 40; + + let Ok(device) = Device::new().await else { + return; + }; + let rows: Vec> = (0..CHUNKS * ROWS_PER_CHUNK) + .map(|row| { + (0..COLS) + .map(|column| (row * COLS + column) as f32 * 0.01) + .collect() + }) + .collect(); + let input = Tensor::new::(&device, &rows); + let mut output = Tensor::splat(&device, 0.0f32, [CHUNKS * ROWS_PER_CHUNK, COLS]); + for chunk in 0..CHUNKS { + let start = chunk * ROWS_PER_CHUNK; + let mut branch = narrow2(&input, 0, start, ROWS_PER_CHUNK); + for _ in 0..OPS_PER_BRANCH { + branch = &branch + 0.25f32; + } + output = output.slice_assign([start..start + ROWS_PER_CHUNK, 0..COLS], &branch); + } + + let (_, kernels) = output.data.materialize(); + assert_eq!( + kernels, 1, + "wide/deep concat should remain one automatic fusion" + ); + let values = output.as_slice::<2, f32>().await.unwrap(); + for (row, source) in rows.iter().enumerate() { + for (column, value) in source.iter().enumerate() { + let expected = value + OPS_PER_BRANCH as f32 * 0.25; + assert!((values[[row, column]] - expected).abs() < 1e-4); + } + } + }); +} + +#[test] +fn live_elementwise_region_has_no_statement_cap() { + pollster::block_on(async { + const STATEMENTS: usize = 20; + let Ok(device) = Device::new().await else { + return; + }; + let input = Tensor::new::(&device, &[1.0, 2.0, 3.0, 4.0]); + let mut current = input; + let mut live = Vec::new(); + for _ in 0..STATEMENTS { + current = ¤t + 1.0f32; + live.push(current.clone()); + } + + let (_, kernels) = current.data.materialize(); + assert_eq!( + kernels, 1, + "all live statements should lower as one legal multi-output region" + ); + for (statement, tensor) in live.iter().enumerate() { + let values = tensor.as_slice::<1, f32>().await.unwrap(); + for (index, base) in [1.0f32, 2.0, 3.0, 4.0].iter().enumerate() { + assert_eq!(values[[index]], base + statement as f32 + 1.0); + } + } + }); +} + /// A staged (divmod) view feeding a single-read elementwise folds into the /// consumer's loads: one dispatch, no gather. The same composition used to /// cost a gather plus the elementwise. @@ -459,8 +637,17 @@ fn conv_im2col_matmul_runs_without_gather() { .is_some_and(|config| config.is_fixed()); for (b, c, h, w, n, kh, kw, implicit) in [ + // N=16 output channels reach the coop kernel through the + // small-side (64, 16) tile, so the implicit-GEMM read fires. ( - 2usize, 8usize, 16usize, 16usize, 16usize, 3usize, 3usize, false, + 2usize, + 8usize, + 16usize, + 16usize, + 16usize, + 3usize, + 3usize, + coop_viable, ), (2, 64, 34, 34, 128, 3, 3, coop_viable), // Coop-tile-unaligned M (1089 divides nothing): the masked-edge @@ -559,3 +746,342 @@ fn deep_branch_chains_collapse_through_cat() { } }); } + +/// Natural-form state updates (dependency-chained elementwise triples whose +/// intermediates are all user-held, the shape an optimizer step takes) must +/// (a) settle into fixed physical buffers via in-place allocation claims and +/// (b) record and replay their flush plan on isomorphic steps. +#[test] +fn natural_form_updates_replay_and_claim_in_place() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + const N: usize = 64; + let mut m = Tensor::new::(&device, &vec![0.5f32; N]); + let mut v = Tensor::new::(&device, &vec![0.25f32; N]); + let mut p = Tensor::new::(&device, &vec![1.0f32; N]); + + let mut previous_ptrs: Option<[usize; 3]> = None; + let mut stable_iterations = 0; + let (mut expected_m, mut expected_v, mut expected_p) = (0.5f32, 0.25f32, 1.0f32); + // Lifecycle: plain resolve, record on the second sighting, then + // replay — so replays (and in-place claim stability) start at the + // third iteration. + for iteration in 0..6 { + let g = Tensor::new::(&device, &vec![0.01f32; N]); + let m2 = &(&m + 0.1f32) + &g; + let v2 = &(&v + 0.2f32) + &g; + let p2 = &(&p - 0.001f32) - &(&m2 + &v2); + expected_m = (expected_m + 0.1) + 0.01; + expected_v = (expected_v + 0.2) + 0.01; + expected_p = expected_p - 0.001 - (expected_m + expected_v); + m = m2; + v = v2; + p = p2; + device.flush(); + + let ptrs = [&m, &v, &p].map(|tensor| { + let (data, _) = tensor.data.materialize(); + std::sync::Arc::as_ptr(data.buffer()) as usize + }); + if iteration >= 2 && previous_ptrs == Some(ptrs) { + stable_iterations += 1; + } + previous_ptrs = Some(ptrs); + } + assert!( + stable_iterations >= 1, + "state buffers never settled in place: updates should claim their \ + dead sources' buffers on isomorphic steps" + ); + + let cache = device.flush_plan_cache(); + assert!( + cache.record_count() >= 1, + "no flush plan was recorded across isomorphic steps" + ); + assert!( + cache.replay_count() >= 1, + "no flush replay fired across isomorphic steps" + ); + + let m_values = m.as_slice::<1, f32>().await.unwrap(); + let v_values = v.as_slice::<1, f32>().await.unwrap(); + let p_values = p.as_slice::<1, f32>().await.unwrap(); + for i in 0..N { + assert!( + (m_values[[i]] - expected_m).abs() < 1e-5, + "replayed m update is wrong at {i}: got {}, expected {expected_m}", + m_values[[i]], + ); + assert!( + (v_values[[i]] - expected_v).abs() < 1e-5, + "replayed v update is wrong at {i}: got {}, expected {expected_v}", + v_values[[i]], + ); + assert!( + (p_values[[i]] - expected_p).abs() < 1e-4, + "replayed p update is wrong at {i}: got {}, expected {expected_p}", + p_values[[i]], + ); + } + }); +} + +#[test] +fn shared_eclass_observations_replay_as_aliases() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let records_before = device.flush_plan_cache().record_count(); + let replays_before = device.flush_plan_cache().replay_count(); + + for iteration in 0..3 { + let input_values = (0..64) + .map(|index| iteration as f32 + index as f32 * 0.25) + .collect::>(); + let input = Tensor::new::(&device, &input_values); + let left = &input * 2.0; + let right = &input * 2.0; + + device.flush(); + let left_values = left.as_slice::<1, f32>().await.unwrap(); + let right_values = right.as_slice::<1, f32>().await.unwrap(); + for (index, &input_value) in input_values.iter().enumerate() { + assert_eq!(left_values[[index]], input_value * 2.0); + assert_eq!(right_values[[index]], input_value * 2.0); + } + } + + assert!(device.flush_plan_cache().record_count() > records_before); + assert!( + device.flush_plan_cache().replay_count() > replays_before, + "a shared e-class observation was recorded as an executable view instead of an alias" + ); + }); +} + +/// Repeated ordinary materialization (the path used by `as_slice`) should +/// share the same reusable-plan lifecycle as batched flushes. Varying input +/// contents verifies that replay rebinds the current graph's buffers instead +/// of retaining data from the recording iteration. +#[test] +fn single_target_materialization_records_and_replays() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let records_before = device.flush_plan_cache().record_count(); + let replays_before = device.flush_plan_cache().replay_count(); + + for iteration in 0..4 { + let input_values = (0..64) + .map(|index| iteration as f32 + index as f32 * 0.25) + .collect::>(); + let input = Tensor::new::(&device, &input_values); + let output = (&input * 2.0f32) + 1.0f32; + let values = output.as_slice::<1, f32>().await.unwrap(); + for (index, &input_value) in input_values.iter().enumerate() { + assert_eq!(values[[index]], input_value * 2.0 + 1.0); + } + drop(values); + drop(output); + drop(input); + assert_eq!( + device.compute_graph().node_count(), + 0, + "single-target replay retained graph nodes after iteration {iteration}", + ); + } + + assert!( + device.flush_plan_cache().record_count() > records_before, + "single-target materialization did not record a reusable plan", + ); + assert!( + device.flush_plan_cache().replay_count() > replays_before, + "single-target materialization did not replay its recorded plan", + ); + + // The same target structure is not replay-safe while a live sibling + // outside its dependency closure still needs the input. Its liveness + // can change allocation claims without changing the target hash. + let input_values = (0..64).map(|index| index as f32).collect::>(); + let input = Tensor::new::(&device, &input_values); + let target = (&input * 2.0f32) + 1.0f32; + let sibling = (&input * 3.0f32) + 7.0f32; + let replays_before_sibling = device.flush_plan_cache().replay_count(); + let target_values = target.as_slice::<1, f32>().await.unwrap(); + assert_eq!( + device.flush_plan_cache().replay_count(), + replays_before_sibling, + "single-target replay ignored a live output outside its closure", + ); + drop(target_values); + let sibling_values = sibling.as_slice::<1, f32>().await.unwrap(); + for (index, &input_value) in input_values.iter().enumerate() { + assert_eq!(sibling_values[[index]], input_value * 3.0 + 7.0); + } + }); +} + +#[test] +fn quantized_materialization_records_and_replays() { + pollster::block_on(async { + const N: usize = 7; + const K: usize = 11; + + let Ok(device) = Device::new().await else { + return; + }; + let records_before = device.flush_plan_cache().record_count(); + let replays_before = device.flush_plan_cache().replay_count(); + + for iteration in 0..3 { + let input_values = (0..K) + .map(|index| iteration as f32 + index as f32 * 0.125) + .collect::>(); + let weight_values = (0..N * K) + .map(|index| (index as f32 + iteration as f32) * 0.03125) + .collect::>(); + let weight_bytes = weight_values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let weight = + crate::QMatrix::from_parts(&device, &weight_bytes, Box::new([N, K]), GgmlType::F32) + .unwrap(); + let input = Tensor::from_slice(&device, [1, K], &input_values); + let output = input.q_mat_mul(&weight); + let values = output.as_slice::<2, f32>().await.unwrap(); + for row in 0..N { + let expected = (0..K) + .map(|column| input_values[column] * weight_values[row * K + column]) + .sum::(); + assert!((values[[0, row]] - expected).abs() < 1e-4); + } + } + + assert!(device.flush_plan_cache().record_count() > records_before); + assert!(device.flush_plan_cache().replay_count() > replays_before); + }); +} + +#[test] +fn dequantization_replay_keeps_raw_and_materialized_bindings_distinct() { + pollster::block_on(async { + const ROWS: usize = 4; + let Ok(device) = Device::new().await else { + return; + }; + let records_before = device.flush_plan_cache().record_count(); + let replays_before = device.flush_plan_cache().replay_count(); + + for _ in 0..3 { + let raw = vec![0u8; ROWS * std::mem::size_of::()]; + let matrix = crate::QMatrix::from_parts( + &device, + &raw, + Box::new([ROWS, BlockQ4K::BLOCK_SIZE]), + GgmlType::Q4K, + ) + .unwrap(); + let dense = matrix.dequantize::(); + let values = dense.as_slice::<2, f32>().await.unwrap(); + assert!(values.as_slice().iter().all(|value| *value == 0.0)); + } + + assert!(device.flush_plan_cache().record_count() > records_before); + assert!(device.flush_plan_cache().replay_count() > replays_before); + }); +} + +/// Plain, encoder-tail, and epilogue-fused QMatMul materialization all cross +/// the unified resolver and must compute the same result. +#[test] +fn qmatmul_materialization_forms_match() { + pollster::block_on(async { + const N: usize = 4; + const K: usize = 8; + + let Ok(plain_device) = Device::new().await else { + return; + }; + let Ok(general_device) = Device::new().await else { + return; + }; + let Ok(tail_device) = Device::new().await else { + return; + }; + + let weights: Vec = (0..N * K).map(|i| 0.1 + i as f32 * 0.05).collect(); + let weight_bytes: Vec = weights + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let make_weight = |device: &Device| { + crate::QMatrix::from_parts( + device, + &weight_bytes, + vec![N, K].into_boxed_slice(), + GgmlType::F32, + ) + .unwrap() + }; + let input = [1.0f32, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0]; + + let plain = Tensor::new::(&plain_device, &[input]) + .q_mat_mul(&make_weight(&plain_device)); + let (_, plain_kernels) = plain.data.materialize(); + assert!(plain_kernels > 0, "plain QMatMul dispatched no kernels"); + + let direct_tail = + Tensor::new::(&tail_device, &[input]).q_mat_mul(&make_weight(&tail_device)); + let (_, tail_kernels, ()) = direct_tail.data.materialize_with_tail(|_, _| ()); + assert!( + tail_kernels > 0, + "encoder-tail QMatMul dispatched no kernels" + ); + + let qmatmul = Tensor::new::(&general_device, &[input]) + .q_mat_mul(&make_weight(&general_device)); + let general = &qmatmul + 0.0f32; + let (_, general_kernels, ()) = general.data.materialize_with_tail(|_, _| ()); + assert!( + general_kernels > 0, + "general-resolver QMatMul dispatched no kernels" + ); + + let plain_values = plain.as_slice::<2, f32>().await.unwrap(); + let tail_values = direct_tail.as_slice::<2, f32>().await.unwrap(); + let general_values = general.as_slice::<2, f32>().await.unwrap(); + for column in 0..N { + let expected = (0..K) + .map(|k| input[k] * weights[column * K + k]) + .sum::(); + let plain_value = plain_values[[0, column]]; + let tail_value = tail_values[[0, column]]; + let general_value = general_values[[0, column]]; + let tolerance = 1e-4 * expected.abs().max(1.0); + assert!( + (plain_value - expected).abs() <= tolerance, + "plain QMatMul mismatch at {column}: got {plain_value}, expected {expected}", + ); + assert!( + (tail_value - expected).abs() <= tolerance, + "encoder-tail QMatMul mismatch at {column}: got {tail_value}, expected {expected}", + ); + assert!( + (general_value - expected).abs() <= tolerance, + "general QMatMul mismatch at {column}: got {general_value}, expected {expected}", + ); + assert!( + (plain_value - tail_value).abs() <= tolerance + && (plain_value - general_value).abs() <= tolerance, + "QMatMul forms differ at {column}: plain {plain_value}, tail {tail_value}, epilogue {general_value}", + ); + } + }); +} diff --git a/fusor-ml/core/src/device.rs b/fusor-ml/core/src/device.rs index e16d9a089..71664ace5 100644 --- a/fusor-ml/core/src/device.rs +++ b/fusor-ml/core/src/device.rs @@ -1,11 +1,11 @@ use std::{ fmt::Debug, - sync::{Arc, OnceLock}, + sync::{Arc, Mutex, OnceLock}, }; use fusor_tile_ir::{CoopMatrixToken, SubgroupToken}; use fusor_tile_ir_kernels::SubgroupConfig; -use fusor_tile_ir_runtime::{BufferPool, KernelCache}; +use fusor_tile_ir_runtime::{BufferPool, FusorConfig, KernelCache}; use wgpu::{BackendOptions, Dx12BackendOptions}; use crate::{ @@ -19,6 +19,56 @@ use web_time::{Duration, Instant}; #[cfg(not(target_arch = "wasm32"))] const GPU_POLL_SPIN_BUDGET: Duration = Duration::from_millis(2); +/// Apple-silicon cooperative-matmul rates (see [`Device::matmul_rates`]). +/// `dram_decibytes_per_ns` is the measured M2 Max bandwidth roof (379.5 +/// GB/s). `mac_per_ns` is the measured simdgroup issue ceiling, 8.9 TF/s: +/// the 5.60 TF/s this used to carry was never a roof at all — the same +/// kernels sustain 7.57 TF/s on 16384x3072x1536 and 7.64 on the 4096-cube, +/// so the model was charging MMA issue above what the hardware bills and the +/// term crowded out every geometry difference it was supposed to arbitrate. +/// +/// The other two have no roof to read and are fitted together against 13 +/// contractions crossed with every legal tile — the eight merged microbench +/// cases (K from 64 to 16,384) plus five standalone `bench_coop_tiles` +/// contractions from 1024^3 to 16384x3072x1536, 65 measured spans. They are +/// what decides the geometry, so they are fitted for ranking, not for +/// absolute span: `workgroup_bytes_per_ns` sets what a wider tile buys and +/// `store_fs_per_element` what it costs, and the values below sit at the +/// centre of the box in which every one of the 13 picks the measured +/// fastest tile (`workgroup_bytes_per_ns` 650..800, `store_fs_per_element` +/// 2,000..4,200). Off the low side the K<=16 shapes take a 128-wide profile +/// and lose 8-22%; off the high side the K>=384 shapes take a 64-wide one +/// and lose 3-7%. +pub(crate) const APPLE_MATMUL_RATES: crate::occupancy::MatmulRates = + crate::occupancy::MatmulRates { + mac_per_ns: 4450, + dram_decibytes_per_ns: 3795, + workgroup_bytes_per_ns: 700, + store_fs_per_element: 4_000, + // Fitted against the direct measurement of the trade: the same + // merged body staged from one pair instead of two, at identical + // tile, split count and grid, runs 2048x64x64 -14.0%, 2048x64x256 + // -11.0%, 2048x256x64 +0.8% and 384x16384x1536 +7.2% (warm-resolve + // span medians, 20 interleaved processes per arm, controls flat). + // Those four shapes bracket the crossover at 4 < k_iterations < 16 + // against the residency credit in T3, which pins this to 103..108. + single_buffered_traffic_pct: 105, + }; + +/// Every other device class. The cost model is a ratio of these rates, so +/// what matters off Apple silicon is their proportion, not their scale: keep +/// the Apple ratios and scale the two absolute roofs down to a conservative +/// mid-range discrete part. Under-stating both roofs equally leaves every +/// comparison unchanged. +pub(crate) const OTHER_MATMUL_RATES: crate::occupancy::MatmulRates = + crate::occupancy::MatmulRates { + mac_per_ns: 2225, + dram_decibytes_per_ns: 1900, + workgroup_bytes_per_ns: 350, + store_fs_per_element: 8_000, + single_buffered_traffic_pct: 105, + }; + #[cfg(not(target_arch = "wasm32"))] fn poll_until_queue_empty(device: &wgpu::Device) -> Result { let start = Instant::now(); @@ -37,10 +87,9 @@ fn poll_until_queue_empty(device: &wgpu::Device) -> Result, ) -> Result { - let desired_adapter_name = std::env::var("WGPU_ADAPTER_NAME") - .ok() - .map(|name| name.to_ascii_lowercase()); + let desired_adapter_name = desired_adapter_name.map(|name| name.to_ascii_lowercase()); let mut adapters = instance.enumerate_adapters(backends).await; if let Some(desired_adapter_name) = desired_adapter_name { @@ -55,7 +104,7 @@ async fn select_adapter( }) .ok_or_else(|| { crate::Error::msg(format!( - "WGPU_ADAPTER_NAME={desired_adapter_name:?} did not match any available adapter" + "adapter name {desired_adapter_name:?} (WGPU_ADAPTER_NAME) did not match any available adapter" )) }); } @@ -136,8 +185,43 @@ fn install_device_diagnostics(device: &wgpu::Device) { })); } +/// One timed dispatch aggregate (per category or per kernel name) from a +/// profiled resolve. +#[derive(Clone, Debug, PartialEq)] +pub struct KernelProfileRow { + pub name: String, + pub count: usize, + pub total_ms: f64, + pub average_us: f64, + pub max_us: f64, +} + +/// GPU timestamp profile of one resolve, recorded when +/// [`FusorConfig::trace_gpu_kernels`] is set and drained with +/// [`Device::take_kernel_profiles`]. +#[derive(Clone, Debug, PartialEq)] +pub struct KernelProfile { + /// `"inside_pass"` when the adapter timestamps individual dispatches, + /// otherwise `"pass_boundary"`. + pub timestamp_mode: &'static str, + pub kernels: usize, + /// Dispatches the GPU did not sample. `accounted_ms` covers + /// `kernels - unmeasured_kernels` dispatches, never all of them silently. + pub unmeasured_kernels: usize, + pub accounted_ms: f64, + /// Wall span from the first sampled begin to the last sampled end, or `None` + /// when any dispatch went unmeasured and the span would not cover the resolve. + pub span_ms: Option, + pub timestamp_period_ns: f64, + /// Per-category aggregates, sorted by total time descending. + pub categories: Vec, + /// The most expensive kernel names, sorted by total time descending. + pub top_names: Vec, +} + struct DeviceInner { device: Arc, + config: Arc, adapter: wgpu::Adapter, /// Cached `adapter.get_info()` / `adapter.limits()`. These are constant for /// the device's lifetime; re-querying them per kernel build (every op, every @@ -148,6 +232,13 @@ struct DeviceInner { queue: Arc, kernel_cache: KernelCache, buffer_pool: BufferPool, + /// First-occurrence cache of recorded dense and quantized materialization + /// plans, replayed by `flush_all_pending`. Lives here beside the kernel + /// cache so it is reachable under the compute-graph write lock. + flush_plan_cache: crate::compute_graph::FlushPlanCache, + /// Structural fusion-plan decisions shared across resolves; templates are + /// matrix-free so nothing here pins buffers or cycles back to this inner. + fusion_plan_store: crate::compute_graph::FusionPlanStore, cooperative_matrix_caps: CooperativeMatrixCaps, compute_graph: OnceLock, /// When set, this device reports `subgroups_supported() == false` so kernel @@ -158,6 +249,30 @@ struct DeviceInner { /// pre-filled with a poison pattern instead of left zeroed, reproducing the /// app's reused buffer pool. A property of the device for the same reason. poison_allocations: bool, + /// GPU kernel profiles recorded by profiled resolves + /// (`FUSOR_TRACE_GPU_KERNELS`), drained by + /// [`Device::take_kernel_profiles`]. + kernel_profiles: Mutex>, + /// Memoized cooperative tile geometry per contraction shape. The scored + /// selection enumerates every table entry against every legal split count + /// (~1.5 us) and is asked five times per matmul per resolve, on a decode + /// path that builds hundreds of dispatches per token. Same class as the + /// cached adapter info above: a device-lifetime cache of a pure function + /// of device state. + coop_tile_memo: Mutex>>, +} + +/// Everything [`crate::matmul::cost::plan_coop_tile`] reads that varies per +/// call; every other input is device state, which the memo's owner pins. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) struct CoopTileKey { + pub(crate) m: u32, + pub(crate) k: u32, + pub(crate) n: u32, + pub(crate) batch: u32, + pub(crate) datatype: crate::DataTypeEnum, + pub(crate) has_epilogues: bool, + pub(crate) probe_group: u32, } impl Debug for DeviceInner { @@ -200,20 +315,26 @@ impl Device { let device = src.device.clone(); let queue = src.queue.clone(); let adapter = src.adapter.clone(); - let kernel_cache = KernelCache::new(device.clone(), &adapter); - let buffer_pool = BufferPool::new(device.clone(), queue.clone()); + let config = src.config.clone(); + let kernel_cache = KernelCache::new(device.clone(), &adapter, config.clone()); + let buffer_pool = BufferPool::new(device.clone(), queue.clone(), &config); let inner = Arc::new(DeviceInner { device, + config, adapter, adapter_info: src.adapter_info.clone(), limits: src.limits.clone(), queue, kernel_cache, buffer_pool, + flush_plan_cache: Default::default(), + fusion_plan_store: Default::default(), cooperative_matrix_caps: src.cooperative_matrix_caps, compute_graph: OnceLock::new(), disable_subgroups, poison_allocations, + kernel_profiles: Default::default(), + coop_tile_memo: Default::default(), }); let device = Device { inner: inner.clone(), @@ -246,6 +367,13 @@ impl Device { } pub async fn new() -> Result { + Self::new_with_config(FusorConfig::from_env()).await + } + + /// Construct a device with an explicit [`FusorConfig`] instead of reading + /// the process environment. + pub async fn new_with_config(config: FusorConfig) -> Result { + let config = Arc::new(config); let dx_compiler = wgpu::Dx12Compiler::from_env().unwrap_or_default(); let backends = wgpu::Backends::from_env().unwrap_or(wgpu::Backends::all()); let instance = wgpu::Instance::new(wgpu::InstanceDescriptor { @@ -259,7 +387,7 @@ impl Device { }, ..wgpu::InstanceDescriptor::new_without_display_handle() }); - let adapter = select_adapter(&instance, backends).await?; + let adapter = select_adapter(&instance, backends, config.adapter_name.as_deref()).await?; let adapter_features = adapter.features(); #[cfg(target_arch = "wasm32")] { @@ -281,7 +409,7 @@ impl Device { if adapter_features.contains(wgpu::Features::SHADER_F16) { required_features |= wgpu::Features::SHADER_F16; } - if std::env::var_os("FUSOR_TRACE_GPU_KERNELS").is_some() { + if config.trace_gpu_kernels { if adapter_features.contains(wgpu::Features::TIMESTAMP_QUERY) { required_features |= wgpu::Features::TIMESTAMP_QUERY; if adapter_features.contains(wgpu::Features::TIMESTAMP_QUERY_INSIDE_PASSES) { @@ -301,6 +429,11 @@ impl Device { // SAFETY: cooperative matrix is an experimental feature that requires opting in features = unsafe { wgpu::ExperimentalFeatures::enabled() }; } + if adapter_features.contains(wgpu::Features::EXPERIMENTAL_WORKGROUP_MEMORY_ALIAS) { + required_features |= wgpu::Features::EXPERIMENTAL_WORKGROUP_MEMORY_ALIAS; + // SAFETY: same experimental opt-in as cooperative matrix. + features = unsafe { wgpu::ExperimentalFeatures::enabled() }; + } features }; #[cfg(target_arch = "wasm32")] @@ -315,9 +448,7 @@ impl Device { required_features, &cooperative_matrix_properties, ); - if std::env::var_os("FUSOR_TRACE_GPU_KERNELS").is_some() - && !cooperative_matrix_properties.is_empty() - { + if config.trace_gpu_kernels && !cooperative_matrix_properties.is_empty() { tracing::info!( "Fusor cooperative matrix properties: {cooperative_matrix_properties:?}" ); @@ -338,31 +469,33 @@ impl Device { let device = Arc::new(device); let queue = Arc::new(queue); - let kernel_cache = KernelCache::new(device.clone(), &adapter); - let buffer_pool = BufferPool::new(device.clone(), queue.clone()); + let kernel_cache = KernelCache::new(device.clone(), &adapter, config.clone()); + let buffer_pool = BufferPool::new(device.clone(), queue.clone(), &config); - // `FUSOR_DISABLE_SUBGROUPS` / `FUSOR_DIRTY_BUFFERS` are construction-time - // defaults for the device flags, so a plain `Device::gpu()` from a repro - // binary reproduces the web path without code changes. Tests instead - // derive `without_subgroups()` / `with_poisoned_allocations()` sibling - // devices explicitly. `var_os` is always `None` on wasm. - let disable_subgroups = std::env::var_os("FUSOR_DISABLE_SUBGROUPS").is_some(); - let poison_allocations = std::env::var_os("FUSOR_DIRTY_BUFFERS").is_some(); + // Capability simulation is explicit through derived test devices; + // production always reflects the adapter's reported subgroup support. + let disable_subgroups = false; + let poison_allocations = false; let adapter_info = adapter.get_info(); let limits = adapter.limits(); let inner = Arc::new(DeviceInner { device, + config, adapter, adapter_info, limits, queue, kernel_cache, buffer_pool, + flush_plan_cache: Default::default(), + fusion_plan_store: Default::default(), cooperative_matrix_caps, compute_graph: OnceLock::new(), disable_subgroups, poison_allocations, + kernel_profiles: Default::default(), + coop_tile_memo: Default::default(), }); let device = Device { @@ -427,14 +560,8 @@ impl Device { /// below this size is treated as cache-resident — re-reading it costs no /// bandwidth — so the reuse-driven tilings (which trade thread-level /// parallelism for explicit reuse) only engage above it. wgpu exposes no - /// cache size, so this is a floor per device class; override with - /// `FUSOR_LAST_LEVEL_CACHE_BYTES` when tuning a specific part. - pub fn last_level_cache_bytes(&self) -> u64 { - if let Ok(value) = std::env::var("FUSOR_LAST_LEVEL_CACHE_BYTES") - && let Ok(parsed) = value.parse::() - { - return parsed; - } + /// cache size, so this is a conservative floor per device class. + pub(crate) fn last_level_cache_bytes(&self) -> u64 { let info = &self.inner.adapter_info; // Apple-silicon system-level cache starts at 8 MiB on the base M1 // and only grows with tier; other GPU classes floor lower (older @@ -446,15 +573,63 @@ impl Device { } } + /// Shader lanes that must be in flight before a dispatch policy may trade + /// thread-level parallelism for per-thread work (register tiling, wider + /// qgemv columns, skipping a fan-out split). wgpu exposes no core count, + /// so this is a conservative per-class floor: base-tier GPUs keep on the + /// order of 16K lanes resident and need ~4x oversubscription to hide + /// memory latency. A conservative under-estimate only makes policies keep + /// MORE parallelism, never less. + pub(crate) fn saturation_lanes(&self) -> u32 { + 64 << 10 + } + + /// The memoized `[bm, bn, bk, row_groups, col_groups]` geometry for one + /// contraction shape, `None` when the shape declines the coop family. + pub(crate) fn coop_tile_memo( + &self, + key: CoopTileKey, + plan: impl FnOnce() -> Option<[u32; 5]>, + ) -> Option<[u32; 5]> { + let mut memo = self + .inner + .coop_tile_memo + .lock() + .expect("coop tile memo poisoned"); + *memo.entry(key).or_insert_with(plan) + } + + /// Physical rates the cooperative-matmul cost model prices its terms in. + /// wgpu exposes no clock, no bandwidth and no core count, so these are + /// per-class values in the same spirit as [`Self::saturation_lanes`] — + /// but unlike a parallelism floor, a rate that is wrong in either + /// direction moves the argmin, so they are anchored on measured roofs + /// where a roof exists and on fitted spans where none does (see + /// [`APPLE_MATMUL_RATES`]). + pub(crate) fn matmul_rates(&self) -> crate::occupancy::MatmulRates { + let info = &self.inner.adapter_info; + if info.backend == wgpu::Backend::Metal && info.name.starts_with("Apple") { + APPLE_MATMUL_RATES + } else { + OTHER_MATMUL_RATES + } + } + + /// The dispatch-sizing policy derived from this device's capabilities. + /// Every "how many workgroups / how much work per thread" decision reads + /// from this one place instead of local constants. + pub(crate) fn dispatch_policy(&self) -> crate::occupancy::DispatchPolicy { + crate::occupancy::DispatchPolicy::from_device(self) + } + pub fn features(&self) -> wgpu::Features { self.inner.device.features() } pub fn subgroups_supported(&self) -> bool { - // A device constructed via `without_subgroups()` (or built with - // `FUSOR_DISABLE_SUBGROUPS` set) reports no subgroups, so kernel - // selection picks the no-subgroup fallbacks. Browser builds also take - // this path because they never request `wgpu::Features::SUBGROUP`. + // A test device constructed via `without_subgroups()` reports no + // subgroups, so selection picks the capability fallback. Browser + // builds also take this path because they do not request the feature. if self.inner.disable_subgroups { return false; } @@ -534,6 +709,15 @@ impl Device { Some(CoopMatrixToken::new_unchecked()) } + /// Proof that the backend can alias workgroup tiles into one byte arena + /// (the Metal workgroup-alias extension): mixed-stride tiles then pack + /// at byte offsets instead of separate typed allocations. + pub(crate) fn byte_arena_token(&self) -> Option { + self.features() + .contains(wgpu::Features::EXPERIMENTAL_WORKGROUP_MEMORY_ALIAS) + .then(fusor_tile_ir::ByteArenaToken::new_unchecked) + } + pub fn wgpu_adapter(&self) -> &wgpu::Adapter { &self.inner.adapter } @@ -564,11 +748,56 @@ impl Device { &self.inner.kernel_cache } + /// The process configuration this device was constructed with. + pub fn config(&self) -> &FusorConfig { + &self.inner.config + } + + /// Every GPU kernel profile recorded since the last call, oldest first. + /// One profile is recorded per profiled resolve when + /// [`FusorConfig::trace_gpu_kernels`] is set; replayed resolves skip + /// profiling. + pub fn take_kernel_profiles(&self) -> Vec { + std::mem::take(&mut self.inner.kernel_profiles.lock().unwrap()) + } + + #[cfg_attr(target_arch = "wasm32", allow(dead_code))] + pub(crate) fn record_kernel_profile(&self, profile: KernelProfile) { + self.inner.kernel_profiles.lock().unwrap().push(profile); + } + + pub(crate) fn flush_plan_cache(&self) -> &crate::compute_graph::FlushPlanCache { + &self.inner.flush_plan_cache + } + + pub(crate) fn fusion_plan_store(&self) -> &crate::compute_graph::FusionPlanStore { + &self.inner.fusion_plan_store + } + /// Reset the initialized flag on all cached buffers. pub fn reset_initialized_buffers(&self) { self.inner.buffer_pool.reset_initialized_buffers(); } + /// Snapshot the cumulative buffer-pool allocation counters (buffers + /// requested / buffers freshly created). Diff two snapshots to measure + /// allocations over a window. + pub fn buffer_pool_counters(&self) -> fusor_tile_ir_runtime::BufferPoolCounters { + self.inner.buffer_pool.counters() + } + + /// Whether the buffer pool holds its own tracked clone of `buffer` in the + /// `(size, usage)` bucket. Used by liveness accounting to enumerate the + /// pool as an expected strong-reference holder. + pub(crate) fn buffer_pool_is_tracked( + &self, + size: u64, + usage: wgpu::BufferUsages, + buffer: &Arc, + ) -> bool { + self.inner.buffer_pool.is_tracked(size, usage, buffer) + } + /// Get or create a buffer of the specified size. Poisoned first when this /// handle was built with [`Device::with_poisoned_allocations`]. pub fn create_buffer(&self, size: u64, usage: wgpu::BufferUsages) -> Arc { @@ -600,6 +829,17 @@ impl Device { .get() .expect("compute_graph should be initialized") } + + /// Resolve every pending lazy tensor now, submitting the work to the GPU + /// without waiting for it or downloading anything. Call at iteration + /// boundaries in training-style loops: it keeps the pending graph (and + /// per-resolve optimizer cost) bounded while leaving the GPU free to run + /// ahead of the host. + pub fn flush(&self) { + if let Some(graph) = self.inner.compute_graph.get() { + graph.flush(); + } + } } #[cfg(test)] @@ -661,6 +901,57 @@ mod dirty_buffer_tests { }); } + /// The pool tracks every buffer it hands out (holding its own strong + /// clone), reports it via `is_tracked` under the exact `(size, usage)` + /// key only, and the allocation counters distinguish fresh creations + /// from pool-cache hits. + #[test] + fn buffer_pool_tracking_and_counters() { + pollster::FutureExt::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + let usage = wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_SRC + | wgpu::BufferUsages::COPY_DST; + let size = 512u64; + + let before = device.buffer_pool_counters(); + let buffer = device.create_buffer(size, usage); + let after = device.buffer_pool_counters(); + assert_eq!(after.requested, before.requested + 1); + assert_eq!(after.created, before.created + 1); + + // Tracked under its exact (size, usage) key, and the pool's own + // clone means a freshly handed-out buffer has strong_count >= 2. + assert!(device.buffer_pool_is_tracked(size, usage, &buffer)); + assert!(Arc::strong_count(&buffer) >= 2); + // Not tracked under a different size or usage. + assert!(!device.buffer_pool_is_tracked(size * 2, usage, &buffer)); + assert!(!device.buffer_pool_is_tracked(size, wgpu::BufferUsages::STORAGE, &buffer)); + // A foreign buffer (same size/usage, allocated outside the pool) + // is not tracked. + let foreign = Arc::new(device.wgpu_device().create_buffer(&wgpu::BufferDescriptor { + label: Some("foreign"), + size, + usage, + mapped_at_creation: false, + })); + assert!(!device.buffer_pool_is_tracked(size, usage, &foreign)); + + // Dropping the handle frees the pooled buffer; the next request + // of the same shape is a cache hit, not a fresh creation. + drop(buffer); + let mid = device.buffer_pool_counters(); + let reused = device.create_buffer(size, usage); + let end = device.buffer_pool_counters(); + assert_eq!(end.requested, mid.requested + 1); + assert_eq!(end.created, mid.created); + assert!(device.buffer_pool_is_tracked(size, usage, &reused)); + }); + } + /// Positive control: a buffer handed out by a poisoned-allocation handle /// must read back as the poison byte, proving the poison actually lands on /// this backend (and is not silently zero-initialized away). diff --git a/fusor-ml/core/src/flash_attention.rs b/fusor-ml/core/src/flash_attention.rs new file mode 100644 index 000000000..d1d605607 --- /dev/null +++ b/fusor-ml/core/src/flash_attention.rs @@ -0,0 +1,737 @@ +//! Recognized attention-pattern execution: the fused cooperative-matrix +//! kernels from `fusor-tile-ir-kernels` behind one execution-graph operation. +//! +//! Each [`AttentionKernel`] kind is one recognizable computational pattern +//! over a scaled-masked score cluster — the fused softmax·v output, the row +//! log-sum-exp, and the three probability-contraction shapes. Route +//! selection lives in pattern recognition; decode shapes (one query row) +//! stay on the attention row program, and shapes no kernel can host run the +//! composed clusters unchanged. + +use std::hash::Hash; +use std::sync::Arc; + +use fusor_tile_ir as tile_ir; +use fusor_tile_ir_kernels as tile_ir_kernels; +use tile_ir_kernels::{ + FlashAttentionLayouts, FlashAttentionShape, FlashBwdLayouts, FlashKvOutputs, FlashMaskLayout, + FlashOperandLayout, FlashRowLayout, +}; + +use crate::{ + DataTypeEnum, Device, TensorData, + compute_graph::NodeIndex, + kernel_selection::CooperativeMatrixKind, + mir::{ + inputs::MirValue, + kernel_backend, + kernel_backend::DirectKernel, + operation::Operation, + workgroup_shape::{Constraint, WorkgroupShape, WorkgroupShapeConstraints}, + }, + row_program::AttentionInputs, +}; + +struct FlashAttentionDirectKernelVariant; + +/// Which fused kernel a recognized attention-pattern cluster lowers to. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum AttentionKernel { + /// `softmax(scale·q·kᵀ [+ mask]) · v`. + Output, + /// Row log-sum-exp of the scores: `m + ln Σ exp(s − m)`. + LogSumExp, + /// `(p ∘ (grad_o·vᵀ − dsum) · scale) · k` with `p = exp(s − lse)`. + GradQ, + /// `(p ∘ (grad_o·vᵀ − dsum) · scale)ᵀ · q`. + GradK, + /// `pᵀ · grad_o`. + GradV, + /// Both KV-side contractions in one dispatch, landing in a tensor whose + /// sequence axis spans `2·kv_len` (dk rows first, dv rows after) — the + /// probability recomputation is shared between them. + GradKV, +} + +/// One recognized attention-pattern cluster bound for a fused kernel. +/// Operand layouts are read from the resolved inputs at kernel-build time, +/// so strided views (transposes, offsets) execute without materialization. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct FlashAttentionOperation { + pub(crate) kind: AttentionKernel, + pub(crate) q: NodeIndex, + pub(crate) k: NodeIndex, + pub(crate) v: Option, + pub(crate) grad_o: Option, + pub(crate) lse: Option, + pub(crate) dsum: Option, + pub(crate) mask: Option, + batch: u32, + heads: u32, + kv_heads: u32, + q_len: u32, + kv_len: u32, + head_dim: u32, + scale: f32, + causal: bool, + datatype: DataTypeEnum, +} + +/// The node set of one matched pattern, in dependency order. +pub(crate) struct AttentionPatternNodes { + pub(crate) q: NodeIndex, + pub(crate) k: NodeIndex, + pub(crate) v: Option, + pub(crate) grad_o: Option, + pub(crate) lse: Option, + pub(crate) dsum: Option, + pub(crate) mask: Option, +} + +impl FlashAttentionOperation { + /// The fused-output route for a recognized forward cluster; `None` + /// falls back to the caller's other attention routes. + pub(crate) fn try_new_output(device: &Device, inputs: &AttentionInputs<'_>) -> Option { + let [batch, heads, q_len, head_dim] = *inputs.q_shape else { + return None; + }; + let [kv_batch, kv_heads, kv_len, kv_dim] = *inputs.k_shape else { + return None; + }; + if inputs.k_shape != inputs.v_shape || kv_batch != batch || kv_dim != head_dim { + return None; + } + match (inputs.mask, inputs.mask_shape, inputs.causal) { + (Some(_), Some(mask_shape), false) if mask_shape == [q_len, kv_len] => {} + (None, None, _) => {} + _ => return None, + } + Self::try_new( + device, + AttentionKernel::Output, + AttentionPatternNodes { + q: inputs.q, + k: inputs.k, + v: Some(inputs.v), + grad_o: None, + lse: None, + dsum: None, + mask: inputs.mask, + }, + [batch, heads, kv_heads, q_len, kv_len, head_dim], + inputs.scale, + inputs.causal, + inputs.input_dtype, + ) + } + + /// Build the operation when the device and shape qualify for `kind`; + /// `None` leaves the composed cluster in place. + #[allow(clippy::too_many_arguments)] + pub(crate) fn try_new( + device: &Device, + kind: AttentionKernel, + nodes: AttentionPatternNodes, + dims: [usize; 6], + scale: f32, + causal: bool, + datatype: DataTypeEnum, + ) -> Option { + if !matches!(datatype, DataTypeEnum::F32 | DataTypeEnum::F16) { + return None; + } + let [batch, heads, kv_heads, q_len, kv_len, head_dim] = dims; + if kv_heads == 0 || heads % kv_heads != 0 { + return None; + } + let operation = Self { + kind, + q: nodes.q, + k: nodes.k, + v: nodes.v, + grad_o: nodes.grad_o, + lse: nodes.lse, + dsum: nodes.dsum, + mask: nodes.mask, + batch: batch.try_into().ok()?, + heads: heads.try_into().ok()?, + kv_heads: kv_heads.try_into().ok()?, + q_len: q_len.try_into().ok()?, + kv_len: kv_len.try_into().ok()?, + head_dim: head_dim.try_into().ok()?, + scale, + causal, + datatype, + }; + if causal && nodes.mask.is_some() { + return None; + } + device.coop_token(CooperativeMatrixKind::F32F32M8N8K8)?; + let subgroups = device.subgroup_config()?; + // The staged tiles must fit the device's workgroup-storage limit + // (16 KB WebGPU default, 32 KB Apple) or pipeline creation fails. + // The analytic bound models the forward kernel's staging + // (`tests/footprint.rs` pins it to the lowered IR); per-kind + // backward footprints are not modelled yet. + let stage = match datatype { + DataTypeEnum::F16 => tile_ir::ScalarElement::F16, + _ => tile_ir::ScalarElement::F32, + }; + if tile_ir_kernels::flash_attention_workgroup_bytes(operation.head_dim, stage) + > u64::from(device.limits().max_compute_workgroup_storage_size) + { + return None; + } + let shape = operation.flash_shape(); + let supported = match kind { + AttentionKernel::Output => { + tile_ir_kernels::flash_attention_supported(&shape, subgroups) + } + AttentionKernel::LogSumExp => { + tile_ir_kernels::flash_attention_bwd_supported(&shape, subgroups) + } + // The streaming contractions recompute probabilities per tile; + // without the causal break they currently lose to the composed + // matmuls, so masked/unmasked shapes keep the composed cluster. + AttentionKernel::GradQ + | AttentionKernel::GradK + | AttentionKernel::GradV + | AttentionKernel::GradKV => { + causal && tile_ir_kernels::flash_attention_bwd_supported(&shape, subgroups) + } + }; + supported.then_some(operation) + } + + fn flash_shape(&self) -> FlashAttentionShape { + FlashAttentionShape { + batch: self.batch, + heads: self.heads, + kv_groups: self.heads / self.kv_heads, + q_len: self.q_len, + kv_len: self.kv_len, + head_dim: self.head_dim, + scale: self.scale, + causal: self.causal, + } + } + + fn out_shape(&self) -> Vec { + let (batch, heads) = (self.batch as usize, self.heads as usize); + let (q_len, kv_len, d) = ( + self.q_len as usize, + self.kv_len as usize, + self.head_dim as usize, + ); + match self.kind { + AttentionKernel::Output | AttentionKernel::GradQ => vec![batch, heads, q_len, d], + AttentionKernel::LogSumExp => vec![batch, heads, q_len], + AttentionKernel::GradK | AttentionKernel::GradV => vec![batch, heads, kv_len, d], + AttentionKernel::GradKV => vec![batch, heads, 2 * kv_len, d], + } + } +} + +/// Element strides of a resolved rank-4 tensor. +fn operand_layout(layout: &crate::Layout) -> Option { + let [batch_stride, head_stride, seq_stride, dim_stride] = *layout.strides() else { + return None; + }; + Some(FlashOperandLayout { + offset: layout.offset().try_into().ok()?, + batch_stride: batch_stride.try_into().ok()?, + head_stride: head_stride.try_into().ok()?, + seq_stride: seq_stride.try_into().ok()?, + dim_stride: dim_stride.try_into().ok()?, + }) +} + +/// Element strides of a resolved rank-3 row statistic. +fn row_layout(layout: &crate::Layout) -> Option { + let [batch_stride, head_stride, seq_stride] = *layout.strides() else { + return None; + }; + Some(FlashRowLayout { + offset: layout.offset().try_into().ok()?, + batch_stride: batch_stride.try_into().ok()?, + head_stride: head_stride.try_into().ok()?, + seq_stride: seq_stride.try_into().ok()?, + }) +} + +/// Element strides of the resolved rank-2 mask. +fn mask_layout(layout: &crate::Layout) -> Option { + let [q_stride, kv_stride] = *layout.strides() else { + return None; + }; + Some(FlashMaskLayout { + offset: layout.offset().try_into().ok()?, + q_stride: q_stride.try_into().ok()?, + kv_stride: kv_stride.try_into().ok()?, + }) +} + +/// A whole-buffer linear view: the kernels compute element indices +/// themselves from the baked operand strides. +fn linear_ref(buffer: Arc) -> Option>> { + let elements: u32 = (buffer.size() / size_of::() as u64).try_into().ok()?; + let layout = tile_ir::Layout::strided( + tile_ir::MemoryLevel::Storage, + tile_ir::Shape::new([elements]), + &[1], + ); + Some(tile_ir::KernelTensorRef::new(buffer, layout)) +} + +struct ResolvedTensors<'a> { + q: &'a TensorData, + k: &'a TensorData, + v: Option<&'a TensorData>, + grad_o: Option<&'a TensorData>, + lse: Option<&'a TensorData>, + dsum: Option<&'a TensorData>, + mask: Option<&'a TensorData>, + output: &'a TensorData, +} + +impl FlashAttentionOperation { + fn split_inputs<'a>(&self, inputs: &'a [MirValue]) -> Option> { + let mut iter = inputs.iter(); + let mut next = || iter.next().and_then(MirValue::as_tensor); + let q = next()?; + let k = next()?; + let v = if self.v.is_some() { + Some(next()?) + } else { + None + }; + let grad_o = if self.grad_o.is_some() { + Some(next()?) + } else { + None + }; + let lse = if self.lse.is_some() { + Some(next()?) + } else { + None + }; + let dsum = if self.dsum.is_some() { + Some(next()?) + } else { + None + }; + let mask = if self.mask.is_some() { + Some(next()?) + } else { + None + }; + let output = next()?; + Some(ResolvedTensors { + q, + k, + v, + grad_o, + lse, + dsum, + mask, + output, + }) + } +} + +impl Operation for FlashAttentionOperation { + fn hash_kernel_fields(&self, state: &mut rustc_hash::FxHasher) { + self.kind.hash(state); + self.batch.hash(state); + self.heads.hash(state); + self.kv_heads.hash(state); + self.q_len.hash(state); + self.kv_len.hash(state); + self.head_dim.hash(state); + self.scale.to_bits().hash(state); + self.causal.hash(state); + self.v.is_some().hash(state); + self.grad_o.is_some().hash(state); + self.lse.is_some().hash(state); + self.dsum.is_some().hash(state); + self.mask.is_some().hash(state); + self.datatype.hash(state); + } + + fn workgroup_shape_constraints(&self, device: &Device) -> WorkgroupShapeConstraints { + let block = device + .subgroup_config() + .map(|subgroups| subgroups.block_for_subgroups(4)) + .unwrap_or(128); + let mut constraints = WorkgroupShapeConstraints::new(); + constraints.add_constraint(0, Constraint::equals(block)); + constraints.add_constraint(1, Constraint::equals(1)); + constraints.add_constraint(2, Constraint::equals(1)); + constraints + } + + fn dispatch_size(&self, _workgroup_shape: &WorkgroupShape, inputs: &[MirValue]) -> [u32; 3] { + let max_per_dim = inputs + .last() + .and_then(MirValue::as_tensor) + .expect("attention output must be a tensor") + .device() + .limits() + .max_compute_workgroups_per_dimension; + let shape = self.flash_shape(); + match self.kind { + AttentionKernel::Output => { + tile_ir_kernels::flash_attention_dispatch(&shape, max_per_dim) + } + AttentionKernel::LogSumExp => tile_ir_kernels::flash_lse_dispatch(&shape, max_per_dim), + AttentionKernel::GradQ => tile_ir_kernels::flash_bwd_q_dispatch(&shape, max_per_dim), + AttentionKernel::GradK | AttentionKernel::GradV | AttentionKernel::GradKV => { + tile_ir_kernels::flash_bwd_kv_dispatch(&shape, max_per_dim) + } + } + } + + fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { + f(self.q); + f(self.k); + for node in [self.v, self.grad_o, self.lse, self.dsum, self.mask] + .into_iter() + .flatten() + { + f(node); + } + } + + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + f(&mut self.q); + f(&mut self.k); + for node in [ + &mut self.v, + &mut self.grad_o, + &mut self.lse, + &mut self.dsum, + &mut self.mask, + ] + .into_iter() + .flatten() + { + f(node); + } + } + + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { + let resolved = |node: NodeIndex| { + nodes + .get_result(node) + .expect("attention inputs must be resolved before kernel launch") + }; + let device = nodes.device(); + let output = TensorData::new_for_shape(&device, &self.out_shape(), self.datatype); + let mut inputs = vec![resolved(self.q).into(), resolved(self.k).into()]; + for node in [self.v, self.grad_o, self.lse, self.dsum, self.mask] + .into_iter() + .flatten() + { + inputs.push(resolved(node).into()); + } + inputs.push(output.into()); + inputs + } + + fn output( + &self, + _nodes: &crate::compute_graph::ComputeGraphInner, + inputs: &[MirValue], + ) -> MirValue { + inputs.last().unwrap().as_tensor().unwrap().clone().into() + } + + fn build_direct_kernel( + &self, + graph: &crate::compute_graph::ComputeGraphInner, + workgroup_shape: &WorkgroupShape, + inputs: &[MirValue], + ) -> Option { + let tensors = self.split_inputs(inputs)?; + let device = graph.device(); + let subgroups = device.subgroup_config()?; + let coop = device.coop_token(CooperativeMatrixKind::F32F32M8N8K8)?; + let byte_arena = device.byte_arena_token(); + let max_per_dim = device.limits().max_compute_workgroups_per_dimension; + let shape = self.flash_shape(); + let dispatch_size = self.dispatch_size(workgroup_shape, inputs); + let cache_key = self.kernel_cache_key_with_dispatch( + kernel_backend::KernelVariantKey::of::(), + Some(workgroup_shape), + dispatch_size, + inputs, + ); + let kind = self.kind; + let datatype = self.datatype; + + let buffer = |data: &TensorData| data.buffer().clone(); + let q_buffer = buffer(tensors.q); + let k_buffer = buffer(tensors.k); + let lq = operand_layout(tensors.q.layout())?; + let lk = operand_layout(tensors.k.layout())?; + let mask_parts = match tensors.mask { + Some(mask) => Some((buffer(mask), mask_layout(mask.layout())?)), + None => None, + }; + let output_buffer = buffer(tensors.output); + + match kind { + AttentionKernel::Output => { + let v = tensors.v?; + let layouts = FlashAttentionLayouts { + q: lq, + k: lk, + v: operand_layout(v.layout())?, + o: operand_layout(tensors.output.layout())?, + }; + let v_buffer = buffer(v); + kernel_backend::run_kernel( + device.kernel_cache(), + self.name(), + cache_key, + dispatch_size, + move |kb| { + if let Some(token) = byte_arena { + kb.program().enable_byte_arena(token); + } + let f32e = match datatype { + DataTypeEnum::F16 => tile_ir::ScalarElement::F16.element(), + _ => tile_ir::ScalarElement::F32.element(), + }; + let q = kb.read(f32e, linear_ref(q_buffer)?); + let k = kb.read(f32e, linear_ref(k_buffer)?); + let v = kb.read(f32e, linear_ref(v_buffer)?); + let mask = match mask_parts { + Some((buffer, layout)) => { + Some((kb.read(f32e, linear_ref(buffer)?), layout)) + } + None => None, + }; + let o = kb.write(f32e, linear_ref(output_buffer)?); + tile_ir_kernels::flash_attention_f32( + kb.program(), + &q, + &k, + &v, + mask.as_ref().map(|(storage, layout)| (storage, *layout)), + &o, + &layouts, + shape, + subgroups, + coop, + max_per_dim, + ) + .then_some(()) + }, + ) + } + AttentionKernel::LogSumExp => { + let lse_layout = row_layout(tensors.output.layout())?; + kernel_backend::run_kernel( + device.kernel_cache(), + self.name(), + cache_key, + dispatch_size, + move |kb| { + if let Some(token) = byte_arena { + kb.program().enable_byte_arena(token); + } + let f32e = match datatype { + DataTypeEnum::F16 => tile_ir::ScalarElement::F16.element(), + _ => tile_ir::ScalarElement::F32.element(), + }; + let q = kb.read(f32e, linear_ref(q_buffer)?); + let k = kb.read(f32e, linear_ref(k_buffer)?); + let mask = match mask_parts { + Some((buffer, layout)) => { + Some((kb.read(f32e, linear_ref(buffer)?), layout)) + } + None => None, + }; + let lse = kb.write(f32e, linear_ref(output_buffer)?); + tile_ir_kernels::flash_lse_f32( + kb.program(), + &q, + &k, + mask.as_ref().map(|(storage, layout)| (storage, *layout)), + &lse, + lq, + lk, + lse_layout, + shape, + subgroups, + coop, + max_per_dim, + ) + .then_some(()) + }, + ) + } + AttentionKernel::GradQ + | AttentionKernel::GradK + | AttentionKernel::GradV + | AttentionKernel::GradKV => { + let grad_o = tensors.grad_o?; + let lse = tensors.lse?; + let needs_dsum = kind != AttentionKernel::GradV; + // The combined output's sequence axis spans 2*kv_len; its + // layout already reflects that from the allocation. + let placeholder = + FlashOperandLayout::contiguous(self.kv_heads, self.kv_len, self.head_dim); + let layouts = FlashBwdLayouts { + q: lq, + k: lk, + v: match tensors.v { + Some(v) => operand_layout(v.layout())?, + None => placeholder, + }, + grad_o: operand_layout(grad_o.layout())?, + lse: row_layout(lse.layout())?, + dsum: match tensors.dsum { + Some(dsum) => row_layout(dsum.layout())?, + None => FlashRowLayout { + offset: 0, + batch_stride: 0, + head_stride: 0, + seq_stride: 0, + }, + }, + out: operand_layout(tensors.output.layout())?, + }; + if needs_dsum && (tensors.v.is_none() || tensors.dsum.is_none()) { + return None; + } + let v_buffer = tensors.v.map(buffer); + let grad_o_buffer = buffer(grad_o); + let lse_buffer = buffer(lse); + let dsum_buffer = tensors.dsum.map(buffer); + kernel_backend::run_kernel( + device.kernel_cache(), + self.name(), + cache_key, + dispatch_size, + move |kb| { + if let Some(token) = byte_arena { + kb.program().enable_byte_arena(token); + } + let f32e = match datatype { + DataTypeEnum::F16 => tile_ir::ScalarElement::F16.element(), + _ => tile_ir::ScalarElement::F32.element(), + }; + let q = kb.read(f32e, linear_ref(q_buffer)?); + let k = kb.read(f32e, linear_ref(k_buffer)?); + let v = match v_buffer { + Some(buffer) => Some(kb.read(f32e, linear_ref(buffer)?)), + None => None, + }; + let grad_o = kb.read(f32e, linear_ref(grad_o_buffer)?); + let lse = kb.read(f32e, linear_ref(lse_buffer)?); + let dsum = match dsum_buffer { + Some(buffer) => Some(kb.read(f32e, linear_ref(buffer)?)), + None => None, + }; + let mask = match mask_parts { + Some((buffer, layout)) => { + Some((kb.read(f32e, linear_ref(buffer)?), layout)) + } + None => None, + }; + let out = kb.write(f32e, linear_ref(output_buffer)?); + let mask = mask.as_ref().map(|(storage, layout)| (storage, *layout)); + match kind { + AttentionKernel::GradQ => tile_ir_kernels::flash_bwd_q_f32( + kb.program(), + &q, + &k, + v.as_ref().expect("grad_q requires values"), + &grad_o, + &lse, + dsum.as_ref().expect("grad_q requires dsum"), + mask, + &out, + &layouts, + shape, + subgroups, + coop, + max_per_dim, + ), + AttentionKernel::GradK => tile_ir_kernels::flash_bwd_kv_f32( + kb.program(), + &q, + &k, + v.as_ref(), + &grad_o, + &lse, + dsum.as_ref(), + mask, + &out, + &layouts, + FlashKvOutputs::Dk, + shape, + subgroups, + coop, + max_per_dim, + ), + AttentionKernel::GradKV => tile_ir_kernels::flash_bwd_kv_f32( + kb.program(), + &q, + &k, + v.as_ref(), + &grad_o, + &lse, + dsum.as_ref(), + mask, + &out, + &layouts, + FlashKvOutputs::Both, + shape, + subgroups, + coop, + max_per_dim, + ), + AttentionKernel::GradV => tile_ir_kernels::flash_bwd_kv_f32( + kb.program(), + &q, + &k, + None, + &grad_o, + &lse, + None, + mask, + &out, + &layouts, + FlashKvOutputs::Dv, + shape, + subgroups, + coop, + max_per_dim, + ), + _ => unreachable!(), + } + .then_some(()) + }, + ) + } + } + } + + fn name(&self) -> String { + let kind = match self.kind { + AttentionKernel::Output => "out", + AttentionKernel::LogSumExp => "lse", + AttentionKernel::GradQ => "dq", + AttentionKernel::GradK => "dk", + AttentionKernel::GradV => "dv", + AttentionKernel::GradKV => "dkv", + }; + format!( + "flash_attention_{kind}_{}x{}x{}x{}x{}", + self.batch, self.heads, self.q_len, self.kv_len, self.head_dim + ) + } +} diff --git a/fusor-ml/core/src/fold.rs b/fusor-ml/core/src/fold.rs new file mode 100644 index 000000000..866fff7af --- /dev/null +++ b/fusor-ml/core/src/fold.rs @@ -0,0 +1,1067 @@ +//! Folds with an explicit loop-carried tuple and an explicit combine. +//! +//! [`ReduceOperation`] is a fold whose carrier is one value and whose combine +//! is one of four built-in operators. That closure is what forces every +//! interesting blocked reduction in this compiler to be hand-written below the +//! IR: split-K lives in the matmul kernel builder, online softmax lives in a +//! `RowStep` matcher, and Welford does not exist at all. They are three +//! instances of one law. +//! +//! A [`FoldOperation`] names the carrier and the combine, so the law becomes a +//! rewrite: +//! +//! > a fold whose `combine` is associative may be split along its axis into +//! > partial folds joined by `combine`. +//! +//! The element expression stays an ordinary [`NaryExpr`], so producer inlining +//! keeps working on it exactly as it does for a reduce. Only the carrier +//! algebra needs a term of its own, and it is closed and small — the same +//! split this codebase already makes between a fusable `expression` and a +//! restricted `post_element_wise` chain. + +use crate::DataTypeEnum; +use crate::compute_graph::NodeIndex; +use crate::nary_wise::{NaryExpr, NaryFunction, NaryOp, NaryScalar, UnaryFunctionChain}; +use crate::reduce::{ReduceFunction, ReduceOp, ReduceOperation}; + +/// One slot of a fold's loop-carried state. +/// +/// A slot is a scalar per row unless it declares a `free_dim`, in which case +/// it carries one value per position of a dimension appended to the fold's +/// output shape. Attention's `Σ p·v` is that: a vector over the head dim +/// whose step reads a *different* tensor (`v`) than the scalar slots do, at +/// the axis coordinate extended by the free coordinate. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CarrierSlot { + pub(crate) name: Option, + pub(crate) datatype: DataTypeEnum, + /// Extent of the dimension this slot ranges over, appended to the output + /// shape. `None` for a scalar slot. + pub(crate) free_dim: Option, + /// A per-element value private to this slot, evaluated at the axis + /// coordinate extended by this slot's free coordinate (`DimIndex(rank)`). + /// Bound in `step` immediately after the shared element — see + /// [`slot_element`]. + pub(crate) element: Option, +} + +impl CarrierSlot { + pub(crate) fn new(name: &str, datatype: DataTypeEnum) -> Self { + Self { + name: Some(name.to_string()), + datatype, + free_dim: None, + element: None, + } + } + + /// A slot ranging over `free_dim`, whose step absorbs `element` read at + /// the axis coordinate extended by the free coordinate. + pub(crate) fn vector( + name: &str, + datatype: DataTypeEnum, + free_dim: usize, + element: NaryExpr, + ) -> Self { + Self { + name: Some(name.to_string()), + datatype, + free_dim: Some(free_dim), + element: Some(element), + } + } + + /// How many values this slot carries per row. + pub(crate) fn width(&self) -> usize { + self.free_dim.unwrap_or(1) + } +} + +/// A private inner axis folded into the element before the element expression +/// runs: attention's `q·k` dot over the head dim. +/// +/// `expression` is evaluated `len` times with the fold coordinate appended to +/// the element's coordinates (`DimIndex(rank)`) and combined with `function`. +/// The folded value is what [`fold_value`] binds inside +/// [`FoldOperation::expression`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ElementFold { + pub(crate) expression: NaryExpr, + pub(crate) len: usize, + pub(crate) function: ReduceFunction, +} + +/// Carrier bodies are ordinary [`NaryExpr`]s. +/// +/// A slot read is `IndexedInput { input_idx: base + k, indices: vec![] }` — +/// the same convention `row_program::slot_expr` already uses for cross-step +/// references, so one expression language serves tensors and carriers alike +/// and every existing pass over `NaryExpr` works on fold bodies unchanged. +/// +/// Binding layout, where `base` is the number of tensor inputs and `n` the +/// carrier width: +/// +/// - `step` sees `[acc_0 .. acc_{n-1}, element]` +/// - `combine` sees `[acc_0 .. acc_{n-1}, rhs_0 .. rhs_{n-1}]` +/// - `init` and `outputs` see `[acc_0 .. acc_{n-1}]` (init reads none of them) +pub(crate) fn slot(base: usize, index: usize) -> NaryExpr { + NaryExpr::IndexedInput { + input_idx: base + index, + indices: Vec::new(), + } +} + +/// Accumulator slot `k`. +pub(crate) fn acc(base: usize, k: usize) -> NaryExpr { + slot(base, k) +} + +/// The element value, valid in `step` only. +pub(crate) fn element(base: usize, width: usize) -> NaryExpr { + slot(base, width) +} + +/// A slot's own per-element value ([`CarrierSlot::element`]), valid in that +/// slot's `step` only. +pub(crate) fn slot_element(base: usize, width: usize) -> NaryExpr { + slot(base, width + 1) +} + +/// The inner fold's value ([`ElementFold`]), valid in +/// [`FoldOperation::expression`] only — it is the element's own private +/// binding space, not the carrier's. +pub(crate) fn fold_value(base: usize) -> NaryExpr { + slot(base, 0) +} + +/// Incoming carrier slot `k`, valid in `combine` only. +pub(crate) fn rhs(base: usize, width: usize, k: usize) -> NaryExpr { + slot(base, width + k) +} + +fn binary(op: NaryOp, lhs: NaryExpr, rhs: NaryExpr, datatype: DataTypeEnum) -> NaryExpr { + NaryExpr::Op { + children: vec![lhs, rhs], + function: NaryFunction::binary(None, op, datatype, datatype, datatype), + } +} + +fn unary(op: NaryOp, value: NaryExpr, datatype: DataTypeEnum) -> NaryExpr { + NaryExpr::Op { + children: vec![value], + function: NaryFunction::unary(None, op, datatype, datatype), + } +} + +/// Whether `expression` reads any slot at or above `index`. +fn reads_slot_from(expression: &NaryExpr, base: usize, index: usize) -> bool { + match expression { + NaryExpr::IndexedInput { input_idx, indices } => { + indices.is_empty() && *input_idx >= base + index + } + NaryExpr::Op { children, .. } => children + .iter() + .any(|child| reads_slot_from(child, base, index)), + _ => false, + } +} + + +/// One finalized output of a fold. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct FoldOutput { + pub(crate) expression: NaryExpr, + pub(crate) datatype: DataTypeEnum, +} + +/// What is known about `combine`, which decides which rewrites may fire. +/// +/// Recorded rather than inferred because for floating point associativity is a +/// *policy*: exact reassociation (max) and error-introducing reassociation +/// (sum) are different permissions, and the online-softmax lift is neither — +/// it is exp's shift-invariance, which introduces fresh correction factors. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum FoldAlgebra { + /// Reassociating changes nothing bit-for-bit (min/max absent NaN). + ExactMonoid, + /// Associative in the reals; reassociating perturbs rounding. + ApproximateMonoid, + /// Associative in the reals, but combining evaluates transcendentals at + /// new arguments. Online softmax lives here. + RescalingMonoid, + /// No associativity claim; the split law must not fire. + Unspecified, +} + +impl FoldAlgebra { + pub(crate) fn splittable_under(self, policy: NumericsPolicy) -> bool { + match self { + FoldAlgebra::ExactMonoid => true, + FoldAlgebra::ApproximateMonoid => policy >= NumericsPolicy::ReassociationPermitted, + FoldAlgebra::RescalingMonoid => policy >= NumericsPolicy::RelativeErrorPermitted, + FoldAlgebra::Unspecified => false, + } + } +} + +/// How much numerical licence a rewrite may take. Ordered by permissiveness. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub(crate) enum NumericsPolicy { + Exact, + ReassociationPermitted, + RelativeErrorPermitted, +} + +/// A fold over one axis with a named carrier and an explicit combine. +/// +/// `expression` is evaluated at every coordinate of `shape` (including `axis`) +/// exactly as in [`ReduceOperation`]; `step` absorbs one such element into the +/// carrier; `combine` joins two carriers; `outputs` finalize. +/// +/// Unlike a row program's phases — which are sequential *independent* +/// reductions, each seeing the previous one's completed value — a fold's slots +/// update **jointly in one pass**: `step[i]` reads the whole running carrier. +/// That is what online softmax needs and what phases cannot express. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct FoldOperation { + pub(crate) inputs: Vec, + pub(crate) expression: NaryExpr, + /// A private inner axis folded into the element before `expression` runs; + /// its value is [`fold_value`]. + pub(crate) element_fold: Option, + pub(crate) shape: Box<[usize]>, + pub(crate) axis: usize, + pub(crate) carrier: Vec, + pub(crate) init: Vec, + pub(crate) step: Vec, + pub(crate) combine: Vec, + pub(crate) outputs: Vec, + pub(crate) algebra: FoldAlgebra, + pub(crate) block: Option, + /// This fold's element is itself a carrier: it folds the partial carriers + /// a blocked fold produced, so its `step` is the original `combine`. + pub(crate) element_is_carrier: bool, +} + +impl FoldOperation { + /// Slot base: carrier slots are numbered after the tensor inputs. + pub(crate) fn base(&self) -> usize { + self.inputs.len() + } + + pub(crate) fn width(&self) -> usize { + self.carrier.len() + } + + pub(crate) fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { + for &input in &self.inputs { + f(input); + } + } + + pub(crate) fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + for input in &mut self.inputs { + f(input); + } + } + + /// Shape of the row grid: the input shape with the folded axis removed. + /// An output that reads a free-dimension slot appends that slot's extent + /// ([`Self::output_shape`]). + pub(crate) fn out_shape(&self) -> Vec { + let mut shape = self.shape.to_vec(); + shape.remove(self.axis); + shape + } + + /// The free dimension an output ranges over: the extent of any + /// free-dimension slot it reads. Two different ones would make the output + /// rank ambiguous, so [`Self::validate`] rejects that. + pub(crate) fn output_free_dim(&self, output: &FoldOutput) -> Option { + let base = self.base(); + self.carrier + .iter() + .enumerate() + .find(|(index, slot)| { + slot.free_dim.is_some() && output.expression.uses_input(base + index) + }) + .and_then(|(_, slot)| slot.free_dim) + } + + /// Shape of output `index`, including its free dimension when it has one. + pub(crate) fn output_shape(&self, output: &FoldOutput) -> Vec { + let mut shape = self.out_shape(); + if let Some(free) = self.output_free_dim(output) { + shape.push(free); + } + shape + } + + /// Whether any slot ranges over a free dimension. + pub(crate) fn has_free_dim(&self) -> bool { + self.carrier.iter().any(|slot| slot.free_dim.is_some()) + } + + /// Byte offsets of each slot inside a flattened carrier record — the + /// scratch layout a blocked fold writes its partials into. + pub(crate) fn slot_offsets(&self) -> (Vec, usize) { + let mut offsets = Vec::with_capacity(self.carrier.len()); + let mut total = 0; + for slot in &self.carrier { + offsets.push(total); + total += slot.width(); + } + (offsets, total) + } + + /// Approximate arithmetic cost of absorbing one element. + pub(crate) fn step_work(&self) -> u128 { + fn work(expression: &NaryExpr) -> u128 { + match expression { + NaryExpr::Op { children, .. } => 1 + children.iter().map(work).sum::(), + NaryExpr::IndexedInput { indices, .. } => { + 1 + indices.iter().map(work).sum::() + } + _ => 1, + } + } + self.step.iter().map(work).sum() + } + + /// Structural hash of everything deciding the kernel. A fold has no + /// `Operation` impl of its own yet, so the interner hashes directly. + pub(crate) fn hash_carrier_fields(&self, hasher: &mut impl std::hash::Hasher) { + use std::hash::Hash; + self.expression.hash(hasher); + if let Some(fold) = &self.element_fold { + fold.expression.hash(hasher); + fold.len.hash(hasher); + fold.function.hash(hasher); + } + self.shape.hash(hasher); + self.axis.hash(hasher); + self.block.hash(hasher); + self.element_is_carrier.hash(hasher); + self.carrier.len().hash(hasher); + for slot in &self.carrier { + slot.datatype.hash(hasher); + slot.free_dim.hash(hasher); + slot.element.hash(hasher); + } + for bodies in [&self.init, &self.step, &self.combine] { + bodies.len().hash(hasher); + for body in bodies { + body.hash(hasher); + } + } + self.outputs.len().hash(hasher); + for output in &self.outputs { + output.datatype.hash(hasher); + output.expression.hash(hasher); + } + } + + /// Structural well-formedness, expressed as slot-range bounds: `init` + /// reads no slot, `step` reads at most the carrier plus the element, and + /// `combine` reads at most two carriers. + pub(crate) fn validate(&self) -> Result<(), String> { + let (base, width) = (self.base(), self.width()); + if width == 0 { + return Err("fold carrier must have at least one slot".into()); + } + for (name, bodies) in [ + ("init", &self.init), + ("step", &self.step), + ("combine", &self.combine), + ] { + if bodies.len() != width { + return Err(format!( + "fold {name} has {} expressions for {width} carrier slots", + bodies.len() + )); + } + } + if self.init.iter().any(|e| reads_slot_from(e, base, 0)) { + return Err("fold init may not read the carrier".into()); + } + let step_limit = if self.element_is_carrier { + width * 2 + } else if self.carrier.iter().any(|slot| slot.element.is_some()) { + // A slot with its own element binds it after the shared one. + width + 2 + } else { + width + 1 + }; + if self.step.iter().any(|e| reads_slot_from(e, base, step_limit)) { + return Err("fold step reads past the accumulator and element".into()); + } + if self + .carrier + .iter() + .any(|slot| slot.element.is_some() != slot.free_dim.is_some()) + { + return Err("a slot's private element and free dimension go together".into()); + } + for output in &self.outputs { + let mut extents = self + .carrier + .iter() + .enumerate() + .filter(|(index, slot)| { + slot.free_dim.is_some() && output.expression.uses_input(base + index) + }) + .filter_map(|(_, slot)| slot.free_dim); + let first = extents.next(); + if extents.any(|extent| Some(extent) != first) { + return Err("an output may not span two different free dimensions".into()); + } + } + if let Some(fold) = &self.element_fold + && fold.len == 0 + { + return Err("an element fold needs a non-empty axis".into()); + } + if self + .combine + .iter() + .any(|e| reads_slot_from(e, base, width * 2)) + { + return Err("fold combine reads past the two carriers".into()); + } + if self.axis >= self.shape.len() { + return Err(format!( + "fold axis {} out of bounds for rank {}", + self.axis, + self.shape.len() + )); + } + Ok(()) + } + + /// Lift a built-in reduction into the general form: a one-slot fold whose + /// step and combine coincide. + pub(crate) fn from_reduce(reduce: &ReduceOperation) -> Self { + let datatype = reduce.function.datatype(); + let base = reduce.inputs.len(); + let op = reduce.function.op; + let join = |lhs: NaryExpr, rhs_expr: NaryExpr| match op { + ReduceOp::Sum => binary(NaryOp::Add, lhs, rhs_expr, datatype), + ReduceOp::Product => binary(NaryOp::Mul, lhs, rhs_expr, datatype), + ReduceOp::Max => binary(NaryOp::Max, lhs, rhs_expr, datatype), + ReduceOp::Min => binary(NaryOp::Min, lhs, rhs_expr, datatype), + }; + let algebra = match op { + ReduceOp::Max | ReduceOp::Min => FoldAlgebra::ExactMonoid, + ReduceOp::Sum | ReduceOp::Product => FoldAlgebra::ApproximateMonoid, + }; + // Finalization is just an output expression, so the post chain wraps + // the carrier read rather than being a separate stage. + let finalize = reduce + .post_element_wise + .functions + .iter() + .fold(acc(base, 0), |value, function| NaryExpr::Op { + children: vec![value], + function: function.clone(), + }); + Self { + inputs: reduce.inputs.clone(), + expression: reduce.expression.clone(), + element_fold: None, + shape: reduce.shape.clone(), + axis: reduce.axis, + carrier: vec![CarrierSlot::new(reduce.function.name(), datatype)], + init: vec![NaryExpr::Scalar(reduce.function.initial_value)], + step: vec![join(acc(base, 0), element(base, 1))], + combine: vec![join(acc(base, 0), rhs(base, 1, 0))], + outputs: vec![FoldOutput { + expression: finalize, + datatype: reduce.out_datatype(), + }], + algebra, + block: None, + element_is_carrier: false, + } + } + + /// The inverse of [`Self::from_reduce`]: recover a built-in reduction when + /// this fold is one. Lets the resolver hold every reduction in the general + /// form while lowering keeps its single-slot path. + pub(crate) fn to_reduce(&self) -> Option { + if self.width() != 1 + || self.block.is_some() + || self.element_is_carrier + || self.element_fold.is_some() + || self.has_free_dim() + { + return None; + } + let base = self.base(); + let datatype = self.carrier[0].datatype; + let NaryExpr::Scalar(initial_value) = self.init[0] else { + return None; + }; + let op = [ + ReduceOp::Sum, + ReduceOp::Product, + ReduceOp::Max, + ReduceOp::Min, + ] + .into_iter() + .find(|&candidate| { + let join = |lhs: NaryExpr, rhs_expr: NaryExpr| match candidate { + ReduceOp::Sum => binary(NaryOp::Add, lhs, rhs_expr, datatype), + ReduceOp::Product => binary(NaryOp::Mul, lhs, rhs_expr, datatype), + ReduceOp::Max => binary(NaryOp::Max, lhs, rhs_expr, datatype), + ReduceOp::Min => binary(NaryOp::Min, lhs, rhs_expr, datatype), + }; + self.step[0] == join(acc(base, 0), element(base, 1)) + && self.combine[0] == join(acc(base, 0), rhs(base, 1, 0)) + })?; + + // Peel finalization back into a post chain: a unary spine bottoming + // out at the carrier read. + let mut functions = Vec::new(); + let mut cursor = &self.outputs.first()?.expression; + let carrier_read = acc(base, 0); + loop { + if *cursor == carrier_read { + break; + } + match cursor { + NaryExpr::Op { children, function } if children.len() == 1 => { + functions.push(function.clone()); + cursor = &children[0]; + } + _ => return None, + } + } + functions.reverse(); + + let mut function = ReduceFunction::new(op, initial_value, datatype); + if let Some(name) = &self.carrier[0].name { + function = function.with_name(name); + } + Some(ReduceOperation { + inputs: self.inputs.clone(), + expression: self.expression.clone(), + shape: self.shape.clone(), + function, + post_element_wise: UnaryFunctionChain::new(functions, datatype), + axis: self.axis, + }) + } + + /// The split law. Blocking `axis` into runs of `factor` turns one fold + /// into a partial fold plus a joining fold over the block index. + /// + /// The joiner's `step` is the original `combine` verbatim: a joining fold's + /// element *is* a whole carrier, and both bodies bind + /// `[acc_0..acc_{n-1}, incoming_0..incoming_{n-1}]`. Mapping the incoming + /// carrier onto a single element would merge the slots. + pub(crate) fn split( + &self, + factor: usize, + policy: NumericsPolicy, + ) -> Result<(FoldOperation, FoldOperation), String> { + if self.block.is_some() { + return Err("fold is already blocked".into()); + } + if factor == 0 { + return Err("split factor must be non-zero".into()); + } + if !self.algebra.splittable_under(policy) { + return Err(format!( + "combine is {:?}, which {policy:?} does not permit splitting", + self.algebra + )); + } + let extent = self.shape[self.axis]; + if extent % factor != 0 { + return Err(format!( + "split factor {factor} does not divide axis extent {extent}" + )); + } + let base = self.base(); + + let mut partial = self.clone(); + partial.block = Some(factor); + // A partial stops at the carrier; finalization belongs to the joiner. + partial.outputs = self + .carrier + .iter() + .enumerate() + .map(|(index, slot)| FoldOutput { + expression: acc(base, index), + datatype: slot.datatype, + }) + .collect(); + + let mut joiner = self.clone(); + joiner.block = None; + joiner.shape = { + let mut shape = self.shape.to_vec(); + shape[self.axis] = extent / factor; + shape.into() + }; + joiner.step = self.combine.clone(); + joiner.element_is_carrier = true; + joiner.outputs = self.outputs.clone(); + + partial.validate()?; + joiner.validate()?; + Ok((partial, joiner)) + } +} + +/// The online-softmax carrier: running max, normalizer, and unnormalized +/// accumulator. +/// +/// ```text +/// (m1,l1,o1) (+) (m2,l2,o2) = (M, e^(m1-M)*l1 + e^(m2-M)*l2, +/// e^(m1-M)*o1 + e^(m2-M)*o2) M = max(m1,m2) +/// ``` +/// +/// A commutative monoid in the reals, but combining evaluates `exp` at +/// arguments depending on both sides, so it is [`FoldAlgebra::RescalingMonoid`] +/// rather than merely approximate. This is the identity flash attention is +/// built on and the reason a carrier must be a tuple. +pub(crate) fn online_softmax_carrier( + inputs: Vec, + expression: NaryExpr, + shape: Box<[usize]>, + axis: usize, + datatype: DataTypeEnum, +) -> FoldOperation { + let base = inputs.len(); + let width = 3; + let neg_inf = NaryScalar::F32(-3.0e38); + let zero = match datatype { + DataTypeEnum::F16 => NaryScalar::F16(half::f16::from_f32(0.0)), + DataTypeEnum::U32 => NaryScalar::U32(0), + DataTypeEnum::F32 => NaryScalar::F32(0.0), + }; + + // step: M = max(m, x); l' = l*e^(m-M) + e^(x-M); o' = o*e^(m-M) + e^(x-M) + let x = element(base, width); + let new_max = binary(NaryOp::Max, acc(base, 0), x.clone(), datatype); + let rescale = unary( + NaryOp::Exp, + binary(NaryOp::Sub, acc(base, 0), new_max.clone(), datatype), + datatype, + ); + let weight = unary( + NaryOp::Exp, + binary(NaryOp::Sub, x, new_max.clone(), datatype), + datatype, + ); + let step_slot = |k: usize| { + binary( + NaryOp::Add, + binary(NaryOp::Mul, acc(base, k), rescale.clone(), datatype), + weight.clone(), + datatype, + ) + }; + + // combine: the same law with the incoming carrier in place of the element. + let joined_max = binary(NaryOp::Max, acc(base, 0), rhs(base, width, 0), datatype); + let alpha = |value: NaryExpr| { + unary( + NaryOp::Exp, + binary(NaryOp::Sub, value, joined_max.clone(), datatype), + datatype, + ) + }; + let combine_slot = |k: usize| { + binary( + NaryOp::Add, + binary( + NaryOp::Mul, + acc(base, k), + alpha(acc(base, 0)), + datatype, + ), + binary( + NaryOp::Mul, + rhs(base, width, k), + alpha(rhs(base, width, 0)), + datatype, + ), + datatype, + ) + }; + + FoldOperation { + inputs, + expression, + shape, + axis, + carrier: vec![ + CarrierSlot::new("max", datatype), + CarrierSlot::new("normalizer", datatype), + CarrierSlot::new("accumulator", datatype), + ], + init: vec![ + NaryExpr::Scalar(neg_inf), + NaryExpr::Scalar(zero), + NaryExpr::Scalar(zero), + ], + step: vec![new_max, step_slot(1), step_slot(2)], + combine: vec![joined_max.clone(), combine_slot(1), combine_slot(2)], + outputs: vec![ + FoldOutput { + expression: binary(NaryOp::Div, acc(base, 2), acc(base, 1), datatype), + datatype, + }, + // The log-sum-exp, free from the same carrier. + FoldOutput { + expression: binary( + NaryOp::Add, + acc(base, 0), + unary(NaryOp::Log, acc(base, 1), datatype), + datatype, + ), + datatype, + }, + ], + algebra: FoldAlgebra::RescalingMonoid, + block: None, + element_is_carrier: false, + } +} + +/// Welford's algorithm as a fold: running count, mean, and sum of squared +/// deviations. The generality check — if the carrier abstraction is right this +/// costs nothing attention-specific, and blocked mean/variance for layer norm +/// falls out of the same split law that blocks attention. +pub(crate) fn welford_carrier( + inputs: Vec, + expression: NaryExpr, + shape: Box<[usize]>, + axis: usize, +) -> FoldOperation { + let datatype = DataTypeEnum::F32; + let base = inputs.len(); + let width = 3; + let zero = NaryExpr::Scalar(NaryScalar::F32(0.0)); + let one = NaryExpr::Scalar(NaryScalar::F32(1.0)); + let x = element(base, width); + + let n_next = binary(NaryOp::Add, acc(base, 0), one, datatype); + let delta = binary(NaryOp::Sub, x.clone(), acc(base, 1), datatype); + let mean_next = binary( + NaryOp::Add, + acc(base, 1), + binary(NaryOp::Div, delta.clone(), n_next.clone(), datatype), + datatype, + ); + let m2_next = binary( + NaryOp::Add, + acc(base, 2), + binary( + NaryOp::Mul, + delta, + binary(NaryOp::Sub, x, mean_next.clone(), datatype), + datatype, + ), + datatype, + ); + + // combine: the pairwise Chan-Golub-LeVeque update. + let n_total = binary(NaryOp::Add, acc(base, 0), rhs(base, width, 0), datatype); + let mean_delta = binary(NaryOp::Sub, rhs(base, width, 1), acc(base, 1), datatype); + let mean_joined = binary( + NaryOp::Add, + acc(base, 1), + binary( + NaryOp::Mul, + mean_delta.clone(), + binary(NaryOp::Div, rhs(base, width, 0), n_total.clone(), datatype), + datatype, + ), + datatype, + ); + let m2_joined = binary( + NaryOp::Add, + binary(NaryOp::Add, acc(base, 2), rhs(base, width, 2), datatype), + binary( + NaryOp::Mul, + binary(NaryOp::Mul, mean_delta.clone(), mean_delta, datatype), + binary( + NaryOp::Div, + binary(NaryOp::Mul, acc(base, 0), rhs(base, width, 0), datatype), + n_total.clone(), + datatype, + ), + datatype, + ), + datatype, + ); + + FoldOperation { + inputs, + expression, + shape, + axis, + carrier: vec![ + CarrierSlot::new("count", datatype), + CarrierSlot::new("mean", datatype), + CarrierSlot::new("m2", datatype), + ], + init: vec![zero.clone(), zero.clone(), zero], + step: vec![n_next, mean_next, m2_next], + combine: vec![n_total.clone(), mean_joined, m2_joined], + outputs: vec![ + FoldOutput { + expression: acc(base, 1), + datatype, + }, + FoldOutput { + expression: binary(NaryOp::Div, acc(base, 2), acc(base, 0), datatype), + datatype, + }, + ], + algebra: FoldAlgebra::ApproximateMonoid, + block: None, + element_is_carrier: false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::nary_wise::UnaryFunctionChain; + + /// Host evaluator over the slot convention, enough to prove the split law + /// on concrete data. Slots are supplied in binding order. + fn eval(expression: &NaryExpr, slots: &[f32]) -> f32 { + match expression { + NaryExpr::IndexedInput { input_idx, indices } => { + assert!(indices.is_empty(), "carrier bodies read scalar slots"); + slots[*input_idx] + } + NaryExpr::Scalar(NaryScalar::F32(value)) => *value, + NaryExpr::Scalar(other) => panic!("unexpected scalar {other:?}"), + NaryExpr::DimIndex(_) => panic!("carrier bodies are coordinate-free"), + NaryExpr::Op { children, function } => { + let args: Vec = children.iter().map(|c| eval(c, slots)).collect(); + match function.op { + NaryOp::Add => args[0] + args[1], + NaryOp::Sub => args[0] - args[1], + NaryOp::Mul => args[0] * args[1], + NaryOp::Div => args[0] / args[1], + NaryOp::Max => args[0].max(args[1]), + NaryOp::Min => args[0].min(args[1]), + NaryOp::Exp => args[0].exp(), + NaryOp::Log => args[0].ln(), + NaryOp::Sqrt => args[0].sqrt(), + ref other => panic!("unexpected op {other:?}"), + } + } + } + } + + fn init_carrier(fold: &FoldOperation) -> Vec { + fold.init.iter().map(|e| eval(e, &[])).collect() + } + + /// One `step`: bindings are `[acc.., element]`. + fn step_once(fold: &FoldOperation, acc: &[f32], value: f32) -> Vec { + let mut slots = acc.to_vec(); + slots.push(value); + fold.step.iter().map(|e| eval(e, &slots)).collect() + } + + /// One `combine`: bindings are `[acc.., rhs..]`. + fn combine_once(fold: &FoldOperation, acc: &[f32], incoming: &[f32]) -> Vec { + let mut slots = acc.to_vec(); + slots.extend_from_slice(incoming); + fold.combine.iter().map(|e| eval(e, &slots)).collect() + } + + fn finish(fold: &FoldOperation, acc: &[f32]) -> Vec { + fold.outputs + .iter() + .map(|out| eval(&out.expression, acc)) + .collect() + } + + fn run_whole(fold: &FoldOperation, data: &[f32]) -> Vec { + let mut acc = init_carrier(fold); + for &value in data { + acc = step_once(fold, &acc, value); + } + finish(fold, &acc) + } + + /// Partial folds per block, joined by the joiner's step (which is the + /// original combine), then finalized. + fn run_split(fold: &FoldOperation, data: &[f32], factor: usize) -> Vec { + let (partial, joiner) = fold + .split(factor, NumericsPolicy::RelativeErrorPermitted) + .expect("split"); + let partials: Vec> = data + .chunks(factor) + .map(|block| { + let mut acc = init_carrier(&partial); + for &value in block { + acc = step_once(&partial, &acc, value); + } + acc + }) + .collect(); + let mut acc = init_carrier(&joiner); + for incoming in &partials { + // A carrier-element fold's step binds two whole carriers. + acc = combine_once(fold, &acc, incoming); + } + finish(&joiner, &acc) + } + + fn parts() -> (Vec, NaryExpr, Box<[usize]>) { + (Vec::new(), NaryExpr::input(0, 2), vec![1usize, 8].into()) + } + + fn reduce_of(function: crate::reduce::ReduceFunction, post: UnaryFunctionChain) -> ReduceOperation { + let (inputs, expression, shape) = parts(); + ReduceOperation { + inputs, + expression, + shape, + function, + post_element_wise: post, + axis: 1, + } + } + + #[test] + fn split_preserves_sum_and_max() { + let data: Vec = (1..=8).map(|v| v as f32).collect(); + let sum = FoldOperation::from_reduce(&reduce_of( + crate::reduce::sum_fn(DataTypeEnum::F32), + UnaryFunctionChain::empty(DataTypeEnum::F32), + )); + sum.validate().unwrap(); + assert_eq!(run_whole(&sum, &data), vec![36.0]); + assert_eq!(run_split(&sum, &data, 4), vec![36.0]); + assert_eq!(run_split(&sum, &data, 2), vec![36.0]); + + let values = vec![3.0, -1.0, 7.5, 2.0, 0.0, 7.4, -9.0, 1.0]; + let max = FoldOperation::from_reduce(&reduce_of( + crate::reduce::max_fn(DataTypeEnum::F32), + UnaryFunctionChain::empty(DataTypeEnum::F32), + )); + assert_eq!(run_whole(&max, &values), vec![7.5]); + assert_eq!(run_split(&max, &values, 4), vec![7.5]); + } + + #[test] + fn split_preserves_online_softmax() { + let (inputs, expression, shape) = parts(); + let fold = online_softmax_carrier(inputs, expression, shape, 1, DataTypeEnum::F32); + fold.validate().unwrap(); + let data = vec![1.0, 3.0, 2.0, -4.0, 0.5, 8.0, 7.0, -2.0]; + + let max = data.iter().cloned().fold(f32::MIN, f32::max); + let expected_lse = max + data.iter().map(|v| (v - max).exp()).sum::().ln(); + + let whole = run_whole(&fold, &data); + assert!((whole[1] - expected_lse).abs() < 1e-5, "lse {whole:?}"); + for factor in [2usize, 4] { + let split = run_split(&fold, &data, factor); + for (a, b) in whole.iter().zip(&split) { + assert!((a - b).abs() < 1e-5, "factor {factor}: {whole:?} vs {split:?}"); + } + } + } + + #[test] + fn split_preserves_welford() { + let (inputs, expression, shape) = parts(); + let fold = welford_carrier(inputs, expression, shape, 1); + fold.validate().unwrap(); + let data = vec![2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0]; + let whole = run_whole(&fold, &data); + assert!((whole[0] - 5.0).abs() < 1e-5, "mean {whole:?}"); + assert!((whole[1] - 4.0).abs() < 1e-5, "variance {whole:?}"); + for factor in [2usize, 4] { + let split = run_split(&fold, &data, factor); + for (a, b) in whole.iter().zip(&split) { + assert!((a - b).abs() < 1e-4, "factor {factor}: {whole:?} vs {split:?}"); + } + } + } + + #[test] + fn numerics_policy_gates_the_split_law() { + let (inputs, expression, shape) = parts(); + let softmax = online_softmax_carrier(inputs, expression, shape, 1, DataTypeEnum::F32); + // The rescaling lift is not licensed by mere reassociation. + assert!(softmax.split(4, NumericsPolicy::ReassociationPermitted).is_err()); + assert!(softmax.split(4, NumericsPolicy::RelativeErrorPermitted).is_ok()); + // Max is exact, so it splits even under the strictest policy. + let max = FoldOperation::from_reduce(&reduce_of( + crate::reduce::max_fn(DataTypeEnum::F32), + UnaryFunctionChain::empty(DataTypeEnum::F32), + )); + assert!(max.split(4, NumericsPolicy::Exact).is_ok()); + } + + #[test] + fn built_in_reductions_round_trip_through_the_general_form() { + let post = UnaryFunctionChain::new( + vec![NaryFunction::unary( + None, + NaryOp::Sqrt, + DataTypeEnum::F32, + DataTypeEnum::F32, + )], + DataTypeEnum::F32, + ); + for function in [ + crate::reduce::sum_fn(DataTypeEnum::F32), + crate::reduce::max_fn(DataTypeEnum::F32), + ] { + for chain in [UnaryFunctionChain::empty(DataTypeEnum::F32), post.clone()] { + let reduce = reduce_of(function.clone(), chain); + let fold = FoldOperation::from_reduce(&reduce); + fold.validate().unwrap(); + let recovered = fold.to_reduce().expect("built-in fold has a reduce form"); + assert_eq!(recovered, reduce, "round trip changed the reduction"); + } + } + } + + #[test] + fn general_folds_have_no_reduce_form() { + let (inputs, expression, shape) = parts(); + assert!( + online_softmax_carrier(inputs, expression, shape, 1, DataTypeEnum::F32) + .to_reduce() + .is_none() + ); + let (inputs, expression, shape) = parts(); + assert!(welford_carrier(inputs, expression, shape, 1).to_reduce().is_none()); + + let sum = FoldOperation::from_reduce(&reduce_of( + crate::reduce::sum_fn(DataTypeEnum::F32), + UnaryFunctionChain::empty(DataTypeEnum::F32), + )); + let (partial, joiner) = sum.split(4, NumericsPolicy::ReassociationPermitted).unwrap(); + assert!(partial.to_reduce().is_none(), "blocked fold is not a reduce"); + assert!(joiner.to_reduce().is_none(), "carrier-element fold is not a reduce"); + } + + #[test] + fn split_rejects_indivisible_factors_and_double_blocking() { + let policy = NumericsPolicy::RelativeErrorPermitted; + let sum = FoldOperation::from_reduce(&reduce_of( + crate::reduce::sum_fn(DataTypeEnum::F32), + UnaryFunctionChain::empty(DataTypeEnum::F32), + )); + assert!(sum.split(3, policy).is_err()); + let (partial, _) = sum.split(4, policy).unwrap(); + assert!(partial.split(2, policy).is_err()); + } +} diff --git a/fusor-ml/core/src/lib.rs b/fusor-ml/core/src/lib.rs index 3a1805a9d..3ba0dc881 100644 --- a/fusor-ml/core/src/lib.rs +++ b/fusor-ml/core/src/lib.rs @@ -5,6 +5,7 @@ pub use device::*; pub use element_wise::CastTensor; pub use fusor_gguf::GgufReadError; +pub use fusor_tile_ir_runtime::FusorConfig; pub use fusor_types::{ Layout, SlidingWindow, StrideSpec, TILE_SIZE, TensorSlice, slice_shape, slice_strides, }; @@ -29,6 +30,8 @@ mod compute_graph; pub use compute_graph::NodeIndex; mod device; mod element_wise; +mod flash_attention; +mod fold; mod index_select; #[doc(hidden)] pub mod kernel_selection; @@ -37,10 +40,12 @@ pub mod matmul; mod mir; mod nary_direct; mod nary_wise; +mod occupancy; mod pair_wise; mod quantized; mod rank; mod reduce; +mod region; mod row_dispatch; mod row_program; mod sampling; diff --git a/fusor-ml/core/src/matmul/coop_gemm.rs b/fusor-ml/core/src/matmul/coop_gemm.rs deleted file mode 100644 index 91efeaeff..000000000 --- a/fusor-ml/core/src/matmul/coop_gemm.rs +++ /dev/null @@ -1,103 +0,0 @@ -use crate::{Device, MatMulOperation, kernel_selection::CooperativeMatrixKind}; - -/// Parameters for cooperative matrix matmul. -#[derive(Debug, Clone, PartialEq, Hash)] -pub struct CoopGemmParams { - pub block_m: u32, - pub block_n: u32, - pub block_k: u32, - pub n_passes: u32, - pub mma_size: u32, - pub wg_threads: u32, - pub(crate) kind: CooperativeMatrixKind, -} - -impl Default for CoopGemmParams { - fn default() -> Self { - Self { - block_m: 128, - block_n: 64, - block_k: 16, - n_passes: 4, - mma_size: 8, - wg_threads: 256, - kind: CooperativeMatrixKind::F32F32M8N8K8, - } - } -} - -impl CoopGemmParams { - pub(crate) fn kind(&self) -> CooperativeMatrixKind { - self.kind - } -} - -pub(super) fn optimal_params( - m: usize, - n: usize, - k: usize, - device: &Device, - kind: CooperativeMatrixKind, -) -> Option { - if !device.cooperative_matrix_caps().supports(kind) - || !device.subgroups_supported() - || device.min_subgroup_size() != device.max_subgroup_size() - || device.limits().max_compute_workgroup_size_x < 64 - { - return None; - } - - let mut params = CoopGemmParams::default(); - if n <= 16 { - params.block_n = 16; - params.n_passes = 1; - } else if n <= 32 { - params.block_n = 32; - params.n_passes = 2; - } - - if m <= 16 { - params.block_m = 16; - params.wg_threads = 64; - } else if m < params.block_m as usize { - params.block_m = 64; - params.wg_threads = 128; - } - - if params.wg_threads > device.limits().max_compute_workgroup_size_x { - return None; - } - - params.kind = kind; - let _ = k; - Some(params) -} - -pub(super) fn workgroup_shape_constraints( - _: &MatMulOperation, - _: &Device, - params: &CoopGemmParams, -) -> crate::mir::workgroup_shape::WorkgroupShapeConstraints { - let mut constraints = crate::mir::workgroup_shape::WorkgroupShapeConstraints::default(); - constraints.add_constraint( - 0, - crate::mir::workgroup_shape::Constraint::Equals(params.wg_threads), - ); - constraints.add_constraint(1, crate::mir::workgroup_shape::Constraint::Equals(1)); - constraints.add_constraint(2, crate::mir::workgroup_shape::Constraint::Equals(1)); - constraints -} - -pub(super) fn dispatch_size( - last_dim_size: usize, - second_to_last_dim_size: usize, - batch_size: usize, - workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, - params: &CoopGemmParams, -) -> [u32; 3] { - [ - (second_to_last_dim_size as u32).div_ceil(params.block_m), - (last_dim_size as u32).div_ceil(params.block_n), - (batch_size as u32).div_ceil(workgroup_shape.z()), - ] -} diff --git a/fusor-ml/core/src/matmul/cost.rs b/fusor-ml/core/src/matmul/cost.rs new file mode 100644 index 000000000..d1181b3e7 --- /dev/null +++ b/fusor-ml/core/src/matmul/cost.rs @@ -0,0 +1,595 @@ +//! General cooperative-tile selection. +//! +//! One scalar additive cost over the joint choice of (tile entry) x (split +//! count) x (staged tile pairs), in integer femtoseconds with physical device +//! rates. A lexicographic +//! key cannot trade split-K scratch against tile geometry — the levels have +//! incommensurate units — so the whole decision surface is a single argmin +//! ([`score_fs`]), with a tie-break chain that is reached and therefore +//! load-bearing. +//! +//! Alignment enters only through what it does to the cost (edge tiles execute +//! padded MACs); the device enters through legality, four measured rates and +//! one occupancy target. Padding is the dominant measured effect — a +//! 33%-padded tile ran ~2x slower than its zero-pad neighbor at N=384, and a +//! 6.7%-padded 128x512 ran ~3x slower than the least-padded choice on the +//! vision shape — and the MAC term is proportional to padded MACs, so it +//! carries that without a rule. +//! +//! Single-buffered profiles are excluded from automatic selection outright: +//! they cannot horizontally merge (the merged body shares one double-buffered +//! tile pair across guarded segments), and their best standalone shape +//! (4096^3) measures 5.57 TF/s against the selected 128x64's 7.64 — even at +//! parity it would not pay for decomposing a merged wave into standalone +//! dispatches (~15% on the vision QKV family). They stay in the kernel table +//! for forced experiments, where `coop_tile_conformance` keeps them verified. +//! +//! Workgroup-memory footprint enters twice. As legality: entries whose +//! `CoopTileEntry::workgroup_bytes` for the contraction's stage element +//! exceed the device's workgroup-storage limit are unselectable (WebGPU's +//! 16 KB default would otherwise fail at pipeline creation). And as +//! residency: shared memory is carved from a per-core pool, so the footprint +//! divides into how many workgroups a core holds at once, and co-resident +//! workgroups cover each other's epilogue drain. That second role is what +//! makes the staging depth a decision rather than a table column — one +//! staged pair loses the load/MMA overlap on every K iteration and buys back +//! a once-per-workgroup drain, so deep-K contractions want two pairs and +//! shallow ones want one. +//! +//! Measured anchors on M2 Max, per-entry minima over the round-robin reps of +//! `ITERS=8 REPS=6 cargo run --release -p fusor-core --example +//! bench_coop_tiles`. f32, `128x64 / 64x64 / 128x128 / 64x16 / 16x64` in ms: +//! 16384x384x384 0.799 / 0.822 / 1.011 / 0.990 / 1.197; 16384x384x1536 +//! 2.705 / 2.893 / 3.456 / 3.823 / 3.996; 16384x1536x384 2.698 / 2.885 / +//! 3.646 / 3.836 / 4.023; 16384x3072x1536 20.41 / 21.16 / 28.30 / 30.58 / +//! 32.44; 1024^3 0.445 / 0.445 / 0.577 / 0.446 / 0.634. The K-deep half of +//! the calibration set: every one of them wants the 128-wide profile, which +//! is what stops the epilogue term from taking 64x64 everywhere. f16 128x64 +//! vs 64x64 is inside 0.3% on the first three and 3-4% the other way on +//! 16384x3072x1536 (19.21 / 18.46) and 4096^3 (17.22 / 16.65), so the two +//! f16 rows the golden table moves onto 64x64 are neutral-to-positive. +//! Earlier anchors from the same harness, still the padding evidence: +//! 4096^3 — 256x256 5.57 TF/s (excluded as unmergeable) / 128x512 5.20; +//! 1944x1280x3840 — 64x64 4.19 (least padded) vs 128x64 2.70 ~= 64x128 +//! 2.61, padded 128x512 1.47; 1000x1024x1024 — 64x64 3.11 / 128x64 2.62, +//! padded 128x512 1.09. +//! +//! Warm-resolve span_ms against the previous calibration of these same terms +//! (both arms one binary behind one env switch, 20 interleaved processes per +//! arm, order flipped every round): 64x2048x64 0.228 -> 0.200, 64x2048x256 +//! 0.617 -> 0.593, 256x2048x64 0.690 -> 0.594, 2048x64x64 0.261 -> 0.220, +//! 2048x256x64 0.738 -> 0.578, 2048x64x256 0.876 -> 0.683, batched 64x64x16 +//! 0.290 -> 0.291, 384x16384x1536 x4 14.46 -> 14.22; softmax control flat at +//! 1.090 -> 1.086. Adding the staging depth then moved the two shapes whose +//! chosen depth changed: 2048x64x64 0.217 -> 0.187 and 2048x64x256 0.671 -> +//! 0.597, every other plan bit-identical and the control at 0.0%. The depth +//! itself was calibrated by forcing it: the same merged body staged from one +//! pair instead of two runs -14.0% / -11.0% on those two, +0.8% on +//! 2048x256x64 (16 K iterations) and +7.2% on 384x16384x1536 (1024), which +//! is the crossover the two constants encode. Every one of those moves is the same decision: six of the +//! eight shapes had been taking a 128-wide or 16-wide profile where the +//! 64x64 one measures fastest, because a flat per-element epilogue rate +//! cannot see that a wider workgroup drains its accumulator more slowly, and +//! a linear occupancy law over-splits a starved grid. + +use std::cmp::Reverse; + +use fusor_tile_ir::ScalarElement; +use fusor_tile_ir_kernels::{CoopTileEntry, coop_tile_entries}; + +use super::variants::CoopTile; +use crate::occupancy::DispatchPolicy; + +/// Cooperative-matrix fragment side; every per-subgroup fragment grid counts +/// whole 8x8 fragments. +const COOP_DIM: u64 = 8; + +/// Storage bindings one matmul segment declares: A, B and its output. The +/// horizontal merger budgets a merged wave with the same count, which is what +/// bounds the group [`tile_probe_group`] scores against. +const MATMUL_SEGMENT_BINDINGS: usize = 3; + +/// The contraction plus how many co-located segments the dispatch that will +/// run it carries. `segments` is 1 for a standalone kernel and +/// `segments.len()` inside `build_merged_matmul_kernel` — the launched grid is +/// `tiles * splits * segments`, never `tiles * splits`. +#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] +pub(crate) struct CoopDispatch { + pub(crate) m: u32, + pub(crate) k: u32, + pub(crate) n: u32, + pub(crate) batch: u32, + pub(crate) segments: u32, + pub(crate) datatype: crate::DataTypeEnum, + pub(crate) has_epilogues: bool, +} + +/// The probe group the TILE is scored against, everywhere and always. +/// +/// Allocation runs before the wave is partitioned, and `hardware_matmul_prep` +/// compares the output's strides for exact equality against the build tile's +/// padding — a tile that moved between `inputs()` and the build does not waste +/// memory, it silently falls back to the generic path. So the tile is scored +/// at a fixed group both can compute, never at the real one: 1 for a +/// contraction that can never merge (epilogues disqualify it at +/// `merge_profile`), otherwise the merger's own maximum group size. +pub(crate) fn tile_probe_group(device: &crate::Device, has_epilogues: bool) -> u32 { + if has_epilogues { + 1 + } else { + (device.nary_direct_input_binding_budget() / MATMUL_SEGMENT_BINDINGS).max(1) as u32 + } +} + +/// The element the kernels stage operand tiles in. Production paths leave +/// `staging` off, so this is the storage element; it sets both the staged +/// byte count and the workgroup-memory footprint the residency term reads. +fn stage_element(datatype: crate::DataTypeEnum) -> ScalarElement { + match datatype { + crate::DataTypeEnum::F16 => ScalarElement::F16, + _ => ScalarElement::F32, + } +} + +/// Cost of running one contraction with one (tile, subgroup split, split +/// count) choice, in integer femtoseconds. +/// +/// Roofline in its literal form: the three issue-side terms sum (MMA issue, +/// threadgroup traffic and the accumulator store all contend for the same +/// per-core issue and load/store slots), DRAM overlaps them (`max`), and the +/// combine dispatch adds behind its barrier. Occupancy scales the issue side +/// only, by the cube root of the residency shortfall against +/// `prefetched_saturation_lanes` — that target is the only role a +/// parallelism floor plays here, never an execution width and never a +/// MAC-equivalent. +/// +/// Integer, not f64: candidates tie exactly (the tile table is built from +/// powers of two and the score is invariant to `n_passes` by construction), so +/// the tie-break chain is load-bearing and must be reached deterministically. +#[allow(clippy::too_many_arguments)] +fn score_fs( + m: u32, + k: u32, + n: u32, + batch: u32, + segments: u32, + entry: &CoopTileEntry, + rg: u32, + cg: u32, + splits: u32, + buffers: u32, + stage: ScalarElement, + policy: &DispatchPolicy, + subgroup_width: u32, +) -> u128 { + let rates = policy.matmul_rates(); + let elem_bytes = stage.byte_size(); + let (bm, bn, bk) = ( + u64::from(entry.tile.bm), + u64::from(entry.tile.bn), + u64::from(entry.tile.bk), + ); + let n_passes = u64::from(entry.n_passes); + let bn_pass = bn / n_passes; + // Per subgroup per kk-step: `tr` A-fragment loads and `tc` B-fragment + // loads feed `tr * tc` MMAs (`coop_load_a_fragments` emits `rows`, + // `coop_load_b_fragments` emits `cols`, `coop_mma_grid` does rows*cols). + let tr = bm / (COOP_DIM * u64::from(rg)); + let tc = bn_pass / (COOP_DIM * u64::from(cg)); + let subgroups = u64::from(rg) * u64::from(cg); + let threads = subgroups * u64::from(subgroup_width); + + let tiles_m = u64::from(m.div_ceil(entry.tile.bm)); + let tiles_n = u64::from(n.div_ceil(entry.tile.bn)); + let tiles_per_segment = tiles_m * tiles_n * u64::from(batch); + let k_iterations = u64::from(k.div_ceil(entry.tile.bk)); + let span_iterations = k_iterations.div_ceil(u64::from(splits)); + let workgroups = tiles_per_segment * u64::from(splits) * u64::from(segments); + let m_padded = tiles_m * bm; + let n_padded = tiles_n * bn; + let per_workgroup = u128::from(workgroups) * u128::from(span_iterations); + + // T1 MMA issue. Per kk-step a workgroup issues `subgroups * tr * tc * 512` + // MACs = `bm * bn_pass * 8`; over `bk/8` kk-steps and `n_passes` passes + // that is `bm * bn * bk` per workgroup per K iteration. + let t_mma = per_workgroup * u128::from(bm * bn * bk) * 1_000_000 + / u128::from(rates.mac_per_ns); + + // T2 threadgroup traffic: cooperative fragment loads plus operand staging, + // at one rate. Their per-load byte ratio varies only 96..128 across the + // whole table, so a fit that separates them is not identifiable. + let fragment_bytes = n_passes * subgroups * (tr + tc) * (bk / COOP_DIM) * 64 * elem_bytes; + let stage_bytes = n_passes * (bm * bk + bk * bn_pass) * elem_bytes; + // One staged pair loses the load/MMA overlap the rates were fitted on. + let overlap_pct = if buffers == 1 { + rates.single_buffered_traffic_pct + } else { + 100 + }; + let t_threadgroup = per_workgroup * u128::from(fragment_bytes + stage_bytes) * 1_000_000 + * u128::from(overlap_pct) + / (u128::from(rates.workgroup_bytes_per_ns) * 100); + + // T3 accumulator zeroing, the cooperative store's fragment shuffles and + // the store itself, over the padded output every workgroup emits. This is + // what makes split-K expensive: every split writes a full padded tile. + // + // Per element AND per subgroup of the emitting workgroup. The epilogue is + // a whole-workgroup drain — every subgroup's accumulator fragments shuffle + // through the one staged tile pair behind the workgroup's own barrier, so + // a wider workgroup serializes more of its output through the same + // threadgroup port and the workgroup cannot retire until the last + // subgroup lands. Measured on 2048x64x256, where every profile does the + // same MACs and emits the same padded output, the implied per-element + // epilogue rate is 32 ps (64x16) and 36 ps (64x64) against 60 ps + // (128x64), 63 ps (128x128) and 72 ps (64x128): the split is exactly the + // 4-vs-8 subgroup count. Nothing else in the table sorts it — 64x64 and + // 64x128 have the same workgroup-memory footprint and land on opposite + // sides, 128x64 and 64x128 have the same tile area and land together. + // + // This is a per-workgroup cost, so it amortizes away down a deep K loop, + // which is why 384x16384x1536 (1024 K-iterations) still takes the widest + // profile while the shallow-K shapes do not. + // + // It is also what a second co-resident workgroup covers. A core's + // shared memory holds `core_workgroup_slots` workgroups of this + // footprint; while one drains its accumulators another issues MMAs, so + // the drain the dispatch actually waits on is the term divided by that + // count. Measured by halving the footprint at fixed tile, splits and + // grid: 2048x64x64 -14.0%, 2048x64x256 -11.0%, against 2048x256x64 + // +0.8% and 384x16384x1536 +7.2% where the deeper K loop makes the + // staging penalty above outweigh it. A fourth root, not a reciprocal: + // the drain is only partly hideable, and the raw ratio predicts a 55% + // swing where the measured pair is 14%. + let slots = policy.core_workgroup_slots(entry.workgroup_bytes_at(stage, buffers)); + let t_store = u128::from(workgroups) + * u128::from(bm * bn) + * u128::from(rates.store_fs_per_element * subgroups) + * 1_000 + / integer_root(u128::from(slots) * 1_000_000_000_000, 4); + + // T4 DRAM: operands once per segment plus every split's padded output. + let dram_bytes = u128::from(segments) + * u128::from(elem_bytes) + * (u128::from(batch) * (u128::from(m) * u128::from(k) + u128::from(k) * u128::from(n)) + + u128::from(splits) * u128::from(batch) * u128::from(m_padded * n_padded)); + let t_dram = dram_bytes * 10_000_000 / u128::from(rates.dram_decibytes_per_ns); + + // T5 combine: reads every partial slice and writes the output, in its own + // barrier-separated dispatch, so it adds rather than overlaps. + let t_combine = if splits > 1 { + u128::from(segments) + * u128::from(splits + 1) + * u128::from(batch) + * u128::from(m_padded * n_padded) + * u128::from(elem_bytes) + * 10_000_000 + / u128::from(rates.dram_decibytes_per_ns) + } else { + 0 + }; + + let issue = t_mma + t_threadgroup + t_store; + let resident = u128::from(workgroups) * u128::from(threads); + let target = u128::from(policy.prefetched_saturation_lanes()); + // A grid short of the floor does not lose issue rate in proportion to the + // lanes it is missing: the lanes it does have keep more of the core's + // issue slots, its threadgroup port and its share of L2 to themselves, so + // each runs faster than it would in a full grid. Measured on the split-K + // sweeps, where the split count is exactly a lane-count dial: 64x2048x64 + // at 64x64 runs 0.202 ms with 20,480 lanes resident and 0.310 with 81,920 + // — the linear law says the starved grid should have been 3.2x slower and + // it is 0.65x. A cube root reproduces that and the 64x2048x256 and + // 256x2048x64 split curves; a linear one over-splits every one of them by + // 2-4x. + let issue_scaled = if resident >= target { + issue + } else { + issue * integer_root(target * 1_000_000_000 / resident.max(1), 3) / 1_000 + }; + issue_scaled.max(t_dram) + t_combine +} + +/// Floor of the `n`th root, by Newton iteration on integers so the argmin +/// stays exactly reproducible across platforms (a floating `powf` is not). +fn integer_root(value: u128, n: u32) -> u128 { + if value < 2 { + return value; + } + let mut x = 1u128 << (value.ilog2() / n + 1); + loop { + let next = ((u128::from(n) - 1) * x + value / x.pow(n - 1)) / u128::from(n); + if next >= x { + return x; + } + x = next; + } +} + +/// Split counts worth scoring: never splitting, plus every divisor of the K +/// loop that leaves at least two iterations per workgroup. The 64 ceiling +/// bounds the candidate count; epilogues make splitting illegal outright +/// (partials carry no epilogue identity). +fn split_candidates(k_iterations: u32, has_epilogues: bool) -> impl Iterator { + let limit = if has_epilogues { + 1 + } else { + (k_iterations / 2).min(64) + }; + (1..=limit.max(1)).filter(move |d| *d == 1 || k_iterations.is_multiple_of(*d)) +} + +/// Staged operand tile pairs worth scoring. Two pairs overlap the next K +/// tile's fill with the current tile's MMAs; one pair halves the workgroup's +/// threadgroup footprint, so a core holds more of them and their epilogue +/// drains cover each other. The split-K partials body is one pair outright: +/// a split grid exists to raise occupancy and a second pair halves it, which +/// measured 55% of wall time on 64x2048x256 at identical tile, splits and +/// grid. +fn staging_depths(splits: u32) -> impl Iterator { + if splits > 1 { 1..=1 } else { 1..=2 } +} + +/// The best split count and staging depth for one fixed geometry on the grid +/// that actually launches. Monotone non-increasing in `dispatch.segments`: every term is +/// linear in the segment count except the occupancy scaling, whose starved +/// branch is independent of it, so a larger group reaches the saturated +/// plateau at an ever-smaller split count. +pub(super) fn plan_coop_splits( + dispatch: CoopDispatch, + tile: CoopTile, + rg: u32, + cg: u32, + policy: &DispatchPolicy, + subgroup_width: u32, +) -> (u32, u32) { + let Some(entry) = coop_tile_entries() + .iter() + .find(|entry| entry.tile.bm == tile.bm && entry.tile.bn == tile.bn && entry.tile.bk == tile.bk) + else { + return (1, 2); + }; + let stage = stage_element(dispatch.datatype); + let k_iterations = dispatch.k.div_ceil(tile.bk); + let mut best: Option<(u128, u32, u32)> = None; + for splits in split_candidates(k_iterations, dispatch.has_epilogues) { + for buffers in staging_depths(splits) { + let score = score_fs( + dispatch.m, + dispatch.k, + dispatch.n, + dispatch.batch, + dispatch.segments.max(1), + entry, + rg, + cg, + splits, + buffers, + stage, + policy, + subgroup_width, + ); + // Ascending candidates with a strict comparison: an exact tie + // keeps the smaller split count (less scratch) and, within a + // split count, the shallower staging (less workgroup memory). + if best.is_none_or(|(best_score, ..)| score < best_score) { + best = Some((score, splits, buffers)); + } + } + } + best.map_or((1, 2), |(_, splits, buffers)| (splits, buffers)) +} + +/// The tile geometry and its subgroup split, scored at `probe_group` with +/// every candidate free to pick its own best split count. `None` routes the +/// contraction to the vector/generic families. +#[allow(clippy::too_many_arguments)] +pub(super) fn plan_coop_tile( + m: u32, + k: u32, + n: u32, + batch: u32, + datatype: crate::DataTypeEnum, + has_epilogues: bool, + probe_group: u32, + policy: &DispatchPolicy, + subgroup_width: u32, +) -> Option<(CoopTile, u32, u32)> { + if m == 0 || n == 0 || k == 0 || batch == 0 { + return None; + } + // The kernels stage operand tiles in the storage element (`staging` stays + // off in production paths). + let stage = stage_element(datatype); + let mut best: Option<((u128, u32, Reverse, Reverse), CoopTile, u32, u32, u128)> = None; + for entry in coop_tile_entries() { + let (bm, bn, bk) = (entry.tile.bm, entry.tile.bn, entry.tile.bk); + if entry.single_buffered { + continue; + } + if entry.workgroup_bytes(stage) > u64::from(policy.max_workgroup_storage_bytes()) { + continue; + } + let (rg, cg) = entry.subgroup_split(); + let threads = rg * cg * subgroup_width; + if threads == 0 || threads > policy.max_workgroup_lanes() { + continue; + } + let k_iterations = k.div_ceil(bk); + let Some(score) = split_candidates(k_iterations, has_epilogues) + .flat_map(|splits| staging_depths(splits).map(move |buffers| (splits, buffers))) + .map(|(splits, buffers)| { + score_fs( + m, + k, + n, + batch, + probe_group.max(1), + entry, + rg, + cg, + splits, + buffers, + stage, + policy, + subgroup_width, + ) + }) + .min() + else { + continue; + }; + // The score is invariant to `n_passes` by construction — a p-pass + // profile does p times the per-workgroup work over a p-times-smaller + // grid — and the recorded sweeps say fewer passes wins at every + // measured shape, so it leads the tie-break. The two Reverse levels + // reproduce the previous selector's tie-breaks and are reached. + let key = (score, entry.n_passes, Reverse(bm), Reverse(bn)); + let padded_macs = u128::from(m.div_ceil(bm)) + * u128::from(n.div_ceil(bn)) + * u128::from(batch) + * u128::from(bm) + * u128::from(bn) + * u128::from(u64::from(k.div_ceil(bk)) * u64::from(bk)); + if best.as_ref().is_none_or(|(best_key, ..)| key < *best_key) { + best = Some((key, CoopTile::new(bm, bn, bk), rg, cg, padded_macs)); + } + } + let (_, tile, rg, cg, padded_macs) = best?; + // Even the best tile may waste more than a quarter of its work on padding + // — degenerate (gemv-shaped) contractions with a tiny M or N. Those belong + // to the vector/generic families; declining here is the routing signal. + // It compares against kernel families this model cannot cost, so it stays + // a guard rather than a term. + let useful_macs = u128::from(m) + * u128::from(n) + * u128::from(batch) + * (u128::from(k.div_ceil(tile.bk)) * u128::from(tile.bk)); + if padded_macs * 4 > useful_macs * 5 { + return None; + } + Some((tile, rg, cg)) +} + +/// Traversal-order parameter for the dense coop grid, decided per shape +/// from the raw swizzle sweeps (`bench_coop_tiles` `SWIZZLE=1,4,8,16`, both +/// element types, round-robin minima; winners with margins): +/// +/// | shape (m,k,n) | f32 | f16 | +/// |--------------------|--------------------|--------------| +/// | 16384,384,1536 | tie | tie | +/// | 16384,1536,384 | sw8 +4.5% | sw1 +4.4% | +/// | 384,16384,1536 | sw8 +4% | sw8 +3.6% | +/// | 16384,384,384 | tie | tie | +/// | 4096,4096,4096 | sw4 +2.7% (sw1 2nd)| sw1 +4.3% | +/// | 16384,3072,1536 | sw1 +7% | sw1 +2.2% | +/// +/// The mechanisms the winners trace: the swizzle exists to share B column +/// slabs across the resident wavefront, and it pays exactly when B is too +/// big to sit in the LLC; K-deep contractions are A-streaming-bound and +/// prefer plain row-major order (the swizzle's M-jumps break the A read +/// stream); halved f16 operands halve cache pressure and shift small-B +/// shapes to row-major too. The one point the rule leaves on the table is +/// 4096-cube f32 (sw4 beats the rule's sw8 by 2.7%) — an isolated winner a +/// dedicated branch would overfit. +pub(super) fn swizzle_group_m( + m: usize, + k: usize, + n: usize, + datatype: crate::DataTypeEnum, +) -> u32 { + let _ = m; + let element_size = datatype.element_size() as u128; + let b_bytes = k as u128 * n as u128 * element_size; + const LLC_CLASS: u128 = 32 << 20; + const SMALL_B: u128 = 4 << 20; + if k >= 2048 && b_bytes <= LLC_CLASS { + // K-deep, cacheable B: streaming A wins (f32 +7%, f16 +2.2%; + // 4096-cube f16 +4.3%). + return 1; + } + if element_size == 2 && b_bytes <= SMALL_B { + // f16 with a small B: row-major wins (+4.4% at 16384x1536x384). + return 1; + } + // B-slab sharing pays when B outsizes the LLC (+4-4.5% at + // 384x16384x1536, both dtypes) and measures as a tie on small grids. + fusor_tile_ir_kernels::DEFAULT_SWIZZLE_GROUP_M +} + +/// The dense matmul family's complete routing for one contraction, pure in +/// device-derived inputs: family variant, cooperative plan, and traversal +/// group. One entry point makes the whole decision surface golden-testable — +/// new scoring terms change this function's output table or they change +/// nothing. +#[cfg(test)] +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct DensePlan { + pub variant: super::variants::DenseMatmulVariant, + /// `(tile, row_groups, col_groups, splits, stage_buffers)`. + pub coop: Option<(CoopTile, u32, u32, u32, u32)>, + pub swizzle_group_m: u32, +} + +#[cfg(test)] +#[allow(clippy::too_many_arguments)] +pub(crate) fn plan_dense_matmul( + m: usize, + k: usize, + n: usize, + batch: u32, + probe_group: u32, + segments: u32, + datatype: crate::DataTypeEnum, + policy: &DispatchPolicy, + max_subgroup_size: u32, + caps: crate::kernel_selection::KernelDeviceCaps, +) -> DensePlan { + let ctx = super::variants::DenseMatmulCtx { + coop_kinds: super::variants::dense_coop_kinds_from_datatype(datatype), + }; + let shape = crate::kernel_selection::KernelShape::new([m, k, n]); + let variant = super::variants::dense_matmul_selector() + .select(shape, &ctx, caps) + .expect("dense matmul selector has a catch-all rule"); + // Exactly how production composes the two stages: the tile at the fixed + // probe group (allocation precedes the merge partition), the split count + // at the grid that actually launches. + let coop = (variant == super::variants::DenseMatmulVariant::Coop) + .then(|| { + let (tile, rg, cg) = plan_coop_tile( + m as u32, + k as u32, + n as u32, + batch, + datatype, + false, + probe_group, + policy, + max_subgroup_size, + )?; + let (splits, buffers) = plan_coop_splits( + CoopDispatch { + m: m as u32, + k: k as u32, + n: n as u32, + batch, + segments, + datatype, + has_epilogues: false, + }, + tile, + rg, + cg, + policy, + max_subgroup_size, + ); + Some((tile, rg, cg, splits, buffers)) + }) + .flatten(); + DensePlan { + variant, + coop, + swizzle_group_m: swizzle_group_m(m, k, n, datatype), + } +} diff --git a/fusor-ml/core/src/matmul/kernel.rs b/fusor-ml/core/src/matmul/kernel.rs index 8f1325d8c..03039d377 100644 --- a/fusor-ml/core/src/matmul/kernel.rs +++ b/fusor-ml/core/src/matmul/kernel.rs @@ -15,13 +15,16 @@ use crate::{ tile_storage_write_with_direct_layout_typed, }, }, + nary_direct::apply_typed_unary_function_chain, nary_wise::{NaryExpr, NaryFunction, NaryOp, NaryScalar, UnaryFunctionChain}, reduce::{ReduceFunction, ReduceOp, ReduceOperation}, tensor::{DataTypeEnum, TensorData}, }; use super::{ - MatMulOperation, MatMulParams, MatrixOperand, coop_gemm, sgemm, sgemv, + MatMulOperation, MatMulParams, MatrixOperand, + cost::CoopDispatch, + sgemm, sgemv, variants::{CoopTile, dense_coop_kinds_from_datatype, select_dense_matmul_params}, }; @@ -29,6 +32,46 @@ fn device_supported(value: Option) -> Result u32 { + if splits <= 1 { + return 1; + } + let slice_bytes = + u64::from(batch_m_padded) * u64::from(n_padded) * datatype.element_size() as u64; + if slice_bytes == 0 { + return 1; + } + let slices = output.buffer().size() / slice_bytes; + splits + .min(u32::try_from(slices.saturating_sub(1)).unwrap_or(u32::MAX)) + .max(1) +} + impl MatMulOperation { pub fn new( datatype: DataTypeEnum, @@ -98,6 +141,30 @@ impl MatMulOperation { matches!(self.datatype, DataTypeEnum::F32 | DataTypeEnum::F16) } + /// The cooperative kernel hosts dtype-preserving unary chains, plus post + /// chains that widen f16 operands into an f32 output (the fused form of + /// matmul-then-cast, which training's mixed-precision backward emits for + /// every weight gradient). The store lands in the chain's output dtype; + /// the in-place epilogue rounds back to the operand dtype before the + /// chain reads it, so the fused result matches the unfused one exactly. + /// Narrowing chains would round the accumulator ahead of the chain, so + /// they keep using the generic fused reduction. + fn coop_epilogues_supported(&self) -> bool { + let post_out = self.post_element_wise.out_datatype(); + self.pre_element_wise.iter().all(|chain| { + chain.input_datatype() == self.datatype && chain.out_datatype() == self.datatype + }) && self.post_element_wise.input_datatype() == self.datatype + && (post_out == self.datatype + || (self.datatype == DataTypeEnum::F16 && post_out == DataTypeEnum::F32)) + } + + fn has_elementwise_epilogues(&self) -> bool { + self.pre_element_wise + .iter() + .any(|chain| !chain.functions.is_empty()) + || !self.post_element_wise.functions.is_empty() + } + /// The contraction in its composed map-reduce form: a multiply over the /// `[batch.., m, n, k]` index space summed along `k`, with the fused /// pre/post chains inlined and accumulation upgraded to f32 (matching @@ -191,52 +258,132 @@ impl MatMulOperation { self.coop_tile(device).is_some() } + /// The contraction as the planner sees it, at the given launched-segment + /// count. `None` when the shape is not a coop contraction at all. + fn coop_dispatch(&self, segments: u32) -> Option { + Some(super::cost::CoopDispatch { + m: self.a.rows().try_into().ok()?, + k: self.a.cols().try_into().ok()?, + n: self.b.cols().try_into().ok()?, + batch: self + .a + .batch_shape() + .iter() + .try_fold(1u32, |acc, &dim| acc.checked_mul(u32::try_from(dim).ok()?))?, + segments, + datatype: self.datatype, + has_epilogues: self.has_elementwise_epilogues(), + }) + } + + /// Split count and staged tile pairs for the dispatch that will actually + /// run this contraction: `(1, _)` for the single-pass body. `segments` is + /// the real launched grid depth — 1 standalone, `segments.len()` for a + /// merged dispatch — never a probe. + fn coop_splits( + &self, + device: &Device, + tile: CoopTile, + rg: u32, + cg: u32, + segments: u32, + ) -> (u32, u32) { + let Some(dispatch) = self.coop_dispatch(segments) else { + return (1, 2); + }; + let (splits, buffers) = super::cost::plan_coop_splits( + dispatch, + tile, + rg, + cg, + &device.dispatch_policy(), + device.max_subgroup_size(), + ); + if device.config().trace_splitk { + let CoopDispatch { m, k, n, batch, .. } = dispatch; + eprintln!( + "matmul_plan name={} m={m} k={k} n={n} batch={batch} segments={segments} \ + tile={}x{}x{} rg={rg} cg={cg} splits={splits} buffers={buffers}", + self.name(), + tile.bm, + tile.bn, + tile.bk, + ); + } + (splits, buffers) + } + /// The tile geometry the cooperative-matrix kernel would run with on /// this device, `None` when any static gate fails and the contraction is /// bound for the generic path. Shapes need not divide the tile: edge /// tiles mask their fills and the output allocation pads to whole tiles. - pub(crate) fn coop_tile(&self, device: &Device) -> Option { + pub(crate) fn coop_tile(&self, device: &Device) -> Option<(CoopTile, u32, u32)> { if !self.can_use_hardware_matmul() || (self.datatype == DataTypeEnum::F16 && !device.f16_supported()) - || !self.pre_element_wise[0].functions.is_empty() - || !self.pre_element_wise[1].functions.is_empty() - || !self.post_element_wise.functions.is_empty() + || !self.coop_epilogues_supported() { return None; } - let MatMulParams::CoopMatMul(params) = &self.parameters else { + let MatMulParams::CoopMatMul = &self.parameters else { return None; }; - device.coop_token(params.kind())?; + let kind = *dense_coop_kinds_from_datatype(self.datatype).first()?; + device.coop_token(kind)?; let subgroup_config = device.subgroup_config()?; if !subgroup_config.is_fixed() { return None; } - let (m, k, n): (u32, u32, u32) = ( - self.a.rows().try_into().ok()?, - self.a.cols().try_into().ok()?, - self.b.cols().try_into().ok()?, - ); - let batch = self - .a - .batch_shape() - .iter() - .try_fold(1u32, |acc, &dim| acc.checked_mul(u32::try_from(dim).ok()?))?; - let limits = device.limits(); - let tile = CoopTile::select( + let CoopDispatch { m, k, n, - limits - .max_compute_workgroup_size_x - .min(limits.max_compute_invocations_per_workgroup), - subgroup_config.max_size(), + batch, + has_epilogues, + .. + } = self.coop_dispatch(1)?; + let limits = device.limits(); + // Memoized on the device: the scored selection is asked once per + // static-viability probe, once per prep, once per allocation and once + // per trace for every matmul in every resolve, and it enumerates the + // whole table against every legal split count each time. + let probe_group = super::cost::tile_probe_group(device, has_epilogues); + let [bm, bn, bk, row_groups, col_groups] = device.coop_tile_memo( + crate::device::CoopTileKey { + m, + k, + n, + batch, + datatype: self.datatype, + has_epilogues, + probe_group, + }, + || { + super::cost::plan_coop_tile( + m, + k, + n, + batch, + self.datatype, + has_epilogues, + probe_group, + &device.dispatch_policy(), + subgroup_config.max_size(), + ) + .map(|(tile, rg, cg)| [tile.bm, tile.bn, tile.bk, rg, cg]) + }, )?; - let total_tiles = m + let tile = CoopTile::new(bm, bn, bk); + // The 1D->3D grid spread plus the kernels' overhang guard cover any + // u32 tile count; the checked math above is the only real bound. A + // per-dimension cap here silently dropped real-vocab lm-head shapes + // (16384x384x32768 = 65536 tiles) onto the generic fallback at ~17x + // the cost. + let _ = m .div_ceil(tile.bm) .checked_mul(n.div_ceil(tile.bn)) .and_then(|tiles| tiles.checked_mul(batch))?; - (total_tiles <= limits.max_compute_workgroups_per_dimension).then_some(tile) + let _ = limits; + Some((tile, row_groups, col_groups)) } /// Row-major strides of the logical output over its padded backing: @@ -255,13 +402,19 @@ impl MatMulOperation { strides.into() } - fn build_hardware_matmul( + /// The shared head of the cooperative-matrix lowering: flatten the + /// operand layouts, validate the contraction geometry, pick the tile, + /// and verify the output allocation carries the tile-padded backing. + /// Used by both the standalone [`Self::build_hardware_matmul`] and the + /// horizontally merged builder ([`build_merged_matmul_kernel`]), so the + /// two agree on every gate by construction. + fn hardware_matmul_prep( &self, device: &Device, input_a: &TensorData, input_b: &TensorData, output: &TensorData, - ) -> Result { + ) -> Result { // Operands with a base map read their producer through it: compose // with the runtime buffer layout, then flatten with the operand's // dim grouping. @@ -309,14 +462,9 @@ impl MatMulOperation { let shape = tile_ir_kernels::DenseMatmulShape { batch, m, k, n }; // Only the cooperative-matrix route stays hand-specialized; gemv - // shapes lower through the generic subgroup-per-output reduce, and - // fused chains lower through the generic tiled reduce. - let tile = device_supported(self.coop_tile(device))?; - let subgroup_config = device_supported(device.subgroup_config())?; - let MatMulParams::CoopMatMul(params) = &self.parameters else { - return Err(kernel_backend::DeviceNotSupported); - }; - let coop = device_supported(device.coop_token(params.kind()))?; + // shapes and dtype-changing fused chains lower through the generic + // row reduction. Dtype-preserving unary chains remain hosted here. + let (tile, row_groups, col_groups) = device_supported(self.coop_tile(device))?; // The store covers whole tiles, so `y` is the padded matrix: rows // padded to `ceil(m / bm) * bm` per batch and columns to @@ -333,7 +481,8 @@ impl MatMulOperation { .checked_mul(m_padded as usize) .and_then(|rows| rows.checked_mul(n_padded as usize)), )?; - let padded_bytes = padded_elements as u64 * self.datatype.element_size() as u64; + let padded_bytes = + padded_elements as u64 * self.post_element_wise.out_datatype().element_size() as u64; if output.layout().offset() != 0 || output.layout().strides() != &*expected_strides || padded_bytes > output.buffer().size() @@ -351,9 +500,99 @@ impl MatMulOperation { &[n_padded, 1], ), }; + Ok(HardwareMatmulPrep { + a_view, + b_view, + y_view, + shape, + tile, + row_groups, + col_groups, + batch_m_padded, + n_padded, + }) + } + + fn build_hardware_matmul( + &self, + device: &Device, + input_a: &TensorData, + input_b: &TensorData, + output: &TensorData, + ) -> Result { + let HardwareMatmulPrep { + a_view, + b_view, + y_view, + shape, + tile, + row_groups, + col_groups, + batch_m_padded, + n_padded, + } = self.hardware_matmul_prep(device, input_a, input_b, output)?; + let subgroup_config = device_supported(device.subgroup_config())?; + let MatMulParams::CoopMatMul = &self.parameters else { + return Err(kernel_backend::DeviceNotSupported); + }; + let kind = *device_supported(dense_coop_kinds_from_datatype(self.datatype).first())?; + let coop = device_supported(device.coop_token(kind))?; let max_wg_per_dim = device.limits().max_compute_workgroups_per_dimension; let datatype = self.datatype; + + let make_epilogue = |label, chain: &UnaryFunctionChain| { + if chain.functions.is_empty() { + return None; + } + let chain = chain.clone(); + Some(tile_ir_kernels::UnaryEpilogue::new(label, move |value| { + apply_typed_unary_function_chain(value, datatype, &chain) + .expect("cooperative matmul epilogue validated before kernel construction") + .0 + })) + }; + let pre_a = make_epilogue("dense_matmul_pre_a", &self.pre_element_wise[0]); + let pre_b = make_epilogue("dense_matmul_pre_b", &self.pre_element_wise[1]); + let post = make_epilogue("dense_matmul_post", &self.post_element_wise); + + // Starved tile grids with a long contraction split K across + // workgroups: partials land in scratch slices of the over-allocated + // output buffer and a combine kernel folds them (sum-reorder-only + // numerics). A weight-gradient shape like 64×2048×64 otherwise runs + // as a single workgroup. + let (standalone_splits, stage_buffers) = + self.coop_splits(device, tile, row_groups, col_groups, 1); + let splits = splits_fitting_allocation( + standalone_splits, + output, + batch_m_padded, + n_padded, + self.datatype, + ); + if splits > 1 + && let Some(kernel) = self.build_split_k_matmul( + device, + input_a, + input_b, + output, + &a_view, + &b_view, + &y_view, + shape, + tile, + row_groups, + col_groups, + subgroup_config, + coop, + batch_m_padded, + n_padded, + splits, + ) + { + return Ok(kernel); + } + let used = std::cell::Cell::new(false); let ir = tile_ir::tile::build(|phase| { let element = match datatype { @@ -361,9 +600,14 @@ impl MatMulOperation { DataTypeEnum::F16 => tile_ir::ElementType::F16, _ => unreachable!("hardware matmul only supports f32/f16"), }; + let out_element = match self.post_element_wise.out_datatype() { + DataTypeEnum::F32 => tile_ir::ElementType::F32, + DataTypeEnum::F16 => tile_ir::ElementType::F16, + _ => unreachable!("hardware matmul only supports f32/f16 outputs"), + }; let a = tile_storage_read_with_direct_layout_typed(phase, element, a_view.clone()); let b = tile_storage_read_with_direct_layout_typed(phase, element, b_view.clone()); - let y = tile_storage_write_with_direct_layout_typed(phase, element, y_view.clone()); + let y = tile_storage_write_with_direct_layout_typed(phase, out_element, y_view.clone()); used.set(tile_ir_kernels::try_batched_coop_matmul( phase, tile_ir_kernels::DenseMatmulTensors { @@ -373,9 +617,9 @@ impl MatMulOperation { }, shape, &tile_ir_kernels::DenseMatmulEpilogues { - pre_a: None, - pre_b: None, - post: None, + pre_a: pre_a.as_ref(), + pre_b: pre_b.as_ref(), + post: post.as_ref(), }, max_wg_per_dim, tile_ir_kernels::DenseCoopMatmulConfig { @@ -386,6 +630,16 @@ impl MatMulOperation { bn: tile.bn, bk: tile.bk, }, + row_groups, + col_groups, + staging: None, + stage_buffers, + swizzle_group_m: super::cost::swizzle_group_m( + self.a.rows(), + self.a.cols(), + self.b.cols(), + self.datatype, + ), }, )); }); @@ -409,7 +663,7 @@ impl MatMulOperation { let cache_key = self.kernel_cache_key_with_dispatch(variant, None, dispatch_size, &inputs); let name = self.name(); - let pipeline = kernel_backend::three_buffer_pipeline_from_ir( + let (pipeline, cached) = kernel_backend::three_buffer_pipeline_from_ir( device.kernel_cache(), &name, cache_key, @@ -420,6 +674,7 @@ impl MatMulOperation { kernel_backend::DirectKernel::from_prepared_three_buffer_pipeline( name, pipeline, + Some(cached), input_a.buffer().clone(), input_b.buffer().clone(), output.buffer().clone(), @@ -427,9 +682,583 @@ impl MatMulOperation { ), ) } + + /// Split-K route for coop matmuls whose tile grid starves the GPU: the + /// partials kernel runs `splits × total_tiles` workgroups, each covering + /// one K-span with the standard coop tile loop and storing an + /// unnormalized partial into one scratch slice of the over-allocated + /// output buffer (slices `1..=splits`, allocated by [`Self::inputs`]); + /// a combine kernel sums the slices into the padded output at slice 0. + /// Keeping the scratch inside the output allocation means every bound + /// buffer stays slot-attributable, so flush-plan recording keeps + /// working. Numerics differ from the single-pass kernel only in + /// summation order. Returns `None` (single-pass coop path proceeds) + /// when the geometry, allocation, or device declines. + #[allow(clippy::too_many_arguments)] + fn build_split_k_matmul( + &self, + device: &Device, + input_a: &TensorData, + input_b: &TensorData, + output: &TensorData, + a_view: &crate::mir::tile_direct::DirectMatrixLayout, + b_view: &crate::mir::tile_direct::DirectMatrixLayout, + y_view: &crate::mir::tile_direct::DirectMatrixLayout, + shape: tile_ir_kernels::DenseMatmulShape, + tile: CoopTile, + row_groups: u32, + col_groups: u32, + subgroup_config: fusor_tile_ir_kernels::SubgroupConfig, + coop: tile_ir::CoopMatrixToken, + batch_m_padded: u32, + n_padded: u32, + splits: u32, + ) -> Option { + let slice_elements = batch_m_padded.checked_mul(n_padded)?; + let total_elements = slice_elements.checked_mul(splits.checked_add(1)?)?; + let required_bytes = total_elements as u64 * self.datatype.element_size() as u64; + // The output allocation must carry the scratch slices; an exact + // allocation (a plan built before the split decision, or an aliased + // buffer) falls back to the single-pass kernel. + if output.buffer().size() < required_bytes { + return None; + } + let scratch_rows = splits.checked_mul(batch_m_padded)?; + let scratch_view = crate::mir::tile_direct::DirectMatrixLayout { + rows: scratch_rows, + cols: n_padded, + offset: slice_elements, + layout: tile_ir::Layout::strided( + tile_ir::MemoryLevel::Storage, + tile_ir::Shape::new([scratch_rows, n_padded]), + &[n_padded, 1], + ), + }; + let element = match self.datatype { + DataTypeEnum::F32 => tile_ir::ElementType::F32, + DataTypeEnum::F16 => tile_ir::ElementType::F16, + _ => return None, + }; + let max_wg_per_dim = device.limits().max_compute_workgroups_per_dimension; + + let used = std::cell::Cell::new(false); + let ir = tile_ir::tile::build(|phase| { + let a = tile_storage_read_with_direct_layout_typed(phase, element, a_view.clone()); + let b = tile_storage_read_with_direct_layout_typed(phase, element, b_view.clone()); + let y = + tile_storage_write_with_direct_layout_typed(phase, element, scratch_view.clone()); + used.set(tile_ir_kernels::try_batched_coop_matmul_split_k( + phase, + tile_ir_kernels::DenseMatmulTensors { + a: &a, + b: &b, + y: &y, + }, + shape, + splits, + max_wg_per_dim, + tile_ir_kernels::DenseCoopMatmulConfig { + coop, + subgroups: subgroup_config, + tile: tile_ir_kernels::DenseCoopMatmulTile { + bm: tile.bm, + bn: tile.bn, + bk: tile.bk, + }, + row_groups, + col_groups, + staging: None, + // The partials body stages one pair; see `staging_depths`. + stage_buffers: 1, + swizzle_group_m: super::cost::swizzle_group_m( + self.a.rows(), + self.a.cols(), + self.b.cols(), + self.datatype, + ), + }, + )); + }); + if !used.get() { + if device.config().trace_splitk { + eprintln!("splitk_declined_by_kernel name={}", self.name()); + } + return None; + } + let dispatch_size = ir.grid; + if dispatch_size.iter().any(|dim| *dim > max_wg_per_dim) { + return None; + } + let inputs = [ + input_a.clone().into(), + input_b.clone().into(), + output.clone().into(), + ]; + let variant = + kernel_backend::KernelVariantKey::with_payload::(|state| { + tile.hash(state); + subgroup_config.hash(state); + splits.hash(state); + 1u64.hash(state); + }); + let cache_key = self.kernel_cache_key_with_dispatch(variant, None, dispatch_size, &inputs); + let name = self.name(); + let (pipeline, cached) = kernel_backend::three_buffer_pipeline_from_ir( + device.kernel_cache(), + &name, + cache_key, + || Some(ir), + )?; + let partials = kernel_backend::DirectKernel::from_prepared_three_buffer_pipeline( + name.clone(), + pipeline, + Some(cached), + input_a.buffer().clone(), + input_b.buffer().clone(), + output.buffer().clone(), + dispatch_size, + ); + + // One read-write view over all `1 + splits` slices: the combine + // reads the partial slices and stores slice 0, through a single + // binding of the shared buffer. + debug_assert_eq!(y_view.offset, 0); + debug_assert_eq!( + y_view.rows, batch_m_padded, + "split-K expects the padded output view" + ); + let all_rows = scratch_rows + batch_m_padded; + let combine_view = crate::mir::tile_direct::DirectMatrixLayout { + rows: all_rows, + cols: n_padded, + offset: 0, + layout: tile_ir::Layout::strided( + tile_ir::MemoryLevel::Storage, + tile_ir::Shape::new([all_rows, n_padded]), + &[n_padded, 1], + ), + }; + let combine_ir = tile_ir::tile::build(|phase| { + let y = + tile_storage_write_with_direct_layout_typed(phase, element, combine_view.clone()); + tile_ir_kernels::split_k_combine( + phase, + &y, + batch_m_padded, + n_padded, + splits, + max_wg_per_dim, + ); + }); + let combine_dispatch = combine_ir.grid; + if combine_dispatch.iter().any(|dim| *dim > max_wg_per_dim) { + return None; + } + let combine_variant = + kernel_backend::KernelVariantKey::with_payload::(|state| { + tile.hash(state); + subgroup_config.hash(state); + splits.hash(state); + 2u64.hash(state); + }); + let combine_key = + self.kernel_cache_key_with_dispatch(combine_variant, None, combine_dispatch, &inputs); + let combine = kernel_backend::dynamic_kernel_from_ir( + device.kernel_cache(), + format!("{name}_split_combine"), + combine_key, + || Some(combine_ir), + [output.buffer().clone()], + combine_dispatch, + )?; + + Some(kernel_backend::DirectKernel::sequence( + name, + vec![partials, combine], + )) + } } struct HardwareMatmulVariant; +struct SplitKMatmulVariant; +struct MergedMatmulVariant; + +/// The horizontal-merge compatibility key of a dense matmul: two matmuls +/// merge into one dispatch only when every field matches, which makes the +/// guarded segment bodies identical up to their storage bindings (same +/// logical shape, tile geometry, workgroup size, split factor, and element +/// type). Only matmuls that will take the cooperative-matrix route produce a +/// key. +#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] +pub(crate) struct MatmulMergeKey { + m: u32, + k: u32, + n: u32, + batch: u32, + /// Whether this profile would split its K loop as a standalone dispatch. + /// The tile and the split count themselves are pure functions of the + /// fields above, so they cannot discriminate two keys these do not; this + /// one derived bool survives only to keep the split-K wave category where + /// it is. + split_candidate: bool, + datatype: DataTypeEnum, +} + +impl MatmulMergeKey { + /// Whether the profile is a split-K candidate, which is the wave category + /// its one consumer wants. + pub(crate) fn splits(&self) -> Option { + self.split_candidate.then_some(1) + } +} + +impl MatMulOperation { + /// See [`MatmulMergeKey`]. `None` = not horizontally mergeable. + pub(crate) fn merge_profile(&self, device: &Device) -> Option { + // Standalone coop kernels host unary epilogues. The guarded merged + // body does not yet carry per-segment epilogue identities/bindings. + if self.has_elementwise_epilogues() { + return None; + } + let (tile, row_groups, col_groups) = self.coop_tile(device)?; + let CoopDispatch { m, k, n, batch, .. } = self.coop_dispatch(1)?; + Some(MatmulMergeKey { + m, + k, + n, + batch, + // Scored at the probe group, not at 1: the category exists to + // keep split-K profiles in their own merged wave, so the question + // is what the merged dispatch will do, and the merged dispatch is + // bounded by exactly this group. + split_candidate: self + .coop_splits( + device, + tile, + row_groups, + col_groups, + super::cost::tile_probe_group(device, false), + ) + .0 + > 1, + datatype: self.datatype, + }) + } +} + +/// One kernel running several independent same-profile dense matmuls (see +/// [`MatmulMergeKey`]): the guarded-segment counterpart of the standalone +/// cooperative-matrix lowering, emitted through +/// [`tile_ir_kernels::try_merged_coop_matmul`]. Split-K profiles produce a +/// two-dispatch sequence — all segments' partials in one kernel, then all +/// combines in another — with each segment's scratch carved out of its own +/// over-allocated output (slot-attributable for flush-plan recording, like +/// the standalone split-K route). +/// +/// Returns `None` when any segment fails the hardware gates (the caller +/// falls back to per-segment kernels and poisons the recording). +pub(crate) fn build_merged_matmul_kernel( + graph: &crate::compute_graph::ComputeGraphInner, + segments: &[MatMulOperation], + segment_inputs: &[Vec], +) -> Option { + let device = graph.device(); + macro_rules! decline { + ($reason:expr) => {{ + if device.config().trace_matmul_merge { + eprintln!("matmul_merge_decline reason={}", $reason); + } + return None; + }}; + } + let first = segments.first()?; + let mut tensors = Vec::with_capacity(segments.len()); + for (op, inputs) in segments.iter().zip(segment_inputs) { + let [input_a, input_b, output] = inputs.as_slice() else { + decline!("input_arity"); + }; + let (Some(input_a), Some(input_b), Some(output)) = + (input_a.as_tensor(), input_b.as_tensor(), output.as_tensor()) + else { + decline!("input_tensors"); + }; + if !op.can_use_hardware_matmul() + || input_a.datatype() != op.datatype + || input_b.datatype() != op.datatype + || output.datatype() != op.datatype + || (op.datatype == DataTypeEnum::F16 && !device.f16_supported()) + { + decline!("datatype_gate"); + } + tensors.push((input_a, input_b, output)); + } + + let mut preps = Vec::with_capacity(segments.len()); + for (op, (input_a, input_b, output)) in segments.iter().zip(&tensors) { + let Ok(prep) = op.hardware_matmul_prep(&device, input_a, input_b, output) else { + decline!(format!("prep {}", op.name())); + }; + preps.push(prep); + } + let tile = preps[0].tile; + let shape = preps[0].shape; + let batch_m_padded = preps[0].batch_m_padded; + let n_padded = preps[0].n_padded; + // The merge key guarantees profile equality; re-verify structurally so a + // drifted caller can never emit mismatched guarded bodies. + if preps.iter().any(|prep| { + prep.tile != tile + || prep.shape.batch != shape.batch + || prep.shape.m != shape.m + || prep.shape.k != shape.k + || prep.shape.n != shape.n + || prep.batch_m_padded != batch_m_padded + || prep.n_padded != n_padded + }) { + decline!("profile_mismatch"); + } + if segments.iter().any(|op| op.datatype != first.datatype) { + decline!("datatype_mismatch"); + } + // The one place the grid that actually launches is known: + // `splits x tiles x segments.len()` workgroups. The tile was scored at a + // fixed probe group (allocation had to precede this partition), so the + // split the real grid wants may exceed the scratch that was allocated; + // clamping is always legal — fewer splits is a correct kernel. + let (merged_splits, stage_buffers) = first.coop_splits( + &device, + tile, + preps[0].row_groups, + preps[0].col_groups, + segments.len() as u32, + ); + let splits = splits_fitting_allocation( + merged_splits, + tensors + .iter() + .map(|(_, _, output)| *output) + .min_by_key(|output| output.buffer().size()) + .expect("merged matmul builds from a non-empty segment list"), + batch_m_padded, + n_padded, + first.datatype, + ); + + let Some(subgroup_config) = device.subgroup_config() else { + decline!("subgroups"); + }; + let MatMulParams::CoopMatMul = &first.parameters else { + decline!("params"); + }; + let Some(&kind) = dense_coop_kinds_from_datatype(first.datatype).first() else { + decline!("coop_kind"); + }; + let Some(coop) = device.coop_token(kind) else { + decline!("coop_token"); + }; + let element = match first.datatype { + DataTypeEnum::F32 => tile_ir::ElementType::F32, + DataTypeEnum::F16 => tile_ir::ElementType::F16, + _ => decline!("element"), + }; + let max_wg_per_dim = device.limits().max_compute_workgroups_per_dimension; + let config = tile_ir_kernels::DenseCoopMatmulConfig { + coop, + subgroups: subgroup_config, + tile: tile_ir_kernels::DenseCoopMatmulTile { + bm: tile.bm, + bn: tile.bn, + bk: tile.bk, + }, + row_groups: preps[0].row_groups, + col_groups: preps[0].col_groups, + staging: None, + stage_buffers, + swizzle_group_m: super::cost::swizzle_group_m( + first.a.rows(), + first.a.cols(), + first.b.cols(), + first.datatype, + ), + }; + + // Split-K segments store partials into scratch slices of their own + // output allocation; verify every allocation carries the slices. + let slice_elements = batch_m_padded.checked_mul(n_padded)?; + if splits > 1 { + let total_elements = slice_elements.checked_mul(splits.checked_add(1)?)?; + let required_bytes = total_elements as u64 * first.datatype.element_size() as u64; + if tensors + .iter() + .any(|(_, _, output)| output.buffer().size() < required_bytes) + { + decline!("scratch_capacity"); + } + } + + let name = if device.config().trace_decode_names { + format!( + "merged_matmul[{}]", + segments + .iter() + .map(|op| op.name()) + .collect::>() + .join("; ") + ) + } else { + format!("merged_matmul_x{}", segments.len()) + }; + + let used = std::cell::Cell::new(false); + let ir = tile_ir::tile::build(|phase| { + let mut storages = Vec::with_capacity(segments.len()); + for prep in &preps { + let a = tile_storage_read_with_direct_layout_typed(phase, element, prep.a_view.clone()); + let b = tile_storage_read_with_direct_layout_typed(phase, element, prep.b_view.clone()); + // Partials land in the scratch slices (`1..=splits`). + let y_view = if splits > 1 { + crate::mir::tile_direct::DirectMatrixLayout { + rows: splits * batch_m_padded, + cols: n_padded, + offset: slice_elements, + layout: tile_ir::Layout::strided( + tile_ir::MemoryLevel::Storage, + tile_ir::Shape::new([splits * batch_m_padded, n_padded]), + &[n_padded, 1], + ), + } + } else { + prep.y_view.clone() + }; + let y = tile_storage_write_with_direct_layout_typed(phase, element, y_view); + storages.push((a, b, y)); + } + let segment_tensors: Vec = storages + .iter() + .map(|(a, b, y)| tile_ir_kernels::DenseMatmulTensors { a, b, y }) + .collect(); + used.set(tile_ir_kernels::try_merged_coop_matmul( + phase, + &segment_tensors, + shape, + splits, + max_wg_per_dim, + config, + )); + }); + if !used.get() { + decline!("tile_ir_declined"); + } + let dispatch_size = ir.grid; + if dispatch_size.iter().any(|dim| *dim > max_wg_per_dim) { + return None; + } + let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + kernel_backend::KernelVariantKey::of::().hash(state); + dispatch_size.hash(state); + tile.hash(state); + subgroup_config.hash(state); + splits.hash(state); + 1u64.hash(state); + crate::compute_graph::resolve::plan_cache::hash_merged_segments( + state, + segments.iter(), + segment_inputs, + ); + }); + let buffers: Vec> = tensors + .iter() + .flat_map(|(input_a, input_b, output)| { + [ + input_a.buffer().clone(), + input_b.buffer().clone(), + output.buffer().clone(), + ] + }) + .collect(); + let Some(main) = kernel_backend::dynamic_kernel_from_ir( + device.kernel_cache(), + name.clone(), + cache_key, + move || Some(ir), + buffers, + dispatch_size, + ) else { + decline!("pipeline"); + }; + if splits <= 1 { + return Some(main); + } + + // The merged combine: every segment's `(1 + splits)`-slice buffer bound + // once read-write, folded by guarded ranges in the same segment order. + let all_rows = (splits + 1) * batch_m_padded; + let combine_ir = tile_ir::tile::build(|phase| { + let storages: Vec = preps + .iter() + .map(|_| { + tile_storage_write_with_direct_layout_typed( + phase, + element, + crate::mir::tile_direct::DirectMatrixLayout { + rows: all_rows, + cols: n_padded, + offset: 0, + layout: tile_ir::Layout::strided( + tile_ir::MemoryLevel::Storage, + tile_ir::Shape::new([all_rows, n_padded]), + &[n_padded, 1], + ), + }, + ) + }) + .collect(); + let ys: Vec<&tile_ir::tile::Storage> = storages.iter().collect(); + tile_ir_kernels::merged_split_k_combine( + phase, + &ys, + batch_m_padded, + n_padded, + splits, + max_wg_per_dim, + ); + }); + let combine_dispatch = combine_ir.grid; + if combine_dispatch.iter().any(|dim| *dim > max_wg_per_dim) { + return None; + } + let combine_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + kernel_backend::KernelVariantKey::of::().hash(state); + combine_dispatch.hash(state); + tile.hash(state); + subgroup_config.hash(state); + splits.hash(state); + 2u64.hash(state); + crate::compute_graph::resolve::plan_cache::hash_merged_segments( + state, + segments.iter(), + segment_inputs, + ); + }); + let combine_buffers: Vec> = tensors + .iter() + .map(|(_, _, output)| output.buffer().clone()) + .collect(); + let Some(combine) = kernel_backend::dynamic_kernel_from_ir( + device.kernel_cache(), + format!("{name}_split_combine"), + combine_key, + move || Some(combine_ir), + combine_buffers, + combine_dispatch, + ) else { + decline!("combine_pipeline"); + }; + Some(kernel_backend::DirectKernel::sequence( + name, + vec![main, combine], + )) +} impl Operation for MatMulOperation { fn hash_kernel_fields(&self, state: &mut FxHasher) { @@ -453,8 +1282,12 @@ impl Operation for MatMulOperation { MatMulParams::MatMul(sgemm_params) => { sgemm::workgroup_shape_constraints(self, device, sgemm_params) } - MatMulParams::CoopMatMul(coop_params) => { - coop_gemm::workgroup_shape_constraints(self, device, coop_params) + // The cooperative kernels carry their own grid and block + // (`ir.grid`); this shape only reaches the generic fused-reduce + // fallback, so it is the fallback's own policy. + MatMulParams::CoopMatMul => { + crate::row_program::RowProgramOperation::from_reduce(&self.as_fused_reduce()) + .workgroup_shape_constraints(device) } } } @@ -493,13 +1326,10 @@ impl Operation for MatMulOperation { workgroup_shape, sgemm_params, ), - MatMulParams::CoopMatMul(coop_params) => coop_gemm::dispatch_size( - last_dim_size, - second_to_last_dim_size, - batch_size, - workgroup_shape, - coop_params, - ), + MatMulParams::CoopMatMul => { + crate::row_program::RowProgramOperation::from_reduce(&self.as_fused_reduce()) + .dispatch_size(workgroup_shape, inputs) + } } } @@ -508,6 +1338,11 @@ impl Operation for MatMulOperation { f(self.second); } + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + f(&mut self.first); + f(&mut self.second); + } + fn inputs( &self, nodes: &crate::compute_graph::ComputeGraphInner, @@ -518,19 +1353,32 @@ impl Operation for MatMulOperation { let datatype = self.post_element_wise.out_datatype(); // The coop kernel stores whole tiles: pad the backing to tile // multiples and view the logical shape over it (consumers never - // read the pad region). Shapes that already divide the tile — and - // anything bound for the generic path — allocate exactly. + // read the pad region). Split-K shapes over-allocate one extra + // padded slice per split for the partials scratch (slice 0 is the + // output; the combine kernel folds slices 1..=splits into it). + // Shapes that already divide the tile — and anything bound for the + // generic path — allocate exactly. let (m, n) = (self.a.rows(), self.b.cols()); - let padded = self.coop_tile(device).and_then(|tile| { + let padded = self.coop_tile(device).and_then(|(tile, rg, cg)| { let m_padded = m.div_ceil(tile.bm as usize) * tile.bm as usize; let n_padded = n.div_ceil(tile.bn as usize) * tile.bn as usize; - (m_padded != m || n_padded != n).then_some((m_padded, n_padded)) + // Allocation predates the merge partition, so it sizes against the + // same fixed probe group the tile was scored at. A merged build of + // a shorter tail chunk may want more splits than this; that build + // clamps, rather than every matmul over-allocating for the worst + // case the way the group-1 sizing did. + let probe = super::cost::tile_probe_group(device, self.has_elementwise_epilogues()); + let slices = self.coop_splits(device, tile, rg, cg, probe).0 as usize + 1; + (slices > 1 || m_padded != m || n_padded != n).then_some((m_padded, n_padded, slices)) }); let output_tensor = match padded { - Some((m_padded, n_padded)) => { + Some((m_padded, n_padded, slices)) => { let batch: usize = self.a.batch_shape().iter().product(); - let backing = - TensorData::new_for_shape(device, &[batch, m_padded, n_padded], datatype); + let backing = TensorData::new_for_shape( + device, + &[slices * batch, m_padded, n_padded], + datatype, + ); TensorData::new_from_parts( device, backing.buffer().clone(), @@ -562,7 +1410,7 @@ impl Operation for MatMulOperation { if self.can_use_hardware_matmul() && input_a.datatype() == self.datatype && input_b.datatype() == self.datatype - && output.datatype() == self.datatype + && output.datatype() == self.post_element_wise.out_datatype() && (self.datatype != DataTypeEnum::F16 || graph.device().f16_supported()) && let Ok(kernel) = self.build_hardware_matmul(&graph.device(), input_a, input_b, output) @@ -572,6 +1420,25 @@ impl Operation for MatMulOperation { // Everything else is the composed contraction's own lowering: the // generic tiled (or serial) fused reduce, identical to what any // unrecognized contraction gets. + if std::env::var_os("FUSOR_TRACE_MATMUL_FALLBACK").is_some() { + tracing::warn!( + "matmul fallback to row reduce: dtype={:?} params={:?} m={} k={} n={} batch={:?} hw={} epi={} f16_dev={} coop_tile={:?} a_dt={:?} b_dt={:?} out_dt={:?}", + self.datatype, + std::mem::discriminant(&self.parameters), + self.a.rows(), + self.a.cols(), + self.b.cols(), + self.a.batch_shape(), + self.can_use_hardware_matmul(), + self.coop_epilogues_supported(), + graph.device().f16_supported(), + self.coop_tile(&graph.device()) + .map(|(tile, ..)| (tile.bm, tile.bn, tile.bk)), + input_a.datatype(), + input_b.datatype(), + output.datatype(), + ); + } let reduce = self.as_fused_reduce(); crate::row_program::RowProgramOperation::from_reduce(&reduce).build_direct_kernel( graph, diff --git a/fusor-ml/core/src/matmul/mod.rs b/fusor-ml/core/src/matmul/mod.rs index 8447ee6b9..8b150c471 100644 --- a/fusor-ml/core/src/matmul/mod.rs +++ b/fusor-ml/core/src/matmul/mod.rs @@ -1,9 +1,8 @@ use crate::{ - Device, Layout, Tensor, compute_graph::NodeIndex, kernel_selection::CooperativeMatrixKind, - nary_wise::UnaryFunctionChain, tensor::DataTypeEnum, + Layout, Tensor, compute_graph::NodeIndex, nary_wise::UnaryFunctionChain, tensor::DataTypeEnum, }; -pub mod coop_gemm; +mod cost; mod kernel; pub mod sgemm; mod sgemm_params; @@ -11,25 +10,25 @@ pub mod sgemv; mod sgemv_params; mod variants; +pub(crate) use kernel::{MatmulMergeKey, build_merged_matmul_kernel}; pub(crate) use variants::CoopTile; -use variants::select_dense_matmul_params; -pub fn get_optimal_params(m: usize, n: usize, k: usize, device: &Device) -> MatMulParams { - select_dense_matmul_params(m, n, k, device, &[CooperativeMatrixKind::F32F32M8N8K8]) -} - -#[derive(Debug, Clone, Hash)] -pub enum MatMulParams { +#[derive(Debug, Clone, PartialEq, Hash)] +pub(crate) enum MatMulParams { Vector(sgemv::SgemvParams), MatMul(sgemm::SgemmParams), - CoopMatMul(coop_gemm::CoopGemmParams), + /// The cooperative-matrix family. Neither geometry nor the matrix + /// kind is a parameter: the scored tile selection derives geometry per + /// kernel build and the kind follows from the datatype, so dispatch, + /// allocation, and kernel agree by construction. + CoopMatMul, } /// An affine relayout between an operand's dims and its node's logical /// space: conv's sliding windows. The kernels concretize it lazily — the /// coop path composes it with the node's runtime buffer layout, the generic /// reduce substitutes it into the load coordinates. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub(crate) struct OperandBaseMap { pub(crate) layout: Layout, pub(crate) base_shape: Box<[usize]>, @@ -51,7 +50,7 @@ impl std::hash::Hash for OperandBaseMap { /// im2col operand keeps the windowed view's dims (mapped onto the node by /// `base_map`), and the kernels divmod the flat matrix coordinates back /// apart per load. -#[derive(Debug, Clone, Hash)] +#[derive(Debug, Clone, PartialEq, Hash)] pub(crate) struct MatrixOperand { pub(crate) shape: Box<[usize]>, pub(crate) batch_dims: usize, @@ -135,7 +134,7 @@ impl MatrixOperand { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub(crate) struct MatMulOperation { pub(crate) datatype: DataTypeEnum, pub(crate) first: NodeIndex, @@ -206,35 +205,6 @@ impl Tensor { })); product.sum(k_dim) } - - /// Matrix multiply with explicit kernel parameters: a tuning/benchmark - /// API. The parameters cannot round-trip through the composed graph, so - /// the operation builds directly against materialized inputs and - /// executes eagerly, returning a fresh leaf tensor. - pub fn mat_mul_with_parameters(&self, other: &Self, parameters: MatMulParams) -> Self { - assert_eq!(self.datatype(), other.datatype()); - self.data.materialize(); - other.data.materialize(); - let operation = MatMulOperation::new( - self.datatype(), - self.key(), - other.key(), - self.shape(), - other.shape(), - Some(parameters), - self.device(), - ); - let output = self - .device() - .compute_graph() - .execute_eager(&operation) - .unwrap_or_else(|| { - panic!( - "mat_mul_with_parameters could not build a kernel for the requested parameters" - ) - }); - Tensor::from(output) - } } #[cfg(test)] @@ -389,73 +359,336 @@ mod selection_tests { ); } + /// Generator for the dense-plan golden table: run with + /// `cargo test -p fusor-core dense_plan_golden -- --nocapture` after a + /// deliberate selection change and paste the printed rows below. #[test] - fn direct_tile_coop_selector_prefers_largest_supported_tile() { - let select = - |m, k, n, max_workgroup_size_x| CoopTile::select(m, k, n, max_workgroup_size_x, 32); - // 4096³ (square) hits Tile128x512 — it has fewer barriers than - // Tile256x256 because it's double-buffered. - assert_eq!( - select(4096, 4096, 4096, 512), - Some(CoopTile::new(128, 512, 16)) - ); - // Shapes where N is divisible by 256 but not 512 — with enough - // tiles — fall to Tile256x256 single-buffer. - assert_eq!( - select(8192, 1024, 4352, 512), - Some(CoopTile::new(256, 256, 16)) - ); - // N=512 doesn't divide 256 on the M side... actually wait, 4096 % 256 == 0. - // For shapes where N is divisible by 512 but M isn't by 256, fall to - // Tile128x512. - assert_eq!( - select(384, 1024, 1024, 512), - Some(CoopTile::new(128, 64, 16)) - ); - // 1024³ doesn't have enough tiles for Tile128x512 OR Tile128x256; - // falls back to Tile128x64 for better parallelism. - assert_eq!( - select(1024, 1024, 1024, 512), - Some(CoopTile::new(128, 64, 16)) - ); - // 8192x256 has tiles_for(128, 256) = 64*1 = 64 — below the threshold, - // so it falls to Tile128x64. - assert_eq!( - select(8192, 1024, 256, 256), - Some(CoopTile::new(128, 64, 16)) - ); - // M=4096, N=1024 gives tiles_for(128, 256) = 32*4 = 128. Below 256. - // Falls to Tile128x64. - assert_eq!( - select(4096, 1024, 1024, 256), - Some(CoopTile::new(128, 64, 16)) - ); - // M=8192, N=512 gives tiles_for(128, 256) = 64*2 = 128 (still <256), - // so falls to Tile128x64. To hit Tile128x256 we need a wider shape: - // 8192x1024 → 64*4 = 256 ✓. + fn dense_plan_golden() { + let policy = apple_policy(64 << 10, 32 << 10); + let shapes: [(usize, usize, usize); 8] = [ + (16384, 384, 1536), + (16384, 1536, 384), + (384, 16384, 1536), + (16384, 384, 384), + (4096, 4096, 4096), + (16384, 3072, 1536), + (64, 2048, 64), + (1, 4096, 4096), + ]; + let mut rows = Vec::new(); + for &(m, k, n) in &shapes { + for datatype in [crate::DataTypeEnum::F32, crate::DataTypeEnum::F16] { + let plan = super::cost::plan_dense_matmul( + m, + k, + n, + 1, + PROBE_GROUP, + 1, + datatype, + &policy, + 32, + caps(true), + ); + rows.push(format!( + "{m}x{k}x{n} {datatype:?} => {:?} tile={:?} groups={:?} splits={:?} \ + buffers={:?} sw={}", + plan.variant, + plan.coop.map(|(tile, ..)| (tile.bm, tile.bn, tile.bk)), + plan.coop.map(|(_, rg, cg, ..)| (rg, cg)), + plan.coop.map(|(.., splits, _)| splits), + plan.coop.map(|(.., buffers)| buffers), + plan.swizzle_group_m + )); + } + } + // The split count sizes the launched grid, so a horizontally merged + // dispatch of `group` same-profile contractions splits less. The tile + // is held at the probe group throughout: allocation precedes the + // partition, so only the split count may move with it. + for group in [2, 4, 8, 16] { + let plan = super::cost::plan_dense_matmul( + 64, + 2048, + 64, + 1, + PROBE_GROUP, + group, + crate::DataTypeEnum::F32, + &policy, + 32, + caps(true), + ); + rows.push(format!( + "64x2048x64 F32 group={group} => splits={:?} buffers={:?}", + plan.coop.map(|(.., splits, _)| splits), + plan.coop.map(|(.., buffers)| buffers) + )); + } + let golden = GOLDEN_PLANS + .trim() + .lines() + .map(str::trim) + .collect::>(); + for row in &rows { + println!("{row}"); + } assert_eq!( - select(8192, 1024, 1024, 256), - Some(CoopTile::new(128, 256, 16)) + rows, golden, + "dense matmul routing changed; regenerate deliberately" ); - // N=128 doesn't divide 256 so Tile128x256/Tile128x512 are out; falls - // back to Tile128x64. + } + + /// The locked routing surface. The `Coop tile=None` row is truthful and + /// interesting: the family selector picks the coop family for the + /// gemv-shaped contraction while the tile scorer's padding gate then + /// declines, which production resolves through the generic fallback — + /// the selector does not consult the tile scorer. + const GOLDEN_PLANS: &str = " + 16384x384x1536 F32 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(1) buffers=Some(2) sw=8 + 16384x384x1536 F16 => Coop tile=Some((64, 64, 16)) groups=Some((2, 2)) splits=Some(1) buffers=Some(2) sw=1 + 16384x1536x384 F32 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(1) buffers=Some(2) sw=8 + 16384x1536x384 F16 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(1) buffers=Some(2) sw=1 + 384x16384x1536 F32 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(2) buffers=Some(1) sw=8 + 384x16384x1536 F16 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(2) buffers=Some(1) sw=8 + 16384x384x384 F32 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(1) buffers=Some(2) sw=8 + 16384x384x384 F16 => Coop tile=Some((64, 64, 16)) groups=Some((2, 2)) splits=Some(1) buffers=Some(2) sw=1 + 4096x4096x4096 F32 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(1) buffers=Some(2) sw=8 + 4096x4096x4096 F16 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(1) buffers=Some(2) sw=1 + 16384x3072x1536 F32 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(1) buffers=Some(2) sw=1 + 16384x3072x1536 F16 => Coop tile=Some((128, 64, 16)) groups=Some((4, 2)) splits=Some(1) buffers=Some(2) sw=1 + 64x2048x64 F32 => Coop tile=Some((64, 64, 16)) groups=Some((2, 2)) splits=Some(32) buffers=Some(1) sw=1 + 64x2048x64 F16 => Coop tile=Some((64, 64, 16)) groups=Some((2, 2)) splits=Some(32) buffers=Some(1) sw=1 + 1x4096x4096 F32 => Coop tile=None groups=None splits=None buffers=None sw=8 + 1x4096x4096 F16 => Coop tile=None groups=None splits=None buffers=None sw=1 + 64x2048x64 F32 group=2 => splits=Some(32) buffers=Some(1) + 64x2048x64 F32 group=4 => splits=Some(32) buffers=Some(1) + 64x2048x64 F32 group=8 => splits=Some(32) buffers=Some(1) + 64x2048x64 F32 group=16 => splits=Some(16) buffers=Some(1) + "; + + /// The group the tile is scored at everywhere in production: the + /// horizontal merger's own maximum, `budget / MATMUL_SEGMENT_BINDINGS`. + const PROBE_GROUP: u32 = 10; + + /// A policy carrying this machine's measured rates, so the selection + /// tests exercise the same decision surface production does. + fn apple_policy( + max_workgroup_lanes: u32, + max_workgroup_storage_bytes: u32, + ) -> crate::occupancy::DispatchPolicy { + crate::occupancy::DispatchPolicy::from_parts( + 64 << 10, + 32, + max_workgroup_lanes, + 8 << 20, + max_workgroup_storage_bytes, + crate::device::APPLE_MATMUL_RATES, + ) + } + + fn select_with_lanes(m: u32, k: u32, n: u32, max_lanes: u32) -> Option { + let policy = apple_policy(max_lanes, 32 << 10); + super::cost::plan_coop_tile( + m, + k, + n, + 1, + crate::DataTypeEnum::F32, + false, + PROBE_GROUP, + &policy, + 32, + ) + .map(|(tile, ..)| tile) + } + + /// At WebGPU's 16 KB default workgroup-storage limit the footprint + /// filter is live: every 64-wide-or-more f32 entry overflows two staged + /// pairs, so the 4096-cube falls back to the narrow 64x16 profile it + /// would never pick at 32 KB, while f16 halves the staged bytes and + /// keeps the profile the full limit chooses. + #[test] + fn footprint_filter_at_16kb_limit() { + let select = |datatype, storage| { + let policy = apple_policy(64 << 10, storage); + super::cost::plan_coop_tile( + 4096, + 4096, + 4096, + 1, + datatype, + false, + PROBE_GROUP, + &policy, + 32, + ) + .map(|(tile, ..)| tile) + }; assert_eq!( - select(1024, 1024, 128, 256), + select(crate::DataTypeEnum::F32, 32 << 10), Some(CoopTile::new(128, 64, 16)) ); assert_eq!( - select(1024, 1024, 1024, 128), - Some(CoopTile::new(64, 64, 16)) + select(crate::DataTypeEnum::F32, 16 << 10), + Some(CoopTile::new(64, 16, 16)) ); - // M=1000 divides nothing: the masked-edge fallback picks the tile - // with the least padded work (all candidates pad M to 1024; the - // preference order breaks the tie toward the biggest tile). assert_eq!( - select(1000, 1024, 1024, 512), + select(crate::DataTypeEnum::F16, 16 << 10), Some(CoopTile::new(128, 64, 16)) ); - // Tiny N inflates every candidate's padding past the waste bound: - // gemv-shaped contractions stay on the generic path. - assert_eq!(select(1000, 1024, 16, 512), None); + } + + /// Every selection must be legal and within the padding bound for every + /// (shape, caps) combination — the properties the scorer guarantees by + /// construction, checked over a deterministic sweep. + #[test] + fn scored_selection_properties() { + let mut lcg = 0x5eed_1234u64; + let mut next = |range: u32| { + lcg = lcg + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((lcg >> 33) as u32) % range + 1 + }; + for _ in 0..2000 { + let (m, k, n) = (next(9000), next(5000), next(9000)); + let max_lanes = [128u32, 256, 512, 1024][(next(4) - 1) as usize]; + let Some(tile) = select_with_lanes(m, k, n, max_lanes) else { + continue; + }; + let entry = fusor_tile_ir_kernels::coop_tile_entries() + .iter() + .find(|entry| entry.tile.bm == tile.bm && entry.tile.bn == tile.bn) + .expect("selected tile must exist in the kernel table"); + let threads = entry.subgroups * 32; + assert!( + threads <= max_lanes, + "m={m} k={k} n={n} lanes={max_lanes}: illegal tile {tile:?}" + ); + let padded = u64::from(m.div_ceil(tile.bm)) + * u64::from(tile.bm) + * u64::from(n.div_ceil(tile.bn)) + * u64::from(tile.bn); + assert!( + padded * 4 <= u64::from(m) * u64::from(n) * 5, + "m={m} k={k} n={n}: padding bound violated by {tile:?}" + ); + } + } + + /// The sgemv bucket table and the sgemm regression tree are measured + /// policies for the non-cooperative fallback families; every cell they + /// can produce must still be structurally legal (kernel divisibility, + /// workgroup lane bounds, shared-memory budget) for every shape. + #[test] + fn fallback_family_params_are_legal_everywhere() { + let mut lcg = 0x0fa1_1bac_c5u64; + let mut next = |range: u32| { + lcg = lcg + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((lcg >> 33) as u32) % range + 1 + }; + for _ in 0..4000 { + let (m, k, n) = ( + next(20_000) as usize, + next(20_000) as usize, + next(20_000) as usize, + ); + + let gemm = crate::matmul::sgemm_params::gemm_parameters(m, n, k); + let (bm, bn, bk) = ( + gemm.block_m_size(), + gemm.block_n_size(), + gemm.block_k_size(), + ); + let (tm, tn) = (gemm.thread_m_size(), gemm.thread_n_size()); + assert!( + bm.is_multiple_of(tm) && bn.is_multiple_of(tn), + "m={m} n={n} k={k}: thread tile must divide the block tile ({gemm:?})" + ); + let lanes = (bm * bn) / (tm * tn); + assert!( + (32..=1024).contains(&lanes), + "m={m} n={n} k={k}: workgroup lanes {lanes} out of range ({gemm:?})" + ); + // A and B staging tiles, doubled when double-buffered, must fit + // Apple's 32 KB workgroup-memory floor. + let buffers = if gemm.double_buffer() { 2 } else { 1 }; + let smem_bytes = u64::from((bm + bn) * bk) * 4 * buffers; + assert!( + smem_bytes <= 32 * 1024, + "m={m} n={n} k={k}: {smem_bytes}B of workgroup memory ({gemm:?})" + ); + + let gemv = crate::matmul::sgemv_params::gemv_parameters(m, n, k); + assert!( + gemv.chunk_size() >= 1 + && matches!(gemv.vector_size(), 1 | 2 | 4) + && (1..=32).contains(&gemv.subgroups_per_workgroup()), + "m={m} n={n} k={k}: illegal gemv params ({gemv:?})" + ); + } } } + +#[cfg(test)] +mod split_k_tests { + //! GPU gates for automatic split-K selection and aligned-span codegen. + + use crate::{Device, Tensor}; + + fn check_dense_split_k(m: usize, k: usize, n: usize) { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let values = |len: usize, freq: f32| -> Vec { + (0..len).map(|i| ((i as f32) * freq).sin()).collect() + }; + let a_data = values(m * k, 0.13); + let b_data = values(k * n, 0.07); + let a = Tensor::from_slice(&device, [m, k], &a_data); + let b = Tensor::from_slice(&device, [k, n], &b_data); + let out = a.mat_mul(&b); + let actual = out.as_slice::<2, f32>().await.unwrap(); + for mi in 0..m { + for ni in 0..n { + let mut acc = 0.0f64; + for ki in 0..k { + acc += a_data[mi * k + ki] as f64 * b_data[ki * n + ni] as f64; + } + let want = acc as f32; + let got = actual[[mi, ni]]; + assert!( + (got - want).abs() < 2e-3 + want.abs() * 1e-3, + "m={m} k={k} n={n} [{mi}, {ni}]: got {got}, expected {want}" + ); + } + } + }); + } + + // The 64×2048×64 weight-gradient shape: K-tiles divide the fan-out, so + // the spans partition K exactly and the K bounds are elided (the vec4 + // staging fast path). + #[test] + fn dense_split_k_elided_bounds() { + check_dense_split_k(64, 2048, 64); + } + + // Ragged K (1000): no useful divisor alignment, the last span overruns + // the logical K extent and the bounds stay live under the dense flag. + #[test] + fn dense_split_k_ragged_k() { + check_dense_split_k(64, 1000, 64); + } + + // Barely past the split gate (k = 520): short trailing spans idle. + #[test] + fn dense_split_k_short_spans() { + check_dense_split_k(64, 520, 64); + } +} + diff --git a/fusor-ml/core/src/matmul/sgemm_params.rs b/fusor-ml/core/src/matmul/sgemm_params.rs index c4c770c59..d8729e644 100644 --- a/fusor-ml/core/src/matmul/sgemm_params.rs +++ b/fusor-ml/core/src/matmul/sgemm_params.rs @@ -1,3 +1,10 @@ +//! Measured selection tree for the generic tiled-matmul fallback. This +//! family only runs where cooperative matrices are unavailable (no-subgroup +//! devices and the WebGPU baseline), so the leaves are measurement, not +//! derivation — rederiving them requires benching those targets. Structural +//! legality of every leaf is pinned by +//! `fallback_family_params_are_legal_everywhere`. + use crate::sgemm::SgemmParams; #[inline] diff --git a/fusor-ml/core/src/matmul/sgemv.rs b/fusor-ml/core/src/matmul/sgemv.rs index 53542e779..17dce80a7 100644 --- a/fusor-ml/core/src/matmul/sgemv.rs +++ b/fusor-ml/core/src/matmul/sgemv.rs @@ -49,7 +49,7 @@ pub(crate) fn workgroup_shape_constraints( constraints } -#[derive(Debug, Clone, Hash)] +#[derive(Debug, Clone, PartialEq, Hash)] pub struct SgemvParams { chunk_size: u32, vector_size: u32, diff --git a/fusor-ml/core/src/matmul/sgemv_params.rs b/fusor-ml/core/src/matmul/sgemv_params.rs index d15f56dea..e1da0b541 100644 --- a/fusor-ml/core/src/matmul/sgemv_params.rs +++ b/fusor-ml/core/src/matmul/sgemv_params.rs @@ -1,3 +1,9 @@ +//! Measured (m, k) bucket table for the vector-family fallback. This family +//! only runs where cooperative matrices are unavailable (no-subgroup devices +//! and the WebGPU baseline), so the cells are measurement, not derivation — +//! rederiving them requires benching those targets. Structural legality of +//! every cell is pinned by `fallback_family_params_are_legal_everywhere`. + use crate::sgemv::SgemvParams; #[inline] diff --git a/fusor-ml/core/src/matmul/variants.rs b/fusor-ml/core/src/matmul/variants.rs index cf91ec12d..78f6ebd0f 100644 --- a/fusor-ml/core/src/matmul/variants.rs +++ b/fusor-ml/core/src/matmul/variants.rs @@ -8,10 +8,9 @@ use crate::{ tensor::DataTypeEnum, }; -use super::{MatMulParams, coop_gemm}; +use super::MatMulParams; pub(super) const DENSE_M: Axis<0> = Axis; -pub(super) const DENSE_K: Axis<1> = Axis; pub(super) const DENSE_N: Axis<2> = Axis; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -40,15 +39,8 @@ pub(super) fn dense_matmul_selector() -> ShapeSelector<3, DenseMatmulCtx, DenseM ShapeSelector::new() .rule( DenseMatmulVariant::Coop, - ShapeRule::new().when(|shape: KernelShape<3>, ctx: &DenseMatmulCtx, caps| { - coop_gemm_params_from_caps( - shape[DENSE_M], - shape[DENSE_N], - shape[DENSE_K], - caps, - ctx.coop_kinds, - ) - .is_some() + ShapeRule::new().when(|_shape: KernelShape<3>, ctx: &DenseMatmulCtx, caps| { + coop_supported(caps, ctx.coop_kinds) }), ) .rule( @@ -76,53 +68,27 @@ pub(super) fn select_dense_matmul_params( .select(shape, &ctx, caps) .expect("dense matmul selector has a catch-all rule") { - DenseMatmulVariant::Coop => MatMulParams::CoopMatMul( - coop_gemm::optimal_params(m, n, k, device, select_coop_kind(caps, coop_kinds)) - .expect("coop selector and coop parameter selection disagree"), - ), + DenseMatmulVariant::Coop => MatMulParams::CoopMatMul, DenseMatmulVariant::Vector => MatMulParams::Vector(gemv_parameters(m, n, k)), DenseMatmulVariant::MatMul => MatMulParams::MatMul(gemm_parameters(m, n, k)), } } -pub(super) fn coop_gemm_params_from_caps( - m: usize, - n: usize, - _k: usize, - caps: KernelDeviceCaps, - coop_kinds: &[CooperativeMatrixKind], -) -> Option { - if !caps.subgroups_supported - || !coop_kinds +/// Whether the cooperative-matrix family is available at all on this +/// device: fixed-width subgroups, a supported cooperative kind, and room +/// for at least the smallest coop workgroup. Geometry is not consulted — +/// the scored tile selection decides it per kernel build, and shapes it +/// declines lower through the generic fused reduction. +pub(super) fn coop_supported(caps: KernelDeviceCaps, coop_kinds: &[CooperativeMatrixKind]) -> bool { + caps.subgroups_supported + && coop_kinds .iter() .any(|kind| caps.cooperative_matrix.supports(*kind)) - || caps.min_subgroup_size != caps.max_subgroup_size - || caps.max_compute_workgroup_size_x < 64 - { - return None; - } - - let mut params = coop_gemm::CoopGemmParams::default(); - if n <= 16 { - params.block_n = 16; - params.n_passes = 1; - } else if n <= 32 { - params.block_n = 32; - params.n_passes = 2; - } - - if m <= 16 { - params.block_m = 16; - params.wg_threads = 64; - } else if m < params.block_m as usize { - params.block_m = 64; - params.wg_threads = 128; - } - - params.kind = select_coop_kind(caps, coop_kinds); - (params.wg_threads <= caps.max_compute_workgroup_size_x).then_some(params) + && caps.min_subgroup_size == caps.max_subgroup_size + && caps.max_compute_workgroup_size_x >= 64 } +#[cfg(test)] pub(super) fn select_coop_kind( caps: KernelDeviceCaps, coop_kinds: &[CooperativeMatrixKind], @@ -134,10 +100,10 @@ pub(super) fn select_coop_kind( .expect("coop selector called with no supported cooperative matrix kind") } -/// (BM, BN, BK) tile dimensions for a cooperative-matrix matmul tile. The -/// `select` helper below returns `Option` (`None` = no coop variant -/// fits the shape); the kernel layer uses the tuple to look up the matching -/// ROW_GROUPS/COL_GROUPS/N_PASSES/BLOCK in its internal table. +/// (BM, BN, BK) tile dimensions for a cooperative-matrix matmul tile. Chosen +/// by [`super::cost::plan_coop_tile`], which also derives the subgroup split +/// the kernel runs with; the kernel layer looks the remaining execution +/// properties (N_PASSES, BLOCK) up by geometry in its own table. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub(crate) struct CoopTile { pub(crate) bm: u32, @@ -150,107 +116,4 @@ impl CoopTile { Self { bm, bn, bk } } - const fn subgroup_groups(self) -> u32 { - match (self.bm, self.bn, self.bk) { - (256, 256, 16) => 8, - (128, 512, 16) => 8, - (128, 256, 16) => 8, - (128, 128, 16) => 16, - (128, 64, 16) => 8, - (64, 128, 16) => 8, - (64, 64, 16) => 4, - _ => 0, - } - } - - fn workgroup_size_supported(self, max_workgroup_size_x: u32, max_subgroup_size: u32) -> bool { - self.subgroup_groups() - .checked_mul(max_subgroup_size) - .is_some_and(|block| block <= max_workgroup_size_x) - } - - /// Pick a cooperative-matrix tile for the given (m, k, n) shape, returning - /// `None` when no coop tile fits. All entries use BK=16 to keep - /// double-buffered workgroup tiles inside Apple's 32 KB limit; the - /// (256, 256, 16) entry runs single-buffered in the inner perf kernel. - /// Heuristic: bigger tiles only fire when (M/BM)*(N/BN) clears a minimum - /// tile count so there's enough work for the GPU. - pub(super) fn select( - m: u32, - k: u32, - n: u32, - max_workgroup_size_x: u32, - max_subgroup_size: u32, - ) -> Option { - let tiles_for = |bm: u32, bn: u32| -> u32 { (m / bm) * (n / bn) }; - if m == 0 || n == 0 || k == 0 { - return None; - } - // Tile256x256 single-buffer has lower memory traffic (sqrt-min) but - // 2× the barriers of Tile128x512 double-buffer; only fires when N - // is divisible by 256 but not by 512. - if m.is_multiple_of(256) - && n.is_multiple_of(256) - && !n.is_multiple_of(512) - && tiles_for(256, 256) >= 256 - { - let tile = Self::new(256, 256, 16); - if tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { - return Some(tile); - } - } - if m.is_multiple_of(128) && n.is_multiple_of(512) && tiles_for(128, 512) >= 256 { - let tile = Self::new(128, 512, 16); - if tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { - return Some(tile); - } - } - if m.is_multiple_of(128) && n.is_multiple_of(256) && tiles_for(128, 256) >= 256 { - let tile = Self::new(128, 256, 16); - if tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { - return Some(tile); - } - } - if m.is_multiple_of(128) && n.is_multiple_of(64) { - let tile = Self::new(128, 64, 16); - if tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { - return Some(tile); - } - } - if m.is_multiple_of(64) && n.is_multiple_of(128) { - let tile = Self::new(64, 128, 16); - if tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { - return Some(tile); - } - } - if m.is_multiple_of(64) && n.is_multiple_of(64) { - let tile = Self::new(64, 64, 16); - if tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { - return Some(tile); - } - } - - // Shapes that divide no tile run with masked edge tiles: pick the - // candidate minimizing padded work, in preference order on ties. - // Selections whose padding inflates the output by more than a - // quarter stay on the generic path — that bound also keeps - // gemv-shaped contractions (tiny M or N) off the tile kernels. - // Candidates stick to geometries the aligned rules already reach - // (the (128, 128) table entry was never selectable and miscomputes). - let mut best: Option<(u64, Self)> = None; - for (bm, bn) in [(128, 64), (64, 128), (64, 64)] { - let tile = Self::new(bm, bn, 16); - if !tile.workgroup_size_supported(max_workgroup_size_x, max_subgroup_size) { - continue; - } - let padded = u64::from(m.div_ceil(bm) * bm) * u64::from(n.div_ceil(bn) * bn); - if padded * 4 > u64::from(m) * u64::from(n) * 5 { - continue; - } - if best.is_none_or(|(best_padded, _)| padded < best_padded) { - best = Some((padded, tile)); - } - } - best.map(|(_, tile)| tile) - } } diff --git a/fusor-ml/core/src/mir/kernel_backend.rs b/fusor-ml/core/src/mir/kernel_backend.rs index 447170764..06ed7bc13 100644 --- a/fusor-ml/core/src/mir/kernel_backend.rs +++ b/fusor-ml/core/src/mir/kernel_backend.rs @@ -1,6 +1,6 @@ pub(crate) use fusor_tile_ir_runtime::{ - DirectKernel, KernelCacheKey, KernelVariantKey, PreparedDirectDispatch, dynamic_kernel_from_ir, - run_direct_kernel, run_kernel, three_buffer_pipeline_from_ir, + DirectKernel, DirectKernelTemplate, KernelCacheKey, KernelVariantKey, PreparedDirectDispatch, + dynamic_kernel_from_ir, run_direct_kernel, run_kernel, three_buffer_pipeline_from_ir, }; /// Marker returned by device-specific direct-kernel builders when the current @@ -16,6 +16,4 @@ impl std::fmt::Display for DeviceNotSupported { impl std::error::Error for DeviceNotSupported {} -pub(crate) mod mirostat; -pub(crate) mod sampling_topk; -pub(crate) mod standard_sampler; +pub(crate) mod sampling; diff --git a/fusor-ml/core/src/mir/kernel_backend/mirostat.rs b/fusor-ml/core/src/mir/kernel_backend/mirostat.rs deleted file mode 100644 index ffce0ee2e..000000000 --- a/fusor-ml/core/src/mir/kernel_backend/mirostat.rs +++ /dev/null @@ -1,128 +0,0 @@ -use std::hash::Hash; - -use crate::{ - Device, - mir::kernel_backend, - sampling::{ - GPU_SAMPLE_RESULT_WORDS, GpuMirostat2Sampler, GpuMirostat2SamplerParams, TOP_K_BLOCK, - row_kernels, - }, - tensor::{DataTypeEnum, TensorData}, -}; -use wgpu::CommandEncoder; - -#[repr(C)] -#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)] -struct Mirostat2Params { - tau: f32, - eta: f32, - random: f32, - _padding: f32, -} - -struct Mirostat2SortedTopKKernelVariant; - -fn mirostat2_params_data(device: &Device, params: GpuMirostat2SamplerParams) -> TensorData { - let params = Mirostat2Params { - tau: params.tau, - eta: params.eta, - random: params.random.clamp(0.0, 0.999_999_94), - _padding: 0.0, - }; - let buffer = device.create_buffer_init( - bytemuck::bytes_of(¶ms), - wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC, - ); - TensorData::new_from_buffer(device, buffer, &[4], DataTypeEnum::F32) -} - -pub(crate) fn sample_from_sorted_top_k_data_with_encoder( - ids: &TensorData, - values: &TensorData, - sampler: &mut GpuMirostat2Sampler, - params: GpuMirostat2SamplerParams, - exactness_flag: Option<&TensorData>, - encoder: Option<&mut CommandEncoder>, -) -> Option { - if ids.datatype() != DataTypeEnum::U32 || values.datatype() != DataTypeEnum::F32 { - return None; - } - if ids.layout().rank() != 1 || values.layout().rank() != 1 { - return None; - } - if let Some(flag) = exactness_flag - && (flag.datatype() != DataTypeEnum::U32 - || flag.layout().rank() != 1 - || flag.layout().shape()[0] == 0 - || !values.device().is_same_device(flag.device())) - { - return None; - } - - let top_k = params - .top_k - .min(ids.layout().shape()[0]) - .min(values.layout().shape()[0]); - if top_k == 0 { - return None; - } - let device = values.device(); - let params = mirostat2_params_data(device, params); - let has_exactness_flag = exactness_flag.is_some(); - let output = TensorData::new_for_shape(device, &[GPU_SAMPLE_RESULT_WORDS], DataTypeEnum::U32); - let meta = row_kernels::SamplerMeta { - top_k: top_k.try_into().ok()?, - ids_offset: ids.layout().offset().try_into().ok()?, - ids_stride: ids.layout().strides()[0].try_into().ok()?, - values_offset: values.layout().offset().try_into().ok()?, - values_stride: values.layout().strides()[0].try_into().ok()?, - has_exactness_flag, - }; - let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { - kernel_backend::KernelVariantKey::of::().hash(state); - TOP_K_BLOCK.hash(state); - top_k.hash(state); - ids.layout().offset().hash(state); - ids.layout().shape().hash(state); - ids.layout().strides().hash(state); - values.layout().offset().hash(state); - values.layout().shape().hash(state); - values.layout().strides().hash(state); - has_exactness_flag.hash(state); - }); - let kernel = kernel_backend::run_kernel( - device.kernel_cache(), - "sample_mirostat2_sorted_top_k_f32", - cache_key, - [1, 1, 1], - |kb| { - row_kernels::mirostat2( - kb, - row_kernels::Mirostat2 { - ids: ids.as_kernel_tensor_ref(), - values: values.as_kernel_tensor_ref(), - state: sampler.state.as_kernel_tensor_ref(), - params: params.as_kernel_tensor_ref(), - output: output.as_kernel_tensor_ref(), - exactness_flag: exactness_flag.map(|t| t.as_kernel_tensor_ref()), - meta, - }, - ) - }, - )?; - - if let Some(encoder) = encoder { - kernel.run(device.kernel_cache(), encoder); - } else { - let mut encoder = - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("sample_mirostat2_sorted_top_k_f32 encoder"), - }); - kernel.run(device.kernel_cache(), &mut encoder); - device.wgpu_queue().submit(Some(encoder.finish())); - } - - Some(output) -} diff --git a/fusor-ml/core/src/mir/kernel_backend/sampling.rs b/fusor-ml/core/src/mir/kernel_backend/sampling.rs new file mode 100644 index 000000000..1dc7340f3 --- /dev/null +++ b/fusor-ml/core/src/mir/kernel_backend/sampling.rs @@ -0,0 +1,591 @@ +use std::hash::Hash; +use std::sync::Arc; + +use fusor_tile_ir as tile_ir; + +use crate::{ + Device, + mir::{kernel_backend, operation::hash_layout}, + sampling::{ + GPU_SAMPLE_RESULT_WORDS, GpuSamplerRequest, TOP_K_BLOCK, TOP_K_CHUNK, + processors::{ + fixed_previous_tokens_data, fixed_previous_tokens_data_with_gpu_tail, + processor_params_data, + }, + row_kernels, + }, + tensor::{DataTypeEnum, TensorData}, +}; +use wgpu::CommandEncoder; + +struct ProveTopKExactKernelVariant; +struct ChunkTopKPairsKernelVariant; +struct MergeSortedChunkTopKPairsKernelVariant; +struct Mirostat2SortedTopKKernelVariant; +struct StandardSamplerSortedTopKKernelVariant; +struct UnfilteredCategoricalSamplerKernelVariant; + +/// Build (or reuse) one sampling kernel and record it. Every stage of the +/// sampling tail binds raw buffers it already holds — no compute-graph +/// access — so they all launch through here: with an encoder the dispatch +/// joins the resolver's submission, without one it takes its own. +fn launch( + device: &Device, + name: &'static str, + cache_key: kernel_backend::KernelCacheKey, + dispatch_size: [u32; 3], + encoder: Option<&mut CommandEncoder>, + body: impl FnOnce(&mut tile_ir::KernelBuilder>) -> Option<()>, +) -> Option<()> { + let kernel = + kernel_backend::run_kernel(device.kernel_cache(), name, cache_key, dispatch_size, body)?; + kernel_backend::run_direct_kernel( + device.kernel_cache(), + device.wgpu_queue(), + &format!("{name} encoder"), + &kernel, + encoder, + ); + Some(()) +} + +/// The four-word parameter buffer every sampler kernel binds. +fn sampler_params_data(device: &Device, words: [f32; 4]) -> TensorData { + let buffer = device.create_buffer_init( + bytemuck::bytes_of(&words), + wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC, + ); + TensorData::new_from_buffer(device, buffer, &[4], DataTypeEnum::F32) +} + +/// True when any dimension of the top-k working set is zero; in that case +/// every top-k kernel short-circuits. +fn top_k_dims_empty(dims: &[usize]) -> bool { + dims.contains(&0) +} + +/// True when the inputs to `top_k_exactness` can't be sharpened any further +/// by running the exactness kernel (no candidates per output, or no top_k). +fn top_k_exactness_ineligible( + top_values_len: usize, + candidate_count: usize, + output_per_chunk: usize, + top_k: usize, +) -> bool { + top_k == 0 || top_values_len < top_k || candidate_count >= output_per_chunk +} + +pub(crate) fn top_k_exactness_flag_data_with_encoder( + top_values: &TensorData, + chunk_values: &TensorData, + chunks: usize, + candidate_count: usize, + output_per_chunk: usize, + top_k: usize, + encoder: Option<&mut CommandEncoder>, +) -> Option { + if top_values.datatype() != DataTypeEnum::F32 + || chunk_values.datatype() != DataTypeEnum::F32 + || top_values.layout().rank() != 1 + || chunk_values.layout().rank() != 1 + || !top_values.device().is_same_device(chunk_values.device()) + { + return None; + } + if top_k_exactness_ineligible( + top_values.layout().shape()[0], + candidate_count, + output_per_chunk, + top_k, + ) { + return None; + } + + let device = top_values.device(); + let flag = TensorData::new_for_shape(device, &[1], DataTypeEnum::U32); + let meta = row_kernels::TopKExactnessMeta { + chunks: chunks.try_into().ok()?, + candidate_count: candidate_count.try_into().ok()?, + output_per_chunk: output_per_chunk.try_into().ok()?, + top_k: top_k.try_into().ok()?, + top_values_offset: top_values.layout().offset().try_into().ok()?, + top_values_stride: top_values.layout().strides()[0].try_into().ok()?, + chunk_values_offset: chunk_values.layout().offset().try_into().ok()?, + chunk_values_stride: chunk_values.layout().strides()[0].try_into().ok()?, + }; + let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + kernel_backend::KernelVariantKey::of::().hash(state); + TOP_K_BLOCK.hash(state); + chunks.hash(state); + candidate_count.hash(state); + output_per_chunk.hash(state); + top_k.hash(state); + hash_layout(state, top_values.layout()); + hash_layout(state, chunk_values.layout()); + }); + launch( + device, + "prove_top_k_exact_f32", + cache_key, + [1, 1, 1], + encoder, + |kb| { + row_kernels::top_k_exactness( + kb, + top_values.as_kernel_tensor_ref(), + chunk_values.as_kernel_tensor_ref(), + flag.as_kernel_tensor_ref(), + meta, + ) + }, + )?; + + Some(flag) +} + +/// Logit-processor settings (temperature scaling and repetition penalty) +/// applied before the top-k reduction. +#[derive(Clone, Copy)] +pub(crate) struct ProcessorSettings { + pub temperature: f32, + pub repetition_penalty: f32, +} + +/// The repetition window the processors score against: the host-side token +/// history, optionally extended by a token still living on the GPU (the +/// previous pending sample, copied into the window on `encoder`). +pub(crate) struct ChunkProcessors<'a> { + pub previous_tokens: &'a [u32], + pub gpu_tail: Option<&'a TensorData>, + pub settings: ProcessorSettings, +} + +pub(crate) fn chunk_top_k_pair_data_with_encoder( + input: &TensorData, + processors: Option>, + candidate_count: usize, + output_per_chunk: usize, + mut encoder: Option<&mut CommandEncoder>, +) -> Option<(TensorData, TensorData)> { + let device = input.device(); + let processors = match processors { + None => None, + Some(processors) => { + let previous_len; + let previous_tokens; + match processors.gpu_tail { + None => { + (previous_tokens, previous_len) = + fixed_previous_tokens_data(device, processors.previous_tokens); + } + Some(gpu_tail) => { + if gpu_tail.datatype() != DataTypeEnum::U32 + || gpu_tail.layout().rank() != 1 + || gpu_tail + .layout() + .shape() + .first() + .copied() + .unwrap_or_default() + == 0 + || !device.is_same_device(gpu_tail.device()) + { + return None; + } + let encoder = encoder.as_deref_mut()?; + (previous_tokens, previous_len) = fixed_previous_tokens_data_with_gpu_tail( + device, + processors.previous_tokens, + gpu_tail, + encoder, + ); + } + } + let params = processor_params_data( + device, + processors.settings.temperature, + processors.settings.repetition_penalty, + previous_len, + ); + Some((previous_tokens, params)) + } + }; + + if input.datatype() != DataTypeEnum::F32 || input.layout().rank() != 1 { + return None; + } + + let input_len = input.layout().shape()[0]; + let chunks = input_len.div_ceil(TOP_K_CHUNK); + let output_len = chunks.checked_mul(output_per_chunk)?; + let ids = TensorData::new_for_shape(device, &[output_len], DataTypeEnum::U32); + let values = TensorData::new_for_shape(device, &[output_len], DataTypeEnum::F32); + if top_k_dims_empty(&[input_len, candidate_count, output_per_chunk]) { + return Some((ids, values)); + } + + let input_offset = input.layout().offset(); + let input_stride = input.layout().strides()[0]; + let has_processors = processors.is_some(); + let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + kernel_backend::KernelVariantKey::of::().hash(state); + TOP_K_BLOCK.hash(state); + TOP_K_CHUNK.hash(state); + input_len.hash(state); + candidate_count.hash(state); + output_per_chunk.hash(state); + input_offset.hash(state); + input_stride.hash(state); + has_processors.hash(state); + }); + + launch( + device, + "chunk_top_k_pairs_f32", + cache_key, + [chunks.try_into().ok()?, 1, 1], + encoder, + |kb| { + row_kernels::top_k_chunk( + kb, + input.as_kernel_tensor_ref(), + ids.as_kernel_tensor_ref(), + values.as_kernel_tensor_ref(), + processors.as_ref().map(|(previous_tokens, params)| { + ( + previous_tokens.as_kernel_tensor_ref(), + params.as_kernel_tensor_ref(), + ) + }), + row_kernels::TopKChunkMeta { + input_len: input_len.try_into().ok()?, + output_per_chunk: output_per_chunk.try_into().ok()?, + input_offset: input_offset.try_into().ok()?, + input_stride: input_stride.try_into().ok()?, + processors: has_processors, + }, + ) + }, + )?; + + Some((ids, values)) +} + +pub(crate) struct MergeSortedChunkTopKParams { + pub chunks: usize, + pub chunk_len: usize, + pub chunk_stride: usize, + pub input_len: usize, + pub k: usize, +} + +pub(crate) fn merge_sorted_chunk_top_k_pair_data_with_encoder( + input_ids: &TensorData, + input_values: &TensorData, + params: MergeSortedChunkTopKParams, + encoder: Option<&mut CommandEncoder>, +) -> Option<(TensorData, TensorData)> { + let MergeSortedChunkTopKParams { + chunks, + chunk_len, + chunk_stride, + input_len, + k, + } = params; + if input_ids.datatype() != DataTypeEnum::U32 || input_values.datatype() != DataTypeEnum::F32 { + return None; + } + if input_ids.layout().rank() != 1 || input_values.layout().rank() != 1 { + return None; + } + let input_ids_len = input_ids.layout().shape()[0]; + let input_values_len = input_values.layout().shape()[0]; + let expected_len = if chunks == 0 { + 0 + } else { + (chunks - 1) + .checked_mul(chunk_stride)? + .checked_add(chunk_len)? + }; + if input_ids_len < expected_len || input_values_len < expected_len { + return None; + } + + let device = input_values.device(); + let output_len = k.min(input_len); + let ids = TensorData::new_for_shape(device, &[output_len], DataTypeEnum::U32); + let values = TensorData::new_for_shape(device, &[output_len], DataTypeEnum::F32); + if top_k_dims_empty(&[chunks, chunk_len, output_len]) { + return Some((ids, values)); + } + + let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + kernel_backend::KernelVariantKey::of::() + .hash(state); + TOP_K_BLOCK.hash(state); + chunks.hash(state); + chunk_len.hash(state); + chunk_stride.hash(state); + input_len.hash(state); + output_len.hash(state); + hash_layout(state, input_ids.layout()); + hash_layout(state, input_values.layout()); + }); + launch( + device, + "merge_sorted_chunk_top_k_pairs_f32", + cache_key, + [1, 1, 1], + encoder, + |kb| { + row_kernels::top_k_merge( + kb, + input_ids.as_kernel_tensor_ref(), + input_values.as_kernel_tensor_ref(), + ids.as_kernel_tensor_ref(), + values.as_kernel_tensor_ref(), + row_kernels::MergeTopKMeta { + chunks: chunks.try_into().ok()?, + chunk_len: chunk_len.try_into().ok()?, + chunk_stride: chunk_stride.try_into().ok()?, + input_len: input_len.try_into().ok()?, + k: output_len.try_into().ok()?, + }, + ) + }, + )?; + + Some((ids, values)) +} + +fn normalized_probability(value: f32, default: f32) -> f32 { + if value.is_finite() { + value.clamp(0.0, 1.0) + } else { + default + } +} + +pub(crate) fn supports_unfiltered_categorical( + input_len: usize, + params: crate::sampling::GpuStandardSamplerParams, +) -> bool { + input_len > 0 + && input_len <= TOP_K_BLOCK as usize + && params.top_k >= input_len + && normalized_probability(params.top_p, 1.0) == 1.0 + && normalized_probability(params.min_p, 0.0) == 0.0 + && params.temperature.is_finite() + && params.random.is_finite() +} + +pub(crate) fn sample_categorical_logits_data_with_encoder( + logits: &TensorData, + params: crate::sampling::GpuStandardSamplerParams, + encoder: Option<&mut CommandEncoder>, +) -> Option { + if logits.datatype() != DataTypeEnum::F32 || logits.layout().rank() != 1 { + return None; + } + let input_len = logits.layout().shape()[0]; + if !supports_unfiltered_categorical(input_len, params) { + return None; + } + + let device = logits.device(); + let params_data = sampler_params_data( + device, + [ + params.random.clamp(0.0, 0.999_999_94), + params.temperature, + 0.0, + 0.0, + ], + ); + let output = TensorData::new_for_shape(device, &[GPU_SAMPLE_RESULT_WORDS], DataTypeEnum::U32); + let block = u32::try_from(input_len.next_power_of_two()).ok()?; + let meta = row_kernels::CategoricalSamplerMeta { + input_len: input_len.try_into().ok()?, + input_offset: logits.layout().offset().try_into().ok()?, + input_stride: logits.layout().strides()[0].try_into().ok()?, + block, + }; + let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + kernel_backend::KernelVariantKey::of::() + .hash(state); + input_len.hash(state); + block.hash(state); + hash_layout(state, logits.layout()); + }); + launch( + device, + "sample_categorical_logits_f32", + cache_key, + [1, 1, 1], + encoder, + |kb| { + row_kernels::categorical_sampler( + kb, + row_kernels::CategoricalSampler { + logits: logits.as_kernel_tensor_ref(), + params: params_data.as_kernel_tensor_ref(), + output: output.as_kernel_tensor_ref(), + meta, + }, + ) + }, + )?; + Some(output) +} + +/// Terminate the top-k tail with the sampler `request` selects. Both kernels +/// read the same sorted candidate pairs and the same optional exactness flag; +/// only the parameter buffer and the selection rule differ. +pub(crate) fn sample_from_sorted_top_k_data_with_encoder( + ids: &TensorData, + values: &TensorData, + request: &mut GpuSamplerRequest<'_>, + exactness_flag: Option<&TensorData>, + encoder: Option<&mut CommandEncoder>, +) -> Option { + if ids.datatype() != DataTypeEnum::U32 || values.datatype() != DataTypeEnum::F32 { + return None; + } + if ids.layout().rank() != 1 || values.layout().rank() != 1 { + return None; + } + if let Some(flag) = exactness_flag + && (flag.datatype() != DataTypeEnum::U32 + || flag.layout().rank() != 1 + || flag.layout().shape()[0] == 0 + || !values.device().is_same_device(flag.device())) + { + return None; + } + + let top_k = request + .top_k() + .min(ids.layout().shape()[0]) + .min(values.layout().shape()[0]); + if top_k == 0 { + return None; + } + let device = values.device(); + let has_exactness_flag = exactness_flag.is_some(); + let output = TensorData::new_for_shape(device, &[GPU_SAMPLE_RESULT_WORDS], DataTypeEnum::U32); + let meta = row_kernels::SamplerMeta { + top_k: top_k.try_into().ok()?, + ids_offset: ids.layout().offset().try_into().ok()?, + ids_stride: ids.layout().strides()[0].try_into().ok()?, + values_offset: values.layout().offset().try_into().ok()?, + values_stride: values.layout().strides()[0].try_into().ok()?, + has_exactness_flag, + }; + let sorted_top_k_cache_key = |variant: kernel_backend::KernelVariantKey| { + kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + variant.hash(state); + TOP_K_BLOCK.hash(state); + top_k.hash(state); + hash_layout(state, ids.layout()); + hash_layout(state, values.layout()); + has_exactness_flag.hash(state); + }) + }; + + match request { + GpuSamplerRequest::Mirostat2 { sampler, params } => { + let params_data = sampler_params_data( + device, + [ + params.tau, + params.eta, + params.random.clamp(0.0, 0.999_999_94), + 0.0, + ], + ); + let variant = + kernel_backend::KernelVariantKey::of::(); + launch( + device, + "sample_mirostat2_sorted_top_k_f32", + sorted_top_k_cache_key(variant), + [1, 1, 1], + encoder, + |kb| { + row_kernels::mirostat2( + kb, + row_kernels::Mirostat2 { + ids: ids.as_kernel_tensor_ref(), + values: values.as_kernel_tensor_ref(), + state: sampler.state.as_kernel_tensor_ref(), + params: params_data.as_kernel_tensor_ref(), + output: output.as_kernel_tensor_ref(), + exactness_flag: exactness_flag.map(|t| t.as_kernel_tensor_ref()), + meta, + }, + ) + }, + )?; + } + GpuSamplerRequest::Standard { params } => { + let params_data = sampler_params_data( + device, + [ + params.random.clamp(0.0, 0.999_999_94), + normalized_probability(params.top_p, 1.0), + normalized_probability(params.min_p, 0.0), + 0.0, + ], + ); + let variant = + kernel_backend::KernelVariantKey::of::(); + launch( + device, + "sample_standard_sorted_top_k_f32", + sorted_top_k_cache_key(variant), + [1, 1, 1], + encoder, + |kb| { + row_kernels::standard_sampler( + kb, + row_kernels::StandardSampler { + ids: ids.as_kernel_tensor_ref(), + values: values.as_kernel_tensor_ref(), + params: params_data.as_kernel_tensor_ref(), + output: output.as_kernel_tensor_ref(), + exactness_flag: exactness_flag.map(|t| t.as_kernel_tensor_ref()), + meta, + }, + ) + }, + )?; + } + } + + Some(output) +} + +#[cfg(test)] +mod selection_tests { + use super::*; + + #[test] + fn empty_dims_short_circuit() { + assert!(top_k_dims_empty(&[0, 4, 4])); + assert!(top_k_dims_empty(&[4, 0, 4])); + assert!(top_k_dims_empty(&[4, 4, 0])); + assert!(!top_k_dims_empty(&[1, 1, 1])); + } + + #[test] + fn exactness_ineligible_matches_old_selector() { + // top_k == 0 → ineligible. + assert!(top_k_exactness_ineligible(1024, 4, 64, 0)); + // top_values_len < top_k → ineligible. + assert!(top_k_exactness_ineligible(100, 4, 64, 200)); + // candidate_count >= output_per_chunk → ineligible. + assert!(top_k_exactness_ineligible(1024, 64, 64, 16)); + // Sized values within the eligible window. + assert!(!top_k_exactness_ineligible(2048, 4, 128, 32)); + } +} diff --git a/fusor-ml/core/src/mir/kernel_backend/sampling_topk.rs b/fusor-ml/core/src/mir/kernel_backend/sampling_topk.rs deleted file mode 100644 index 76a6c11c6..000000000 --- a/fusor-ml/core/src/mir/kernel_backend/sampling_topk.rs +++ /dev/null @@ -1,413 +0,0 @@ -use std::hash::Hash; - -use crate::{ - mir::kernel_backend, - sampling::{ - TOP_K_BLOCK, TOP_K_CHUNK, - processors::{ - fixed_previous_tokens_data, fixed_previous_tokens_data_with_gpu_tail, - processor_params_data, - }, - row_kernels, - }, - tensor::{DataTypeEnum, TensorData}, -}; -use wgpu::CommandEncoder; - -struct ProveTopKExactKernelVariant; -struct ChunkTopKPairsKernelVariant; -struct MergeSortedChunkTopKPairsKernelVariant; - -/// True when any dimension of the top-k working set is zero; in that case -/// every top-k kernel short-circuits. -fn top_k_dims_empty(dims: &[usize]) -> bool { - dims.contains(&0) -} - -/// True when the inputs to `top_k_exactness` can't be sharpened any further -/// by running the exactness kernel (no candidates per output, or no top_k). -fn top_k_exactness_ineligible( - top_values_len: usize, - candidate_count: usize, - output_per_chunk: usize, - top_k: usize, -) -> bool { - top_k == 0 || top_values_len < top_k || candidate_count >= output_per_chunk -} - -pub(crate) fn top_k_exactness_flag_data_with_encoder( - top_values: &TensorData, - chunk_values: &TensorData, - chunks: usize, - candidate_count: usize, - output_per_chunk: usize, - top_k: usize, - encoder: Option<&mut CommandEncoder>, -) -> Option { - if top_values.datatype() != DataTypeEnum::F32 - || chunk_values.datatype() != DataTypeEnum::F32 - || top_values.layout().rank() != 1 - || chunk_values.layout().rank() != 1 - || !top_values.device().is_same_device(chunk_values.device()) - { - return None; - } - if top_k_exactness_ineligible( - top_values.layout().shape()[0], - candidate_count, - output_per_chunk, - top_k, - ) { - return None; - } - - let device = top_values.device(); - let flag = TensorData::new_for_shape(device, &[1], DataTypeEnum::U32); - let meta = row_kernels::TopKExactnessMeta { - chunks: chunks.try_into().ok()?, - candidate_count: candidate_count.try_into().ok()?, - output_per_chunk: output_per_chunk.try_into().ok()?, - top_k: top_k.try_into().ok()?, - top_values_offset: top_values.layout().offset().try_into().ok()?, - top_values_stride: top_values.layout().strides()[0].try_into().ok()?, - chunk_values_offset: chunk_values.layout().offset().try_into().ok()?, - chunk_values_stride: chunk_values.layout().strides()[0].try_into().ok()?, - }; - let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { - kernel_backend::KernelVariantKey::of::().hash(state); - TOP_K_BLOCK.hash(state); - chunks.hash(state); - candidate_count.hash(state); - output_per_chunk.hash(state); - top_k.hash(state); - top_values.layout().offset().hash(state); - top_values.layout().shape().hash(state); - top_values.layout().strides().hash(state); - chunk_values.layout().offset().hash(state); - chunk_values.layout().shape().hash(state); - chunk_values.layout().strides().hash(state); - }); - let kernel = kernel_backend::run_kernel( - device.kernel_cache(), - "prove_top_k_exact_f32", - cache_key, - [1, 1, 1], - |kb| { - row_kernels::top_k_exactness( - kb, - top_values.as_kernel_tensor_ref(), - chunk_values.as_kernel_tensor_ref(), - flag.as_kernel_tensor_ref(), - meta, - ) - }, - )?; - - kernel_backend::run_direct_kernel( - device.kernel_cache(), - device.wgpu_queue(), - "prove_top_k_exact_f32 encoder", - &kernel, - encoder, - ); - - Some(flag) -} - -pub(crate) fn chunk_top_k_pair_data_with_encoder( - input: &TensorData, - candidate_count: usize, - output_per_chunk: usize, - encoder: Option<&mut CommandEncoder>, -) -> Option<(TensorData, TensorData)> { - chunk_top_k_pair_data_inner_with_encoder( - input, - candidate_count, - output_per_chunk, - None, - encoder, - ) -} - -/// Logit-processor settings (temperature scaling and repetition penalty) -/// applied before the top-k reduction. -#[derive(Clone, Copy)] -pub(crate) struct ProcessorSettings { - pub temperature: f32, - pub repetition_penalty: f32, -} - -pub(crate) fn chunk_top_k_pair_data_with_processors_with_encoder( - input: &TensorData, - previous_tokens: &[u32], - settings: ProcessorSettings, - candidate_count: usize, - output_per_chunk: usize, - encoder: Option<&mut CommandEncoder>, -) -> Option<(TensorData, TensorData)> { - let device = input.device(); - let (previous_tokens, previous_len) = fixed_previous_tokens_data(device, previous_tokens); - let params = processor_params_data( - device, - settings.temperature, - settings.repetition_penalty, - previous_len, - ); - chunk_top_k_pair_data_inner_with_encoder( - input, - candidate_count, - output_per_chunk, - Some((&previous_tokens, ¶ms)), - encoder, - ) -} - -pub(crate) fn chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder( - input: &TensorData, - previous_tokens: &[u32], - gpu_tail: Option<&TensorData>, - settings: ProcessorSettings, - candidate_count: usize, - output_per_chunk: usize, - encoder: Option<&mut CommandEncoder>, -) -> Option<(TensorData, TensorData)> { - let Some(gpu_tail) = gpu_tail else { - return chunk_top_k_pair_data_with_processors_with_encoder( - input, - previous_tokens, - settings, - candidate_count, - output_per_chunk, - encoder, - ); - }; - if gpu_tail.datatype() != DataTypeEnum::U32 - || gpu_tail.layout().rank() != 1 - || gpu_tail - .layout() - .shape() - .first() - .copied() - .unwrap_or_default() - == 0 - || !input.device().is_same_device(gpu_tail.device()) - { - return None; - } - - let device = input.device(); - let encoder = encoder?; - let (previous_tokens, previous_len) = - fixed_previous_tokens_data_with_gpu_tail(device, previous_tokens, gpu_tail, encoder); - let params = processor_params_data( - device, - settings.temperature, - settings.repetition_penalty, - previous_len, - ); - chunk_top_k_pair_data_inner_with_encoder( - input, - candidate_count, - output_per_chunk, - Some((&previous_tokens, ¶ms)), - Some(encoder), - ) -} - -fn chunk_top_k_pair_data_inner_with_encoder( - input: &TensorData, - candidate_count: usize, - output_per_chunk: usize, - processors: Option<(&TensorData, &TensorData)>, - encoder: Option<&mut CommandEncoder>, -) -> Option<(TensorData, TensorData)> { - if input.datatype() != DataTypeEnum::F32 || input.layout().rank() != 1 { - return None; - } - - let input_len = input.layout().shape()[0]; - let chunks = input_len.div_ceil(TOP_K_CHUNK); - let output_len = chunks.checked_mul(output_per_chunk)?; - let device = input.device(); - let ids = TensorData::new_for_shape(device, &[output_len], DataTypeEnum::U32); - let values = TensorData::new_for_shape(device, &[output_len], DataTypeEnum::F32); - if top_k_dims_empty(&[input_len, candidate_count, output_per_chunk]) { - return Some((ids, values)); - } - - let input_offset = input.layout().offset(); - let input_stride = input.layout().strides()[0]; - let has_processors = processors.is_some(); - let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { - kernel_backend::KernelVariantKey::of::().hash(state); - TOP_K_BLOCK.hash(state); - TOP_K_CHUNK.hash(state); - input_len.hash(state); - candidate_count.hash(state); - output_per_chunk.hash(state); - input_offset.hash(state); - input_stride.hash(state); - has_processors.hash(state); - }); - - let kernel = kernel_backend::run_kernel( - device.kernel_cache(), - "chunk_top_k_pairs_f32", - cache_key, - [chunks.try_into().ok()?, 1, 1], - |kb| { - row_kernels::top_k_chunk( - kb, - input.as_kernel_tensor_ref(), - ids.as_kernel_tensor_ref(), - values.as_kernel_tensor_ref(), - processors.map(|(previous_tokens, params)| { - ( - previous_tokens.as_kernel_tensor_ref(), - params.as_kernel_tensor_ref(), - ) - }), - row_kernels::TopKChunkMeta { - input_len: input_len.try_into().ok()?, - output_per_chunk: output_per_chunk.try_into().ok()?, - input_offset: input_offset.try_into().ok()?, - input_stride: input_stride.try_into().ok()?, - processors: has_processors, - }, - ) - }, - )?; - - kernel_backend::run_direct_kernel( - device.kernel_cache(), - device.wgpu_queue(), - "chunk_top_k_pairs_f32 encoder", - &kernel, - encoder, - ); - - Some((ids, values)) -} - -pub(crate) struct MergeSortedChunkTopKParams { - pub chunks: usize, - pub chunk_len: usize, - pub chunk_stride: usize, - pub input_len: usize, - pub k: usize, -} - -pub(crate) fn merge_sorted_chunk_top_k_pair_data_with_encoder( - input_ids: &TensorData, - input_values: &TensorData, - params: MergeSortedChunkTopKParams, - encoder: Option<&mut CommandEncoder>, -) -> Option<(TensorData, TensorData)> { - let MergeSortedChunkTopKParams { - chunks, - chunk_len, - chunk_stride, - input_len, - k, - } = params; - if input_ids.datatype() != DataTypeEnum::U32 || input_values.datatype() != DataTypeEnum::F32 { - return None; - } - if input_ids.layout().rank() != 1 || input_values.layout().rank() != 1 { - return None; - } - let input_ids_len = input_ids.layout().shape()[0]; - let input_values_len = input_values.layout().shape()[0]; - let expected_len = if chunks == 0 { - 0 - } else { - (chunks - 1) - .checked_mul(chunk_stride)? - .checked_add(chunk_len)? - }; - if input_ids_len < expected_len || input_values_len < expected_len { - return None; - } - - let device = input_values.device(); - let output_len = k.min(input_len); - let ids = TensorData::new_for_shape(device, &[output_len], DataTypeEnum::U32); - let values = TensorData::new_for_shape(device, &[output_len], DataTypeEnum::F32); - if top_k_dims_empty(&[chunks, chunk_len, output_len]) { - return Some((ids, values)); - } - - let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { - kernel_backend::KernelVariantKey::of::() - .hash(state); - TOP_K_BLOCK.hash(state); - chunks.hash(state); - chunk_len.hash(state); - chunk_stride.hash(state); - input_len.hash(state); - output_len.hash(state); - input_ids.layout().offset().hash(state); - input_ids.layout().shape().hash(state); - input_ids.layout().strides().hash(state); - input_values.layout().offset().hash(state); - input_values.layout().shape().hash(state); - input_values.layout().strides().hash(state); - }); - let kernel = kernel_backend::run_kernel( - device.kernel_cache(), - "merge_sorted_chunk_top_k_pairs_f32", - cache_key, - [1, 1, 1], - |kb| { - row_kernels::top_k_merge( - kb, - input_ids.as_kernel_tensor_ref(), - input_values.as_kernel_tensor_ref(), - ids.as_kernel_tensor_ref(), - values.as_kernel_tensor_ref(), - row_kernels::MergeTopKMeta { - chunks: chunks.try_into().ok()?, - chunk_len: chunk_len.try_into().ok()?, - chunk_stride: chunk_stride.try_into().ok()?, - input_len: input_len.try_into().ok()?, - k: output_len.try_into().ok()?, - }, - ) - }, - )?; - - kernel_backend::run_direct_kernel( - device.kernel_cache(), - device.wgpu_queue(), - "merge_sorted_chunk_top_k_pairs_f32 encoder", - &kernel, - encoder, - ); - - Some((ids, values)) -} - -#[cfg(test)] -mod selection_tests { - use super::*; - - #[test] - fn empty_dims_short_circuit() { - assert!(top_k_dims_empty(&[0, 4, 4])); - assert!(top_k_dims_empty(&[4, 0, 4])); - assert!(top_k_dims_empty(&[4, 4, 0])); - assert!(!top_k_dims_empty(&[1, 1, 1])); - } - - #[test] - fn exactness_ineligible_matches_old_selector() { - // top_k == 0 → ineligible. - assert!(top_k_exactness_ineligible(1024, 4, 64, 0)); - // top_values_len < top_k → ineligible. - assert!(top_k_exactness_ineligible(100, 4, 64, 200)); - // candidate_count >= output_per_chunk → ineligible. - assert!(top_k_exactness_ineligible(1024, 64, 64, 16)); - // Sized values within the eligible window. - assert!(!top_k_exactness_ineligible(2048, 4, 128, 32)); - } -} diff --git a/fusor-ml/core/src/mir/kernel_backend/standard_sampler.rs b/fusor-ml/core/src/mir/kernel_backend/standard_sampler.rs deleted file mode 100644 index 2d1249b40..000000000 --- a/fusor-ml/core/src/mir/kernel_backend/standard_sampler.rs +++ /dev/null @@ -1,132 +0,0 @@ -use std::hash::Hash; - -use crate::{ - Device, - mir::kernel_backend, - sampling::{GPU_SAMPLE_RESULT_WORDS, GpuStandardSamplerParams, TOP_K_BLOCK, row_kernels}, - tensor::{DataTypeEnum, TensorData}, -}; -use wgpu::CommandEncoder; - -#[repr(C)] -#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)] -struct StandardSamplerParams { - random: f32, - top_p: f32, - min_p: f32, - _padding: f32, -} - -struct StandardSamplerSortedTopKKernelVariant; - -fn normalized_probability(value: f32, default: f32) -> f32 { - if value.is_finite() { - value.clamp(0.0, 1.0) - } else { - default - } -} - -fn standard_sampler_params_data(device: &Device, params: GpuStandardSamplerParams) -> TensorData { - let params = StandardSamplerParams { - random: params.random.clamp(0.0, 0.999_999_94), - top_p: normalized_probability(params.top_p, 1.0), - min_p: normalized_probability(params.min_p, 0.0), - _padding: 0.0, - }; - let buffer = device.create_buffer_init( - bytemuck::bytes_of(¶ms), - wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC, - ); - TensorData::new_from_buffer(device, buffer, &[4], DataTypeEnum::F32) -} - -pub(crate) fn sample_from_sorted_top_k_data_with_encoder( - ids: &TensorData, - values: &TensorData, - params: GpuStandardSamplerParams, - exactness_flag: Option<&TensorData>, - encoder: Option<&mut CommandEncoder>, -) -> Option { - if ids.datatype() != DataTypeEnum::U32 || values.datatype() != DataTypeEnum::F32 { - return None; - } - if ids.layout().rank() != 1 || values.layout().rank() != 1 { - return None; - } - if let Some(flag) = exactness_flag - && (flag.datatype() != DataTypeEnum::U32 - || flag.layout().rank() != 1 - || flag.layout().shape()[0] == 0 - || !values.device().is_same_device(flag.device())) - { - return None; - } - - let top_k = params - .top_k - .min(ids.layout().shape()[0]) - .min(values.layout().shape()[0]); - if top_k == 0 { - return None; - } - let device = values.device(); - let params = standard_sampler_params_data(device, params); - let has_exactness_flag = exactness_flag.is_some(); - let output = TensorData::new_for_shape(device, &[GPU_SAMPLE_RESULT_WORDS], DataTypeEnum::U32); - let meta = row_kernels::SamplerMeta { - top_k: top_k.try_into().ok()?, - ids_offset: ids.layout().offset().try_into().ok()?, - ids_stride: ids.layout().strides()[0].try_into().ok()?, - values_offset: values.layout().offset().try_into().ok()?, - values_stride: values.layout().strides()[0].try_into().ok()?, - has_exactness_flag, - }; - let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { - kernel_backend::KernelVariantKey::of::() - .hash(state); - TOP_K_BLOCK.hash(state); - top_k.hash(state); - ids.layout().offset().hash(state); - ids.layout().shape().hash(state); - ids.layout().strides().hash(state); - values.layout().offset().hash(state); - values.layout().shape().hash(state); - values.layout().strides().hash(state); - has_exactness_flag.hash(state); - }); - let kernel = kernel_backend::run_kernel( - device.kernel_cache(), - "sample_standard_sorted_top_k_f32", - cache_key, - [1, 1, 1], - |kb| { - row_kernels::standard_sampler( - kb, - row_kernels::StandardSampler { - ids: ids.as_kernel_tensor_ref(), - values: values.as_kernel_tensor_ref(), - params: params.as_kernel_tensor_ref(), - output: output.as_kernel_tensor_ref(), - exactness_flag: exactness_flag.map(|t| t.as_kernel_tensor_ref()), - meta, - }, - ) - }, - )?; - - if let Some(encoder) = encoder { - kernel.run(device.kernel_cache(), encoder); - } else { - let mut encoder = - device - .wgpu_device() - .create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("sample_standard_sorted_top_k_f32 encoder"), - }); - kernel.run(device.kernel_cache(), &mut encoder); - device.wgpu_queue().submit(Some(encoder.finish())); - } - - Some(output) -} diff --git a/fusor-ml/core/src/mir/operation.rs b/fusor-ml/core/src/mir/operation.rs index 82a9cc741..696c9a3ed 100644 --- a/fusor-ml/core/src/mir/operation.rs +++ b/fusor-ml/core/src/mir/operation.rs @@ -14,13 +14,74 @@ use super::{ workgroup_shape::{WorkgroupShape, WorkgroupShapeConstraints}, }; -pub(crate) trait Operation: Debug + 'static { +/// The complete direct-kernel lowering for one operation. +/// +/// Most operations produce exactly one kernel. Some valid operations produce +/// no dispatch for an empty output, while others require an ordered sequence +/// of kernels. Keeping those cases in the operation interface lets executors +/// treat every operation uniformly. +pub(crate) struct DirectKernelPlan { + kernels: Vec, +} + +impl DirectKernelPlan { + pub(crate) fn empty() -> Self { + Self { + kernels: Vec::new(), + } + } + + pub(crate) fn single(kernel: DirectKernel) -> Self { + Self { + kernels: vec![kernel], + } + } + + pub(crate) fn many(kernels: Vec) -> Self { + Self { kernels } + } + + pub(crate) fn into_kernels(self) -> Vec { + self.kernels + } +} + +#[derive(Debug, Clone)] +pub(crate) struct DirectKernelLoweringError { + operation: String, +} + +impl DirectKernelLoweringError { + pub(crate) fn new(operation: String) -> Self { + Self { operation } + } +} + +impl std::fmt::Display for DirectKernelLoweringError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "operation did not provide a direct kernel plan: {}", + self.operation + ) + } +} + +impl std::error::Error for DirectKernelLoweringError {} + +pub(crate) trait Operation: Debug + Send + Sync + 'static { fn workgroup_shape_constraints(&self, device: &Device) -> WorkgroupShapeConstraints; fn dispatch_size(&self, workgroup_shape: &WorkgroupShape, inputs: &[MirValue]) -> [u32; 3]; fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)); + /// Visit the same dependency slots as [`Self::visit_dependencies`], in + /// the same order, so callers can rebind them. Materializing an operation + /// that was interned from a structurally identical instance rewrites + /// every slot through this visitor. + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)); + fn inputs(&self, nodes: &ComputeGraphInner) -> Vec; fn output(&self, nodes: &ComputeGraphInner, inputs: &[MirValue]) -> MirValue; @@ -32,6 +93,21 @@ pub(crate) trait Operation: Debug + 'static { inputs: &[MirValue], ) -> Option; + /// Lower this operation to its complete zero-, one-, or many-kernel plan. + /// Operations with a singular lowering only implement + /// [`Self::build_direct_kernel`]; multi-dispatch and empty-output + /// operations override this method. + fn build_direct_kernel_plan( + &self, + nodes: &ComputeGraphInner, + workgroup_shape: &WorkgroupShape, + inputs: &[MirValue], + ) -> Result { + self.build_direct_kernel(nodes, workgroup_shape, inputs) + .map(DirectKernelPlan::single) + .ok_or_else(|| DirectKernelLoweringError::new(self.name())) + } + fn name(&self) -> String; /// Hash structural operation fields that affect generated kernel IR. @@ -51,7 +127,7 @@ pub(crate) trait Operation: Debug + 'static { kernel_backend::KernelCacheKey::from_hash_inputs(|hasher| { // Version the shared key layout so future changes cannot silently // collide with cache entries produced by an older hash recipe. - 1u64.hash(hasher); + 2u64.hash(hasher); variant.hash(hasher); TypeId::of::().hash(hasher); self.hash_kernel_fields(hasher); @@ -67,7 +143,7 @@ pub(crate) trait Operation: Debug + 'static { } } -fn hash_mir_value(state: &mut FxHasher, value: &MirValue) { +pub(crate) fn hash_mir_value(state: &mut FxHasher, value: &MirValue) { std::mem::discriminant(value).hash(state); match value { MirValue::QMatrix(matrix) => { @@ -85,7 +161,7 @@ fn hash_mir_value(state: &mut FxHasher, value: &MirValue) { } } -fn hash_layout(state: &mut FxHasher, layout: &crate::Layout) { +pub(crate) fn hash_layout(state: &mut FxHasher, layout: &crate::Layout) { layout.offset().hash(state); layout.shape().hash(state); layout.strides().hash(state); diff --git a/fusor-ml/core/src/nary_direct.rs b/fusor-ml/core/src/nary_direct.rs index 245507df8..a3e9751dc 100644 --- a/fusor-ml/core/src/nary_direct.rs +++ b/fusor-ml/core/src/nary_direct.rs @@ -17,7 +17,7 @@ use crate::{ visit_tiled::MaybeQData, }; -const BLOCK: usize = 256; +const BLOCK: usize = crate::occupancy::FULL_WORKGROUP_LANES as usize; const SMALL_BLOCK: usize = 1; struct NaryDirectKernelVariant; @@ -79,7 +79,7 @@ fn build_nary_direct_kernel_with_output_index( let total_elements = total_elements(&operation.shape)?; let plan = plan_nary_tiling(operation, &graph.device(), &values, output_index); if let Some(plan) = &plan - && std::env::var_os("FUSOR_TRACE_REDUCE_TILED").is_some() + && graph.device().config().trace_reduce_tiled { eprintln!( "nary_tiled dim={} invariant={:?} threads={} shape={:?}", @@ -104,6 +104,7 @@ fn build_nary_direct_kernel_with_output_index( if let Some(plan) = &plan { plan.dim.hash(state); plan.invariant.hash(state); + plan.tm.hash(state); } }); let cache_key = operation.kernel_cache_key_with_dispatch( @@ -112,7 +113,7 @@ fn build_nary_direct_kernel_with_output_index( dispatch_size, inputs, ); - let name = if std::env::var_os("FUSOR_TRACE_DECODE_NAMES").is_some() { + let name = if graph.device().config().trace_decode_names { operation.name() } else { format!("nary_direct_out_{output_index}") @@ -140,23 +141,282 @@ fn build_nary_direct_kernel_with_output_index( ) } -/// Outputs per thread along the tiled dim of a reuse-tiled elementwise -/// kernel. -const NARY_TM: u32 = 4; -/// Floor on post-tiling thread count: trading threads for register reuse -/// must leave the device saturated. -const MIN_TILED_THREADS: u32 = 65536; +struct MergedRegionKernelVariant; + +/// Bind a value tile into a register so later statements can reuse it +/// without re-evaluating. Statement values are cast to a storable dtype +/// before binding, so `Bool` never reaches here. +/// For every tensor value across all segments (flattened), the index of its +/// first occurrence — the cross-segment binding-sharing pattern. +fn cross_segment_alias_classes<'a>(segments: impl Iterator) -> Vec { + // Sharing equality must match the declare-time dedup exactly: same + // buffer, same datatype, same layout. + let mut seen: Vec<(usize, Option<(DataTypeEnum, crate::Layout)>)> = Vec::new(); + let mut classes = Vec::new(); + for segment in segments { + for value in segment { + let key = match value { + MaybeQData::Tensor(tensor) => ( + std::sync::Arc::as_ptr(tensor.buffer()) as usize, + Some((tensor.datatype(), tensor.layout().clone())), + ), + MaybeQData::QMatrix(matrix) => { + (std::sync::Arc::as_ptr(matrix.buffer()) as usize, None) + } + }; + let class = seen + .iter() + .position(|entry| *entry == key) + .unwrap_or(seen.len()); + classes.push(class); + seen.push(key); + } + } + classes +} + +fn bind_value_tile(program: &mut tile_ir::tile::TileBlock<'_>, value: ValueTile) -> ValueTile { + match value { + ValueTile::F32(tile) => ValueTile::F32(program.bind(tile)), + ValueTile::F16(tile) => ValueTile::F16(program.bind(tile)), + ValueTile::U32(tile) => ValueTile::U32(program.bind(tile)), + ValueTile::Bool(_) => unreachable!("region statements are cast to storable dtypes"), + } +} + +/// One kernel executing several independent multi-output regions: like +/// [`build_merged_nary_kernel`], each segment owns a contiguous range of +/// workgroups behind a uniform guard, but a segment's body is a statement +/// chain — statement values stay in registers (the `extras` slots of +/// [`eval_nary_expr`]) and every externally-live statement stores to its own +/// output binding. Segment values are inputs-then-outputs. +pub(crate) fn build_merged_region_kernel( + graph: &crate::compute_graph::ComputeGraphInner, + segments: &[crate::region::ElementwiseRegionOperation], + segment_inputs: &[Vec], +) -> Option { + let device = graph.device(); + let max_per_dim = device.limits().max_compute_workgroups_per_dimension; + + struct Segment { + values: Vec, + /// Per output (relative index): the input slot whose buffer the + /// output writes in place, bound once as read-write. + fold: Vec>, + elements: u32, + base: u32, + groups: u32, + } + let mut prepared = Vec::with_capacity(segments.len()); + let mut total_groups = 0u32; + for (op, inputs) in segments.iter().zip(segment_inputs) { + let values = inputs + .iter() + .map(|input| MaybeQData::try_from(input.clone()).ok()) + .collect::>>()?; + if values + .iter() + .any(|value| matches!(value, MaybeQData::QMatrix(_))) + { + return None; + } + if values.iter().any(|value| { + matches!(value, MaybeQData::Tensor(tensor) + if tensor.datatype() == DataTypeEnum::F16 && !device.f16_supported()) + }) { + return None; + } + debug_assert_eq!(values.len(), op.inputs.len() + op.output_count()); + let input_count = op.inputs.len(); + let buffer_of = |value: &MaybeQData| match value { + MaybeQData::Tensor(tensor) => Some(std::sync::Arc::as_ptr(tensor.buffer()) as usize), + _ => None, + }; + // An output sharing an input's buffer must bind it exactly once, + // read-write: wgpu rejects one buffer bound read-only and read-write + // in the same dispatch. + let fold: Vec> = values[input_count..] + .iter() + .map(|output| { + let output_ptr = buffer_of(output)?; + values[..input_count] + .iter() + .position(|input| buffer_of(input) == Some(output_ptr)) + }) + .collect(); + let elements = total_elements(&op.shape)?; + let groups = elements.div_ceil(BLOCK as u32); + prepared.push(Segment { + values, + fold, + elements, + base: total_groups, + groups, + }); + total_groups = total_groups.checked_add(groups)?; + } + + let dispatch_size = crate::visit_tiled::distribute_workgroups(total_groups, max_per_dim); + let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + kernel_backend::KernelVariantKey::of::().hash(state); + dispatch_size.hash(state); + segments.len().hash(state); + for ((op, inputs), segment) in segments.iter().zip(segment_inputs).zip(&prepared) { + op.hash_kernel_fields(state); + segment.fold.hash(state); + inputs.len().hash(state); + for input in inputs { + crate::mir::operation::hash_mir_value(state, input); + } + } + // Shared read-only inputs bind once across segments; the sharing + // pattern changes the generated bindings, so it keys the kernel. + for class in + cross_segment_alias_classes(prepared.iter().map(|segment| segment.values.as_slice())) + { + class.hash(state); + } + }); + let name = if device.config().trace_decode_names { + format!( + "merged_region[{}]", + segments + .iter() + .map(|op| op.name()) + .collect::>() + .join("; ") + ) + } else { + format!("merged_region_x{}", segments.len()) + }; + + kernel_backend::run_kernel( + device.kernel_cache(), + name, + cache_key, + dispatch_size, + move |kb| { + let mut declared = Vec::with_capacity(prepared.len()); + // Read-only inputs shared across segments (a learning-rate + // tensor read by every optimizer segment) bind once: wgpu + // rejects one buffer bound at several slots of a dispatch, and + // one binding is cheaper anyway. + let mut shared_reads: Vec<(usize, DataTypeEnum, crate::Layout, Storage2, TensorMeta)> = + Vec::new(); + for (op, segment) in segments.iter().zip(&prepared) { + let input_count = op.inputs.len(); + let folded_inputs: rustc_hash::FxHashSet = + segment.fold.iter().flatten().copied().collect(); + let mut storages: Vec = Vec::with_capacity(segment.values.len()); + let mut metas: Vec = Vec::with_capacity(segment.values.len()); + for (binding, value) in segment.values.iter().enumerate() { + if binding >= input_count + && let Some(source) = segment.fold[binding - input_count] + { + storages.push(storages[source].clone()); + metas.push(metas[source].clone()); + continue; + } + let write = binding >= input_count || folded_inputs.contains(&binding); + if !write && let MaybeQData::Tensor(tensor) = value { + let ptr = std::sync::Arc::as_ptr(tensor.buffer()) as usize; + let datatype = tensor.datatype(); + // Only reads through the identical view share a + // binding: the same buffer read through different + // layouts needs its own metadata. + if let Some((.., storage, meta)) = + shared_reads.iter().find(|(p, d, layout, _, _)| { + *p == ptr && *d == datatype && layout == tensor.layout() + }) + { + storages.push(storage.clone()); + metas.push(meta.clone()); + continue; + } + let (storage, meta) = declare_value(kb, value, false)?; + shared_reads.push(( + ptr, + datatype, + tensor.layout().clone(), + storage.clone(), + meta.clone(), + )); + storages.push(storage); + metas.push(meta); + continue; + } + let (storage, meta) = declare_value(kb, value, write)?; + storages.push(storage); + metas.push(meta); + } + declared.push((storages, metas)); + } + + kb.program() + .program_grid(BLOCK as u32, dispatch_size, |program| { + let lane = program.lane(); + let group = program.bind(linear_group(program, dispatch_size)); + for (op, (segment, (storages, metas))) in + segments.iter().zip(prepared.iter().zip(&declared)) + { + let in_segment = group.clone().ge(segment.base) + & group.clone().lt(segment.base + segment.groups); + program.if_then(in_segment, |program| { + let flat = (group.clone() - segment.base) * BLOCK as u32 + lane.clone(); + let in_bounds = flat.clone().lt(segment.elements); + let dims = output_dims_from_flat(flat, &op.shape); + let input_count = op.inputs.len(); + let (input_storages, output_storages) = storages.split_at(input_count); + let (input_metas, output_metas) = metas.split_at(input_count); + let mut extras: Vec<(ValueTile, DataTypeEnum)> = Vec::new(); + let mut out_idx = 0usize; + for statement in &op.statements { + let (value, _) = eval_nary_expr( + program, + &statement.expression, + &dims, + input_storages, + input_metas, + in_bounds.clone(), + &extras, + ); + let value = + bind_value_tile(program, value.cast_to(statement.datatype)); + if statement.output.is_some() { + let index = layout_index(&output_metas[out_idx], &dims); + output_storages[out_idx].store( + program, + index, + value.clone(), + in_bounds.clone(), + ); + out_idx += 1; + } + extras.push((value, statement.datatype)); + } + }); + } + }); + Some(()) + }, + ) +} /// A register-reuse tiling for an elementwise kernel: each thread covers -/// `NARY_TM` outputs along `dim`, loading the inputs that are invariant -/// along `dim` once instead of per output. +/// `tm` outputs along `dim`, loading the inputs that are invariant along +/// `dim` once instead of per output. struct NaryTilePlan { dim: usize, /// Per input: invariant along `dim` (its loads hoist out of the run). invariant: Vec, /// Per input: index-space dim read by each input dimension. dims: Vec>, - /// The output shape with `dim` divided by `NARY_TM`. + /// Outputs per thread along `dim` + /// ([`crate::occupancy::DispatchPolicy::work_per_thread`]); carried in + /// the plan so the planner's arithmetic and the kernel builder's loop + /// trip count cannot diverge. + tm: u32, + /// The output shape with `dim` divided by `tm`. thread_shape: Vec, total_threads: u32, } @@ -194,17 +454,21 @@ fn plan_nary_tiling( .collect::>()?; let access = crate::access_analysis::InputAccesses::collect(&operation.expression, input_count, &metas)?; + let policy = device.dispatch_policy(); + let tm = policy.work_per_thread(crate::occupancy::RegPressure::ElementwiseFew); let mut best: Option<(u64, usize)> = None; for dim in 0..rank.saturating_sub(1) { - if operation.shape[dim] < NARY_TM as usize { + if operation.shape[dim] < tm as usize { continue; } let invariant_bytes: u64 = (0..input_count) .filter(|&i| !access.depends_on(i, dim)) .map(|i| input_allocation_bytes(&metas[i], &values[i])) .sum(); - if invariant_bytes < device.last_level_cache_bytes() { + if policy.cache_resident(invariant_bytes) { + // Cache-resident re-reads are free; the tiling would only cost + // thread-level parallelism. continue; } if best @@ -220,15 +484,16 @@ fn plan_nary_tiling( .collect(); let mut thread_shape = operation.shape.to_vec(); - thread_shape[dim] = thread_shape[dim].div_ceil(NARY_TM as usize); + thread_shape[dim] = thread_shape[dim].div_ceil(tm as usize); let total_threads = total_elements(&thread_shape)?; - if total_threads < MIN_TILED_THREADS { + if !policy.tiling_leaves_saturated(total_threads) { return None; } Some(NaryTilePlan { dim, invariant, dims: access.dims, + tm, thread_shape, total_threads, }) @@ -264,7 +529,7 @@ fn build_nary_tiled_ir( let flat = group * BLOCK as u32 + lane; let in_bounds = flat.clone().lt(total_threads); let mut coords = output_dims_from_flat(flat, &plan.thread_shape); - let base = program.bind(coords[plan.dim].clone() * NARY_TM); + let base = program.bind(coords[plan.dim].clone() * plan.tm); // Invariant loads hoist out of the per-output run; the base // coordinate is always in range for an in-bounds thread. @@ -287,7 +552,7 @@ fn build_nary_tiled_ir( }) .collect(); - for j in 0..NARY_TM { + for j in 0..plan.tm { let coord = base.clone() + j; let in_bounds_j = in_bounds.clone() & coord.clone().lt(extent); coords[plan.dim] = coord; @@ -356,6 +621,21 @@ pub(crate) enum ValueTile { } impl ValueTile { + fn from_typed_tile(value: tile_ir::tile::Tile, datatype: DataTypeEnum) -> Self { + match datatype { + DataTypeEnum::F32 => Self::F32(value), + DataTypeEnum::F16 => Self::F16(value), + DataTypeEnum::U32 => Self::U32(value), + } + } + + fn into_typed_tile(self, datatype: DataTypeEnum) -> tile_ir::tile::Tile { + match self.cast_to(datatype) { + Self::F32(value) | Self::F16(value) | Self::U32(value) => value, + Self::Bool(_) => unreachable!("tensor datatypes never store bool tiles"), + } + } + pub(crate) fn cast_to(self, target: DataTypeEnum) -> Self { match (self, target) { (Self::F32(v), DataTypeEnum::F32) => Self::F32(v), @@ -445,6 +725,7 @@ fn bool_as_u32(value: tile_ir::tile::Mask) -> tile_ir::tile::Tile { tile_ir::tile::Tile::select(value, 1u32.into(), 0u32.into()) } +#[derive(Clone)] pub(crate) enum Storage2 { F32(tile_ir::tile::Storage), F16(tile_ir::tile::Storage), @@ -775,6 +1056,14 @@ fn emit_function(function: &NaryFunction, values: &mut [(ValueTile, DataTypeEnum .0 .clone() .binary(tile_ir::TileBinaryOp::Pow, values[1].0.clone()), + NaryOp::Max => values[0] + .0 + .clone() + .binary(tile_ir::TileBinaryOp::Max, values[1].0.clone()), + NaryOp::Min => values[0] + .0 + .clone() + .binary(tile_ir::TileBinaryOp::Min, values[1].0.clone()), NaryOp::Neg => values[0].0.clone().unary(tile_ir::TileUnaryOp::Neg), NaryOp::Cast => values[0].0.clone().cast_to(function.output_type), NaryOp::Select => match values[1].0.clone().cast_to(function.output_type) { @@ -849,6 +1138,31 @@ fn emit_function(function: &NaryFunction, values: &mut [(ValueTile, DataTypeEnum values[1].0.clone(), function.output_type, ), + NaryOp::Less => values[0].0.clone().compare( + tile_ir::TileCompareOp::Lt, + values[1].0.clone(), + function.output_type, + ), + NaryOp::Equal => values[0].0.clone().compare( + tile_ir::TileCompareOp::Eq, + values[1].0.clone(), + function.output_type, + ), + NaryOp::NotEqual => values[0].0.clone().compare( + tile_ir::TileCompareOp::Ne, + values[1].0.clone(), + function.output_type, + ), + NaryOp::Greater => values[0].0.clone().compare( + tile_ir::TileCompareOp::Gt, + values[1].0.clone(), + function.output_type, + ), + NaryOp::GreaterEqual => values[0].0.clone().compare( + tile_ir::TileCompareOp::Ge, + values[1].0.clone(), + function.output_type, + ), NaryOp::AddConst(scalar) => values[0].0.clone().binary( tile_ir::TileBinaryOp::Add, tile_literal(scalar).cast_to(values[0].1), @@ -983,6 +1297,32 @@ pub(crate) fn apply_unary_function_chain( Some((value.into_f32(), value_ty)) } +/// Apply a unary chain while preserving its declared output tile type. +/// Cooperative dense matmul uses this for f16 as well as f32 staging/store +/// epilogues; the older helper above intentionally normalizes to f32 for row +/// reduction callers. +pub(crate) fn apply_typed_unary_function_chain( + value: tile_ir::tile::Tile, + value_ty: DataTypeEnum, + chain: &UnaryFunctionChain, +) -> Option<(tile_ir::tile::Tile, DataTypeEnum)> { + if chain.input_datatype() != value_ty { + return None; + } + + let mut value = ValueTile::from_typed_tile(value, value_ty); + let mut value_ty = value_ty; + for function in &chain.functions { + if function.input_types.as_slice() != [value_ty] { + return None; + } + let mut values = [(value, value_ty)]; + value = emit_function(function, &mut values); + value_ty = function.output_type; + } + Some((value.into_typed_tile(value_ty), value_ty)) +} + pub(crate) fn apply_single_input_elementwise_expr( value: tile_ir::tile::Tile, value_ty: DataTypeEnum, @@ -1051,23 +1391,37 @@ pub(crate) fn output_dims_from_flat( flat: tile_ir::tile::Tile, shape: &[usize], ) -> Vec { - (0..shape.len()) - .map(|axis| { - let dim = shape[axis] as u32; - if dim == 1 { - return tile_u32(0); - } - let divisor = shape[axis + 1..] - .iter() - .fold(1u32, |acc, dim| acc.saturating_mul(*dim as u32)); - let quotient = if divisor == 1 { - flat.clone() - } else { - flat.clone() / tile_u32(divisor) - }; - quotient % tile_u32(dim) - }) - .collect() + // Peel innermost-out with a running quotient rather than dividing the + // flat index by each axis' suffix product. Load-bearing, not a style + // choice: the Apple Metal compiler miscompiles chains of u32 div/mod by + // large non-power-of-two constants (e.g. `flat / 393216` for a + // [64,6,256,256] delinearize) once the grid needs a second dispatch + // dimension — stores land at wild addresses, which with unchecked + // shaders corrupts arbitrary GPU memory (observed M2 Max, macOS 26; + // reproduced with stock wgpu + trivial WGSL). The peeled form only ever + // divides by a single dimension extent, on a quotient already reduced by + // the inner axes, which compiles correctly at any size. + // The outermost non-trivial axis takes the raw quotient without a `%`: + // for in-bounds lanes it is already < dim, out-of-bounds lanes never + // touch memory (stores are branch-masked, masked loads clamp their index + // to 0), and a trailing non-power-of-two `%` re-triggers the miscompile + // (observed with `% 48` on a [48,6,256,256] delinearize). + let mut dims = vec![tile_u32(0); shape.len()]; + let mut rest = flat; + for axis in (0..shape.len()).rev() { + let dim = shape[axis] as u32; + if dim == 1 { + continue; + } + if shape[..axis].iter().any(|&outer| outer != 1) { + dims[axis] = rest.clone() % tile_u32(dim); + rest = rest / tile_u32(dim); + } else { + dims[axis] = rest; + break; + } + } + dims } pub(crate) fn layout_index( diff --git a/fusor-ml/core/src/nary_wise.rs b/fusor-ml/core/src/nary_wise.rs index 859799931..a06e99478 100644 --- a/fusor-ml/core/src/nary_wise.rs +++ b/fusor-ml/core/src/nary_wise.rs @@ -86,6 +86,16 @@ pub(crate) enum NaryOp { /// index-dependent masking (e.g. composed causal attention compares the /// kv position against the query position). LessEqual, + /// Binary `a < b` comparison producing 1/0 in the output type. + Less, + /// Binary `a == b` comparison producing 1/0 in the output type. + Equal, + /// Binary `a != b` comparison producing 1/0 in the output type. + NotEqual, + /// Binary `a > b` comparison producing 1/0 in the output type. + Greater, + /// Binary `a >= b` comparison producing 1/0 in the output type. + GreaterEqual, AddConst(NaryScalar), SubConst(NaryScalar), RSubConst(NaryScalar), @@ -102,6 +112,14 @@ pub(crate) enum NaryOp { LessEqualConst(NaryScalar), GreaterConst(NaryScalar), GreaterEqualConst(NaryScalar), + /// Binary maximum. `ReduceOp::Max` predates this and lowers separately; + /// the fold algebra needs a two-value max as an expression. + /// + /// Appended, not inserted: these discriminants are hashed into kernel + /// cache keys, so reordering invalidates every cached plan. + Max, + /// Binary minimum. Appended for the same reason. + Min, } /// A function that can be applied in the expression tree. @@ -165,7 +183,7 @@ impl NaryFunction { /// A chain of unary functions used for pre/post processing in reduce/matmul/dequantize. /// Each function takes a single input and produces a single output; the chain is applied sequentially. -#[derive(Clone, Debug, Hash)] +#[derive(Clone, Debug, PartialEq, Eq, Hash)] pub(crate) struct UnaryFunctionChain { input_datatype: DataTypeEnum, pub(crate) functions: Vec, @@ -295,27 +313,6 @@ impl NaryExpr { } } - pub(crate) fn remap_inputs(&self, mapping: &[usize]) -> NaryExpr { - match self { - NaryExpr::Op { children, function } => NaryExpr::Op { - children: children - .iter() - .map(|child| child.remap_inputs(mapping)) - .collect(), - function: function.clone(), - }, - NaryExpr::IndexedInput { input_idx, indices } => NaryExpr::IndexedInput { - input_idx: mapping[*input_idx], - indices: indices - .iter() - .map(|index| index.remap_inputs(mapping)) - .collect(), - }, - NaryExpr::DimIndex(dim) => NaryExpr::DimIndex(*dim), - NaryExpr::Scalar(value) => NaryExpr::Scalar(*value), - } - } - /// Create a select expression (ternary operator) /// Semantics: condition != 0 ? on_true : on_false pub fn select( @@ -462,7 +459,7 @@ impl NaryExpr { /// N-ary operation combining multiple inputs with arbitrary operations. /// Can fuse chains of element-wise and pair-wise operations into a single kernel. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct ElementwiseOperation { /// Input tensors (leaves of expression tree) pub(crate) inputs: Vec, @@ -571,6 +568,12 @@ impl Operation for ElementwiseOperation { } } + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + for input in &mut self.inputs { + f(input); + } + } + fn inputs(&self, nodes: &ComputeGraphInner) -> Vec { let mut mir_inputs: Vec = self .inputs diff --git a/fusor-ml/core/src/occupancy.rs b/fusor-ml/core/src/occupancy.rs new file mode 100644 index 000000000..7bed353da --- /dev/null +++ b/fusor-ml/core/src/occupancy.rs @@ -0,0 +1,279 @@ +//! Device-derived dispatch sizing. +//! +//! Every "how many workgroups / how wide a workgroup / how much work per +//! thread" decision in kernel dispatch reads from [`DispatchPolicy`] instead +//! of a local constant. The policy is derived from device capabilities plus +//! one calibrated parallelism floor ([`crate::Device::saturation_lanes`]), +//! so a policy value is never a shape rule: shapes enter only as the +//! arguments of the predicate being asked. + +use crate::Device; + +/// One full-width committed workgroup. WebGPU guarantees at least 256 +/// invocations per workgroup, so this is exact on every conformant device; +/// [`DispatchPolicy::preferred_workgroup_lanes`] is its runtime clamp for +/// devices reporting less. Kept as a `const` because several kernel builders +/// take the block size as a const generic. +pub(crate) const FULL_WORKGROUP_LANES: u32 = 256; + +/// Register-pressure class of a kernel body, used to pick how many outputs +/// one thread computes when a tiling trades thread-level parallelism for +/// register reuse. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum RegPressure { + /// A handful of values live across the reuse loop (elementwise bodies). + ElementwiseFew, +} + +/// Dispatch-sizing policy for one device. Cheap to construct (`Copy` data +/// gathered from cached device state); build it on demand via +/// [`crate::Device::dispatch_policy`]. +#[derive(Clone, Copy, Debug)] +pub(crate) struct DispatchPolicy { + /// Lanes that must stay in flight before any policy trades parallelism + /// for per-thread work. See [`crate::Device::saturation_lanes`]. + saturation_lanes: u32, + /// Fixed subgroup width when known; 32 as a documented fallback — every + /// policy derived from it is a floor, and 32 is the narrowest width on + /// hardware fusor targets, so a wrong guess only keeps more parallelism. + subgroup_width: u32, + /// Hardware cap on one workgroup's invocations. + max_workgroup_lanes: u32, + last_level_cache_bytes: u64, + /// Hardware cap on one workgroup's shared-memory bytes + /// (`max_compute_workgroup_storage_size`: 16 KB WebGPU baseline, 32 KB + /// on Apple silicon). + max_workgroup_storage_bytes: u32, + /// Physical rates the matmul cost model prices its terms in. See the + /// accessors for units; all four are per-class floors derived in + /// [`crate::Device::matmul_rates`]. + rates: MatmulRates, +} + +/// The four device rates the cooperative-matmul cost model needs to express +/// MMA issue, threadgroup traffic, output stores and DRAM traffic in one +/// commensurable unit (femtoseconds). Integers so the argmin is exact and +/// bit-reproducible — candidate scores tie often, the tile table being built +/// from powers of two. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct MatmulRates { + /// Peak MAC issue in MACs per nanosecond. + pub(crate) mac_per_ns: u64, + /// Achievable DRAM bandwidth in bytes per nanosecond x10 (one decimal). + pub(crate) dram_decibytes_per_ns: u64, + /// Achievable threadgroup-memory bandwidth in bytes per nanosecond, + /// covering both operand staging and cooperative fragment loads. + pub(crate) workgroup_bytes_per_ns: u64, + /// Accumulator zeroing, the cooperative store's fragment shuffles and the + /// store itself, in femtoseconds per padded output element a workgroup + /// emits *per subgroup in that workgroup* — the drain is barrier-gated + /// across the whole workgroup, so its per-element cost tracks the + /// workgroup's width. See the T3 term in `matmul::cost`. + pub(crate) store_fs_per_element: u64, + /// Threadgroup traffic of a single-buffered body as a percentage of the + /// same traffic double-buffered. A single-buffered kernel cannot overlap + /// the next K iteration's staging with the current iteration's MMAs, and + /// the other three rates are all fitted on double-buffered bodies, so a + /// body that stages from one tile pair pays this on its staging — every + /// K iteration, which is why the penalty grows with the K loop while what + /// it buys ([`Self::core_workgroup_slots`]) does not. + pub(crate) single_buffered_traffic_pct: u64, +} + +impl DispatchPolicy { + pub(crate) fn from_device(device: &Device) -> Self { + let limits = device.limits(); + let subgroup_width = if device.subgroups_supported() { + device.max_subgroup_size().max(1) + } else { + 32 + }; + Self::from_parts( + device.saturation_lanes(), + subgroup_width, + limits + .max_compute_workgroup_size_x + .min(limits.max_compute_invocations_per_workgroup), + device.last_level_cache_bytes(), + limits.max_compute_workgroup_storage_size, + device.matmul_rates(), + ) + } + + pub(crate) fn from_parts( + saturation_lanes: u32, + subgroup_width: u32, + max_workgroup_lanes: u32, + last_level_cache_bytes: u64, + max_workgroup_storage_bytes: u32, + rates: MatmulRates, + ) -> Self { + Self { + saturation_lanes: saturation_lanes.max(1), + subgroup_width: subgroup_width.max(1), + max_workgroup_lanes: max_workgroup_lanes.max(1), + last_level_cache_bytes, + max_workgroup_storage_bytes, + rates, + } + } + + /// Physical rates for the cooperative-matmul cost model. + pub(crate) fn matmul_rates(&self) -> MatmulRates { + self.rates + } + + /// The workgroup width for full-width dispatches. + pub(crate) fn preferred_workgroup_lanes(&self) -> u32 { + FULL_WORKGROUP_LANES.min(self.max_workgroup_lanes) + } + + /// Hardware cap on one workgroup's invocations. + pub(crate) fn max_workgroup_lanes(&self) -> u32 { + self.max_workgroup_lanes + } + + /// Hardware cap on one workgroup's shared-memory bytes. + pub(crate) fn max_workgroup_storage_bytes(&self) -> u32 { + self.max_workgroup_storage_bytes + } + + /// The device-parallelism floor (see [`crate::Device::saturation_lanes`]). + pub(crate) fn saturation_lanes(&self) -> u32 { + self.saturation_lanes + } + + /// The parallelism floor for a body that prefetches its own operands. + /// [`Self::saturation_lanes`] carries a ~4x oversubscription factor + /// because a latency-exposed kernel can only cover a DRAM round trip by + /// having other warps ready to issue. A double-buffered cooperative-matmul + /// body issues the next K tile's loads before running the current tile's + /// MMAs, so it covers that round trip in software and reaches peak issue + /// at half the residency. + pub(crate) fn prefetched_saturation_lanes(&self) -> u32 { + (self.saturation_lanes / 2).max(1) + } + + /// How many workgroups of the given threadgroup-memory footprint a core + /// can hold at once. Shared memory is carved from a per-core pool the + /// same size as the per-workgroup cap, so this is a plain division — + /// and it is the only place a kernel's footprint costs anything beyond + /// legality. Co-resident workgroups cover each other's epilogue drain, + /// which is why the T3 term in `matmul::cost` is divided by it. + pub(crate) fn core_workgroup_slots(&self, workgroup_bytes: u64) -> u64 { + (u64::from(self.max_workgroup_storage_bytes) / workgroup_bytes.max(1)).max(1) + } + + /// Smallest workgroup worth a subgroup-accelerated whole-block + /// reduction: two subgroups — one subgroup makes the cross-subgroup + /// combine tree degenerate. + pub(crate) fn min_reduction_lanes(&self) -> u32 { + (2 * self.subgroup_width).min(self.preferred_workgroup_lanes()) + } + + /// A natural one-workgroup-per-row dispatch leaves the device idle, so a + /// fan-out-plus-combine split pays for its combine kernel. + pub(crate) fn should_split_for_occupancy(&self, natural_wgs: u32, wg_lanes: u32) -> bool { + (natural_wgs as u64) * (wg_lanes as u64) < self.saturation_lanes as u64 + } + + /// Register tiling may trade threads for per-thread work only when the + /// post-tiling thread count still saturates the device. + pub(crate) fn tiling_leaves_saturated(&self, total_threads: u32) -> bool { + total_threads >= self.saturation_lanes + } + + /// Outputs one thread computes when a register-reuse tiling engages. + /// Doubling past this halves thread count below the saturation floor at + /// the engagement point; halving it buys too little load amortization. + pub(crate) fn work_per_thread(&self, class: RegPressure) -> u32 { + match class { + RegPressure::ElementwiseFew => 4, + } + } + + /// Monomorphization buckets for dynamic-axis row kernels, smallest + /// first: powers of two from half a full workgroup up to the hardware + /// cap, at most four buckets to bound compile count. + pub(crate) fn dynamic_block_buckets(&self) -> impl Iterator + use<> { + let start = self.preferred_workgroup_lanes() / 2; + let max = self.max_workgroup_lanes; + (0..4u32) + .map(move |i| start << i) + .filter(move |&b| b >= 1 && b <= max) + } + + /// Data strictly below the cache watermark is treated as cache-resident: + /// re-reads cost no bandwidth, so reuse-driven tilings should not + /// engage. Strict comparison preserves the pre-policy gate exactly. + pub(crate) fn cache_resident(&self, bytes: u64) -> bool { + bytes < self.last_level_cache_bytes + } + + /// Horizontal-merge ceiling for elementwise segments: the smallest op + /// that could take the register-reuse tiled path must stay unmerged, so + /// the bound is exactly the tiled path's engagement element count. + pub(crate) fn merge_elements_bound(&self) -> usize { + self.saturation_lanes as usize * self.work_per_thread(RegPressure::ElementwiseFew) as usize + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Apple-silicon Metal: 32-wide subgroups, 1024-lane workgroups. + fn apple() -> DispatchPolicy { + DispatchPolicy::from_parts(64 << 10, 32, 1024, 8 << 20, 32 << 10, crate::device::APPLE_MATMUL_RATES) + } + + /// WebGPU baseline limits: 256-lane workgroups. + fn webgpu_baseline() -> DispatchPolicy { + DispatchPolicy::from_parts(64 << 10, 32, 256, 4 << 20, 16 << 10, crate::device::APPLE_MATMUL_RATES) + } + + /// The derived values must reproduce the constants the kernels were + /// tuned with before the policy existed. If one of these fails, a + /// dispatch policy silently moved — decide deliberately, then update + /// both the derivation and this pin. + #[test] + fn derived_values_match_legacy_constants() { + let p = apple(); + assert_eq!(p.preferred_workgroup_lanes(), 256); // ex-BLOCK + assert_eq!(p.min_reduction_lanes(), 64); // ex-MIN_STATIC_BLOCK + assert_eq!(p.merge_elements_bound(), 262_144); // ex-MAX_MERGED_NARY_ELEMENTS + assert_eq!( + p.dynamic_block_buckets().collect::>(), + vec![128, 256, 512, 1024] // ex-ROW_DYNAMIC_BLOCKS + ); + // ex-MIN_TILED_THREADS + assert!(p.tiling_leaves_saturated(65_536)); + assert!(!p.tiling_leaves_saturated(65_535)); + // ex-SPLIT_ROWS_TARGET: rows < 256 at 256-lane blocks + assert!(p.should_split_for_occupancy(255, 256)); + assert!(!p.should_split_for_occupancy(256, 256)); + // ex-NARY_TM + assert_eq!(p.work_per_thread(RegPressure::ElementwiseFew), 4); + } + + #[test] + fn baseline_device_clamps() { + let p = webgpu_baseline(); + assert_eq!(p.preferred_workgroup_lanes(), 256); + assert_eq!( + p.dynamic_block_buckets().collect::>(), + vec![128, 256] + ); + } + + #[test] + fn no_subgroup_device_floors() { + // Subgroup width falls back to 32 → same reduction floor. + let p = DispatchPolicy::from_parts(64 << 10, 32, 512, 4 << 20, 32 << 10, crate::device::APPLE_MATMUL_RATES); + assert_eq!(p.min_reduction_lanes(), 64); + assert_eq!( + p.dynamic_block_buckets().collect::>(), + vec![128, 256, 512] + ); + } +} diff --git a/fusor-ml/core/src/pair_wise.rs b/fusor-ml/core/src/pair_wise.rs index 32066a2f6..99b6038de 100644 --- a/fusor-ml/core/src/pair_wise.rs +++ b/fusor-ml/core/src/pair_wise.rs @@ -1,7 +1,7 @@ use std::ops::{Add, Div, Mul, Sub}; use crate::{ - Tensor, + DataType, Tensor, nary_wise::{NaryFunction, NaryOp}, }; @@ -81,3 +81,51 @@ macro_rules! impl_pairwise_method { } impl_pairwise_method!(pow, NaryOp::Pow, "pow", pow_, |a, b| a.pow(&b)); + +/// Emit a tensor-tensor comparison method producing 1/0 in the output type +/// `D`, mirroring the scalar comparisons in `element_wise`. +macro_rules! impl_pairwise_cmp { + ($(#[$meta:meta])* $method:ident, $nary_op:expr, $op_name:literal) => { + impl Tensor { + $(#[$meta])* + pub fn $method(&self, other: &Self) -> Tensor { + assert_eq!(self.datatype(), other.datatype()); + self.binary_nary( + other, + NaryFunction::binary( + Some($op_name.to_string()), + $nary_op, + self.datatype(), + other.datatype(), + D::DATA_TYPE, + ), + ) + } + } + }; +} + +impl_pairwise_cmp!( + /// Element-wise `self == other` returning 1 for true and 0 for false. + eq_tensor, NaryOp::Equal, "eq" +); +impl_pairwise_cmp!( + /// Element-wise `self != other` returning 1 for true and 0 for false. + ne_tensor, NaryOp::NotEqual, "ne" +); +impl_pairwise_cmp!( + /// Element-wise `self < other` returning 1 for true and 0 for false. + lt_tensor, NaryOp::Less, "lt" +); +impl_pairwise_cmp!( + /// Element-wise `self <= other` returning 1 for true and 0 for false. + lte_tensor, NaryOp::LessEqual, "lte" +); +impl_pairwise_cmp!( + /// Element-wise `self > other` returning 1 for true and 0 for false. + gt_tensor, NaryOp::Greater, "gt" +); +impl_pairwise_cmp!( + /// Element-wise `self >= other` returning 1 for true and 0 for false. + gte_tensor, NaryOp::GreaterEqual, "gte" +); diff --git a/fusor-ml/core/src/quantized/dequantize.rs b/fusor-ml/core/src/quantized/dequantize.rs index 858814760..a263cba90 100644 --- a/fusor-ml/core/src/quantized/dequantize.rs +++ b/fusor-ml/core/src/quantized/dequantize.rs @@ -18,7 +18,7 @@ use crate::{ use super::{QMatrix, QMatrixStorageLayout}; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub(crate) struct DequantizeOperation { pub(crate) matrix: QMatrix, pub(crate) datatype: DataTypeEnum, @@ -110,6 +110,8 @@ impl Operation for DequantizeOperation { fn visit_dependencies(&self, _: &mut dyn FnMut(crate::compute_graph::NodeIndex)) {} + fn visit_dependencies_mut(&mut self, _: &mut dyn FnMut(&mut crate::compute_graph::NodeIndex)) {} + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { let shape = &self.matrix.shape; let output_tensor = TensorData::new_for_shape(&nodes.device(), shape, self.datatype); diff --git a/fusor-ml/core/src/quantized/embedding.rs b/fusor-ml/core/src/quantized/embedding.rs index 457b4ebc3..e2be7d253 100644 --- a/fusor-ml/core/src/quantized/embedding.rs +++ b/fusor-ml/core/src/quantized/embedding.rs @@ -85,7 +85,7 @@ fn datatype_element(datatype: DataTypeEnum) -> Option { }) } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub(crate) struct QEmbeddingOperation { pub(crate) indexes: NodeIndex, pub(crate) matrix: QMatrix, @@ -222,6 +222,10 @@ impl Operation for QEmbeddingOperation { f(self.indexes); } + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + f(&mut self.indexes); + } + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { let indexes = nodes .get_result(self.indexes) diff --git a/fusor-ml/core/src/quantized/matmul/kernel.rs b/fusor-ml/core/src/quantized/matmul/kernel.rs index d104e93ea..f8ffbe46a 100644 --- a/fusor-ml/core/src/quantized/matmul/kernel.rs +++ b/fusor-ml/core/src/quantized/matmul/kernel.rs @@ -13,33 +13,6 @@ enum QmatmulDirectTokens { } impl QMatMulOperation { - /// Lower this operation to its kernel plan. Recognition and epilogue - /// fusion only build operations the direct paths can lower (see - /// `supports_elementwise_epilogue_fusion`), so a `None` from every path - /// here is an invariant violation, not a recoverable state. - pub(crate) fn build_direct_kernels( - &self, - graph: &crate::compute_graph::ComputeGraphInner, - workgroup_shape: &crate::mir::workgroup_shape::WorkgroupShape, - inputs: &[MirValue], - ) -> Result { - if inputs - .last() - .and_then(MirValue::as_tensor) - .is_some_and(|output| output.layout().shape().contains(&0)) - { - return Ok(QMatMulKernelPlan::EmptyOutput); - } - - if let Some(kernels) = self.build_m_padded_kernels(graph, inputs) { - return Ok(QMatMulKernelPlan::Kernels(kernels)); - } - if let Some(kernel) = self.build_direct_kernel(graph, workgroup_shape, inputs) { - return Ok(QMatMulKernelPlan::Kernels(vec![kernel])); - } - Err(QMatMulLoweringError::new(self.name())) - } - /// Lower an M-padded matmul: copy the activation into a zero-padded /// scratch tensor (the same kernel a `resize` view lowers to) and run the /// matmul over the padded views. The output buffer's slack rows were @@ -437,71 +410,20 @@ impl QMatMulOperation { } ^ if f16_storage { 0xF16F_0001u64 } else { 0 } ^ accumulator_offsets_identity; - let fast_dispatch_size = if use_workgroup_qmatmul { - // The workgroup-tiled kernel computes its own grid inside - // `tile::build`; skip the pre-built-pipeline fast path. - None - } else { - match variant { - QMatmulPath::Q5SmallSingleRow | QMatmulPath::SingleRow => { - let qgemv_cols_per_workgroup = - qgemv_cols_per_workgroup_for_direct(format, k, n); - let qgemv_workgroups = n.div_ceil(qgemv_cols_per_workgroup); - let [dispatch_x, _] = split_workgroups_2d(qgemv_workgroups, max_workgroups)?; - qmatmul_workgroups_x = dispatch_x; - Some([ - qmatmul_workgroups_x, - qgemv_workgroups.div_ceil(qmatmul_workgroups_x), - 1, - ]) - } - // The IR-build fallback (cached=false catch-all) is the only - // path that defers the dispatch to the IR builder; every - // tile-aligned coop variant has a precomputed `[n/BN, m/BM, 1]`. - QMatmulPath::Tile { - cached: false, - tile, - } if tile == CoopTile::new(64, 64, QMATMUL_COOP_BK) => None, - QMatmulPath::Workgroup => None, - QMatmulPath::Q8Wide(tile) | QMatmulPath::Tile { tile, .. } => { - Some([n / tile.bn, m / tile.bm, 1]) - } - } - }; - let kernel_name = kernel_name.into(); - // The pre-built-pipeline fast path can only be reused when there's no - // epilogue attached — otherwise the cached pipeline encodes the wrong - // (no-epilogue) kernel. Skip the fast path entirely when fusing. - if pre_extra_tensors.is_empty() - && post_extra_tensors.is_empty() - && !has_custom_accumulator_offsets - && pre_epilogue_with_extras.is_none() - && post_epilogue_with_extras.is_none() - && let Some(dispatch_size) = fast_dispatch_size + // The one qgemv geometry decision: the dispatch here and the kernel + // body below both consume this value, so they cannot disagree. + let qgemv_shape = tile_ir_kernels::qgemv_selected_shape(format, k, n); + if !use_workgroup_qmatmul + && matches!( + variant, + QMatmulPath::Q5SmallSingleRow | QMatmulPath::SingleRow + ) { - if dispatch_size.iter().any(|dim| *dim > max_workgroups) { - return None; - } - let pipeline_key = QMatMulDirectPipelineKey::new( - matrix.datatype(), - matrix.storage_layout(), - crate::quantized::QMatMulShape { m, k, n: matrix_n }, - subgroup_size_range, - dispatch_size, - input.layout(), - output.layout(), - ); - if let Some(kernel) = cached_qmatmul_direct_kernel( - &kernel_name, - matrix, - &pipeline_key, - input, - output, - dispatch_size, - ) { - return Some(kernel); - } + let qgemv_workgroups = n.div_ceil(qgemv_shape.cols_per_workgroup()); + let [dispatch_x, _] = split_workgroups_2d(qgemv_workgroups, max_workgroups)?; + qmatmul_workgroups_x = dispatch_x; } + let kernel_name = kernel_name.into(); let pre_with_extras_for_ir = pre_epilogue_with_extras.clone(); let post_with_extras_for_ir = post_epilogue_with_extras.clone(); let post_accumulator_offsets_for_ir = post_accumulator_offsets.to_vec(); @@ -519,25 +441,15 @@ impl QMatMulOperation { y_view, ); let epilogues = tile_ir_kernels::QmatmulEpilogues::default(); - if m == 1 { - tile_ir_kernels::qgemv_workgroup_storage_f16_with_epilogue( - phase, - &a, - &b, - &y, - &epilogues, - max_workgroups, - ); - } else { - tile_ir_kernels::qmatmul_workgroup_storage_f16_with_epilogues( - phase, - &a, - &b, - &y, - &epilogues, - max_workgroups, - ); - } + tile_ir_kernels::qmatmul_workgroup_with_epilogues( + phase, + &a, + &b, + &y, + tile_ir::ScalarElement::F16, + &epilogues, + max_workgroups, + ); return; } let a = tile_storage_read_with_direct_layout(phase, a_view); @@ -605,47 +517,20 @@ impl QMatMulOperation { }, }; if use_workgroup_qmatmul { - if m == 1 { - if use_f16_workgroup_tiles { - tile_ir_kernels::qgemv_workgroup_f16_with_epilogue( - phase, - &a, - &b, - &y, - &epilogues, - max_workgroups, - ); - } else { - tile_ir_kernels::qgemv_workgroup_with_epilogue( - phase, - &a, - &b, - &y, - &epilogues, - max_workgroups, - ); - } + let staging = if use_f16_workgroup_tiles { + tile_ir::ScalarElement::F16 } else { - if use_f16_workgroup_tiles { - tile_ir_kernels::qmatmul_workgroup_f16_with_epilogues( - phase, - &a, - &b, - &y, - &epilogues, - max_workgroups, - ); - } else { - tile_ir_kernels::qmatmul_workgroup_with_epilogues( - phase, - &a, - &b, - &y, - &epilogues, - max_workgroups, - ); - } - } + tile_ir::ScalarElement::F32 + }; + tile_ir_kernels::qmatmul_workgroup_with_epilogues( + phase, + &a, + &b, + &y, + staging, + &epilogues, + max_workgroups, + ); return; } // Map the selected variant to its cooperative tile dimensions. @@ -663,6 +548,7 @@ impl QMatMulOperation { &y, qmatmul_workgroups_x, subgroups, + qgemv_shape, &epilogues, ); return; @@ -685,16 +571,6 @@ impl QMatMulOperation { if dispatch_size.iter().any(|dim| *dim > max_workgroups) { return None; } - let pipeline_key = QMatMulDirectPipelineKey::new_with_epilogue( - matrix.datatype(), - matrix.storage_layout(), - crate::quantized::QMatMulShape { m, k, n: matrix_n }, - epilogue_identity, - subgroup_size_range, - dispatch_size, - input.layout(), - output.layout(), - ); let cache_key = qmatmul_direct_cache_key::( |state| { variant.hash(state); @@ -723,11 +599,9 @@ impl QMatMulOperation { ); qmatmul_direct_kernel_from_ir( device, - kernel_name.clone(), kernel_name, cache_key, matrix, - pipeline_key, input, pre_extra_tensors, post_extra_tensors, @@ -752,6 +626,32 @@ fn hash_qmatmul_epilogue(state: &mut FxHasher, epilogue: &Option Result { + if inputs + .last() + .and_then(MirValue::as_tensor) + .is_some_and(|output| output.layout().shape().contains(&0)) + { + return Ok(DirectKernelPlan::empty()); + } + + if let Some(kernels) = self.build_m_padded_kernels(graph, inputs) { + return Ok(DirectKernelPlan::many(kernels)); + } + if let Some(kernel) = self.build_direct_kernel(graph, workgroup_shape, inputs) { + return Ok(DirectKernelPlan::single(kernel)); + } + Err(DirectKernelLoweringError::new(self.name())) + } + fn hash_kernel_fields(&self, state: &mut FxHasher) { self.input_datatype.hash(state); self.in_shape.hash(state); @@ -816,6 +716,21 @@ impl Operation for QMatMulOperation { } } + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + f(&mut self.input); + for epilogue in [ + &mut self.pre_element_wise_expr, + &mut self.post_element_wise_expr, + ] + .into_iter() + .flatten() + { + for extra in &mut epilogue.extras { + f(extra); + } + } + } + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { let m_pad = self.m_pad_target(KernelDeviceCaps::from_device(&nodes.device())); let base = diff --git a/fusor-ml/core/src/quantized/matmul/mod.rs b/fusor-ml/core/src/quantized/matmul/mod.rs index 96fd9bb95..b2d918f77 100644 --- a/fusor-ml/core/src/quantized/matmul/mod.rs +++ b/fusor-ml/core/src/quantized/matmul/mod.rs @@ -1,4 +1,4 @@ -use std::{fmt, hash::Hash}; +use std::hash::Hash; use crate::{ DataTypeEnum, Device, Layout, Tensor, TensorData, @@ -12,7 +12,7 @@ use crate::{ inputs::MirValue, kernel_backend, kernel_backend::DirectKernel, - operation::Operation, + operation::{DirectKernelLoweringError, DirectKernelPlan, Operation}, tile_direct::{ flatten_matrix_layout, tile_storage_read_with_direct_layout, tile_storage_read_with_direct_layout_typed, tile_storage_write_with_direct_layout, @@ -28,7 +28,7 @@ use fusor_tile_ir as tile_ir; use fusor_tile_ir_kernels as tile_ir_kernels; use rustc_hash::FxHasher; -use super::{QMatMulDirectPipelineKey, QMatrix, QMatrixStorageLayout}; +use super::{QMatrix, QMatrixStorageLayout}; mod kernel; #[cfg(test)] @@ -167,7 +167,9 @@ fn qgemv_subgroup_supported( n: u32, caps: KernelDeviceCaps, ) -> bool { - let subgroups = tile_ir_kernels::qgemv_subgroups_per_workgroup_for_shape(format, k, n); + // The subgroup count the emitted kernel will actually use — not the + // per-format default, which disagrees in the shape-specialized regimes. + let subgroups = tile_ir_kernels::qgemv_selected_shape(format, k, n).subgroups; subgroup_runtime_supported(caps) && subgroup_workgroup_size_supported(caps, subgroups) } @@ -418,7 +420,7 @@ pub(crate) struct ElementwiseEpilogue { pub(crate) output_datatype: DataTypeEnum, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub(crate) struct QMatMulOperation { pub(crate) input_datatype: DataTypeEnum, pub(crate) input: NodeIndex, @@ -591,48 +593,6 @@ pub(crate) struct DirectKernelTensors<'a> { pub output: &'a TensorData, } -pub(crate) enum QMatMulKernelPlan { - EmptyOutput, - Kernels(Vec), -} - -impl QMatMulKernelPlan { - pub(crate) fn dispatch_count(&self) -> usize { - match self { - Self::EmptyOutput => 0, - Self::Kernels(kernels) => kernels.len(), - } - } - - pub(crate) fn into_kernels(self) -> Vec { - match self { - Self::EmptyOutput => Vec::new(), - Self::Kernels(kernels) => kernels, - } - } -} - -#[derive(Debug, Clone)] -pub(crate) struct QMatMulLoweringError { - operation: String, -} - -impl QMatMulLoweringError { - fn new(operation: String) -> Self { - Self { operation } - } -} - -impl fmt::Display for QMatMulLoweringError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "QMatMul lowering produced no kernel plan for {}", - self.operation - ) - } -} - pub(crate) struct DirectKernelChains<'a> { pub pre_expr: Option<&'a ElementwiseEpilogue>, pub post_expr: Option<&'a ElementwiseEpilogue>, @@ -700,9 +660,8 @@ fn qmatmul_qgemv_dispatch_supported( n: u32, max_workgroups_per_dimension: u32, ) -> bool { - let qgemv_cols_per_workgroup = qgemv_cols_per_workgroup_for_direct(format, k, n); - let qgemv_workgroups = n.div_ceil(qgemv_cols_per_workgroup); - split_workgroups_2d(qgemv_workgroups, max_workgroups_per_dimension).is_some() + let cols = tile_ir_kernels::qgemv_selected_shape(format, k, n).cols_per_workgroup(); + split_workgroups_2d(n.div_ceil(cols), max_workgroups_per_dimension).is_some() } #[allow(clippy::too_many_arguments)] @@ -781,40 +740,12 @@ fn qmatrix_direct_quant_format(matrix: &QMatrix) -> Option Option { - let pipeline = matrix - .direct_pipeline_cache() - .write() - .get(pipeline_key) - .cloned()?; - Some( - kernel_backend::DirectKernel::from_prepared_three_buffer_pipeline( - kernel_name.to_owned(), - pipeline, - input.buffer().clone(), - matrix.buffer().clone(), - output.buffer().clone(), - dispatch_size, - ), - ) -} - #[allow(clippy::too_many_arguments)] fn qmatmul_direct_kernel_from_ir( device: &Device, - cached_kernel_name: String, kernel_name: String, cache_key: kernel_backend::KernelCacheKey, matrix: &QMatrix, - pipeline_key: QMatMulDirectPipelineKey, input: &TensorData, pre_extra_tensors: &[&TensorData], post_extra_tensors: &[&TensorData], @@ -845,31 +776,17 @@ fn qmatmul_direct_kernel_from_ir( dispatch_size, ); } - if let Some(kernel) = cached_qmatmul_direct_kernel( - &cached_kernel_name, - matrix, - &pipeline_key, - input, - output, - dispatch_size, - ) { - return Some(kernel); - } - let pipeline = kernel_backend::three_buffer_pipeline_from_ir( + let (pipeline, cached) = kernel_backend::three_buffer_pipeline_from_ir( device.kernel_cache(), &kernel_name, cache_key, build_ir, )?; - let pipeline = matrix - .direct_pipeline_cache() - .write() - .get_or_insert(pipeline_key, || pipeline.clone()) - .clone(); Some( kernel_backend::DirectKernel::from_prepared_three_buffer_pipeline( kernel_name, pipeline, + Some(cached), input.buffer().clone(), matrix.buffer().clone(), output.buffer().clone(), @@ -892,46 +809,6 @@ fn effective_qmatmul_max_workgroups_per_dimension(limits: &wgpu::Limits) -> u32 limits.max_compute_workgroups_per_dimension.max(1) } -/// Output columns per workgroup for the direct qgemv path, by (format, K, N). -fn qgemv_cols_per_workgroup_for_direct(format: tile_ir::GgmlQuantFormat, k: u32, n: u32) -> u32 { - // Q4K specializations. - if format.is_q4k_family() && k <= 4096 && (4096..8192).contains(&n) { - return 4; // was Q4KSmallWide4 - } - if format.is_q4k_family() && k <= 4096 && n >= 8192 { - return 8; // was Q4KSmallWide8 - } - if format.is_q4k_family() && n <= 4096 && k > 4096 { - return 8; // was Q4KLargeNarrow8 - } - // Q6K specializations. - if format.is_q6k_family() && k <= 4096 && n >= 8192 { - return 8; // was Q6KSmallWide8 - } - if format.is_q6k_family() && n <= 4096 && k > 4096 { - return 4; // was Q6KLargeNarrow4 - } - // Q8_0 wide. - if format.is_q8_0_family() && k <= 1024 && n >= 8192 { - return 32; // was Q8WideAccelerated32 - } - // FormatAccelerated: Q5_0 mid (K,N both 2048..=4096), Q4K/Q6K general, - // or Q5_0 large (K*N >= 4M). Delegates to the format-aware helper. - let q5_mid = - format.is_q5_0_family() && (2048..=4096).contains(&k) && (2048..=4096).contains(&n); - let q5_large = format.is_q5_0_family() - && (k as u64) - .checked_mul(n as u64) - .is_some_and(|elements| elements >= 4 * 1024 * 1024); - if q5_mid || format.is_q4k_family() || format.is_q6k_family() || q5_large { - return tile_ir_kernels::qgemv_cols_per_workgroup_for_shape(format, k, n); - } - if format.is_q5_0_family() && k <= 1024 && n <= 4096 { - return 8; // was Q5Small8 - } - 4 // was Default4 -} - fn qmatmul_m_pad_target_for_caps(m: usize, n: usize, caps: KernelDeviceCaps) -> Option { if !qmatmul_coop_supported(caps) { return None; diff --git a/fusor-ml/core/src/quantized/matmul/tests.rs b/fusor-ml/core/src/quantized/matmul/tests.rs index 1ea09b838..2f2bb011b 100644 --- a/fusor-ml/core/src/quantized/matmul/tests.rs +++ b/fusor-ml/core/src/quantized/matmul/tests.rs @@ -538,6 +538,67 @@ mod tests { }); } + // Q6K in the large regime (K <= 4096, 8192 <= N <= 16384): the dispatch + // ladder and the kernel's own geometry must agree on columns per + // workgroup, or the launched grid covers only part of the output. Checks + // the high columns that an under-dispatch leaves unwritten. + #[test] + fn q6k_large_single_row_qgemv_covers_all_columns() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + if !device.subgroups_supported() { + return; + } + + let weight_shape = [8192usize, 2048usize]; + let blocks_per_row = weight_shape[1] / BlockQ6K::BLOCK_SIZE; + let raw_bytes = patterned_q6k_bytes(weight_shape); + let matrix = + QMatrix::from_parts(&device, &raw_bytes, weight_shape.into(), GgmlType::Q6K) + .unwrap(); + let blocks: &[BlockQ6K] = bytemuck::cast_slice(&raw_bytes); + let input_values = (0..weight_shape[1]) + .map(|index| { + let bucket = (index.wrapping_mul(37).wrapping_add(11)) % 101; + (bucket as f32 - 50.0) * 0.0025 + }) + .collect::>(); + // Two runs: the first builds the kernel (whose own grid is + // authoritative), the second takes the cached-pipeline fast + // path whose dispatch the core-side ladder computes. + for run in 0..2 { + let input = Tensor::from_slice::(&device, [1, weight_shape[1]], &input_values); + let result = input.q_mat_mul(&matrix).as_slice::<2, f32>().await.unwrap(); + + assert_eq!(result.shape(), &[1, weight_shape[0]]); + for col in [0usize, 1, 4095, 4096, 6000, 8190, 8191] { + let expected = (0..blocks_per_row) + .map(|block_col| { + let block = &blocks[col * blocks_per_row + block_col]; + let weights = block.dequantize(); + weights + .as_ref() + .iter() + .enumerate() + .map(|(offset, weight)| { + input_values[block_col * BlockQ6K::BLOCK_SIZE + offset] + * *weight + }) + .sum::() + }) + .sum::(); + let actual = result[[0, col]]; + assert!( + (actual - expected).abs() <= 1e-2_f32.max(expected.abs() * 1.0e-4), + "run={run} col={col} actual={actual} expected={expected}" + ); + } + } + }); + } + #[test] fn q4k_large_single_row_qgemv_handles_tail_columns_with_subgroups() { pollster::block_on(async { diff --git a/fusor-ml/core/src/quantized/mod.rs b/fusor-ml/core/src/quantized/mod.rs index a63d3f20e..450e3a393 100644 --- a/fusor-ml/core/src/quantized/mod.rs +++ b/fusor-ml/core/src/quantized/mod.rs @@ -1,125 +1,15 @@ -use std::{borrow::Cow, mem::size_of, num::NonZeroUsize, sync::Arc}; +use std::{borrow::Cow, mem::size_of, sync::Arc}; -use crate::{Device, Layout}; +use crate::Device; use fusor_gguf::{ BlockQ4_0, BlockQ4K, BlockQ5_0, BlockQ5K, BlockQ6K, BlockQ8_0, GgmlType, GgufBlock, GgufMetadata, GgufReadError, GgufTensorMetadata, }; -use lru::LruCache; -use parking_lot::RwLock; -use rustc_hash::FxBuildHasher; pub(crate) mod dequantize; pub(crate) mod embedding; pub(crate) mod matmul; -const QMATRIX_DIRECT_PIPELINE_CACHE_SIZE: usize = 16; - -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -pub(crate) struct QMatMulDirectPipelineKey { - format: u8, - storage_layout: QMatrixStorageLayout, - m: u32, - k: u32, - n: u32, - // Structural hash of attached epilogues and lowering-only choices. - // Zero for plain qmatmul; non-zero values disambiguate kernels whose - // bindings/dispatch are otherwise identical. - epilogue_identity: u64, - subgroup_size_range: [u32; 2], - dispatch_size: [u32; 3], - input_layout: QMatMulDirectLayoutKey, - output_layout: QMatMulDirectLayoutKey, -} - -#[derive(Clone, Copy)] -pub(crate) struct QMatMulShape { - pub m: u32, - pub k: u32, - pub n: u32, -} - -impl QMatMulDirectPipelineKey { - pub(crate) fn new( - format: GgmlType, - storage_layout: QMatrixStorageLayout, - shape: QMatMulShape, - subgroup_size_range: [u32; 2], - dispatch_size: [u32; 3], - input_layout: &Layout, - output_layout: &Layout, - ) -> Self { - Self::new_with_epilogue( - format, - storage_layout, - shape, - 0, - subgroup_size_range, - dispatch_size, - input_layout, - output_layout, - ) - } - - #[allow(clippy::too_many_arguments)] - pub(crate) fn new_with_epilogue( - format: GgmlType, - storage_layout: QMatrixStorageLayout, - shape: QMatMulShape, - epilogue_identity: u64, - subgroup_size_range: [u32; 2], - dispatch_size: [u32; 3], - input_layout: &Layout, - output_layout: &Layout, - ) -> Self { - let QMatMulShape { m, k, n } = shape; - Self { - format: format as u8, - storage_layout, - m, - k, - n, - epilogue_identity, - subgroup_size_range, - dispatch_size, - input_layout: QMatMulDirectLayoutKey::new(input_layout), - output_layout: QMatMulDirectLayoutKey::new(output_layout), - } - } -} - -#[derive(Clone, Debug, PartialEq, Eq, Hash)] -enum QMatMulDirectLayoutKey { - Rank2 { - offset: usize, - shape: [usize; 2], - strides: [usize; 2], - }, - General { - offset: usize, - shape: Box<[usize]>, - strides: Box<[usize]>, - }, -} - -impl QMatMulDirectLayoutKey { - fn new(layout: &Layout) -> Self { - if layout.shape().len() == 2 && layout.strides().len() == 2 { - Self::Rank2 { - offset: layout.offset(), - shape: [layout.shape()[0], layout.shape()[1]], - strides: [layout.strides()[0], layout.strides()[1]], - } - } else { - Self::General { - offset: layout.offset(), - shape: layout.shape().into(), - strides: layout.strides().into(), - } - } - } -} - fn padded_copy_size(size: u64) -> u64 { let align_mask = wgpu::COPY_BUFFER_ALIGNMENT - 1; ((size + align_mask) & !align_mask).max(wgpu::COPY_BUFFER_ALIGNMENT) @@ -209,23 +99,6 @@ pub(crate) enum QMatrixStorageLayout { GpuF32Scales, } -fn env_bool(value: &str) -> Option { - match value.to_ascii_lowercase().as_str() { - "1" | "true" | "yes" | "on" | "native" => Some(true), - "0" | "false" | "no" | "off" | "expanded" | "f32" => Some(false), - _ => None, - } -} - -fn native_half_scale_storage_enabled( - shader_f16_supported: bool, - env_override: Option<&str>, -) -> bool { - env_override - .and_then(env_bool) - .unwrap_or(shader_f16_supported) -} - fn qmatrix_storage_layout_for_parts( ty: GgmlType, shader_f16_supported: bool, @@ -235,20 +108,12 @@ fn qmatrix_storage_layout_for_parts( let _ = shader_f16_supported; false }; - qmatrix_storage_layout_for_parts_with_env( - ty, - shader_f16_supported, - std::env::var("FUSOR_Q_NATIVE") - .ok() - .or_else(|| std::env::var("FUSOR_Q4K_NATIVE").ok()) - .as_deref(), - ) + qmatrix_storage_layout_for_parts_with_capability(ty, shader_f16_supported) } -fn qmatrix_storage_layout_for_parts_with_env( +fn qmatrix_storage_layout_for_parts_with_capability( ty: GgmlType, shader_f16_supported: bool, - env_override: Option<&str>, ) -> QMatrixStorageLayout { // Q6K is always stored with f32 scales. Its native block is 210 bytes — // not a multiple of 4 — so blocks past the first are not word-aligned, which @@ -263,7 +128,7 @@ fn qmatrix_storage_layout_for_parts_with_env( if matches!( ty, GgmlType::Q4_0 | GgmlType::Q5_0 | GgmlType::Q8_0 | GgmlType::Q4K | GgmlType::Q5K - ) && native_half_scale_storage_enabled(shader_f16_supported, env_override) + ) && shader_f16_supported { QMatrixStorageLayout::Native } else if matches!( @@ -283,8 +148,6 @@ pub struct QMatrix { buffer: Arc, datatype: GgmlType, storage_layout: QMatrixStorageLayout, - direct_pipeline_cache: - Arc>>, } impl std::fmt::Debug for QMatrix { @@ -378,10 +241,6 @@ impl QMatrix { buffer, datatype, storage_layout, - direct_pipeline_cache: Arc::new(RwLock::new(LruCache::with_hasher( - NonZeroUsize::new(QMATRIX_DIRECT_PIPELINE_CACHE_SIZE).unwrap(), - Default::default(), - ))), }) } @@ -497,10 +356,6 @@ impl QMatrix { buffer, datatype, storage_layout, - direct_pipeline_cache: Arc::new(RwLock::new(LruCache::with_hasher( - NonZeroUsize::new(QMATRIX_DIRECT_PIPELINE_CACHE_SIZE).unwrap(), - Default::default(), - ))), }) } @@ -508,12 +363,6 @@ impl QMatrix { &self.buffer } - pub(crate) fn direct_pipeline_cache( - &self, - ) -> &RwLock> { - &self.direct_pipeline_cache - } - pub fn shape(&self) -> &[usize] { &self.shape } @@ -545,32 +394,20 @@ mod tests { GgmlType::Q5K, ] { assert_eq!( - qmatrix_storage_layout_for_parts_with_env(ty, true, None), + qmatrix_storage_layout_for_parts_with_capability(ty, true), QMatrixStorageLayout::Native ); assert_eq!( - qmatrix_storage_layout_for_parts_with_env(ty, false, None), + qmatrix_storage_layout_for_parts_with_capability(ty, false), QMatrixStorageLayout::GpuF32Scales ); } assert_eq!( - qmatrix_storage_layout_for_parts_with_env(GgmlType::Q6K, true, None), - QMatrixStorageLayout::GpuF32Scales - ); - assert_eq!( - qmatrix_storage_layout_for_parts_with_env(GgmlType::Q6K, false, None), + qmatrix_storage_layout_for_parts_with_capability(GgmlType::Q6K, true), QMatrixStorageLayout::GpuF32Scales ); - } - - #[test] - fn q4k_storage_env_can_force_native_or_f32_expanded() { - assert_eq!( - qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4K, false, Some("1")), - QMatrixStorageLayout::Native - ); assert_eq!( - qmatrix_storage_layout_for_parts_with_env(GgmlType::Q4K, true, Some("0")), + qmatrix_storage_layout_for_parts_with_capability(GgmlType::Q6K, false), QMatrixStorageLayout::GpuF32Scales ); } diff --git a/fusor-ml/core/src/reduce.rs b/fusor-ml/core/src/reduce.rs index de97aa444..94bcad14b 100644 --- a/fusor-ml/core/src/reduce.rs +++ b/fusor-ml/core/src/reduce.rs @@ -43,7 +43,7 @@ fn unsqueeze_dim(tensor: &Tensor, dim_idx: usize) -> Tensor { /// `input(0, rank)`; the resolver widens it by inlining upstream elementwise /// expressions, so composed map-reduce clusters (contractions included) lower /// as a single kernel without materializing the intermediate. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct ReduceOperation { /// Producer inputs referenced by `expression`. pub(crate) inputs: Vec, @@ -135,6 +135,12 @@ impl Operation for ReduceOperation { } } + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + for input in &mut self.inputs { + f(input); + } + } + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { let mut mir_inputs: Vec = self .inputs @@ -190,7 +196,7 @@ impl Operation for ReduceOperation { } } -#[derive(Clone, Debug, Hash)] +#[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct ReduceFunction { pub(crate) name: Option, pub(crate) op: ReduceOp, @@ -199,7 +205,7 @@ pub struct ReduceFunction { } impl ReduceFunction { - fn new(op: ReduceOp, initial_value: NaryScalar, datatype: DataTypeEnum) -> Self { + pub(crate) fn new(op: ReduceOp, initial_value: NaryScalar, datatype: DataTypeEnum) -> Self { Self { name: None, op, diff --git a/fusor-ml/core/src/region.rs b/fusor-ml/core/src/region.rs new file mode 100644 index 000000000..f1005fa80 --- /dev/null +++ b/fusor-ml/core/src/region.rs @@ -0,0 +1,193 @@ +//! Multi-output elementwise regions. +//! +//! A region is a topologically-ordered list of elementwise statements fused +//! into one kernel body: statement `k` may read statement `j < k`'s value as +//! a register instead of a materialized tensor, and every statement whose +//! value is externally live (a flush target or user-held node) writes its +//! own output buffer. This generalizes the sole-consumer nary fusion rule: +//! a producer with several consumers still fuses when *all* of its consumers +//! land in the same region — external liveness is satisfied by emitting the +//! value as one of the region's outputs rather than by blocking fusion. +//! +//! Register reads use the `extras` slot-overflow convention of +//! [`crate::nary_direct::eval_nary_expr`]: statement `j`'s value is read as +//! `NaryExpr::IndexedInput { input_idx: inputs.len() + j, indices: [] }`. +//! +//! Regions exist only inside the resolver between optimization and lowering: +//! the inner compute graph never contains one, so the flush fingerprint +//! recipe is unaffected by their formation. + +use rustc_hash::FxHasher; +use std::hash::Hash; + +use crate::DataTypeEnum; +use crate::compute_graph::NodeIndex; +use crate::nary_wise::{ElementwiseOperation, NaryExpr}; + +#[derive(Debug, Clone)] +pub(crate) struct RegionStatement { + /// Slots `0..inputs.len()` read region inputs (elementwise or + /// custom-indexed, e.g. folded broadcast views). Slot `inputs.len() + j` + /// reads statement `j`'s register value (elementwise only). + pub(crate) expression: NaryExpr, + pub(crate) datatype: DataTypeEnum, + /// `Some(inner)` = externally live: stored to its own output buffer and + /// cached under that inner-graph node. `None` = register-only. + pub(crate) output: Option, +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct InputReadSummary { + pub(crate) last_reader: Option, + pub(crate) identity_only: bool, +} + +#[derive(Debug, Clone)] +pub(crate) struct ElementwiseRegionOperation { + /// External producers, deduplicated; order defines input slots. + pub(crate) inputs: Vec, + /// Topologically ordered; the last statement is the region's sink and + /// always has `output: Some(_)`. + pub(crate) statements: Vec, + /// Shared index space: every statement evaluates over this shape. + pub(crate) shape: Box<[usize]>, +} + +impl ElementwiseRegionOperation { + /// A single elementwise operation as a one-statement region. + pub(crate) fn from_nary(op: ElementwiseOperation, node: NodeIndex) -> Self { + Self { + inputs: op.inputs, + shape: op.shape, + statements: vec![RegionStatement { + expression: op.expression, + datatype: op.output_datatype, + output: Some(node), + }], + } + } + + /// The inverse of [`Self::from_nary`], for lowering a lone + /// single-statement region through the standalone elementwise path. + pub(crate) fn into_nary(self) -> Option { + let mut statements = self.statements; + if statements.len() != 1 { + return None; + } + let statement = statements.pop().expect("length checked"); + statement.output?; + Some(ElementwiseOperation { + inputs: self.inputs, + expression: statement.expression, + shape: self.shape, + output_datatype: statement.datatype, + }) + } + + pub(crate) fn output_count(&self) -> usize { + self.statements + .iter() + .filter(|statement| statement.output.is_some()) + .count() + } + + /// Storage bindings a merged kernel declares for this region. + pub(crate) fn binding_count(&self) -> usize { + self.inputs.len() + self.output_count() + } + + pub(crate) fn visit_dependencies(&self, f: &mut dyn FnMut(NodeIndex)) { + for &input in &self.inputs { + f(input); + } + } + + pub(crate) fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + for input in &mut self.inputs { + f(input); + } + } + + pub(crate) fn name(&self) -> String { + format!("region_x{}", self.statements.len()) + } + + /// Per input slot: the last statement index that reads it from memory + /// and whether every such read is identity-indexed (element `i` of the + /// input read exactly at output coordinate `i`). An output statement may + /// write its buffer over an input's only when all reads of that input + /// are identity (threads own disjoint elements, and each thread loads + /// before it stores) and none happen after the writing statement. + pub(crate) fn input_read_summary(&self) -> Vec { + let rank = self.shape.len(); + let mut summary = vec![ + InputReadSummary { + last_reader: None, + identity_only: true, + }; + self.inputs.len() + ]; + for (position, statement) in self.statements.iter().enumerate() { + Self::scan_reads( + &statement.expression, + self.inputs.len(), + rank, + position, + false, + &mut summary, + ); + } + summary + } + + fn scan_reads( + expr: &NaryExpr, + input_count: usize, + rank: usize, + position: usize, + in_index: bool, + summary: &mut [InputReadSummary], + ) { + match expr { + NaryExpr::Op { children, .. } => { + for child in children { + Self::scan_reads(child, input_count, rank, position, in_index, summary); + } + } + NaryExpr::IndexedInput { input_idx, indices } => { + if *input_idx < input_count { + let entry = &mut summary[*input_idx]; + entry.last_reader = Some( + entry + .last_reader + .map_or(position, |last| last.max(position)), + ); + let identity = !in_index + && indices.len() == rank + && indices.iter().enumerate().all( + |(dim, index)| matches!(index, NaryExpr::DimIndex(d) if *d == dim), + ); + if !identity { + entry.identity_only = false; + } + } + for index in indices { + Self::scan_reads(index, input_count, rank, position, true, summary); + } + } + NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => {} + } + } + + /// Hash every field that affects the generated kernel body. + pub(crate) fn hash_kernel_fields(&self, state: &mut FxHasher) { + self.inputs.len().hash(state); + self.shape.hash(state); + self.statements.len().hash(state); + for statement in &self.statements { + statement.expression.hash(state); + statement.datatype.hash(state); + statement.output.is_some().hash(state); + } + } +} diff --git a/fusor-ml/core/src/row_program.rs b/fusor-ml/core/src/row_program.rs index 03fc02fa8..789397ae9 100644 --- a/fusor-ml/core/src/row_program.rs +++ b/fusor-ml/core/src/row_program.rs @@ -27,7 +27,7 @@ use tile_ir::{ }; use crate::{ - compute_graph::{GraphOperation, NodeIndex}, + compute_graph::NodeIndex, mir::{ inputs::MirValue, kernel_backend::{self, DirectKernel}, @@ -44,38 +44,90 @@ use crate::{ visit_tiled::{MaybeQData, distribute_workgroups}, }; -const BLOCK: u32 = 256; - -/// Below this many rows a long axis is split across workgroups (one tile -/// each) with a combine kernel folding the spans — decode has too few rows -/// to fill the device with one workgroup per row. -const SPLIT_ROWS_TARGET: u32 = 256; - -/// Workgroup buckets for dynamic-axis row programs, smallest first. The -/// kernel monomorphizes per bucket; the active axis length rides in the -/// params input. -const ROW_DYNAMIC_BLOCKS: [u32; 4] = [128, 256, 512, 1024]; - /// One reduction phase: `expression` (over the external inputs and the /// slots of earlier phases) folded along the row axis, then `post_chain` /// applied to the combined value once per row. -#[derive(Debug, Clone, Hash)] +#[derive(Debug, Clone, PartialEq, Hash)] pub(crate) struct RowReduce { pub(crate) expression: NaryExpr, - pub(crate) function: ReduceFunction, + pub(crate) combine: RowCombine, pub(crate) post_chain: UnaryFunctionChain, } +/// How a row phase folds its element into the slot. +/// +/// `BuiltIn` is the four closed operators, which have subgroup intrinsics and +/// keep the existing emission byte-for-byte. `General` names the step as an +/// expression over the accumulator and the element, which is what lets a +/// carrier like online softmax be *stated* instead of pattern-matched back out +/// of a fixed step sequence. +#[derive(Debug, Clone, PartialEq, Hash)] +pub(crate) enum RowCombine { + BuiltIn(ReduceFunction), + /// `step` reads the accumulator as slot `acc_slot` and the element as + /// slot `element_slot`, using the same `slot_expr` convention as every + /// other cross-step reference. + General { + init: NaryScalar, + /// `acc (+) element`, evaluated once per axis position. + step: NaryExpr, + /// `acc (+) rhs` over two accumulators, evaluated when partials from + /// different lanes merge. For a monoid whose step and combine + /// coincide (sum, max) this is the same expression. + combine: NaryExpr, + /// Slot index the accumulator is bound to during evaluation; the + /// element (or incoming partial) takes `acc_slot + 1`. Emission + /// appends both to the live slot list, so these are the next two + /// indices after the phases already computed. + acc_slot: usize, + datatype: DataTypeEnum, + }, +} + +impl RowCombine { + pub(crate) fn datatype(&self) -> DataTypeEnum { + match self { + RowCombine::BuiltIn(function) => function.datatype(), + RowCombine::General { datatype, .. } => *datatype, + } + } + + pub(crate) fn initial_value(&self) -> NaryScalar { + match self { + RowCombine::BuiltIn(function) => function.initial_value, + RowCombine::General { init, .. } => *init, + } + } + + /// The closed operator, when there is one. Emission needs it for the + /// cross-lane group reduction: subgroup intrinsics are per-operator, so a + /// general combine can only span lanes once tile-ir grows a group reduce + /// parameterized by an expression. + pub(crate) fn built_in(&self) -> Option<&ReduceFunction> { + match self { + RowCombine::BuiltIn(function) => Some(function), + RowCombine::General { .. } => None, + } + } + + pub(crate) fn name(&self) -> &str { + match self { + RowCombine::BuiltIn(function) => function.name(), + RowCombine::General { .. } => "fold", + } + } +} + /// The private inner axis of an element phase: the expression is evaluated /// `len` times with the fold coordinate as `DimIndex(rank)` and accumulated /// with `function` — an inline dot product per axis position. -#[derive(Debug, Clone, Hash)] +#[derive(Debug, Clone, PartialEq, Hash)] pub(crate) struct RowFold { pub(crate) len: usize, pub(crate) function: ReduceFunction, } -#[derive(Debug, Clone, Hash)] +#[derive(Debug, Clone, PartialEq, Hash)] pub(crate) enum RowOutput { /// One output element per index-space position; output shape == `shape`. Map(NaryExpr), @@ -95,7 +147,7 @@ pub(crate) enum RowOutput { /// Dynamic-axis configuration: the kernel is compiled for the `block` /// capacity bucket and reads the active axis length from a trailing u32 /// params input, so per-token axis growth (the KV cache) reuses one kernel. -#[derive(Debug, Clone, Hash)] +#[derive(Debug, Clone, PartialEq, Hash)] pub(crate) struct DynamicAxis { /// Workgroup size and per-tile capacity; axis lengths beyond it stream /// through the online tile loop. @@ -112,7 +164,7 @@ pub(crate) struct DynamicAxis { /// One ordered row-program step. All non-output steps produce slots for later /// expressions; the final step must be [`RowStep::Output`] and defines the /// tensor write contract. -#[derive(Debug, Clone, Hash)] +#[derive(Debug, Clone, PartialEq, Hash)] pub(crate) enum RowStep { /// Fold over the row axis into a per-row scalar slot. Reduce(RowReduce), @@ -190,8 +242,13 @@ fn match_online_softmax<'a>(steps: &'a [RowStep], input_count: usize) -> Option< else { return None; }; - if max_phase.function.op != ReduceOp::Max - || sum_phase.function.op != ReduceOp::Sum + let (Some(max_function), Some(sum_function)) = + (max_phase.combine.built_in(), sum_phase.combine.built_in()) + else { + return None; + }; + if max_function.op != ReduceOp::Max + || sum_function.op != ReduceOp::Sum || !max_phase.post_chain.functions.is_empty() || !sum_phase.post_chain.functions.is_empty() { @@ -233,7 +290,7 @@ fn match_online_softmax<'a>(steps: &'a [RowStep], input_count: usize) -> Option< Some(OnlineSoftmax { score: (score, score_fold), scaled, - max_identity: max_phase.function.initial_value, + max_identity: max_function.initial_value, weight, }) } @@ -243,7 +300,7 @@ fn match_online_softmax<'a>(steps: &'a [RowStep], input_count: usize) -> Option< /// value for non-output steps (a per-row scalar for reduce steps, a /// per-element value for element steps). The final output step never produces /// a reusable slot. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub(crate) struct RowProgramOperation { pub(crate) inputs: Vec, /// The full row-parallel index space (including the axis). @@ -266,7 +323,7 @@ impl RowProgramOperation { steps: vec![ RowStep::Reduce(RowReduce { expression: reduce.expression.clone(), - function: reduce.function.clone(), + combine: RowCombine::BuiltIn(reduce.function.clone()), post_chain: reduce.post_element_wise.clone(), }), RowStep::Output(RowOutput::Scalar(scalar_slot)), @@ -276,6 +333,71 @@ impl RowProgramOperation { } } + /// Re-express every built-in phase as the equivalent general combine. + /// + /// Sum, Product, Max and Min are monoids whose step and combine coincide, + /// so each becomes the same binary expression over the accumulator and the + /// incoming value. Used by `FUSOR_SPIKE_GENERAL_COMBINE` to validate the + /// general emission path against the closed-operator one. + pub(crate) fn as_general_combines(&self) -> Option { + // Dynamic-axis programs lower through the bespoke online-softmax + // streaming path, which pattern-matches the closed `Max`/`Sum` phases + // via `match_online_softmax`. That coupling is exactly what a general + // combine is meant to retire, but until it does, re-expressing those + // phases would break the matcher rather than exercise the new path. + if self.dynamic_axis.is_some() { + return None; + } + let input_count = self.inputs.len(); + let mut steps = self.steps.clone(); + let mut rewrote = false; + for (phase, step) in steps.iter_mut().enumerate() { + let RowStep::Reduce(reduce) = step else { + continue; + }; + let Some(function) = reduce.combine.built_in() else { + continue; + }; + let datatype = function.datatype(); + let op = match function.op { + ReduceOp::Sum => NaryOp::Add, + ReduceOp::Product => NaryOp::Mul, + ReduceOp::Max => NaryOp::Max, + ReduceOp::Min => NaryOp::Min, + }; + let body = |acc_index: usize| NaryExpr::Op { + children: vec![ + slot_expr(input_count, acc_index), + slot_expr(input_count, acc_index + 1), + ], + function: NaryFunction::binary(None, op, datatype, datatype, datatype), + }; + reduce.combine = RowCombine::General { + init: function.initial_value, + step: body(phase), + combine: body(phase), + acc_slot: phase, + datatype, + }; + rewrote = true; + } + rewrote.then(|| Self { + steps, + ..self.clone() + }) + } + + /// Whether any phase folds with a general combine. Emission needs a + /// closed operator for the cross-lane group reduction — subgroup + /// intrinsics are per-operator — so until tile-ir grows a group reduce + /// parameterized by an expression, these programs have no direct kernel. + pub(crate) fn has_general_combine(&self) -> bool { + self.phase_steps().iter().any(|step| match step { + RowStep::Reduce(reduce) => reduce.combine.built_in().is_none(), + _ => false, + }) + } + fn output_step(&self) -> &RowOutput { match self.steps.last() { Some(RowStep::Output(output)) => output, @@ -325,10 +447,22 @@ impl RowProgramOperation { fn block(&self, device: &crate::Device) -> u32 { match &self.dynamic_axis { Some(dynamic) => dynamic.block, - None => device.limits().max_compute_workgroup_size_x.min(BLOCK), + None => static_axis_block(device, self.shape[self.axis]), } } + /// Whether this program is a plain chunked-map program the horizontal + /// merge pass can host: static axis, reduce-only phases, and a map or + /// per-row-scalar output (no free-dim reduce, no staged elements). + pub(crate) fn mergeable_chunked_map(&self) -> bool { + self.dynamic_axis.is_none() + && self + .phase_steps() + .iter() + .all(|step| matches!(step, RowStep::Reduce(_))) + && matches!(self.output_step(), RowOutput::Map(_) | RowOutput::Scalar(_)) + } + fn uses_custom_indexing_for_input(&self, input_idx: usize) -> bool { let output_uses_custom_indexing = match self.output_step() { RowOutput::Map(expr) | RowOutput::Scalar(expr) => { @@ -399,6 +533,12 @@ impl Operation for RowProgramOperation { } } + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + for input in &mut self.inputs { + f(input); + } + } + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { let mut mir_inputs: Vec = self .inputs @@ -456,12 +596,6 @@ impl Operation for RowProgramOperation { } } -impl GraphOperation for RowProgramOperation { - fn category(&self) -> &'static str { - "row_program" - } -} - struct RowProgramKernelVariant; /// The shared cache-key recipe, except inputs with a dynamic-axis dimension @@ -523,6 +657,50 @@ fn row_program_cache_key( }) } +/// Wrap a raw tile as a typed value for slot binding. +fn typed_tile(value: Tile, datatype: DataTypeEnum) -> ValueTile { + match datatype { + DataTypeEnum::F32 => ValueTile::F32(value), + DataTypeEnum::F16 => ValueTile::F16(value), + DataTypeEnum::U32 => ValueTile::U32(value), + } +} + +/// Evaluate a carrier body that reads the accumulator and one incoming value. +/// +/// The two are appended to the live slot list, so a body built for phase `p` +/// reads them as slots `p` and `p + 1` — the same `slot_expr` convention every +/// other cross-step reference uses. +#[allow(clippy::too_many_arguments)] +fn emit_carrier_body( + program: &mut tile_ir::tile::TileBlock<'_>, + expression: &NaryExpr, + storages: &[crate::nary_direct::Storage2], + metas: &[crate::nary_direct::TensorMeta], + slots: &[(ValueTile, DataTypeEnum)], + acc: Tile, + incoming: Tile, + datatype: DataTypeEnum, +) -> Tile { + // Storages are passed through so a carrier body indexes slots on the same + // base as every other row-program expression: `slot_expr(input_count, k)` + // resolves to `extras[k]`. A carrier body reads no tensor, but giving it a + // different base would be a silent index footgun. + let mut extended = slots.to_vec(); + extended.push((typed_tile(acc, datatype), datatype)); + extended.push((typed_tile(incoming, datatype), datatype)); + let (value, _) = eval_nary_expr( + program, + expression, + &[], + storages, + metas, + tile_ir::tile::Mask::from(true), + &extended, + ); + raw_tile(value.cast_to(datatype)) +} + fn raw_tile(value: ValueTile) -> Tile { match value { ValueTile::F32(tile) | ValueTile::F16(tile) | ValueTile::U32(tile) => tile, @@ -530,6 +708,190 @@ fn raw_tile(value: ValueTile) -> Tile { } } +/// A chunked map program whose axis fits one chunk evaluates every phase and +/// the output at the same coordinates, so each distinct tensor read can be +/// loaded once per lane and kept in a register across the phases instead of +/// being re-read from storage per phase (softmax otherwise reads its input +/// once for the max, again for the exp sum, and a third time for the output). +struct StagedReads { + /// The original `IndexedInput` leaves, evaluated once each, in order. + probes: Vec, + /// Phase reduce expressions rewritten over the staged probe slots + /// (`0..probes.len()`); phase slots follow at `probes.len() + p`. + phases: Vec, + /// The output expression, rewritten the same way. + output: NaryExpr, +} + +/// Whether `expr` references any phase slot (directly or through an index +/// expression). Reads whose indices depend on a slot cannot be staged before +/// the phases run. +fn uses_any_slot(expr: &NaryExpr, input_count: usize) -> bool { + match expr { + NaryExpr::Op { children, .. } => children + .iter() + .any(|child| uses_any_slot(child, input_count)), + NaryExpr::IndexedInput { input_idx, indices } => { + *input_idx >= input_count + || indices + .iter() + .any(|index| uses_any_slot(index, input_count)) + } + NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => false, + } +} + +/// Collect the distinct tensor-read leaves of `expr` into `probes`, +/// returning `false` when a read cannot be staged. +fn collect_probe_reads(expr: &NaryExpr, input_count: usize, probes: &mut Vec) -> bool { + match expr { + NaryExpr::Op { children, .. } => children + .iter() + .all(|child| collect_probe_reads(child, input_count, probes)), + NaryExpr::IndexedInput { input_idx, indices } => { + if *input_idx >= input_count { + // A phase-slot reference: stays a slot in the rewrite. + return indices.is_empty(); + } + if indices + .iter() + .any(|index| uses_any_slot(index, input_count)) + { + return false; + } + if !probes.iter().any(|probe| probe == expr) { + probes.push(expr.clone()); + } + true + } + NaryExpr::DimIndex(_) | NaryExpr::Scalar(_) => true, + } +} + +/// Rewrite `expr` over the staged slot space: tensor-read leaves become bare +/// slot references `0..probes.len()`; phase slots shift up by `probes.len()`. +fn rewrite_staged(expr: &NaryExpr, input_count: usize, probes: &[NaryExpr]) -> NaryExpr { + match expr { + NaryExpr::Op { children, function } => NaryExpr::Op { + children: children + .iter() + .map(|child| rewrite_staged(child, input_count, probes)) + .collect(), + function: function.clone(), + }, + NaryExpr::IndexedInput { input_idx, .. } if *input_idx >= input_count => { + NaryExpr::IndexedInput { + input_idx: probes.len() + (input_idx - input_count), + indices: vec![], + } + } + leaf @ NaryExpr::IndexedInput { .. } => { + let slot = probes + .iter() + .position(|probe| probe == leaf) + .expect("every stageable tensor read was collected"); + NaryExpr::IndexedInput { + input_idx: slot, + indices: vec![], + } + } + other => other.clone(), + } +} + +/// Build the staged rewrite for a program whose axis fits the per-lane +/// register budget: `None` when any read resists staging (slot-dependent +/// indices, non-reduce phases, reducing outputs) or there is nothing to stage. +fn stage_chunk_reads( + phase_steps: &[RowStep], + output: &RowOutput, + input_count: usize, +) -> Option { + let mut phase_sources = Vec::with_capacity(phase_steps.len()); + for step in phase_steps { + let RowStep::Reduce(reduce) = step else { + return None; + }; + phase_sources.push(&reduce.expression); + } + let output_expr = match output { + RowOutput::Map(expr) | RowOutput::Scalar(expr) => expr, + RowOutput::Reduce { .. } => return None, + }; + let mut probes = Vec::new(); + for expr in phase_sources.iter().copied().chain([output_expr]) { + if !collect_probe_reads(expr, input_count, &mut probes) { + return None; + } + } + if probes.is_empty() { + return None; + } + Some(StagedReads { + phases: phase_sources + .iter() + .map(|expr| rewrite_staged(expr, input_count, &probes)) + .collect(), + output: rewrite_staged(output_expr, input_count, &probes), + probes, + }) +} + +/// Evaluate every probe once at `coords` and bind the values to registers. +fn stage_probe_values( + program: &mut tile_ir::tile::TileBlock<'_>, + probes: &[NaryExpr], + coords: &[Tile], + storages: &[crate::nary_direct::Storage2], + metas: &[crate::nary_direct::TensorMeta], + active: Mask, +) -> Vec<(ValueTile, DataTypeEnum)> { + probes + .iter() + .map(|probe| { + let (value, ty) = + eval_nary_expr(program, probe, coords, storages, metas, active.clone(), &[]); + let bound = match value { + ValueTile::F32(tile) => ValueTile::F32(program.bind(tile)), + ValueTile::F16(tile) => ValueTile::F16(program.bind(tile)), + ValueTile::U32(tile) => ValueTile::U32(program.bind(tile)), + ValueTile::Bool(mask) => ValueTile::Bool(mask), + }; + (bound, ty) + }) + .collect() +} + +/// The staged reads of every chunk the lane owns, outer index = chunk. A lane +/// group narrower than the axis strides it `chunks` times, so staging holds +/// one register set per stride — the phases and the output then run entirely +/// out of registers however many strides there are. +type StagedChunks = Vec>; + +/// Evaluate the probes of every chunk the lane owns. +#[allow(clippy::too_many_arguments)] +fn stage_chunk_values( + program: &mut tile_ir::tile::TileBlock<'_>, + probes: &[NaryExpr], + chunks: u32, + k: u32, + k_group: u32, + lane: &Tile, + in_bounds: &Mask, + full_coords: impl Fn(Tile) -> Vec, + storages: &[crate::nary_direct::Storage2], + metas: &[crate::nary_direct::TensorMeta], +) -> StagedChunks { + (0..chunks) + .map(|chunk| { + let k_index = lane.clone() + tile_u32(chunk * k_group); + let active = in_bounds.clone() & k_index.clone().lt(k); + let coords = full_coords(k_index); + stage_probe_values(program, probes, &coords, storages, metas, active) + }) + .collect() +} + fn f32_literal(value: f32) -> Tile { Tile::literal(tile_ir::TileLiteral::f32(value)) } @@ -567,12 +929,197 @@ fn tile_reduce_op(op: ReduceOp) -> tile_ir::TileReduceOp { } } +/// The fixed-subgroup proof used to accelerate whole-workgroup reductions: +/// `(token, subgroup width)` when the device reports one subgroup size. +type FixedSubgroups = Option<(tile_ir::tile::SubgroupToken, u32)>; + +fn fixed_subgroups(device: &crate::Device) -> FixedSubgroups { + device + .subgroup_config() + .filter(|config| config.is_fixed()) + .map(|config| (config.token(), config.max_size())) +} + +/// Workgroup width for a static-axis (chunked-map) row program. +/// +/// With a fixed subgroup width the per-row reduction is a subgroup collective +/// over a lane group narrower than the workgroup ([`lane_group_width`]), so +/// the workgroup is always full width and packs as many rows as fit. Without +/// one every reduction is a shared-memory tree whose barrier count grows with +/// the group, so the workgroup still tracks the axis: a k=64 reduce runs +/// 64-lane workgroups rather than paying a 256-wide tree per row. +fn static_axis_block(device: &crate::Device, axis_len: usize) -> u32 { + let policy = device.dispatch_policy(); + let max_block = policy.preferred_workgroup_lanes(); + if fixed_subgroups(device).is_some() { + return max_block; + } + let k = u32::try_from(axis_len).unwrap_or(max_block); + k.max(1) + .next_power_of_two() + .clamp(policy.min_reduction_lanes().min(max_block), max_block) +} + +/// Lanes per row for a chunked-map program: the workgroup splits into +/// `block / group` contiguous groups, each owning one row. +/// +/// Narrowing the group packs more rows per workgroup and, once the group is +/// exactly one subgroup, turns the per-row reduction into a barrier-free +/// subgroup collective — a k=64 softmax goes from one 64-lane workgroup per +/// row with four threadgroup barriers to eight rows per full workgroup with +/// none. The floor on narrowing is the per-thread work budget: each lane then +/// strides the axis `k / group` times. +fn lane_group_width( + policy: &crate::occupancy::DispatchPolicy, + subgroups: FixedSubgroups, + k: u32, + block: u32, +) -> u32 { + let mut group = k.next_power_of_two().min(block).max(1); + if !block.is_multiple_of(group) { + return block; + } + let Some((_, subgroup_width)) = subgroups else { + return group; + }; + let budget = policy.work_per_thread(crate::occupancy::RegPressure::ElementwiseFew); + while group > subgroup_width && k.div_ceil(group / 2) <= budget { + group /= 2; + } + group +} + +/// Split a workgroup's lanes into `(row within the workgroup, position in the +/// row's axis span)`. +/// +/// When the lane group is subgroup-aligned the split is derived from the +/// subgroup builtins, matching how [`emit_group_reduce`] folds the groups: +/// the mapping from `local_invocation_index` onto subgroups is implementation +/// defined, so a row must not be defined by lane index. The shared-memory +/// tree, in contrast, derives group membership from the lane index itself. +fn split_lane_groups( + program: &mut tile_ir::tile::TileBlock<'_>, + subgroups: FixedSubgroups, + full_lane: Tile, + k_group: u32, + block: u32, +) -> (Tile, Tile) { + match subgroups { + Some((token, subgroup_width)) + if k_group.is_multiple_of(subgroup_width) && block.is_multiple_of(k_group) => + { + let subgroup_id = token.subgroup_id(program); + let subgroup_lane = token.subgroup_lane(program); + let per_group = k_group / subgroup_width; + if per_group == 1 { + (program.bind(subgroup_id), program.bind(subgroup_lane)) + } else { + let row = program.bind(subgroup_id.clone() / per_group); + let lane = program.bind(subgroup_id % per_group * subgroup_width + subgroup_lane); + (row, lane) + } + } + _ => ( + program.bind(full_lane.clone() / k_group), + program.bind(full_lane % k_group), + ), + } +} + +/// The per-row-group reduction: subgroup-accelerated whenever the group is +/// subgroup-aligned on a fixed-subgroup device (no barrier at all when it is +/// one subgroup), the shared-memory tree otherwise. A 1-wide group is the +/// lane's own value (the k=1 bias-grad sum), skipping the scratch round-trip +/// entirely. +/// Cross-lane reduction for a row phase, dispatching on the combine. +/// +/// Built-ins keep the subgroup/tree path unchanged. A general combine has no +/// per-operator intrinsic, so it stages through workgroup memory via +/// [`tile_ir::tile::TileBlock::group_reduce_with`]. +#[allow(clippy::too_many_arguments)] +fn emit_group_reduce_combined( + program: &mut tile_ir::tile::TileBlock<'_>, + subgroups: FixedSubgroups, + combine: &RowCombine, + storages: &[crate::nary_direct::Storage2], + metas: &[crate::nary_direct::TensorMeta], + slots: &[(ValueTile, DataTypeEnum)], + datatype: DataTypeEnum, + op: tile_ir::TileReduceOp, + group_size: u32, + block: u32, + value: Tile, +) -> Tile { + match combine { + RowCombine::BuiltIn(_) => { + emit_group_reduce(program, subgroups, op, group_size, block, value) + } + RowCombine::General { + combine: combine_expr, + .. + } => { + let slots = slots.to_vec(); + program.group_reduce_with(group_size, value, |program, acc, incoming| { + emit_carrier_body( + program, + combine_expr, + storages, + metas, + &slots, + acc, + incoming, + datatype, + ) + }) + } + } +} + +fn emit_group_reduce( + program: &mut tile_ir::tile::TileBlock<'_>, + subgroups: FixedSubgroups, + op: tile_ir::TileReduceOp, + group_size: u32, + block: u32, + value: Tile, +) -> Tile { + if group_size == 1 { + return program.bind(value); + } + match subgroups { + Some((token, subgroup_size)) + if group_size.is_multiple_of(subgroup_size) && block.is_multiple_of(group_size) => + { + token.group_reduce(program, op, subgroup_size, group_size, value) + } + _ => program.group_reduce(op, group_size, value), + } +} + fn build_row_program_kernel( operation: &RowProgramOperation, graph: &crate::compute_graph::ComputeGraphInner, workgroup_shape: &WorkgroupShape, inputs: &[MirValue], ) -> Option { + // Validation spike: re-express every built-in phase as the equivalent + // general combine, so the general path runs on real reductions and must + // reproduce the closed-operator results bit for bit. + let rewritten; + let operation = if graph.device().config().spike_general_combine + && let Some(general) = operation.as_general_combines() + { + rewritten = general; + &rewritten + } else { + operation + }; + // Reject before emitting rather than panicking inside it: a general + // combine is a capability gap in the cross-lane reduction, and the caller + // already surfaces a missing direct kernel as a lowering error. + if operation.has_general_combine() && !graph.device().config().spike_general_combine { + return None; + } let (output, producers) = inputs.split_last()?; let output = output.as_tensor()?.clone(); let (params, producers) = if operation.dynamic_axis.is_some() { @@ -620,7 +1167,10 @@ fn build_row_program_kernel( let splits: u32 = match free_dim_out { Some(free) if lanes_own_axis - && rows < SPLIT_ROWS_TARGET + && graph + .device() + .dispatch_policy() + .should_split_for_occupancy(rows, block) && tiles > 1 && tiles <= block && (free as u32 + 2) <= block => @@ -629,9 +1179,23 @@ fn build_row_program_kernel( } _ => 1, }; + // Chunked programs pack several rows per workgroup: lanes split into + // `block / k_group` groups of `k_group` lanes, each group owning one row, + // reduced with per-group reductions. Without packing a k=64 reduce runs + // one row per 256-lane workgroup with 75% of the lanes idle (and a k=1 + // bias-grad sum runs one workgroup per output scalar). + let subgroups = fixed_subgroups(&graph.device()); + let k_group: u32 = if lanes_own_axis || splits > 1 { + block + } else { + lane_group_width(&graph.device().dispatch_policy(), subgroups, k, block) + }; + let rows_per_workgroup = block / k_group; + let dispatch_rows = rows.div_ceil(rows_per_workgroup); + let max_dispatch_dim = graph.device().limits().max_compute_workgroups_per_dimension; let dispatch_spec = crate::row_dispatch::RowDispatchSpec::distributed( - rows.saturating_mul(splits), + dispatch_rows.saturating_mul(splits), block, max_dispatch_dim, ); @@ -640,6 +1204,16 @@ fn build_row_program_kernel( row_program_cache_key(operation, workgroup_shape, dispatch_size, inputs, 1, splits); let input_count = operation.inputs.len(); + let stage_budget = graph + .device() + .dispatch_policy() + .work_per_thread(crate::occupancy::RegPressure::ElementwiseFew); + // Chunked-map programs whose axis fits the per-lane register budget stage + // each distinct tensor read once per chunk and evaluate every phase (and + // the output) from the registers instead of re-reading storage per phase. + let staged_reads = (!lanes_own_axis && k.div_ceil(k_group) <= stage_budget) + .then(|| stage_chunk_reads(&phase_steps, &output_kind, input_count)) + .flatten(); let axis_bound_dim = operation .dynamic_axis .as_ref() @@ -660,7 +1234,7 @@ fn build_row_program_kernel( )) }); let online_max_identity = match phase_steps.get(1) { - Some(RowStep::Reduce(reduce)) => Some(reduce.function.initial_value), + Some(RowStep::Reduce(reduce)) => Some(reduce.combine.initial_value()), _ => None, }; let combine = if splits > 1 { @@ -707,8 +1281,14 @@ fn build_row_program_kernel( j_active.clone(), ))); let masked_m = Tile::select(j_active.clone(), m_j.clone(), max_identity); - let global_max = - program.group_reduce(tile_reduce_op(ReduceOp::Max), block, masked_m); + let global_max = emit_group_reduce( + program, + subgroups, + tile_reduce_op(ReduceOp::Max), + block, + block, + masked_m, + ); let global_max = program.bind(global_max); let l_index = layout_index( &scratch_meta, @@ -721,8 +1301,14 @@ fn build_row_program_kernel( l_j * (m_j - global_max.clone()).exp(), f32_literal(0.0), ); - let denom = - program.group_reduce(tile_reduce_op(ReduceOp::Sum), block, weighted); + let denom = emit_group_reduce( + program, + subgroups, + tile_reduce_op(ReduceOp::Sum), + block, + block, + weighted, + ); let denom = program.bind(denom); // Lane = one free-dim position: rescale and fold the @@ -810,15 +1396,26 @@ fn build_row_program_kernel( .collect(); crate::row_dispatch::emit_row_grid(phase_handle, dispatch_spec, |program, ctx| { - let lane = ctx.lane; let wg_flat = ctx.row; - let (row_flat, split_idx) = if splits > 1 { + // With packing, `lane` becomes the position inside the row's + // `k_group`-wide lane group; every later axis index and the + // scalar-store lane test use it unchanged. + let (row_flat, split_idx, lane) = if rows_per_workgroup > 1 { + let (row_local, lane) = + split_lane_groups(program, subgroups, ctx.lane, k_group, block); + ( + program.bind(wg_flat * rows_per_workgroup + row_local), + tile_u32(0), + lane, + ) + } else if splits > 1 { ( program.bind(wg_flat.clone() / splits), program.bind(wg_flat % splits), + ctx.lane, ) } else { - (wg_flat, tile_u32(0)) + (wg_flat, tile_u32(0), ctx.lane) }; let in_bounds = row_flat.clone().lt(rows); let row_dims = output_dims_from_flat(row_flat.clone(), &row_shape); @@ -916,7 +1513,13 @@ fn build_row_program_kernel( len, function: fold_fn, }) => { - let fold_dtype = fold_fn.datatype(); + let fold_wire_dtype = fold_fn.datatype(); + // f32 accumulation for half-precision dots, + // rounded once below. + let fold_dtype = match fold_wire_dtype { + DataTypeEnum::F16 => DataTypeEnum::F32, + other => other, + }; let identity = Tile::literal(tile_literal_for( fold_fn.initial_value, fold_dtype, @@ -941,7 +1544,14 @@ fn build_row_program_kernel( [fold_acc.binary(reduce_op.binary(), value)] }, ); - folded + if fold_dtype == fold_wire_dtype { + folded + } else { + raw_tile( + ValueTile::F32(program.bind(folded)) + .cast_to(fold_wire_dtype), + ) + } } None => { let (value, _) = eval_nary_expr( @@ -975,9 +1585,12 @@ fn build_row_program_kernel( raw_tile(scaled.cast_to(DataTypeEnum::F32)), max_identity(), )); - let tile_max = program.group_reduce( + let tile_max = emit_group_reduce( + program, + subgroups, tile_reduce_op(ReduceOp::Max), block, + block, scaled.clone(), ); let old_max = program.load_local(&running_max); @@ -990,9 +1603,12 @@ fn build_row_program_kernel( (scaled - new_max).exp(), f32_literal(0.0), )); - let tile_sum = program.group_reduce( + let tile_sum = emit_group_reduce( + program, + subgroups, tile_reduce_op(ReduceOp::Sum), block, + block, prob.clone(), ); let sum = program.load_local(&running_sum); @@ -1083,81 +1699,237 @@ fn build_row_program_kernel( return; } - // Chunked map program: lanes stride the axis for each phase's - // fold; full-shape intermediates referenced by several phases - // are recomputed per phase — the same trade every multi-pass - // normalization kernel makes. - let chunks = k.div_ceil(block); - for phase in &phase_steps { + // Chunked map program: lanes stride the axis by their lane + // group's width for each phase's fold. Programs within the + // register budget stage each distinct tensor read per chunk up + // front (`staged_reads`) and unroll the stride; longer axes + // roll the stride into a loop and re-evaluate per phase — the + // same trade every multi-pass normalization kernel makes. + let chunks = k.div_ceil(k_group); + let staged_values = staged_reads.as_ref().map(|staged| { + stage_chunk_values( + program, + &staged.probes, + chunks, + k, + k_group, + &lane, + &in_bounds, + &full_coords, + &storages, + &metas, + ) + }); + for (phase_index, phase) in phase_steps.iter().enumerate() { let RowStep::Reduce(reduce) = phase else { unreachable!("element phases require a dynamic-axis row program") }; - let phase_dtype = reduce.function.datatype(); + let function_dtype = reduce.combine.datatype(); + // Half-precision folds accumulate in f32 and round once + // after the group reduction — the same accumulator policy + // as the matmul kernels and the composed fused reduce. + let phase_dtype = match function_dtype { + DataTypeEnum::F16 => DataTypeEnum::F32, + other => other, + }; let identity = || { - Tile::literal(tile_literal_for(reduce.function.initial_value, phase_dtype)) + Tile::literal(tile_literal_for(reduce.combine.initial_value(), phase_dtype)) }; - let reduce_op = tile_reduce_op(reduce.function.op); - let [partial] = program.fold( - tile_ir::tile::range(chunks), - [identity()], - |program, chunk, [acc]| { - let k_index = chunk * block + lane.clone(); - let active = in_bounds.clone() & k_index.clone().lt(k); - let coords = full_coords(k_index); - let (value, _) = eval_nary_expr( + // Only the built-in branch consults this; both dispatch + // sites below match on the combine kind first, so a + // general combine never reads it. + let reduce_op = reduce + .combine + .built_in() + .map(|function| tile_reduce_op(function.op)) + .unwrap_or(tile_ir::TileReduceOp::Sum); + // Built-ins keep the closed-operator path exactly; a + // general combine evaluates its own body. The kernel + // builder rejects general combines up front today, so this + // is inert for every program that currently lowers. + let combine_kind = &reduce.combine; + let accumulate = |program: &mut tile_ir::tile::TileBlock<'_>, + acc: Tile, + value: Tile, + slots: &[(ValueTile, DataTypeEnum)]| + -> Tile { + match combine_kind { + RowCombine::BuiltIn(_) => acc.binary(reduce_op.binary(), value), + RowCombine::General { step, .. } => { + emit_carrier_body( + program, step, &storages, &metas, slots, acc, value, + phase_dtype, + ) + } + } + }; + let chunk_value = |program: &mut tile_ir::tile::TileBlock<'_>, + k_index: Tile, + chunk_reads: Option<&Vec<(ValueTile, DataTypeEnum)>>, + slots: &[(ValueTile, DataTypeEnum)]| + -> Tile { + let active = in_bounds.clone() & k_index.clone().lt(k); + let coords = full_coords(k_index); + let (value, _) = match (&staged_reads, chunk_reads) { + (Some(staged), Some(chunk_reads)) => { + let mut extras = chunk_reads.clone(); + extras.extend(slots.iter().cloned()); + eval_nary_expr( + program, + &staged.phases[phase_index], + &coords, + &[], + &[], + active.clone(), + &extras, + ) + } + _ => eval_nary_expr( program, &reduce.expression, &coords, &storages, &metas, active.clone(), - &slots, + slots, + ), + }; + let value = raw_tile(value.cast_to(phase_dtype)); + Tile::select(active, value, identity()) + }; + let partial = match &staged_values { + Some(staged_values) => { + let mut acc = identity(); + for (chunk, chunk_reads) in staged_values.iter().enumerate() { + let k_index = lane.clone() + tile_u32(chunk as u32 * k_group); + let masked = + chunk_value(program, k_index, Some(chunk_reads), &slots); + acc = accumulate(program, acc, masked, &slots); + } + acc + } + None => { + let [partial] = program.fold( + tile_ir::tile::range(chunks), + [identity()], + |program, chunk, [acc]| { + let k_index = chunk * k_group + lane.clone(); + let masked = chunk_value(program, k_index, None, &slots); + [accumulate(program, acc, masked, &slots)] + }, ); - let value = raw_tile(value.cast_to(phase_dtype)); - let masked = Tile::select(active, value, identity()); - [acc.binary(reduce_op.binary(), masked)] - }, - ); - // The workgroup reduction broadcasts: every lane reads the - // combined value, so later phases can use it directly. - let combined = program.group_reduce(reduce_op, block, partial); + partial + } + }; + // The per-group workgroup reduction broadcasts: every lane + // in the row's `k_group`-wide lane group reads the combined + // value, so later phases can use it directly. With one row + // per workgroup `k_group == block` and this is the full + // workgroup reduction. + let combined = + emit_group_reduce_combined( + program, + subgroups, + combine_kind, + &storages, + &metas, + &slots, + phase_dtype, + reduce_op, + k_group, + block, + partial, + ); + let combined = if phase_dtype == function_dtype { + combined + } else { + raw_tile(ValueTile::F32(program.bind(combined)).cast_to(function_dtype)) + }; let (combined, combined_ty) = - apply_unary_function_chain(combined, phase_dtype, &reduce.post_chain) + apply_unary_function_chain(combined, function_dtype, &reduce.post_chain) .expect("validated row program post chain"); let scalar = ValueTile::F32(program.bind(combined)).cast_to(combined_ty); slots.push((scalar, combined_ty)); } + let eval_output = |program: &mut tile_ir::tile::TileBlock<'_>, + output_expr: &NaryExpr, + coords: &[Tile], + active: Mask, + chunk_reads: Option<&Vec<(ValueTile, DataTypeEnum)>>, + slots: &[(ValueTile, DataTypeEnum)]| + -> ValueTile { + let (value, _) = match (&staged_reads, chunk_reads) { + (Some(staged), Some(chunk_reads)) => { + let mut extras = chunk_reads.clone(); + extras.extend(slots.iter().cloned()); + eval_nary_expr( + program, + &staged.output, + coords, + &[], + &[], + active, + &extras, + ) + } + _ => eval_nary_expr( + program, + output_expr, + coords, + &storages, + &metas, + active, + slots, + ), + }; + value + }; match &output_kind { RowOutput::Map(output_expr) => { - program.loop_range(chunks, |program, chunk| { - let k_index = chunk * block + lane.clone(); + let store_chunk = |program: &mut tile_ir::tile::TileBlock<'_>, + k_index: Tile, + chunk_reads: Option< + &Vec<(ValueTile, DataTypeEnum)>, + >| { let active = in_bounds.clone() & k_index.clone().lt(k); let coords = full_coords(k_index); - let (value, _) = eval_nary_expr( + let value = eval_output( program, output_expr, &coords, - &storages, - &metas, active.clone(), + chunk_reads, &slots, ); let value = value.cast_to(output_dtype); let output_index = layout_index(&output_meta, &coords); output_storage.store(program, output_index, value, active); - }); + }; + match &staged_values { + Some(staged_values) => { + for (chunk, chunk_reads) in staged_values.iter().enumerate() { + let k_index = lane.clone() + tile_u32(chunk as u32 * k_group); + store_chunk(program, k_index, Some(chunk_reads)); + } + } + None => program.loop_range(chunks, |program, chunk| { + store_chunk(program, chunk * k_group + lane.clone(), None); + }), + } } RowOutput::Scalar(output_expr) => { let active = in_bounds & lane.eq(0u32); let coords = full_coords(tile_u32(0)); - let (value, _) = eval_nary_expr( + // Only lane zero of the group stores, and chunk zero is + // the chunk it staged position zero of the axis in. + let chunk_reads = staged_values.as_ref().map(|values| &values[0]); + let value = eval_output( program, output_expr, &coords, - &storages, - &metas, active.clone(), + chunk_reads, &slots, ); let value = value.cast_to(output_dtype); @@ -1182,6 +1954,446 @@ fn build_row_program_kernel( } } +struct MergedRowProgramKernelVariant; + +/// One kernel executing several independent chunked-map row programs: each +/// segment owns a contiguous range of workgroups guarded by a uniform +/// linear-workgroup-id range compare. The guard condition depends only on +/// the workgroup id, so the per-segment workgroup reductions stay in +/// workgroup-uniform control flow. Callers gate segments through +/// [`RowProgramOperation::mergeable_chunked_map`] and bound total bindings. +pub(crate) fn build_merged_row_program_kernel( + graph: &crate::compute_graph::ComputeGraphInner, + segments: &[RowProgramOperation], + segment_inputs: &[Vec], +) -> Option { + use std::hash::Hash; + let device = graph.device(); + let max_per_dim = device.limits().max_compute_workgroups_per_dimension; + // The shared workgroup width is the widest any segment would pick on its + // own (`static_axis_block`), so a merge never widens a segment's lane + // groups relative to its unmerged lowering. + let policy = device.dispatch_policy(); + let subgroups = fixed_subgroups(&device); + let block = segments + .iter() + .map(|op| static_axis_block(&device, op.shape[op.axis])) + .max()?; + + struct Segment { + values: Vec, + output: MaybeQData, + rows: u32, + k: u32, + k_group: u32, + rows_per_workgroup: u32, + row_shape: Vec, + base: u32, + groups: u32, + } + let mut prepared = Vec::with_capacity(segments.len()); + let mut total_groups = 0u32; + for (op, inputs) in segments.iter().zip(segment_inputs) { + if !op.mergeable_chunked_map() { + return None; + } + let (output, producers) = inputs.split_last()?; + let output = output.as_tensor()?.clone(); + let values = producers + .iter() + .map(|input| MaybeQData::try_from(input.clone()).ok()) + .collect::>>()?; + if values + .iter() + .any(|value| matches!(value, MaybeQData::QMatrix(_))) + { + return None; + } + if !device.f16_supported() + && (output.datatype() == DataTypeEnum::F16 + || values.iter().any(|value| { + matches!(value, MaybeQData::Tensor(tensor) + if tensor.datatype() == DataTypeEnum::F16) + })) + { + return None; + } + let rows: u32 = op.rows().try_into().ok()?; + let k: u32 = op.shape[op.axis].try_into().ok()?; + let k_group = lane_group_width(&policy, subgroups, k, block); + let rows_per_workgroup = block / k_group; + let groups = rows.div_ceil(rows_per_workgroup); + prepared.push(Segment { + values, + output: MaybeQData::Tensor(output), + rows, + k, + k_group, + rows_per_workgroup, + row_shape: op.row_shape(), + base: total_groups, + groups, + }); + total_groups = total_groups.checked_add(groups)?; + } + + let dispatch_size = distribute_workgroups(total_groups, max_per_dim); + let stage_budget = policy.work_per_thread(crate::occupancy::RegPressure::ElementwiseFew); + // Segments within the per-lane register budget stage each distinct tensor + // read per chunk (see `stage_chunk_reads`); the rewrite is deterministic + // from the segment's steps, so it stays out of the cache key. + let staged_segments: Vec> = segments + .iter() + .zip(&prepared) + .map(|(op, segment)| { + (segment.k.div_ceil(segment.k_group) <= stage_budget) + .then(|| stage_chunk_reads(op.phase_steps(), op.output_step(), op.inputs.len())) + .flatten() + }) + .collect(); + let cache_key = kernel_backend::KernelCacheKey::from_hash_inputs(|state| { + kernel_backend::KernelVariantKey::of::().hash(state); + dispatch_size.hash(state); + block.hash(state); + crate::compute_graph::resolve::plan_cache::hash_merged_segments( + state, + segments.iter(), + segment_inputs, + ); + }); + let name = if device.config().trace_decode_names { + format!( + "merged_row_program[{}]", + segments + .iter() + .map(|op| op.name()) + .collect::>() + .join("; ") + ) + } else { + format!("merged_row_program_x{}", segments.len()) + }; + + kernel_backend::run_kernel( + device.kernel_cache(), + name, + cache_key, + dispatch_size, + move |kb| { + let mut declared = Vec::with_capacity(prepared.len()); + for segment in &prepared { + let mut storages = Vec::with_capacity(segment.values.len()); + let mut metas = Vec::with_capacity(segment.values.len()); + for value in &segment.values { + let (storage, meta) = declare_value(kb, value, false)?; + storages.push(storage); + metas.push(meta); + } + let output = declare_value(kb, &segment.output, true)?; + declared.push((storages, metas, output)); + } + + kb.program().program_grid(block, dispatch_size, |program| { + let full_lane = program.lane(); + let group = program.bind(crate::nary_direct::linear_group(program, dispatch_size)); + for ((op, staged_reads), (segment, (storages, metas, output))) in segments + .iter() + .zip(&staged_segments) + .zip(prepared.iter().zip(&declared)) + { + let in_segment = group.clone().ge(segment.base) + & group.clone().lt(segment.base + segment.groups); + program.if_then(in_segment, |program| { + let (output_storage, output_meta) = output; + let wg_local = program.bind(group.clone() - segment.base); + let (row_flat, lane) = if segment.rows_per_workgroup > 1 { + let (row_local, lane) = split_lane_groups( + program, + subgroups, + full_lane.clone(), + segment.k_group, + block, + ); + ( + program.bind( + wg_local.clone() * segment.rows_per_workgroup + row_local, + ), + lane, + ) + } else { + (wg_local, full_lane.clone()) + }; + let in_bounds = row_flat.clone().lt(segment.rows); + let row_dims = output_dims_from_flat(row_flat.clone(), &segment.row_shape); + let axis = op.axis; + let rank = op.shape.len(); + let full_coords = |k_index: Tile| -> Vec { + let mut coords = Vec::with_capacity(rank); + let mut row_dim = 0; + for dim in 0..rank { + if dim == axis { + coords.push(k_index.clone()); + } else { + coords.push(row_dims[row_dim].clone()); + row_dim += 1; + } + } + coords + }; + + let k = segment.k; + let k_group = segment.k_group; + let chunks = k.div_ceil(k_group); + let staged_values = staged_reads.as_ref().map(|staged| { + stage_chunk_values( + program, + &staged.probes, + chunks, + k, + k_group, + &lane, + &in_bounds, + &full_coords, + storages, + metas, + ) + }); + let mut slots: Vec<(ValueTile, DataTypeEnum)> = Vec::new(); + for (phase_index, phase) in op.phase_steps().iter().enumerate() { + let RowStep::Reduce(reduce) = phase else { + unreachable!("merged row programs are reduce-phase only") + }; + let function_dtype = reduce.combine.datatype(); + // Same f32 accumulation policy as the standalone + // reduce phases above. + let phase_dtype = match function_dtype { + DataTypeEnum::F16 => DataTypeEnum::F32, + other => other, + }; + let identity = || { + Tile::literal(tile_literal_for( + reduce.combine.initial_value(), + phase_dtype, + )) + }; + // Only the built-in branch consults this; both dispatch + // sites below match on the combine kind first, so a + // general combine never reads it. + let reduce_op = reduce + .combine + .built_in() + .map(|function| tile_reduce_op(function.op)) + .unwrap_or(tile_ir::TileReduceOp::Sum); + // Built-ins keep the closed-operator path exactly; a + // general combine evaluates its own body. The kernel + // builder rejects general combines up front today, so this + // is inert for every program that currently lowers. + let combine_kind = &reduce.combine; + let accumulate = |program: &mut tile_ir::tile::TileBlock<'_>, + acc: Tile, + value: Tile, + slots: &[(ValueTile, DataTypeEnum)]| + -> Tile { + match combine_kind { + RowCombine::BuiltIn(_) => acc.binary(reduce_op.binary(), value), + RowCombine::General { step, .. } => { + emit_carrier_body( + program, step, &storages, &metas, slots, acc, value, + phase_dtype, + ) + } + } + }; + let chunk_value = + |program: &mut tile_ir::tile::TileBlock<'_>, + k_index: Tile, + chunk_reads: Option<&Vec<(ValueTile, DataTypeEnum)>>, + slots: &[(ValueTile, DataTypeEnum)]| + -> Tile { + let active = in_bounds.clone() & k_index.clone().lt(k); + let coords = full_coords(k_index); + let (value, _) = match (&staged_reads, chunk_reads) { + (Some(staged), Some(chunk_reads)) => { + let mut extras = chunk_reads.clone(); + extras.extend(slots.iter().cloned()); + eval_nary_expr( + program, + &staged.phases[phase_index], + &coords, + &[], + &[], + active.clone(), + &extras, + ) + } + _ => eval_nary_expr( + program, + &reduce.expression, + &coords, + storages, + metas, + active.clone(), + slots, + ), + }; + let value = raw_tile(value.cast_to(phase_dtype)); + Tile::select(active, value, identity()) + }; + let partial = match &staged_values { + Some(staged_values) => { + let mut acc = identity(); + for (chunk, chunk_reads) in staged_values.iter().enumerate() { + let k_index = + lane.clone() + tile_u32(chunk as u32 * k_group); + let masked = chunk_value( + program, + k_index, + Some(chunk_reads), + &slots, + ); + acc = accumulate(program, acc, masked, &slots); + } + acc + } + None => { + let [partial] = program.fold( + tile_ir::tile::range(chunks), + [identity()], + |program, chunk, [acc]| { + let k_index = chunk * k_group + lane.clone(); + let masked = + chunk_value(program, k_index, None, &slots); + [accumulate(program, acc, masked, &slots)] + }, + ); + partial + } + }; + let combined = emit_group_reduce( + program, + subgroups, + reduce_op, + segment.k_group, + block, + partial, + ); + let combined = if phase_dtype == function_dtype { + combined + } else { + raw_tile( + ValueTile::F32(program.bind(combined)).cast_to(function_dtype), + ) + }; + let (combined, combined_ty) = apply_unary_function_chain( + combined, + function_dtype, + &reduce.post_chain, + ) + .expect("validated row program post chain"); + let scalar = + ValueTile::F32(program.bind(combined)).cast_to(combined_ty); + slots.push((scalar, combined_ty)); + } + + let output_dtype = op.output_datatype; + let eval_output = |program: &mut tile_ir::tile::TileBlock<'_>, + output_expr: &NaryExpr, + coords: &[Tile], + active: Mask, + chunk_reads: Option<&Vec<(ValueTile, DataTypeEnum)>>, + slots: &[(ValueTile, DataTypeEnum)]| + -> ValueTile { + let (value, _) = match (&staged_reads, chunk_reads) { + (Some(staged), Some(chunk_reads)) => { + let mut extras = chunk_reads.clone(); + extras.extend(slots.iter().cloned()); + eval_nary_expr( + program, + &staged.output, + coords, + &[], + &[], + active, + &extras, + ) + } + _ => eval_nary_expr( + program, + output_expr, + coords, + storages, + metas, + active, + slots, + ), + }; + value + }; + match op.output_step() { + RowOutput::Map(output_expr) => { + let store_chunk = |program: &mut tile_ir::tile::TileBlock<'_>, + k_index: Tile, + chunk_reads: Option< + &Vec<(ValueTile, DataTypeEnum)>, + >| { + let active = in_bounds.clone() & k_index.clone().lt(k); + let coords = full_coords(k_index); + let value = eval_output( + program, + output_expr, + &coords, + active.clone(), + chunk_reads, + &slots, + ); + let value = value.cast_to(output_dtype); + let output_index = layout_index(output_meta, &coords); + output_storage.store(program, output_index, value, active); + }; + match &staged_values { + Some(staged_values) => { + for (chunk, chunk_reads) in staged_values.iter().enumerate() + { + let k_index = + lane.clone() + tile_u32(chunk as u32 * k_group); + store_chunk(program, k_index, Some(chunk_reads)); + } + } + None => program.loop_range(chunks, |program, chunk| { + store_chunk(program, chunk * k_group + lane.clone(), None); + }), + } + } + RowOutput::Scalar(output_expr) => { + let active = in_bounds.clone() & lane.clone().eq(0u32); + let coords = full_coords(tile_u32(0)); + // Lane zero of the group stores, and chunk zero + // is where it staged axis position zero. + let chunk_reads = staged_values.as_ref().map(|values| &values[0]); + let value = eval_output( + program, + output_expr, + &coords, + active.clone(), + chunk_reads, + &slots, + ); + let value = value.cast_to(output_dtype); + let output_index = layout_index(output_meta, &row_dims); + output_storage.store(program, output_index, value, active); + } + RowOutput::Reduce { .. } => { + unreachable!("merged row programs never have a reducing output") + } + } + }); + } + }); + Some(()) + }, + ) +} + /// Scaled dot-product attention as a row program over the KV axis: an /// element phase stages the q·k score per lane via an inline head-dim fold, /// a max phase and an exp-sum phase form the softmax statistics, a second @@ -1252,16 +2464,15 @@ pub(crate) fn attention_row_program( if mask.is_some() != mask_shape.is_some() { return None; } - let limits = device.limits(); + let policy = device.dispatch_policy(); // The workgroup bucket is one axis tile: small tiles let the split // lowering fan decode across workgroups and stream longer axes through - // the online loop with good occupancy. - let needed = head_dim.max(256) as u32; - let block = ROW_DYNAMIC_BLOCKS.iter().copied().find(|&candidate| { - candidate >= needed - && candidate <= limits.max_compute_workgroup_size_x - && candidate <= limits.max_compute_invocations_per_workgroup - })?; + // the online loop with good occupancy. The floor is one full-width + // workgroup; the kernel monomorphizes per bucket. + let needed = head_dim.max(policy.preferred_workgroup_lanes() as usize) as u32; + let block = policy + .dynamic_block_buckets() + .find(|&candidate| candidate >= needed)?; let groups = num_heads / num_kv_heads; let f32 = DataTypeEnum::F32; @@ -1353,12 +2564,12 @@ pub(crate) fn attention_row_program( }, RowStep::Reduce(RowReduce { expression: scaled_score(), - function: max_fn(f32), + combine: RowCombine::BuiltIn(max_fn(f32)), post_chain: UnaryFunctionChain::empty(f32), }), RowStep::Reduce(RowReduce { expression: shifted_exp(), - function: sum_fn(f32), + combine: RowCombine::BuiltIn(sum_fn(f32)), post_chain: UnaryFunctionChain::empty(f32), }), RowStep::Element { diff --git a/fusor-ml/core/src/sampling/mirostat.rs b/fusor-ml/core/src/sampling/mirostat.rs deleted file mode 100644 index f81ef8c80..000000000 --- a/fusor-ml/core/src/sampling/mirostat.rs +++ /dev/null @@ -1 +0,0 @@ -pub(super) use crate::mir::kernel_backend::mirostat::sample_from_sorted_top_k_data_with_encoder; diff --git a/fusor-ml/core/src/sampling/mod.rs b/fusor-ml/core/src/sampling/mod.rs index 4003dd348..14c8b9e18 100644 --- a/fusor-ml/core/src/sampling/mod.rs +++ b/fusor-ml/core/src/sampling/mod.rs @@ -1,20 +1,17 @@ use crate::{Device, Tensor, tensor::TensorData}; -mod mirostat; mod pipeline; pub(crate) mod processors; pub(crate) mod row_kernels; -mod standard_sampler; -mod topk; #[cfg(test)] mod tests; -pub(crate) use pipeline::{GpuSamplerRequest, sample_token_pending, sample_token_to_host}; -pub(crate) use topk::{ +pub(crate) use crate::mir::kernel_backend::sampling::{ MergeSortedChunkTopKParams, chunk_top_k_pair_data_with_encoder, merge_sorted_chunk_top_k_pair_data_with_encoder, }; +pub(crate) use pipeline::{GpuSamplerRequest, sample_token_pending, sample_token_to_host}; pub(crate) const TOP_K_BLOCK: u32 = 256; pub(crate) const TOP_K_CHUNK: usize = TOP_K_BLOCK as usize; @@ -51,10 +48,9 @@ impl PendingGpuSampledToken { pub async fn read_token(self) -> Result, wgpu::BufferAsyncError> { // Diagnostic: time the wait for GPU completion + token readback (the // per-token sync the pending decode path stalls on). On wasm this is - // always on; native gates on the FUSOR_TRACE_* env vars. - let trace = cfg!(target_arch = "wasm32") - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some(); + // always on; native gates on the trace config flags. + let config = self.token.device().config(); + let trace = cfg!(target_arch = "wasm32") || config.trace_decode || config.trace_sampler; let await_start = trace.then(web_time::Instant::now); self.receiver.await.map_err(|_| wgpu::BufferAsyncError)??; if let Some(start) = await_start { @@ -83,10 +79,9 @@ impl PendingGpuSampledToken { } } -pub(crate) fn min_top_k_candidates_per_chunk() -> usize { - std::env::var("FUSOR_TOP_K_MIN_CANDIDATES_PER_CHUNK") - .ok() - .and_then(|value| value.parse::().ok()) +pub(crate) fn min_top_k_candidates_per_chunk(config: &fusor_tile_ir_runtime::FusorConfig) -> usize { + config + .top_k_min_candidates_per_chunk .filter(|value| *value > 0) .unwrap_or(DEFAULT_MIN_TOP_K_CANDIDATES_PER_CHUNK) .min(TOP_K_CHUNK) diff --git a/fusor-ml/core/src/sampling/pipeline.rs b/fusor-ml/core/src/sampling/pipeline.rs index 889d500ac..018358c3c 100644 --- a/fusor-ml/core/src/sampling/pipeline.rs +++ b/fusor-ml/core/src/sampling/pipeline.rs @@ -1,5 +1,11 @@ use crate::{ Layout, Tensor, + mir::kernel_backend::sampling::{ + ChunkProcessors, MergeSortedChunkTopKParams, ProcessorSettings, + chunk_top_k_pair_data_with_encoder, merge_sorted_chunk_top_k_pair_data_with_encoder, + sample_categorical_logits_data_with_encoder, sample_from_sorted_top_k_data_with_encoder, + supports_unfiltered_categorical, top_k_exactness_flag_data_with_encoder, + }, tensor::{DataTypeEnum, LazyTensorData, TensorData}, }; use web_time::Instant; @@ -10,12 +16,6 @@ use super::{ GPU_SAMPLE_RESULT_WORDS, GPU_SAMPLE_STATUS_INVALID, GPU_SAMPLE_STATUS_RETRY_NEEDED, GPU_SAMPLE_STATUS_SAMPLED, GpuMirostat2Sampler, GpuMirostat2SamplerParams, GpuStandardSamplerParams, PendingGpuSampledToken, TOP_K_CHUNK, min_top_k_candidates_per_chunk, - mirostat::sample_from_sorted_top_k_data_with_encoder, - standard_sampler::sample_from_sorted_top_k_data_with_encoder as sample_standard_from_sorted_top_k_data_with_encoder, - topk::{ - ProcessorSettings, chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder, - merge_sorted_chunk_top_k_pair_data_with_encoder, top_k_exactness_flag_data_with_encoder, - }, }; /// Which sampler kernel terminates the top-k tail, along with its parameters. @@ -31,7 +31,7 @@ pub(crate) enum GpuSamplerRequest<'a> { } impl GpuSamplerRequest<'_> { - fn top_k(&self) -> usize { + pub(crate) fn top_k(&self) -> usize { match self { Self::Mirostat2 { params, .. } => params.top_k, Self::Standard { params } => params.top_k, @@ -50,32 +50,6 @@ impl GpuSamplerRequest<'_> { }, } } - - fn encode_sample( - &mut self, - ids: &TensorData, - values: &TensorData, - exactness_flag: Option<&TensorData>, - encoder: &mut CommandEncoder, - ) -> Option { - match self { - Self::Mirostat2 { sampler, params } => sample_from_sorted_top_k_data_with_encoder( - ids, - values, - sampler, - *params, - exactness_flag, - Some(encoder), - ), - Self::Standard { params } => sample_standard_from_sorted_top_k_data_with_encoder( - ids, - values, - *params, - exactness_flag, - Some(encoder), - ), - } - } } #[derive(Clone, Copy)] @@ -86,10 +60,19 @@ struct SampleAttemptDims { input_len: usize, } -fn initial_sampler_candidate_count(top_k: usize, chunks: usize) -> usize { +struct EncodedSampleAttempt { + output: TensorData, + debug_top_k: Option<(TensorData, TensorData)>, +} + +fn initial_sampler_candidate_count( + config: &fusor_tile_ir_runtime::FusorConfig, + top_k: usize, + chunks: usize, +) -> usize { top_k .div_ceil(chunks) - .max(min_top_k_candidates_per_chunk()) + .max(min_top_k_candidates_per_chunk(config)) .min(top_k) .min(TOP_K_CHUNK) } @@ -109,16 +92,13 @@ fn next_sampler_candidate_count(candidate_count: usize, top_k: usize) -> usize { .min(TOP_K_CHUNK) } -fn sampler_trace_enabled() -> bool { - cfg!(target_arch = "wasm32") - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_SAMPLER").is_some() +fn sampler_trace_enabled(config: &fusor_tile_ir_runtime::FusorConfig) -> bool { + cfg!(target_arch = "wasm32") || config.trace_decode || config.trace_sampler } -/// Encode one full sampling attempt — processed chunk top-k, merge, -/// optional exactness proof, and the sampler kernel — into `encoder`. -/// Returns the `[status, token]` output along with the sorted top-k -/// ids/values (kept for `FUSOR_DEBUG_SAMPLER` dumps). +/// Encode one full sampling attempt into `encoder`. Complete, unfiltered +/// Standard requests over one workgroup take a direct categorical kernel; +/// every other request retains the processed chunk-top-k pipeline. /// /// This runs inside the resolver tail while the graph lock is held, so it /// must only touch raw buffers — no compute-graph access. @@ -129,23 +109,37 @@ fn encode_sample_attempt( request: &mut GpuSamplerRequest<'_>, dims: SampleAttemptDims, encoder: &mut CommandEncoder, -) -> Option<(TensorData, TensorData, TensorData)> { +) -> Option { + if previous_tokens.is_empty() + && previous_gpu_token.is_none() + && let GpuSamplerRequest::Standard { params } = request + && supports_unfiltered_categorical(dims.input_len, *params) + && let Some(output) = + sample_categorical_logits_data_with_encoder(logits, *params, Some(encoder)) + { + return Some(EncodedSampleAttempt { + output, + debug_top_k: None, + }); + } + let output_per_chunk = sampler_output_per_chunk(dims.candidate_count); - let (chunk_ids, chunk_values) = - chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder( - logits, + let (chunk_ids, chunk_values) = chunk_top_k_pair_data_with_encoder( + logits, + Some(ChunkProcessors { previous_tokens, - previous_gpu_token, - request.processor_settings(), - dims.candidate_count, - output_per_chunk, - Some(encoder), - )?; + gpu_tail: previous_gpu_token, + settings: request.processor_settings(), + }), + dims.candidate_count, + output_per_chunk, + Some(encoder), + )?; let (ids, values) = merge_sorted_chunk_top_k_pair_data_with_encoder( &chunk_ids, &chunk_values, - crate::sampling::topk::MergeSortedChunkTopKParams { + MergeSortedChunkTopKParams { chunks: dims.chunks, chunk_len: dims.candidate_count, chunk_stride: output_per_chunk, @@ -170,8 +164,17 @@ fn encode_sample_attempt( None }; - let output = request.encode_sample(&ids, &values, exactness_flag.as_ref(), encoder)?; - Some((output, ids, values)) + let output = sample_from_sorted_top_k_data_with_encoder( + &ids, + &values, + request, + exactness_flag.as_ref(), + Some(encoder), + )?; + Some(EncodedSampleAttempt { + output, + debug_top_k: Some((ids, values)), + }) } fn encode_token_download( @@ -250,9 +253,10 @@ pub(crate) async fn sample_token_to_host( } let chunks = input_len.div_ceil(TOP_K_CHUNK); - let mut candidate_count = initial_sampler_candidate_count(top_k, chunks); - let trace = sampler_trace_enabled(); - let debug_dump = std::env::var_os("FUSOR_DEBUG_SAMPLER").is_some(); + let config = logits.device.config(); + let mut candidate_count = initial_sampler_candidate_count(config, top_k, chunks); + let trace = sampler_trace_enabled(config); + let debug_dump = config.debug_sampler; let (logits_data, _, mut attempt) = logits.materialize_with_tail(|logits_data, encoder| { let attempt = encode_sample_attempt( @@ -268,7 +272,7 @@ pub(crate) async fn sample_token_to_host( }, encoder, )?; - let download = encode_token_download(&attempt.0, encoder, "sampled token download"); + let download = encode_token_download(&attempt.output, encoder, "sampled token download"); Some((attempt, download)) }); let device = logits_data.device().clone(); @@ -276,9 +280,10 @@ pub(crate) async fn sample_token_to_host( let mut attempt_index = 0usize; loop { attempt_index += 1; - let Some(((_, ids, values), download)) = attempt else { + let Some((attempt_data, download)) = attempt else { return Ok(None); }; + let debug_top_k = attempt_data.debug_top_k; let (status, token) = read_sample_result(&device, download, trace).await?; match status { GPU_SAMPLE_STATUS_SAMPLED => { @@ -302,7 +307,7 @@ pub(crate) async fn sample_token_to_host( "sampler_trace invalid attempt={attempt_index} top_k={top_k} chunks={chunks} candidate_count={candidate_count} status={status}" ); } - if debug_dump { + if debug_dump && let Some((ids, values)) = debug_top_k { debug_dump_invalid_sample(&logits_data, &ids, &values, previous_tokens).await; } return Ok(None); @@ -336,7 +341,7 @@ pub(crate) async fn sample_token_to_host( ) .map(|attempt| { let download = - encode_token_download(&attempt.0, &mut encoder, "sampled token download"); + encode_token_download(&attempt.output, &mut encoder, "sampled token download"); (attempt, download) }); device.wgpu_queue().submit(Some(encoder.finish())); @@ -386,8 +391,9 @@ pub(crate) fn sample_token_pending( dims, encoder, )?; - let download = encode_token_download(&attempt.0, encoder, "pending sampled token download"); - Some((attempt.0, download)) + let download = + encode_token_download(&attempt.output, encoder, "pending sampled token download"); + Some((attempt.output, download)) }); let (output, download) = tail?; diff --git a/fusor-ml/core/src/sampling/row_kernels.rs b/fusor-ml/core/src/sampling/row_kernels.rs index 183f19139..6dcb1d22d 100644 --- a/fusor-ml/core/src/sampling/row_kernels.rs +++ b/fusor-ml/core/src/sampling/row_kernels.rs @@ -56,6 +56,14 @@ pub(crate) struct SamplerMeta { pub(crate) has_exactness_flag: bool, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct CategoricalSamplerMeta { + pub(crate) input_len: u32, + pub(crate) input_offset: u32, + pub(crate) input_stride: u32, + pub(crate) block: u32, +} + fn u32t(value: u32) -> Tile { Tile::u32(value) } @@ -515,6 +523,208 @@ fn store_sample_result(program: &mut TileBlock<'_>, output: &Storage, status: u3 program.store(output.at(1), token, Mask::all()); } +pub(crate) struct CategoricalSampler { + pub(crate) logits: KernelTensorRef, + pub(crate) params: KernelTensorRef, + pub(crate) output: KernelTensorRef, + pub(crate) meta: CategoricalSamplerMeta, +} + +/// Sample directly from one complete, unfiltered logits row. This avoids the +/// separate chunk-sort, merge, and sampler dispatches used by the general +/// top-k path. The in-workgroup sort preserves that path's sampling order: +/// processed logit descending, with token id descending for ties. +pub(crate) fn categorical_sampler( + kb: &mut KernelBuilder, + spec: CategoricalSampler, +) -> Option<()> { + let CategoricalSampler { + logits, + params, + output, + meta, + } = spec; + if meta.input_len == 0 + || meta.block == 0 + || !meta.block.is_power_of_two() + || meta.input_len > meta.block + || meta.block > TOP_K_BLOCK + { + return None; + } + + let logits = kb.read(ElementType::F32, logits); + let params = kb.read(ElementType::F32, params); + let output = kb.write(ElementType::U32, output); + let phase = kb.program(); + let scratch_values = phase.alloc_workgroup_array(ScalarElement::F32, meta.block); + let scratch_ids = phase.alloc_workgroup_array(ScalarElement::U32, meta.block); + let weights = phase.alloc_workgroup_array(ScalarElement::F32, meta.block); + + emit_row_grid( + phase, + RowDispatchSpec::single(meta.block), + |program, ctx| { + let lane = ctx.lane; + let sort_current_value = program.private(ElementType::F32); + let sort_current_id = program.private(ElementType::U32); + let sort_partner_value = program.private(ElementType::F32); + let sort_partner_id = program.private(ElementType::U32); + let active = lane.clone().lt(u32t(meta.input_len)); + let input_index = index1(meta.input_offset, meta.input_stride, lane.clone()); + let raw = program.load( + logits.at(input_index), + active.clone(), + TileLiteral::f32(NEG_MAX_F32), + ); + let scaled = program.private(ElementType::F32); + program.store_local(&scaled, raw.clone()); + let temperature = load_param_f32(program, ¶ms, 1); + program.if_then(temperature.clone().ne(f32t(0.0)), |program| { + program.store_local(&scaled, raw.clone() / temperature.clone()); + }); + let scaled = program.load_local(&scaled); + let valid = active & is_finite(raw) & is_finite(scaled.clone()); + let value = Tile::select(valid.clone(), scaled, f32t(NEG_MAX_F32)); + let id = Tile::select(valid, lane.clone(), u32t(u32::MAX)); + program.store_workgroup(&scratch_values, lane.clone(), value); + program.store_workgroup(&scratch_ids, lane.clone(), id); + program.workgroup_barrier(); + + let mut size = 2; + while size <= meta.block { + let mut stride = size / 2; + while stride > 0 { + let partner = lane.clone() ^ stride; + let lower_lane = (lane.clone() & stride).eq(u32t(0)); + program.if_then(lower_lane, |program| { + let current_value = program.load_workgroup(&scratch_values, lane.clone()); + let current_id = program.load_workgroup(&scratch_ids, lane.clone()); + let partner_value = + program.load_workgroup(&scratch_values, partner.clone()); + let partner_id = program.load_workgroup(&scratch_ids, partner.clone()); + program.store_local(&sort_current_value, current_value); + program.store_local(&sort_current_id, current_id); + program.store_local(&sort_partner_value, partner_value); + program.store_local(&sort_partner_id, partner_id); + + let current_value = program.load_local(&sort_current_value); + let current_id = program.load_local(&sort_current_id); + let partner_value = program.load_local(&sort_partner_value); + let partner_id = program.load_local(&sort_partner_id); + let descending = (lane.clone() & size).eq(u32t(0)); + let partner_better = better_candidate( + partner_value.clone(), + partner_id.clone(), + current_value.clone(), + current_id.clone(), + ); + let current_better = better_candidate( + current_value.clone(), + current_id.clone(), + partner_value.clone(), + partner_id.clone(), + ); + let ascending = descending.clone().eq(Tile::bool(false)); + let should_swap = + (descending & partner_better) | (ascending & current_better); + program.if_then(should_swap, |program| { + let current_value = program.load_local(&sort_current_value); + let current_id = program.load_local(&sort_current_id); + let partner_value = program.load_local(&sort_partner_value); + let partner_id = program.load_local(&sort_partner_id); + program.store_workgroup( + &scratch_values, + lane.clone(), + partner_value.clone(), + ); + program.store_workgroup(&scratch_ids, lane.clone(), partner_id.clone()); + program.store_workgroup( + &scratch_values, + partner.clone(), + current_value.clone(), + ); + program.store_workgroup(&scratch_ids, partner, current_id); + }); + }); + program.workgroup_barrier(); + stride /= 2; + } + size *= 2; + } + + let max_value = program.load_workgroup(&scratch_values, u32t(0)); + let sorted_value = program.load_workgroup(&scratch_values, lane.clone()); + let sorted_id = program.load_workgroup(&scratch_ids, lane.clone()); + let valid = sorted_id.ne(u32t(u32::MAX)); + let weight = Tile::select( + valid, + (sorted_value - max_value).unary(TileUnaryOp::Exp), + f32t(0.0), + ); + program.store_workgroup(&weights, lane.clone(), weight.clone()); + program.workgroup_barrier(); + let total = program.reduce_sum(weight); + let total = program.bind(total); + + program.if_then(first_lane(&lane), |program| { + let first_id = program.load_workgroup(&scratch_ids, u32t(0)); + let total_invalid = is_finite(total.clone()).eq(Tile::bool(false)) + | total.clone().le(f32t(0.0)) + | first_id.clone().eq(u32t(u32::MAX)); + program.if_then(total_invalid, |program| { + store_sample_result(program, &output, GPU_SAMPLE_STATUS_INVALID, u32t(0)); + program.return_(); + }); + + let random = load_param_f32(program, ¶ms, 0); + let cutoff = program.private(ElementType::U32); + let cutoff_sum = program.private(ElementType::F32); + program.store_local(&cutoff, u32t(meta.input_len)); + program.store_local(&cutoff_sum, f32t(0.0)); + program.fold_state( + u32t(0), + |_, index| index.ge(u32t(meta.input_len)), + |program, index| { + let weight = program.load_workgroup(&weights, index.clone()); + let next = program.load_local(&cutoff_sum) + weight; + program.store_local(&cutoff_sum, next.clone()); + program.if_then(next.ge(total.clone()), |program| { + program.store_local(&cutoff, index.clone() + u32t(1)); + program.break_loop(); + }); + index + u32t(1) + }, + ); + let cutoff_sum = program.load_local(&cutoff_sum).max(f32t(1.0e-20)); + let threshold = random * cutoff_sum; + let cumulative = program.private(ElementType::F32); + let selected = program.private(ElementType::U32); + program.store_local(&cumulative, f32t(0.0)); + program.store_local(&selected, first_id); + program.fold_state( + u32t(0), + |program, index| index.ge(program.load_local(&cutoff)), + |program, index| { + let weight = program.load_workgroup(&weights, index.clone()); + let next = program.load_local(&cumulative) + weight.clone(); + program.if_then(next.clone().ge(threshold.clone()), |program| { + let id = program.load_workgroup(&scratch_ids, index.clone()); + program.store_local(&selected, id); + program.break_loop(); + }); + program.store_local(&cumulative, next); + index + u32t(1) + }, + ); + let selected = program.load_local(&selected); + store_sample_result(program, &output, GPU_SAMPLE_STATUS_SAMPLED, selected); + }); + }, + ); + Some(()) +} + fn emit_sampler_guards( program: &mut TileBlock<'_>, lane: &Tile, diff --git a/fusor-ml/core/src/sampling/standard_sampler.rs b/fusor-ml/core/src/sampling/standard_sampler.rs deleted file mode 100644 index 68415764f..000000000 --- a/fusor-ml/core/src/sampling/standard_sampler.rs +++ /dev/null @@ -1 +0,0 @@ -pub(super) use crate::mir::kernel_backend::standard_sampler::sample_from_sorted_top_k_data_with_encoder; diff --git a/fusor-ml/core/src/sampling/tests.rs b/fusor-ml/core/src/sampling/tests.rs index 9c77305a0..e157123b8 100644 --- a/fusor-ml/core/src/sampling/tests.rs +++ b/fusor-ml/core/src/sampling/tests.rs @@ -3,16 +3,29 @@ use std::mem::size_of; use crate::{DataTypeEnum, Device, Tensor, TensorData, quantized::QMatrix}; use fusor_gguf::{BlockQ4_0, GgmlType}; -use crate::mir::kernel_backend::sampling_topk::chunk_top_k_pair_data_with_processors_with_encoder; +use crate::mir::kernel_backend::sampling::{ + ChunkProcessors, ProcessorSettings, chunk_top_k_pair_data_with_encoder, + sample_categorical_logits_data_with_encoder, sample_from_sorted_top_k_data_with_encoder, + supports_unfiltered_categorical, +}; use super::{ - GPU_SAMPLE_STATUS_RETRY_NEEDED, GPU_SAMPLE_STATUS_SAMPLED, GpuMirostat2Sampler, - GpuMirostat2SamplerParams, GpuSamplerRequest, GpuStandardSamplerParams, - mirostat::sample_from_sorted_top_k_data_with_encoder, sample_token_to_host, - standard_sampler::sample_from_sorted_top_k_data_with_encoder as sample_standard_from_sorted_top_k_data_with_encoder, - topk::ProcessorSettings, + GPU_SAMPLE_STATUS_INVALID, GPU_SAMPLE_STATUS_RETRY_NEEDED, GPU_SAMPLE_STATUS_SAMPLED, + GpuMirostat2Sampler, GpuMirostat2SamplerParams, GpuSamplerRequest, GpuStandardSamplerParams, + sample_token_to_host, }; +fn unfiltered_params(input_len: usize, temperature: f32, random: f32) -> GpuStandardSamplerParams { + GpuStandardSamplerParams { + top_k: input_len, + temperature, + repetition_penalty: 1.0, + top_p: 1.0, + min_p: 0.0, + random, + } +} + #[test] fn top_k_pairs_match_cpu_sorted_order() { pollster::block_on(async { @@ -83,13 +96,16 @@ fn processed_chunk_top_k_applies_temperature_and_repetition_penalty() { ); let data = TensorData::new_from_buffer(&device, buffer, &[values.len()], DataTypeEnum::F32); let previous_tokens = [0, 3, 9]; - let (ids, logits) = chunk_top_k_pair_data_with_processors_with_encoder( + let (ids, logits) = chunk_top_k_pair_data_with_encoder( &data, - &previous_tokens, - ProcessorSettings { - temperature: 0.5, - repetition_penalty: 2.0, - }, + Some(ChunkProcessors { + previous_tokens: &previous_tokens, + gpu_tail: None, + settings: ProcessorSettings { + temperature: 0.5, + repetition_penalty: 2.0, + }, + }), 5, 5, None, @@ -135,6 +151,176 @@ fn processed_chunk_top_k_applies_temperature_and_repetition_penalty() { }); } +#[test] +fn unfiltered_categorical_fast_path_eligibility_is_strict() { + let eligible = unfiltered_params(65, 0.8, 0.25); + assert!(supports_unfiltered_categorical(65, eligible)); + + let mut filtered = eligible; + filtered.top_k = 64; + assert!(!supports_unfiltered_categorical(65, filtered)); + filtered = eligible; + filtered.top_p = 0.99; + assert!(!supports_unfiltered_categorical(65, filtered)); + filtered = eligible; + filtered.min_p = 0.01; + assert!(!supports_unfiltered_categorical(65, filtered)); + filtered = eligible; + filtered.temperature = f32::NAN; + assert!(!supports_unfiltered_categorical(65, filtered)); + filtered = eligible; + filtered.random = f32::NAN; + assert!(!supports_unfiltered_categorical(65, filtered)); + assert!(!supports_unfiltered_categorical(0, eligible)); + assert!(!supports_unfiltered_categorical(257, eligible)); +} + +#[test] +fn backend_unfiltered_categorical_matches_generic_boundaries() { + pollster::block_on(async { + let device = Device::new().await.unwrap().with_poisoned_allocations(); + let logits = vec![0.0; 65]; + let data = TensorData::new_from_buffer( + &device, + device.create_buffer_init( + bytemuck::cast_slice(&logits), + wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_SRC + | wgpu::BufferUsages::COPY_DST, + ), + &[logits.len()], + DataTypeEnum::F32, + ); + let generic_logits = Tensor::new(&device, logits.as_slice()); + let first_boundary = 1.0f32 / logits.len() as f32; + let after_first_boundary = f32::from_bits(first_boundary.to_bits() + 1); + let last_boundary = 64.0f32 / logits.len() as f32; + for random in [ + 0.0, + first_boundary, + after_first_boundary, + 0.5, + last_boundary, + 0.999_999, + ] { + let params = unfiltered_params(logits.len(), 0.8, random); + let expected = generic_logits + .sample_standard_token(&[u32::MAX], params) + .await + .unwrap(); + let output = sample_categorical_logits_data_with_encoder(&data, params, None).unwrap(); + let result = Tensor::from(output).as_slice::<1, u32>().await.unwrap(); + assert_eq!(result.as_slice()[0], GPU_SAMPLE_STATUS_SAMPLED); + assert_eq!( + result.as_slice()[1], + expected, + "categorical boundary mismatch for random={random}" + ); + if random == 0.0 { + assert_eq!(expected, 64, "ties must prefer the larger token id"); + } + } + }); +} + +#[test] +fn unfiltered_categorical_matches_generic_seeded_order() { + pollster::block_on(async { + let device = Device::new().await.unwrap().with_poisoned_allocations(); + let logits = [0.0, 2.0, 2.0, -1.0, 2.0, f32::NAN, 1.0]; + let tensor = Tensor::new(&device, logits.as_slice()); + + for random in [0.0, 0.2, 0.5, 0.9, 0.999_999] { + let params = unfiltered_params(logits.len(), 0.8, random); + let fast = tensor.sample_standard_token(&[], params).await.unwrap(); + // A non-matching previous token forces the generic processed-top-k + // route without changing any processed logits. + let generic = tensor + .sample_standard_token(&[u32::MAX], params) + .await + .unwrap(); + assert_eq!(fast, generic, "seeded sampler mismatch for random={random}"); + if random == 0.0 { + assert_eq!(fast, 4, "ties must prefer the larger token id"); + } + } + }); +} + +#[test] +fn backend_unfiltered_categorical_handles_temperature_and_invalid_logits() { + pollster::block_on(async { + let device = Device::new().await.unwrap(); + let logits = [f32::NAN, -3.0, 0.25, 2.0, f32::INFINITY, -1.0, 1.5]; + let data = TensorData::new_from_buffer( + &device, + device.create_buffer_init( + bytemuck::cast_slice(&logits), + wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_SRC + | wgpu::BufferUsages::COPY_DST, + ), + &[logits.len()], + DataTypeEnum::F32, + ); + let generic_logits = Tensor::new(&device, logits.as_slice()); + for (temperature, random) in [(0.5, 0.2), (0.5, 0.8), (2.0, 0.6)] { + let params = unfiltered_params(logits.len(), temperature, random); + let expected = generic_logits + .sample_standard_token(&[u32::MAX], params) + .await + .unwrap(); + let output = sample_categorical_logits_data_with_encoder(&data, params, None).unwrap(); + let result = Tensor::from(output).as_slice::<1, u32>().await.unwrap(); + assert_eq!( + result.as_slice(), + &[GPU_SAMPLE_STATUS_SAMPLED, expected], + "categorical mismatch for temperature={temperature} random={random}" + ); + } + + let invalid_logits = [f32::NAN, f32::INFINITY, f32::NEG_INFINITY]; + let invalid = TensorData::new_from_buffer( + &device, + device.create_buffer_init( + bytemuck::cast_slice(&invalid_logits), + wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_SRC + | wgpu::BufferUsages::COPY_DST, + ), + &[invalid_logits.len()], + DataTypeEnum::F32, + ); + let output = sample_categorical_logits_data_with_encoder( + &invalid, + unfiltered_params(invalid_logits.len(), 1.0, 0.5), + None, + ) + .unwrap(); + let result = Tensor::from(output).as_slice::<1, u32>().await.unwrap(); + assert_eq!(result.as_slice(), &[GPU_SAMPLE_STATUS_INVALID, 0]); + }); +} + +#[test] +fn pending_unfiltered_categorical_exposes_the_sampled_gpu_token() { + pollster::block_on(async { + let device = Device::new().await.unwrap().with_poisoned_allocations(); + let logits = vec![0.0; 65]; + let params = unfiltered_params(logits.len(), 0.8, 0.0); + let expected = 64; + let tensor = Tensor::new(&device, logits.as_slice()).sin(); + + let pending = tensor + .sample_standard_token_pending(&[], None, params) + .expect("full-vocabulary categorical sampling should stay pending"); + let dependent = pending.token_tensor() + 7u32; + let dependent = dependent.as_slice::<1, u32>().await.unwrap(); + assert_eq!(dependent.as_slice(), &[expected + 7]); + assert_eq!(pending.read_token().await.unwrap(), Some(expected)); + }); +} + fn cpu_mirostat2_selected_token(values: &[f32], mu: f32, params: GpuMirostat2SamplerParams) -> u32 { let mut top = values .iter() @@ -258,8 +444,10 @@ fn backend_mirostat2_sampler_uses_full_top_k_without_surprise_cutoff() { let output = sample_from_sorted_top_k_data_with_encoder( &ids_data, &values_data, - &mut sampler, - params, + &mut GpuSamplerRequest::Mirostat2 { + sampler: &mut sampler, + params, + }, None, None, ) @@ -302,10 +490,10 @@ fn backend_standard_sampler_applies_top_p_on_gpu() { random: 0.8, }; - let output = sample_standard_from_sorted_top_k_data_with_encoder( + let output = sample_from_sorted_top_k_data_with_encoder( &ids_data, &values_data, - params, + &mut GpuSamplerRequest::Standard { params }, None, None, ) @@ -354,8 +542,10 @@ fn backend_mirostat2_sampler_matches_cpu_reference_for_sorted_top_k() { let output = sample_from_sorted_top_k_data_with_encoder( &ids_data, &values_data, - &mut sampler, - params, + &mut GpuSamplerRequest::Mirostat2 { + sampler: &mut sampler, + params, + }, None, None, ) @@ -412,8 +602,10 @@ fn retry_status_does_not_mutate_mirostat_state() { let output = sample_from_sorted_top_k_data_with_encoder( &ids_data, &values_data, - &mut sampler, - params, + &mut GpuSamplerRequest::Mirostat2 { + sampler: &mut sampler, + params, + }, Some(&exactness_flag), None, ) diff --git a/fusor-ml/core/src/sampling/topk.rs b/fusor-ml/core/src/sampling/topk.rs deleted file mode 100644 index 95c4a55d6..000000000 --- a/fusor-ml/core/src/sampling/topk.rs +++ /dev/null @@ -1,9 +0,0 @@ -pub(crate) use crate::mir::kernel_backend::sampling_topk::{ - MergeSortedChunkTopKParams, chunk_top_k_pair_data_with_encoder, - merge_sorted_chunk_top_k_pair_data_with_encoder, -}; - -pub(super) use crate::mir::kernel_backend::sampling_topk::{ - ProcessorSettings, chunk_top_k_pair_data_with_processors_and_gpu_tail_with_encoder, - top_k_exactness_flag_data_with_encoder, -}; diff --git a/fusor-ml/core/src/slice_assign.rs b/fusor-ml/core/src/slice_assign.rs index b25720d67..6a5c43f56 100644 --- a/fusor-ml/core/src/slice_assign.rs +++ b/fusor-ml/core/src/slice_assign.rs @@ -137,6 +137,11 @@ impl Operation for SliceAssignOperation { f(self.input); } + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + f(&mut self.value); + f(&mut self.input); + } + fn inputs(&self, nodes: &ComputeGraphInner) -> Vec { // Pass the ORIGINAL input tensor (not sliced) and the value tensor let input = nodes.get_cached_result(self.input).unwrap(); diff --git a/fusor-ml/core/src/tensor/eager_data.rs b/fusor-ml/core/src/tensor/eager_data.rs index 31a7f96e0..3cccc2c7e 100644 --- a/fusor-ml/core/src/tensor/eager_data.rs +++ b/fusor-ml/core/src/tensor/eager_data.rs @@ -123,6 +123,21 @@ impl TensorData { Self::new_from_buffer(device, buffer, shape, D::DATA_TYPE) } + pub(crate) fn new_from_slice( + device: &Device, + data: &[D], + shape: &[usize], + ) -> Self { + let buffer = device.create_buffer_init( + bytemuck::cast_slice(data), + wgpu::BufferUsages::STORAGE + | wgpu::BufferUsages::COPY_SRC + | wgpu::BufferUsages::COPY_DST, + ); + + Self::new_from_buffer(device, buffer, shape, D::DATA_TYPE) + } + pub fn slice(&self, ranges: &[Range]) -> Self { let layout = self.info.layout.slice(ranges); Self { diff --git a/fusor-ml/core/src/tensor/mod.rs b/fusor-ml/core/src/tensor/mod.rs index 03d2de2fa..5b97bf4d6 100644 --- a/fusor-ml/core/src/tensor/mod.rs +++ b/fusor-ml/core/src/tensor/mod.rs @@ -8,7 +8,7 @@ use tabbycat::Graph; use wgpu::COPY_BUFFER_ALIGNMENT; use crate::{ - Device, ReduceFunction, ReduceOperation, + Device, Layout, ReduceFunction, ReduceOperation, compute_graph::NodeIndex, nary_wise::{ElementwiseOperation, NaryExpr, NaryFunction}, quantized::QMatrix, @@ -102,7 +102,9 @@ impl Tensor { shape.iter().product::(), "Data length must match shape" ); - Tensor::new_inner(device, data.iter(), shape) + Self::from_parts(LazyTensorData::new(TensorData::new_from_slice( + device, data, shape, + ))) } /// Allocate a concrete tensor backing for `shape` without uploading @@ -150,28 +152,55 @@ impl Tensor { pub(crate) async fn as_slice_from_tensor_data( tensor: &TensorData, ) -> Result, wgpu::BufferAsyncError> { + let device = tensor.device.wgpu_device(); + let queue = tensor.device.wgpu_queue(); + let mut encoder = + device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); + let download = Self::enqueue_download::(tensor, &mut encoder); + queue.submit(Some(encoder.finish())); + Self::map_download(tensor, download).await + } + + fn enqueue_download( + tensor: &TensorData, + encoder: &mut wgpu::CommandEncoder, + ) -> (wgpu::Buffer, Layout) { assert_eq!(tensor.datatype(), D::DATA_TYPE); assert_eq!(tensor.layout().shape().len(), R); let buffer = tensor.buffer(); let device = tensor.device.wgpu_device(); - let queue = tensor.device.wgpu_queue(); - let size = buffer.size(); - - // Create a staging buffer for reading + let layout = tensor.layout(); + let element_size = tensor.datatype().element_size() as u64; + let source_offset = layout.offset() as u64 * element_size; + let compact_size = padded_tensor_size(layout.num_elements() as u64 * element_size); + let dense_strides = Layout::continuous_strides(layout.shape()); + let can_copy_compact = layout.strides() == dense_strides.as_ref() + && source_offset.is_multiple_of(COPY_BUFFER_ALIGNMENT) + && source_offset + compact_size <= buffer.size(); + let (source_offset, size, download_layout) = if can_copy_compact { + ( + source_offset, + compact_size, + Layout::contiguous(layout.shape()), + ) + } else { + (0, buffer.size(), layout.clone()) + }; let download = device.create_buffer(&wgpu::BufferDescriptor { size, usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, mapped_at_creation: false, label: None, }); + encoder.copy_buffer_to_buffer(buffer, source_offset, &download, 0, size); + (download, download_layout) + } - // Copy data to staging buffer - let mut encoder = - device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); - encoder.copy_buffer_to_buffer(buffer, 0, &download, 0, size); - queue.submit(Some(encoder.finish())); - - // Map the staging buffer using map_async which correctly uses WasmNotSend + async fn map_download( + tensor: &TensorData, + download: (wgpu::Buffer, Layout), + ) -> Result, wgpu::BufferAsyncError> { + let (download, layout) = download; let (sender, receiver) = futures_channel::oneshot::channel(); download .slice(..) @@ -185,10 +214,7 @@ impl Tensor { // Get the mapped view let view = download.slice(..).get_mapped_range(); - Ok(TensorSlice::new( - MappedBuffer { view }, - tensor.layout().clone(), - )) + Ok(TensorSlice::new(MappedBuffer { view }, layout)) } /// Synchronously dispatch and wait for GPU completion using device.poll(). @@ -255,6 +281,11 @@ impl Tensor { count } + /// Whether fully resolving this tensor takes exactly `N` kernel calls + pub fn resolves_in(&self) -> bool { + self.count_kernels_to_resolve() == N + } + pub async fn as_slice( &self, ) -> Result, wgpu::BufferAsyncError> { @@ -262,12 +293,14 @@ impl Tensor { self.assert_datatype::(); #[cfg(not(target_arch = "wasm32"))] let start_time = std::time::Instant::now(); - let (tensor, _) = self.data.materialize(); + let (tensor, _, download) = self + .data + .materialize_with_tail(Self::enqueue_download::); #[cfg(not(target_arch = "wasm32"))] tracing::trace!("Materialized tensor in {:?}", start_time.elapsed()); #[cfg(not(target_arch = "wasm32"))] let start_time = std::time::Instant::now(); - let out = Self::as_slice_from_tensor_data(&tensor).await; + let out = Self::map_download(&tensor, download).await; #[cfg(not(target_arch = "wasm32"))] tracing::trace!("Downloaded tensor in {:?}", start_time.elapsed()); out diff --git a/fusor-ml/core/src/tensor/sampling.rs b/fusor-ml/core/src/tensor/sampling.rs index 5d7a14643..1565343d4 100644 --- a/fusor-ml/core/src/tensor/sampling.rs +++ b/fusor-ml/core/src/tensor/sampling.rs @@ -100,7 +100,9 @@ impl Tensor { let chunks = input_len.div_ceil(crate::top_k::TOP_K_CHUNK); let mut candidate_count = k .div_ceil(chunks) - .max(crate::top_k::min_top_k_candidates_per_chunk()) + .max(crate::top_k::min_top_k_candidates_per_chunk( + self.device().config(), + )) .min(k) .min(crate::top_k::TOP_K_CHUNK); @@ -206,6 +208,7 @@ fn encode_top_k_attempt( }; let (chunk_ids, chunk_values) = crate::top_k::chunk_top_k_pair_data_with_encoder( input, + None, candidate_count, output_per_chunk, Some(encoder), diff --git a/fusor-ml/core/src/view.rs b/fusor-ml/core/src/view.rs index 03411b7aa..24ab93a5d 100644 --- a/fusor-ml/core/src/view.rs +++ b/fusor-ml/core/src/view.rs @@ -22,7 +22,7 @@ const BLOCKSIZE: u32 = 256; /// logical value space). `defined` is a prefix box of `layout.shape()`: /// coordinates with `coord_i < defined[i]` for every axis read data; /// anything outside reads `fill`. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq)] pub(crate) struct ViewStage { pub(crate) layout: Layout, pub(crate) input_shape: Box<[usize]>, @@ -133,7 +133,7 @@ impl ViewStage { /// /// `input` is never itself a view node: view-over-view always collapses /// into one node at construction. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq)] pub(crate) struct ViewOperation { pub(crate) input: NodeIndex, /// Innermost first; never empty. @@ -499,36 +499,37 @@ pub(crate) fn row_major_indices_from_flat( flat: NaryExpr, shape: &[usize], ) -> Option> { - let mut indices = Vec::with_capacity(shape.len()); - for axis in 0..shape.len() { - let divisor = shape[axis + 1..] - .iter() - .try_fold(1u32, |acc, dim| acc.checked_mul((*dim).try_into().ok()?))?; + // Peel innermost-out with a running quotient instead of dividing by each + // axis' suffix product — the Apple Metal compiler miscompiles u32 + // div/mod chains with large non-power-of-two constants (see + // output_dims_from_flat). The outermost non-trivial axis keeps the raw + // quotient, matching the previous form. + let mut indices = vec![NaryExpr::scalar(NaryScalar::U32(0)); shape.len()]; + let mut rest = flat; + for axis in (0..shape.len()).rev() { let dim = u32::try_from(shape[axis]).ok()?; - let quotient = if divisor == 1 { - flat.clone() - } else { - NaryExpr::unary_op( - flat.clone(), - "div_const", - NaryOp::DivConst(NaryScalar::U32(divisor)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - }; - indices.push(if dim == 1 { - NaryExpr::scalar(NaryScalar::U32(0)) - } else if axis == 0 { - quotient - } else { - NaryExpr::unary_op( - quotient, - "rem_const", - NaryOp::RemConst(NaryScalar::U32(dim)), - DataTypeEnum::U32, - DataTypeEnum::U32, - ) - }); + if dim == 1 { + continue; + } + let outer_nontrivial = shape[..axis].iter().any(|&outer| outer != 1); + if !outer_nontrivial { + indices[axis] = rest; + break; + } + indices[axis] = NaryExpr::unary_op( + rest.clone(), + "rem_const", + NaryOp::RemConst(NaryScalar::U32(dim)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); + rest = NaryExpr::unary_op( + rest, + "div_const", + NaryOp::DivConst(NaryScalar::U32(dim)), + DataTypeEnum::U32, + DataTypeEnum::U32, + ); } Some(indices) } @@ -571,6 +572,10 @@ impl Operation for ViewOperation { f(self.input); } + fn visit_dependencies_mut(&mut self, f: &mut dyn FnMut(&mut NodeIndex)) { + f(&mut self.input); + } + fn inputs(&self, nodes: &crate::compute_graph::ComputeGraphInner) -> Vec { let input = nodes.get_cached_result(self.input).unwrap().clone(); let output = TensorData::new_for_shape(input.device(), self.shape(), input.datatype()); diff --git a/fusor-ml/core/src/visit_tiled.rs b/fusor-ml/core/src/visit_tiled.rs index 1fb708b5c..b57b7ade4 100644 --- a/fusor-ml/core/src/visit_tiled.rs +++ b/fusor-ml/core/src/visit_tiled.rs @@ -82,10 +82,16 @@ pub(crate) fn titled_map_workgroup_size_constraints( pub(crate) fn distribute_workgroups(total_workgroups: u32, max_per_dim: u32) -> [u32; 3] { let max_per_dim = max_per_dim.max(1); - let x = total_workgroups.min(max_per_dim); - let remaining = total_workgroups.div_ceil(x.max(1)); - let y = remaining.min(max_per_dim).max(1); - let z = total_workgroups.div_ceil(x.max(1) * y).max(1); + if total_workgroups <= max_per_dim { + return [total_workgroups, 1, 1]; + } + // Pick the slab count first, then size X to the slab. Saturating X instead + // leaves the last slab nearly empty (122,880 groups would launch + // [65535, 2, 1] = 131,070), and every extra workgroup still runs the + // kernel prologue and the in-range compares before falling through. + let y = total_workgroups.div_ceil(max_per_dim).min(max_per_dim); + let x = total_workgroups.div_ceil(y).min(max_per_dim); + let z = total_workgroups.div_ceil(x.saturating_mul(y)).max(1); [x, y, z] } @@ -103,3 +109,26 @@ pub(crate) fn titled_map_dispatch_size( distribute_workgroups(total_workgroups, max_per_dim) } + +#[test] +fn distribute_workgroups_covers_without_slack() { + const MAX: u32 = 0xFFFF; + for total in (0..3_000_000) + .step_by(1409) + .chain([0, 1, MAX, MAX + 1, 122_880, u32::MAX]) + { + let [x, y, z] = distribute_workgroups(total, MAX); + assert!( + x <= MAX && y <= MAX && z <= MAX, + "{total} exceeds the limit" + ); + let launched = x as u64 * y as u64 * z as u64; + assert!(launched >= total as u64, "{total} is not covered"); + // Every launched group past `total` still runs the kernel prologue, so + // the slack must stay below one slab rather than up to a full X extent. + assert!( + total > MAX * MAX || (launched - total as u64) < y as u64, + "{total} launches {launched}" + ); + } +} diff --git a/fusor-ml/core/tests/attention.rs b/fusor-ml/core/tests/attention.rs index 2aee33924..9eab67a72 100644 --- a/fusor-ml/core/tests/attention.rs +++ b/fusor-ml/core/tests/attention.rs @@ -113,13 +113,12 @@ fn check_attention(case: AttentionCase, tolerance: f32) { .map(|data| Tensor::from_slice(&device, [q_len, kv_len], data.as_slice())); let out = if causal { - q.flash_attention_causal(&k, &v, scale) + q.attention_causal(&k, &v, scale) } else { - q.flash_attention(&k, &v, scale, mask_tensor.as_ref()) + q.attention(&k, &v, scale, mask_tensor.as_ref()) }; - assert_eq!( - out.count_kernels_to_resolve(), - 1, + assert!( + out.resolves_in::<1>(), "attention should lower as one row-program kernel" ); let actual = out.as_slice::<4, f32>().await.unwrap(); @@ -279,7 +278,7 @@ fn attention_f16_io() { let k = Tensor::from_slice(&device, [1, heads, kv_len, head_dim], &to_f16(&k_data)); let v = Tensor::from_slice(&device, [1, heads, kv_len, head_dim], &to_f16(&v_data)); - let out = q.flash_attention(&k, &v, scale, None); + let out = q.attention(&k, &v, scale, None); let actual = out.as_slice::<4, half::f16>().await.unwrap(); let case = AttentionCase { diff --git a/fusor-ml/core/tests/conv_repro.rs b/fusor-ml/core/tests/conv_repro.rs new file mode 100644 index 000000000..f1c89db90 --- /dev/null +++ b/fusor-ml/core/tests/conv_repro.rs @@ -0,0 +1,137 @@ +//! Temporary diagnostic: conv im2col case-2 repro (m=2048, k=576, n=128). +//! Runs the windowed (implicit-GEMM) matmul and the same contraction as a +//! plain dense matmul, many rounds, full-output check against a CPU +//! reference. Not part of the permanent suite. + +use fusor_core::{Device, StrideSpec, Tensor}; + +fn cpu_reference( + b: usize, + c: usize, + h: usize, + w: usize, + n: usize, + kh: usize, + kw: usize, + input_host: &[f32], + weight_host: &[f32], +) -> Vec { + let (oh, ow) = (h - kh + 1, w - kw + 1); + let (m, k) = (b * oh * ow, c * kh * kw); + let mut out = vec![0f32; m * n]; + for mi in 0..m { + let (bi, rest) = (mi / (oh * ow), mi % (oh * ow)); + let (ohi, owi) = (rest / ow, rest % ow); + for ni in 0..n { + let mut acc = 0f32; + for ci in 0..c { + for khi in 0..kh { + for kwi in 0..kw { + let iv = input_host[((bi * c + ci) * h + ohi + khi) * w + owi + kwi]; + let wv = weight_host[ni * k + (ci * kh + khi) * kw + kwi]; + acc += iv * wv; + } + } + } + out[mi * n + ni] = acc; + } + } + out +} + +fn check(round: usize, label: &str, result: &[f32], expected: &[f32], n: usize) -> usize { + let mut mismatches = 0usize; + let mut worst = 0f32; + let mut first = None; + for (i, (&got, &exp)) in result.iter().zip(expected).enumerate() { + let err = (got - exp).abs(); + if err > 1e-3 * exp.abs().max(1.0) { + mismatches += 1; + if err > worst { + worst = err; + } + if first.is_none() { + first = Some((i / n, i % n, got, exp)); + } + } + } + if let Some((mi, ni, got, exp)) = first { + println!( + "round {round} {label}: {mismatches} mismatches, worst {worst:.6}, first [{mi},{ni}] got {got} exp {exp}" + ); + } else { + println!("round {round} {label}: clean"); + } + mismatches +} + +#[test] +fn conv_case2_repro() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (b, c, h, w, n, kh, kw) = (2usize, 64, 34, 34, 128, 3, 3); + let (oh, ow) = (h - kh + 1, w - kw + 1); + let (m, k) = (b * oh * ow, c * kh * kw); + let input_host: Vec = (0..b * c * h * w).map(|i| (i % 13) as f32 * 0.1).collect(); + let weight_host: Vec = (0..n * k).map(|i| (i % 7) as f32 * 0.01).collect(); + let expected = cpu_reference(b, c, h, w, n, kh, kw, &input_host, &weight_host); + + // Host-side im2col for the plain dense variant. + let mut a_dense = vec![0f32; m * k]; + for mi in 0..m { + let (bi, rest) = (mi / (oh * ow), mi % (oh * ow)); + let (ohi, owi) = (rest / ow, rest % ow); + for ci in 0..c { + for khi in 0..kh { + for kwi in 0..kw { + a_dense[mi * k + (ci * kh + khi) * kw + kwi] = + input_host[((bi * c + ci) * h + ohi + khi) * w + owi + kwi]; + } + } + } + } + + let mut windowed_bad = 0usize; + let mut dense_bad = 0usize; + for round in 0..10 { + let input = Tensor::from_slice(&device, [b, c, h, w], &input_host); + let weight = Tensor::from_slice(&device, [n, k], &weight_host); + let windows = input.restride([ + StrideSpec::dim(0, b), + StrideSpec::dim_with(2, oh, 1), + StrideSpec::dim_with(3, ow, 1), + StrideSpec::dim(1, c), + StrideSpec::dim(2, kh), + StrideSpec::dim(3, kw), + ]); + let a = windows.reshape([m, k]); + let b_mat = weight.restride([StrideSpec::dim(1, k), StrideSpec::dim(0, n)]); + let out = a.mat_mul(&b_mat); + let result = out.as_slice::<2, f32>().await.unwrap(); + let flat: Vec = (0..m) + .flat_map(|mi| (0..n).map(move |ni| (mi, ni))) + .map(|(mi, ni)| result[[mi, ni]]) + .collect(); + if check(round, "windowed", &flat, &expected, n) > 0 { + windowed_bad += 1; + } + + let a_t = Tensor::from_slice(&device, [m, k], &a_dense); + let weight2 = Tensor::from_slice(&device, [n, k], &weight_host); + let b_mat2 = weight2.restride([StrideSpec::dim(1, k), StrideSpec::dim(0, n)]); + let out2 = a_t.mat_mul(&b_mat2); + let result2 = out2.as_slice::<2, f32>().await.unwrap(); + let flat2: Vec = (0..m) + .flat_map(|mi| (0..n).map(move |ni| (mi, ni))) + .map(|(mi, ni)| result2[[mi, ni]]) + .collect(); + if check(round, "dense ", &flat2, &expected, n) > 0 { + dense_bad += 1; + } + } + println!("SUMMARY windowed_bad={windowed_bad}/10 dense_bad={dense_bad}/10"); + assert!(windowed_bad == 0 && dense_bad == 0); + }); +} diff --git a/fusor-ml/core/tests/coop_determinism.rs b/fusor-ml/core/tests/coop_determinism.rs new file mode 100644 index 000000000..7360266b6 --- /dev/null +++ b/fusor-ml/core/tests/coop_determinism.rs @@ -0,0 +1,46 @@ +//! The cooperative pair-loop matmul must be bit-deterministic run to run. +//! +//! Guards the expression-sinking class of miscompile: naga backends must +//! force-bake `CooperativeLoad` at its `Emit` point (like `Load`), or +//! single-use simdgroup loads inline into the deferred accumulator chain and +//! sink past the workgroup barrier, racing the next tile refill. That bug +//! produced nondeterministic 8x8-fragment corruption at a rate proportional +//! to workgroups x k-iterations — invisible to small conformance shapes, so +//! this test uses a grid large enough to expose it (it failed 10/10 rounds +//! against the broken backend). + +use fusor_core::{Device, StrideSpec, Tensor}; + +#[test] +fn coop_matmul_is_run_to_run_deterministic() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (m, k, n) = (2048usize, 576, 128); + let a_host: Vec = (0..m * k).map(|i| (i as f32 * 0.137).sin() * 0.5).collect(); + let weight_host: Vec = (0..n * k).map(|i| (i as f32 * 0.071).sin() * 0.1).collect(); + let a = Tensor::from_slice(&device, [m, k], &a_host); + let weight = Tensor::from_slice(&device, [n, k], &weight_host); + // The transposed-B restride keeps B staging on the strided path, + // matching the configuration that exposed the original race. + let b_mat = weight.restride([StrideSpec::dim(1, k), StrideSpec::dim(0, n)]); + + let reference = a.mat_mul(&b_mat).as_slice::<2, f32>().await.unwrap(); + for round in 0..3 { + let repeat = a.mat_mul(&b_mat).as_slice::<2, f32>().await.unwrap(); + let mut diffs = 0usize; + for mi in 0..m { + for ni in 0..n { + if (repeat[[mi, ni]] - reference[[mi, ni]]).abs() > 1e-6 { + diffs += 1; + } + } + } + assert_eq!( + diffs, 0, + "round {round}: {diffs} elements differ between identical dispatches" + ); + } + }); +} diff --git a/fusor-ml/core/tests/coop_tile_conformance.rs b/fusor-ml/core/tests/coop_tile_conformance.rs new file mode 100644 index 000000000..ece961130 --- /dev/null +++ b/fusor-ml/core/tests/coop_tile_conformance.rs @@ -0,0 +1,265 @@ +//! Cooperative-matrix conformance over every tile-table geometry: through +//! the selector on shapes derived from each entry, and with each entry's +//! tile forced through the kernel builder. + +use fusor_core::{Device, Layout, Tensor}; +use fusor_tile_ir::{CoopMatrixToken, ScalarElement, Shape, SubgroupToken, tile}; +use fusor_tile_ir_kernels::{ + CoopTileEntry, DEFAULT_SWIZZLE_GROUP_M, DenseCoopMatmulConfig, DenseMatmulEpilogues, + DenseMatmulShape, DenseMatmulTensors, SubgroupConfig, coop_tile_entries, + try_batched_coop_matmul, +}; + +fn values(len: usize, freq: f32) -> Vec { + (0..len).map(|i| ((i as f32) * freq).sin()).collect() +} + +/// f64-accumulated reference for `a[batch, m, k] @ b[batch, k, n]`. +fn cpu_matmul(a: &[f32], b: &[f32], batch: usize, m: usize, k: usize, n: usize) -> Vec { + let mut out = vec![0.0f32; batch * m * n]; + for bi in 0..batch { + for mi in 0..m { + for ni in 0..n { + let mut acc = 0.0f64; + for ki in 0..k { + let a_val = a[(bi * m + mi) * k + ki] as f64; + let b_val = b[(bi * k + ki) * n + ni] as f64; + acc += a_val * b_val; + } + out[(bi * m + mi) * n + ni] = acc as f32; + } + } + } + out +} + +async fn check_automatic( + device: &Device, + batch: usize, + m: usize, + k: usize, + n: usize, + transpose_b: bool, +) { + let a_data = values(batch * m * k, 0.13); + let b_data = values(batch * k * n, 0.07); + + let a = Tensor::from_slice(device, [batch, m, k], &a_data); + let b = if transpose_b { + let b_t_data: Vec = (0..batch * n * k) + .map(|i| { + let (bi, rest) = (i / (n * k), i % (n * k)); + let (ni, ki) = (rest / k, rest % k); + b_data[(bi * k + ki) * n + ni] + }) + .collect(); + let b_t = Tensor::from_slice(device, [batch, n, k], &b_t_data); + b_t.restride_layout(Layout::from_parts( + 0, + vec![batch, k, n].into(), + vec![n * k, 1, k].into(), + )) + } else { + Tensor::from_slice(device, [batch, k, n], &b_data) + }; + + let out = a.mat_mul(&b); + let actual = out.as_slice::<3, f32>().await.unwrap(); + let expected = cpu_matmul(&a_data, &b_data, batch, m, k, n); + for bi in 0..batch { + for mi in 0..m { + for ni in 0..n { + let want = expected[(bi * m + mi) * n + ni]; + let got = actual[[bi, mi, ni]]; + assert!( + (got - want).abs() < 1e-3 + want.abs() * 1e-3, + "batch={batch} m={m} k={k} n={n} transpose_b={transpose_b} \ + [{bi}, {mi}, {ni}]: got {got}, expected {want}", + ); + } + } + } +} + +/// Build one tile geometry's kernel directly and check its whole logical +/// output against the host reference. The output buffer covers whole tiles, +/// so the rows and columns past `m` and `n` hold pad values the comparison +/// skips. +fn check_forced(device: &Device, entry: &CoopTileEntry, m: u32, k: u32, n: u32) { + let tile = entry.tile; + let (row_groups, col_groups) = entry.subgroup_split(); + let a_data = values((m * k) as usize, 0.13); + let b_data = values((k * n) as usize, 0.07); + let m_pad = m.div_ceil(tile.bm) * tile.bm; + let n_pad = n.div_ceil(tile.bn) * tile.bn; + let a_buf = device.create_buffer_init( + bytemuck::cast_slice(&a_data), + wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST, + ); + let b_buf = device.create_buffer_init( + bytemuck::cast_slice(&b_data), + wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST, + ); + let y_buf = device.create_buffer( + u64::from(m_pad) * u64::from(n_pad) * 4, + wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + ); + + let element = ScalarElement::F32.element(); + let mut built = false; + let ir = tile::build(|program| { + let a = program.storage_read(element, Shape::new([m, k])); + let b = program.storage_read(element, Shape::new([k, n])); + let y = program.storage_write(element, Shape::new([m_pad, n_pad])); + built = try_batched_coop_matmul( + program, + DenseMatmulTensors { + a: &a, + b: &b, + y: &y, + }, + DenseMatmulShape { batch: 1, m, k, n }, + &DenseMatmulEpilogues::empty(), + 65535, + DenseCoopMatmulConfig { + coop: CoopMatrixToken::new_unchecked(), + subgroups: SubgroupConfig::fixed(SubgroupToken::new_unchecked(), 32), + tile, + row_groups, + col_groups, + staging: None, + stage_buffers: 2, + swizzle_group_m: DEFAULT_SWIZZLE_GROUP_M, + }, + ); + }); + assert!(built, "{tile:?} declined {m}x{k}x{n}"); + let grid = ir.grid; + let kernel = ir + .lower_to_naga() + .unwrap_or_else(|error| panic!("{tile:?} lowering failed: {error}")); + + let wgpu_device = device.wgpu_device(); + let module = unsafe { + wgpu_device.create_shader_module_trusted( + wgpu::ShaderModuleDescriptor { + label: None, + source: wgpu::ShaderSource::Naga(std::borrow::Cow::Owned(kernel.module().clone())), + }, + wgpu::ShaderRuntimeChecks::unchecked(), + ) + }; + let pipeline = wgpu_device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: None, + module: &module, + entry_point: Some("main"), + compilation_options: wgpu::PipelineCompilationOptions { + zero_initialize_workgroup_memory: false, + ..Default::default() + }, + cache: None, + }); + let bind_group = wgpu_device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &pipeline.get_bind_group_layout(0), + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: a_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: b_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: y_buf.as_entire_binding(), + }, + ], + }); + let staging = wgpu_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: y_buf.size(), + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut encoder = wgpu_device.create_command_encoder(&Default::default()); + { + let mut pass = encoder.begin_compute_pass(&Default::default()); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &bind_group, &[]); + pass.dispatch_workgroups(grid[0], grid[1], grid[2]); + } + encoder.copy_buffer_to_buffer(&y_buf, 0, &staging, 0, staging.size()); + device.wgpu_queue().submit([encoder.finish()]); + let slice = staging.slice(..); + let (sender, receiver) = std::sync::mpsc::channel(); + slice.map_async(wgpu::MapMode::Read, move |result| { + sender.send(result).unwrap() + }); + device.poll_wait(); + receiver.recv().unwrap().unwrap(); + let view = slice.get_mapped_range(); + let actual: &[f32] = bytemuck::cast_slice(&view); + + let expected = cpu_matmul(&a_data, &b_data, 1, m as usize, k as usize, n as usize); + for mi in 0..m as usize { + for ni in 0..n as usize { + let want = expected[mi * n as usize + ni]; + let got = actual[mi * n_pad as usize + ni]; + assert!( + (got - want).abs() < 1e-3 + want.abs() * 1e-3, + "{tile:?} m={m} k={k} n={n} [{mi}, {ni}]: got {got}, expected {want}", + ); + } + } +} + +/// The scorer orders least-padded first and fewest passes next, and padded +/// MACs are monotone in tile extent, so every multi-pass entry — and the +/// single-buffered one selection skips outright — is unreachable through the +/// automatic sweep below. The 128x128 profile shipped an all-zero miscompile +/// behind exactly that gap, so each entry's kernel is driven directly here, +/// tile-aligned and with every extent past its tile. +#[test] +fn forced_coop_tiles_compute_their_own_geometry_correctly() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + if device.fixed_width_subgroup_size() != Some(32) + || !device + .features() + .contains(wgpu::Features::EXPERIMENTAL_COOPERATIVE_MATRIX) + { + return; + } + for entry in coop_tile_entries() { + let tile = entry.tile; + check_forced(&device, entry, tile.bm, 48, tile.bn); + check_forced(&device, entry, tile.bm + 7, 30, tile.bn + 5); + } + }); +} + +#[test] +fn automatic_coop_selection_computes_table_derived_shapes_correctly() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let geometries: Vec<(u32, u32)> = coop_tile_entries() + .iter() + .map(|entry| (entry.tile.bm, entry.tile.bn)) + .collect(); + for (bm, bn) in geometries { + let (m, n) = (2 * bm as usize, 2 * bn as usize); + // Aligned, all edges masked at once, batched, and transposed-B. + check_automatic(&device, 1, m, 64, n, false).await; + check_automatic(&device, 1, m - 13, 50, n.saturating_sub(9).max(1), false).await; + check_automatic(&device, 3, m, 64, n, false).await; + check_automatic(&device, 2, m, 64, n, true).await; + } + }); +} diff --git a/fusor-ml/core/tests/decode_optimizer.rs b/fusor-ml/core/tests/decode_optimizer.rs new file mode 100644 index 000000000..944317133 --- /dev/null +++ b/fusor-ml/core/tests/decode_optimizer.rs @@ -0,0 +1,59 @@ +//! Decode regression for the default aggressive optimizer. + +use fusor_core::{Device, QMatrix, Tensor}; +use fusor_gguf::GgmlType; + +const N: usize = 4; +const K: usize = 8; +const QMATMULS: usize = 16; + +fn weight(device: &Device) -> (QMatrix, Vec) { + let values = (0..N * K) + .map(|i| 0.1 + i as f32 * 0.05) + .collect::>(); + let bytes = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let matrix = + QMatrix::from_parts(device, &bytes, vec![N, K].into_boxed_slice(), GgmlType::F32).unwrap(); + (matrix, values) +} + +#[test] +fn decode_runs_the_full_optimizer_by_default() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (weight, weights) = weight(&device); + let input = [1.0f32, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0]; + let activation = Tensor::from_slice(&device, [1, K], &input); + let outputs = (0..QMATMULS) + .map(|_| activation.q_mat_mul(&weight)) + .collect::>(); + let mut total = outputs[0].clone(); + for output in &outputs[1..] { + total = &total + output; + } + + assert!( + total.count_kernels_to_resolve() < QMATMULS * 2 - 1, + "decode should take the automatic planning and fusion path", + ); + + let actual = total.as_slice::<2, f32>().await.unwrap(); + for column in 0..N { + let one_matmul = (0..K) + .map(|k| input[k] * weights[column * K + k]) + .sum::(); + let expected = one_matmul * QMATMULS as f32; + let tolerance = 1e-4 * expected.abs().max(1.0); + assert!( + (actual[[0, column]] - expected).abs() <= tolerance, + "column {column}: got {}, expected {expected}", + actual[[0, column]], + ); + } + }); +} diff --git a/fusor-ml/core/tests/dense_repro.rs b/fusor-ml/core/tests/dense_repro.rs new file mode 100644 index 000000000..dd2fc1f8f --- /dev/null +++ b/fusor-ml/core/tests/dense_repro.rs @@ -0,0 +1,55 @@ +//! Temporary diagnostic: plain dense 2048x576x128 matmul, many rounds, +//! full-output check. Not part of the permanent suite. + +use fusor_core::{Device, StrideSpec, Tensor}; + +#[test] +fn dense_2048_576_128_repro() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (m, k, n) = (2048usize, 576, 128); + let a_host: Vec = (0..m * k).map(|i| (i % 13) as f32 * 0.1).collect(); + let weight_host: Vec = (0..n * k).map(|i| (i % 7) as f32 * 0.01).collect(); + let mut expected = vec![0f32; m * n]; + for mi in 0..m { + for ni in 0..n { + let mut acc = 0f32; + for ki in 0..k { + acc += a_host[mi * k + ki] * weight_host[ni * k + ki]; + } + expected[mi * n + ni] = acc; + } + } + + let mut bad = 0usize; + for round in 0..10 { + let a = Tensor::from_slice(&device, [m, k], &a_host); + let weight = Tensor::from_slice(&device, [n, k], &weight_host); + let b_mat = weight.restride([StrideSpec::dim(1, k), StrideSpec::dim(0, n)]); + let out = a.mat_mul(&b_mat); + let result = out.as_slice::<2, f32>().await.unwrap(); + let mut mismatches = 0usize; + let mut first = None; + for mi in 0..m { + for ni in 0..n { + let got = result[[mi, ni]]; + let exp = expected[mi * n + ni]; + if (got - exp).abs() > 1e-3 * exp.abs().max(1.0) { + mismatches += 1; + if first.is_none() { + first = Some((mi, ni, got, exp)); + } + } + } + } + println!("round {round}: mismatches={mismatches} first={first:?}"); + if mismatches > 0 { + bad += 1; + } + } + println!("SUMMARY bad={bad}/10"); + assert_eq!(bad, 0); + }); +} diff --git a/fusor-ml/core/tests/fused_reduce.rs b/fusor-ml/core/tests/fused_reduce.rs index f819f6fe7..287b06570 100644 --- a/fusor-ml/core/tests/fused_reduce.rs +++ b/fusor-ml/core/tests/fused_reduce.rs @@ -43,9 +43,8 @@ fn broadcast_composed_contraction_fuses_to_single_kernel() { let (a3, b3) = broadcast_factors(&device, m, n, k, &a_values, &b_values); let out = (&a3 * &b3).sum(2); - assert_eq!( - out.count_kernels_to_resolve(), - 1, + assert!( + out.resolves_in::<1>(), "broadcast-composed contraction must fuse into one map-reduce kernel" ); @@ -80,7 +79,7 @@ fn small_composed_contraction_fuses_through_serial_path() { let (a3, b3) = broadcast_factors(&device, m, n, k, &a_values, &b_values); let out = (&a3 * &b3).sum(2); - assert_eq!(out.count_kernels_to_resolve(), 1, "small fused contraction"); + assert!(out.resolves_in::<1>(), "small fused contraction"); let slice = out.as_slice::<2, f32>().await.unwrap(); for row in 0..m { @@ -114,7 +113,7 @@ fn max_contraction_fuses_with_masked_tiles() { let (a3, b3) = broadcast_factors(&device, m, n, k, &a_values, &b_values); let out = (&a3 * &b3).max(2); - assert_eq!(out.count_kernels_to_resolve(), 1, "fused max contraction"); + assert!(out.resolves_in::<1>(), "fused max contraction"); let slice = out.as_slice::<2, f32>().await.unwrap(); for row in [0usize, 39, m - 1] { @@ -168,7 +167,7 @@ fn quantized_weighted_reduce_fuses_to_single_kernel() { let wd = w.dequantize::(); let xb = x.reshape([1, k]).broadcast_as([n, k]); let out = (&wd * &xb).sum(1); - assert_eq!(out.count_kernels_to_resolve(), 1, "fused quantized reduce"); + assert!(out.resolves_in::<1>(), "fused quantized reduce"); let slice = out.as_slice::<1, f32>().await.unwrap(); for row in [0usize, 17, n - 1] { @@ -192,9 +191,8 @@ fn weighted_sum_fuses_to_single_kernel() { }; // `sum_k w[k] * x[m, k]`: the k-dependent inputs share no (row, col) - // pair. Lowers serially by default; with the cache threshold forced - // down (FUSOR_LAST_LEVEL_CACHE_BYTES) it takes the 1D register - // tiling — correct either way. + // pair. The automatic dispatch policy decides whether to use the 1D + // register tiling; both legal lowerings are reference-correct. let (m, k) = (512usize, 64usize); let x_values = pattern(m * k, 0.13); let w_values = pattern(k, 0.07); @@ -203,7 +201,7 @@ fn weighted_sum_fuses_to_single_kernel() { let w2 = w.reshape([1, k]).broadcast_as([m, k]); let out = (&x * &w2).sum(1); - assert_eq!(out.count_kernels_to_resolve(), 1, "fused weighted sum"); + assert!(out.resolves_in::<1>(), "fused weighted sum"); let slice = out.as_slice::<1, f32>().await.unwrap(); for row in [0usize, 100, 255, m - 1] { @@ -235,7 +233,7 @@ fn broadcast_table_elementwise_reuses_invariant_loads() { let t3 = t.reshape([1, s, h]).broadcast_as([b, s, h]); let out = &x * &t3; - assert_eq!(out.count_kernels_to_resolve(), 1, "broadcast table apply"); + assert!(out.resolves_in::<1>(), "broadcast table apply"); let slice = out.as_slice::<3, f32>().await.unwrap(); for batch in 0..b { @@ -269,9 +267,8 @@ fn contraction_with_k_independent_factor_fuses() { let c3 = c.reshape([m, n, 1]).broadcast_as([m, n, k]); let out = (&(&a3 * &b3) * &c3).sum(2); - assert_eq!( - out.count_kernels_to_resolve(), - 1, + assert!( + out.resolves_in::<1>(), "contraction with k-independent factor" ); @@ -292,3 +289,107 @@ fn contraction_with_k_independent_factor_fuses() { } }); } + +#[test] +fn reshaped_elementwise_producer_folds_into_reduce() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // The broadcast-gradient pattern: an elementwise product reshaped + // flat and then summed. The reshape view folds into the reduce and + // the producer inlines through the composed index expressions, so + // the whole chain is one dispatch. + let (b, s, h) = (4usize, 8usize, 16usize); + let x_values = pattern(b * s * h, 0.13); + let y_values = pattern(b * s * h, 0.29); + let x = Tensor::from_slice(&device, [b, s, h], &x_values); + let y = Tensor::from_slice(&device, [b, s, h], &y_values); + + let out = (&x * &y).reshape([b * s, h]).sum(0); + assert!( + out.resolves_in::<1>(), + "reshape + sum over an exclusive producer must fuse" + ); + + let slice = out.as_slice::<1, f32>().await.unwrap(); + for col in 0..h { + let expected: f32 = (0..b * s) + .map(|row| x_values[row * h + col] * y_values[row * h + col]) + .sum(); + let actual = slice[[col]]; + assert!( + (actual - expected).abs() < 1e-3, + "[{col}]: got {actual}, expected {expected}" + ); + } + }); +} + +#[test] +fn unary_chain_across_keepdim_view_fuses_into_reduce() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // `sum_keepdim(x, 1) / k` applies the division at the keepdim'd + // `[m, 1]` shape, behind the unsqueeze view. The chain rewrites into + // a reduce of its own (indices substituted into the row dims, the + // axis appended) instead of a separate scalar dispatch. + let (m, k) = (64usize, 48usize); + let x_values = pattern(m * k, 0.23); + let x = Tensor::from_slice(&device, [m, k], &x_values); + + let mean = &x.sum_keepdim(1) / (k as f32); + assert!( + mean.resolves_in::<1>(), + "post-keepdim unary chain must fold into the reduce" + ); + + let slice = mean.as_slice::<2, f32>().await.unwrap(); + for row in [0usize, 17, m - 1] { + let expected: f32 = (0..k).map(|col| x_values[row * k + col]).sum::() / k as f32; + let actual = slice[[row, 0]]; + assert!( + (actual - expected).abs() < 1e-4, + "[{row}]: got {actual}, expected {expected}" + ); + } + }); +} + +#[test] +fn unit_axis_reduce_collapses_into_consumer() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // Reducing a size-1 axis is the identity: the node collapses to an + // elementwise gather and inlines into its (binary) consumer, which a + // reduce post chain could never absorb. + let (m, k) = (32usize, 24usize); + let x_values = pattern(m * k, 0.11); + let y_values = pattern(m * k, 0.31); + let x = Tensor::from_slice(&device, [1, m, k], &x_values); + let y = Tensor::from_slice(&device, [m, k], &y_values); + + let out = &x.sum(0) + &y; + assert!( + out.resolves_in::<1>(), + "size-1-axis reduce must collapse into the consumer" + ); + + let slice = out.as_slice::<2, f32>().await.unwrap(); + for (row, col) in [(0usize, 0usize), (13, 7), (m - 1, k - 1)] { + let expected = x_values[row * k + col] + y_values[row * k + col]; + let actual = slice[[row, col]]; + assert!( + (actual - expected).abs() < 1e-5, + "[{row}, {col}]: got {actual}, expected {expected}" + ); + } + }); +} diff --git a/fusor-ml/core/tests/recognition.rs b/fusor-ml/core/tests/recognition.rs index c0cc06d15..3a0f4fe45 100644 --- a/fusor-ml/core/tests/recognition.rs +++ b/fusor-ml/core/tests/recognition.rs @@ -52,7 +52,7 @@ fn composed_dense_matmul_resolves_to_single_kernel() { let a = Tensor::new::(&device, &[[1.0f32, 2.0], [3.0, 4.0]]); let b = Tensor::new::(&device, &[[5.0f32, 6.0], [7.0, 8.0]]); let out = a.mat_mul(&b); - assert_eq!(out.count_kernels_to_resolve(), 1, "dense matmul"); + assert!(out.resolves_in::<1>(), "dense matmul"); let slice = out.as_slice::<2, f32>().await.unwrap(); assert_eq!(slice[[0, 0]], 19.0); assert_eq!(slice[[0, 1]], 22.0); @@ -68,7 +68,7 @@ fn composed_dense_matmul_resolves_to_single_kernel() { &[[[1.0f32, 0.0], [0.0, 1.0]], [[2.0, 0.0], [0.0, 2.0]]], ); let out = a.mat_mul(&b); - assert_eq!(out.count_kernels_to_resolve(), 1, "batched matmul"); + assert!(out.resolves_in::<1>(), "batched matmul"); let slice = out.as_slice::<3, f32>().await.unwrap(); assert_eq!(slice[[0, 0, 0]], 1.0); assert_eq!(slice[[1, 0, 0]], 10.0); @@ -76,6 +76,48 @@ fn composed_dense_matmul_resolves_to_single_kernel() { }); } +#[test] +fn cooperative_dense_matmul_keeps_unary_chains_in_one_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + const M: usize = 61; + const K: usize = 63; + const N: usize = 59; + let a_data: Vec = (0..M * K) + .map(|index| (index as f32 * 0.013).sin() * 0.1) + .collect(); + let b_data: Vec = (0..K * N) + .map(|index| (index as f32 * 0.017).cos() * 0.1) + .collect(); + let a = Tensor::from_slice(&device, [M, K], &a_data); + let b = Tensor::from_slice(&device, [K, N], &b_data); + let matmul = (-a).mat_mul(&b.sin()); + let out = matmul.cos() + 1.0; + assert!( + out.resolves_in::<1>(), + "cooperative matmul should host both pre and post chains" + ); + + let actual = out.as_slice::<2, f32>().await.unwrap(); + for row in 0..M { + for col in 0..N { + let mut sum = 0.0f32; + for k in 0..K { + sum += -a_data[row * K + k] * b_data[k * N + col].sin(); + } + let expected = sum.cos() + 1.0; + let got = actual[[row, col]]; + assert!( + (got - expected).abs() < 1e-3, + "[{row}, {col}]: got {got}, expected {expected}" + ); + } + } + }); +} + #[test] fn composed_qmatmul_resolves_to_single_kernel() { pollster::block_on(async { @@ -87,7 +129,7 @@ fn composed_qmatmul_resolves_to_single_kernel() { let w = f32_weight(&device, 4, K); let x = Tensor::new::(&device, &[[1.0f32, 0.5, -1.0, 2.0, 0.0, 1.0, -0.5, 3.0]]); let out = x.q_mat_mul(&w); - assert_eq!(out.count_kernels_to_resolve(), 1, "bare quantized matmul"); + assert!(out.resolves_in::<1>(), "bare quantized matmul"); }); } @@ -107,11 +149,7 @@ fn composed_q8_qmatmul_with_epilogue_resolves_to_single_kernel() { let x = Tensor::from_slice(&device, [1, K], &x_values); let bias = Tensor::from_slice(&device, [1, N], &bias_values); let out = x.q_mat_mul(&w) + &bias; - assert_eq!( - out.count_kernels_to_resolve(), - 1, - "q8_0 matmul + bias epilogue" - ); + assert!(out.resolves_in::<1>(), "q8_0 matmul + bias epilogue"); let slice = out.as_slice::<2, f32>().await.unwrap(); for col in 0..N { @@ -172,7 +210,7 @@ fn composed_softmax_resolves_to_fused_kernel() { let values: Vec = (0..256).map(|i| ((i as f32) * 0.1).sin()).collect(); let x = Tensor::from_slice(&device, [2, 128], &values); let out = x.softmax(1); - assert_eq!(out.count_kernels_to_resolve(), 1, "single-pass softmax"); + assert!(out.resolves_in::<1>(), "single-pass softmax"); let slice = out.as_slice::<2, f32>().await.unwrap(); for row in 0..2 { @@ -205,7 +243,7 @@ fn composed_rms_norm_with_bias_resolves_to_fused_kernel() { let weight = Tensor::from_slice(&device, [4], &[0.5f32, 1.0, 1.5, 2.0]); let bias = Tensor::from_slice(&device, [4], &[0.1f32, -0.2, 0.3, -0.4]); let out = x.rms_norm_fused(&weight, Some(&bias), 1e-5); - assert_eq!(out.count_kernels_to_resolve(), 1, "rms norm with bias"); + assert!(out.resolves_in::<1>(), "rms norm with bias"); let slice = out.as_slice::<2, f32>().await.unwrap(); let mean_square = (1.0 + 4.0 + 9.0 + 16.0) / 4.0; @@ -224,7 +262,7 @@ fn composed_rms_norm_with_bias_resolves_to_fused_kernel() { } #[test] -fn composed_attention_resolves_to_flash_kernel() { +fn composed_attention_resolves_to_attention_kernel() { pollster::block_on(async { let Ok(device) = Device::new().await else { return; @@ -244,11 +282,10 @@ fn composed_attention_resolves_to_flash_kernel() { let v = Tensor::from_slice(&device, [1, kv_heads, kv_len, d], &v_data); let scale = 1.0 / (d as f32).sqrt(); - let out = q.flash_attention(&k, &v, scale, None); - assert_eq!( - out.count_kernels_to_resolve(), - 1, - "gqa decode attention should recognize as one fused flash kernel" + let out = q.attention(&k, &v, scale, None); + assert!( + out.resolves_in::<1>(), + "gqa decode attention should recognize as one fused attention kernel" ); let slice = out.as_slice::<4, f32>().await.unwrap(); diff --git a/fusor-ml/core/tests/row_fusion.rs b/fusor-ml/core/tests/row_fusion.rs new file mode 100644 index 000000000..fea35b1ed --- /dev/null +++ b/fusor-ml/core/tests/row_fusion.rs @@ -0,0 +1,99 @@ +//! Aggressive row-fusion regression. Interleaved-view scalar absorption is +//! part of the one default optimizer path for every graph size. + +use fusor_core::{Device, Tensor}; + +fn pattern(len: usize, scale: f32) -> Vec { + (0..len).map(|i| ((i as f32) * scale).sin()).collect() +} + +#[test] +fn keepdim_scalar_chain_sandwich_fuses_to_single_kernel() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + + // Layer-norm-forward shape: `x - broadcast(sum_keepdim(x) / h)`. + // The `sum_keepdim` unsqueeze sits *between* the scalar division and + // the reduce, so absorbing the mean as a row phase requires the + // interleaved-view scalar walk. + let (b, s, h) = (4usize, 8usize, 64usize); + let x_values = pattern(b * s * h, 0.17); + let x = Tensor::from_slice(&device, [b, s, h], &x_values); + + let mean = &x.sum_keepdim(2) / (h as f32); + let centered = &x - &mean.broadcast_as([b, s, h]); + assert!( + centered.resolves_in::<1>(), + "keepdim mean chain must fold into one row program" + ); + + let slice = centered.as_slice::<3, f32>().await.unwrap(); + for (batch, row) in [(0usize, 0usize), (1, 5), (b - 1, s - 1)] { + let base = batch * s * h + row * h; + let mean: f32 = (0..h).map(|col| x_values[base + col]).sum::() / h as f32; + for col in [0usize, 31, h - 1] { + let expected = x_values[base + col] - mean; + let actual = slice[[batch, row, col]]; + assert!( + (actual - expected).abs() < 1e-4, + "[{batch}, {row}, {col}]: got {actual}, expected {expected}" + ); + } + } + }); +} + +/// Lane groups narrower than the workgroup make the axis stride in chunks and +/// pin one row per group of lanes, so every axis length exercises a different +/// (group width, chunk count, masked tail) triple. Sweep the widths around the +/// subgroup width and the chunking thresholds against a host reference. +#[test] +fn row_reductions_match_host_across_axis_widths() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + // The no-subgroup sibling takes the shared-memory reduction tree, the + // only path the web build has. + sweep_axis_widths(&device).await; + sweep_axis_widths(&device.without_subgroups()).await; + }); +} + +async fn sweep_axis_widths(device: &Device) { + const ROWS: usize = 37; + for k in [ + 1usize, 2, 3, 8, 16, 31, 32, 33, 63, 64, 65, 100, 128, 129, 200, 256, 257, 384, 512, 1000, + ] { + let values = pattern(ROWS * k, 0.31); + let x = Tensor::from_slice(device, [ROWS, k], &values); + + let sums = x.sum(1).as_slice::<1, f32>().await.unwrap(); + let softmax = x.softmax_last_dim().as_slice::<2, f32>().await.unwrap(); + for row in [0usize, 1, ROWS / 2, ROWS - 1] { + let base = row * k; + let span = &values[base..base + k]; + + let expected: f32 = span.iter().sum(); + let actual = sums[[row]]; + let tolerance = 1e-4 * expected.abs().max(1.0); + assert!( + (actual - expected).abs() < tolerance, + "k={k} row={row}: sum {actual}, expected {expected}" + ); + + let max = span.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let denominator: f32 = span.iter().map(|v| (v - max).exp()).sum(); + for col in [0usize, k / 2, k - 1] { + let expected = (span[col] - max).exp() / denominator; + let actual = softmax[[row, col]]; + assert!( + (actual - expected).abs() < 1e-5, + "k={k} row={row} col={col}: softmax {actual}, expected {expected}" + ); + } + } + } +} diff --git a/fusor-ml/core/tests/small_tile_matmul.rs b/fusor-ml/core/tests/small_tile_matmul.rs new file mode 100644 index 000000000..c92436147 --- /dev/null +++ b/fusor-ml/core/tests/small_tile_matmul.rs @@ -0,0 +1,178 @@ +//! Correctness gates for the small-side cooperative-matrix tiles +//! ((64, 16) and (16, 64)): contractions with a 16-wide (or 16-padded) +//! M or N side — the attention head_dim family and narrow-vocab lm_head +//! shapes — now route to the coop kernel instead of the generic reduce. +//! Each case A/Bs the GPU result against a CPU reference over exact-multiple +//! shapes, masked-edge shapes (M, N, and K edges), batched forms, and a +//! transposed-B operand. The (128, 128) precedent showed a tile table entry +//! can satisfy the fragment-size rules yet miscompute, so every new geometry +//! is exercised here. + +use fusor_core::{Device, Layout, Tensor}; + +fn values(len: usize, freq: f32) -> Vec { + (0..len).map(|i| ((i as f32) * freq).sin()).collect() +} + +/// f64-accumulated reference for `a[batch, m, k] @ b[batch, k, n]`. +fn cpu_matmul(a: &[f32], b: &[f32], batch: usize, m: usize, k: usize, n: usize) -> Vec { + let mut out = vec![0.0f32; batch * m * n]; + for bi in 0..batch { + for mi in 0..m { + for ni in 0..n { + let mut acc = 0.0f64; + for ki in 0..k { + let a_val = a[(bi * m + mi) * k + ki] as f64; + let b_val = b[(bi * k + ki) * n + ni] as f64; + acc += a_val * b_val; + } + out[(bi * m + mi) * n + ni] = acc as f32; + } + } + } + out +} + +fn check_matmul(batch: usize, m: usize, k: usize, n: usize, transpose_b: bool) { + check_matmul_tol(batch, m, k, n, transpose_b, 1e-3); +} + +fn check_matmul_tol(batch: usize, m: usize, k: usize, n: usize, transpose_b: bool, tol: f32) { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let a_data = values(batch * m * k, 0.13); + let b_data = values(batch * k * n, 0.07); + + let a = Tensor::from_slice(&device, [batch, m, k], &a_data); + let b = if transpose_b { + // Store B in [batch, n, k] order and view it as [batch, k, n]: + // the transposed-operand geometry every dA/dB backward matmul + // reads through. + let b_t_data: Vec = (0..batch * n * k) + .map(|i| { + let (bi, rest) = (i / (n * k), i % (n * k)); + let (ni, ki) = (rest / k, rest % k); + b_data[(bi * k + ki) * n + ni] + }) + .collect(); + let b_t = Tensor::from_slice(&device, [batch, n, k], &b_t_data); + b_t.restride_layout(Layout::from_parts( + 0, + vec![batch, k, n].into(), + vec![n * k, 1, k].into(), + )) + } else { + Tensor::from_slice(&device, [batch, k, n], &b_data) + }; + + let out = a.mat_mul(&b); + assert!( + out.resolves_in::<1>(), + "batch={batch} m={m} k={k} n={n} transpose_b={transpose_b}: \ + the contraction should resolve to one matmul kernel" + ); + let actual = out.as_slice::<3, f32>().await.unwrap(); + let expected = cpu_matmul(&a_data, &b_data, batch, m, k, n); + for bi in 0..batch { + for mi in 0..m { + for ni in 0..n { + let want = expected[(bi * m + mi) * n + ni]; + let got = actual[[bi, mi, ni]]; + assert!( + (got - want).abs() < tol + want.abs() * 1e-3, + "batch={batch} m={m} k={k} n={n} transpose_b={transpose_b} \ + [{bi}, {mi}, {ni}]: got {got}, expected {want}" + ); + } + } + } + }); +} + +// (64, 16): exact-multiple shapes — the P@V / dQ / dV attention family. +#[test] +fn small_tile_64x16_exact() { + check_matmul(4, 64, 64, 16, false); +} + +// (64, 16): every edge masked at once (M 120→128, N 14→16, K=50 tail tile). +#[test] +fn small_tile_64x16_masked_edges() { + check_matmul(3, 120, 50, 14, false); +} + +// (64, 16): N-edge padding on the vocab-65 lm_head geometry (65 → 80). +#[test] +fn small_tile_64x16_vocab_edge() { + check_matmul(1, 192, 64, 65, false); +} + +// (64, 16): transposed-B operand (the dS @ Kᵀ-style backward reads). +#[test] +fn small_tile_64x16_transposed_b() { + check_matmul(2, 64, 64, 16, true); +} + +// (16, 64): exact-multiple shapes — the Qᵀ@dS attention family. +#[test] +fn small_tile_16x64_exact() { + check_matmul(4, 16, 64, 64, false); +} + +// (16, 64): M-edge padding (65 → 80) with a masked K tail. +#[test] +fn small_tile_16x64_masked_edges() { + check_matmul(2, 65, 50, 64, false); +} + +// (16, 64): transposed-B operand. +#[test] +fn small_tile_16x64_transposed_b() { + check_matmul(2, 16, 64, 64, true); +} + +// (16, 64): long-K single-tile-grid shape (the lm_head gradient geometry); +// with K = 2048 this also routes through the split-K partials + combine +// sequence. +#[test] +fn small_tile_16x64_long_k() { + check_matmul(1, 65, 2048, 64, false); +} + +// Split-K A/B gates: starved tile grids with a long contraction run the +// partials + combine two-kernel sequence; the result must match the +// f64-accumulated reference within sum-reorder tolerance. + +// The 64×2048×64 weight-gradient shape: one (64, 64) tile, 16 K-spans. +#[test] +fn split_k_weight_grad_square() { + check_matmul_tol(1, 64, 2048, 64, false, 2e-3); +} + +// The 64×2048×256 shape: two (64, 128) tiles across N, 16 K-spans. +#[test] +fn split_k_weight_grad_wide() { + check_matmul_tol(1, 64, 2048, 256, false, 2e-3); +} + +// K not a multiple of the span width: the last split's tiles read past the +// logical K extent and must fill zero (k = 1000 → 63 K-tiles over 16 spans). +#[test] +fn split_k_ragged_k() { + check_matmul_tol(1, 64, 1000, 64, false, 2e-3); +} + +// Barely past the split gate (k = 520 → 33 K-tiles, the last spans idle), +// with a transposed-B operand. +#[test] +fn split_k_short_spans_transposed_b() { + check_matmul_tol(1, 64, 520, 64, true, 2e-3); +} + +// Batched split-K: scratch rows interleave (split, batch) correctly. +#[test] +fn split_k_batched() { + check_matmul_tol(3, 64, 640, 64, false, 2e-3); +} diff --git a/fusor-ml/core/tests/windowed_repro.rs b/fusor-ml/core/tests/windowed_repro.rs new file mode 100644 index 000000000..1644490b2 --- /dev/null +++ b/fusor-ml/core/tests/windowed_repro.rs @@ -0,0 +1,72 @@ +//! Temporary diagnostic: windowed (im2col) 2048x576x128 matmul only. + +use fusor_core::{Device, StrideSpec, Tensor}; + +#[test] +fn windowed_2048_576_128_repro() { + pollster::block_on(async { + let Ok(device) = Device::new().await else { + return; + }; + let (b, c, h, w, n, kh, kw) = (2usize, 64, 34, 34, 128, 3, 3); + let (oh, ow) = (h - kh + 1, w - kw + 1); + let (m, k) = (b * oh * ow, c * kh * kw); + let input_host: Vec = (0..b * c * h * w).map(|i| (i % 13) as f32 * 0.1).collect(); + let weight_host: Vec = (0..n * k).map(|i| (i % 7) as f32 * 0.01).collect(); + let mut expected = vec![0f32; m * n]; + for mi in 0..m { + let (bi, rest) = (mi / (oh * ow), mi % (oh * ow)); + let (ohi, owi) = (rest / ow, rest % ow); + for ni in 0..n { + let mut acc = 0f32; + for ci in 0..c { + for khi in 0..kh { + for kwi in 0..kw { + acc += input_host[((bi * c + ci) * h + ohi + khi) * w + owi + kwi] + * weight_host[ni * k + (ci * kh + khi) * kw + kwi]; + } + } + } + expected[mi * n + ni] = acc; + } + } + + let mut bad = 0usize; + for round in 0..10 { + let input = Tensor::from_slice(&device, [b, c, h, w], &input_host); + let weight = Tensor::from_slice(&device, [n, k], &weight_host); + let windows = input.restride([ + StrideSpec::dim(0, b), + StrideSpec::dim_with(2, oh, 1), + StrideSpec::dim_with(3, ow, 1), + StrideSpec::dim(1, c), + StrideSpec::dim(2, kh), + StrideSpec::dim(3, kw), + ]); + let a = windows.reshape([m, k]); + let b_mat = weight.restride([StrideSpec::dim(1, k), StrideSpec::dim(0, n)]); + let out = a.mat_mul(&b_mat); + let result = out.as_slice::<2, f32>().await.unwrap(); + let mut mismatches = 0usize; + let mut first = None; + for mi in 0..m { + for ni in 0..n { + let got = result[[mi, ni]]; + let exp = expected[mi * n + ni]; + if (got - exp).abs() > 1e-3 * exp.abs().max(1.0) { + mismatches += 1; + if first.is_none() { + first = Some((mi, ni, got, exp)); + } + } + } + } + println!("round {round}: mismatches={mismatches} first={first:?}"); + if mismatches > 0 { + bad += 1; + } + } + println!("SUMMARY bad={bad}/10"); + assert_eq!(bad, 0); + }); +} diff --git a/fusor-ml/cpu/src/cast.rs b/fusor-ml/cpu/src/cast.rs index 3e8f2b01a..4cc1dfc49 100644 --- a/fusor-ml/cpu/src/cast.rs +++ b/fusor-ml/cpu/src/cast.rs @@ -87,6 +87,13 @@ impl CastTo for f64 { } } +impl CastTo for u32 { + #[inline(always)] + fn cast(self) -> half::f16 { + half::f16::from_f32(self as f32) + } +} + /// Cast a tensor from one element type to another pub(crate) fn cast_tensor( input: &ConcreteTensor, diff --git a/fusor-ml/cpu/src/comparison.rs b/fusor-ml/cpu/src/comparison.rs index eadaedf71..c10271afc 100644 --- a/fusor-ml/cpu/src/comparison.rs +++ b/fusor-ml/cpu/src/comparison.rs @@ -175,6 +175,45 @@ macro_rules! impl_all_comparisons { impl_all_comparisons!(f32, f64, i8, i16, i32, i64, u8, u16, u32, u64); +// f16 comparisons through the single-lane scalar wrapper, mirroring the f16 +// pairwise ops: pulp has no native f16 SIMD. +macro_rules! impl_f16_comparison_op { + ($op:ty, $cmp:expr) => { + impl SimdBinaryOp for $op { + #[inline(always)] + fn apply_simd_vec( + _simd: S, + a: crate::F16Scalar, + b: crate::F16Scalar, + ) -> crate::F16Scalar { + let cmp: fn(half::f16, half::f16) -> bool = $cmp; + crate::F16Scalar(if cmp(a.0, b.0) { + half::f16::ONE + } else { + half::f16::ZERO + }) + } + + #[inline(always)] + fn apply_scalar(a: half::f16, b: half::f16) -> half::f16 { + let cmp: fn(half::f16, half::f16) -> bool = $cmp; + if cmp(a, b) { + half::f16::ONE + } else { + half::f16::ZERO + } + } + } + }; +} + +impl_f16_comparison_op!(EqOp, |a, b| a == b); +impl_f16_comparison_op!(NeOp, |a, b| a != b); +impl_f16_comparison_op!(LtOp, |a, b| a < b); +impl_f16_comparison_op!(LteOp, |a, b| a <= b); +impl_f16_comparison_op!(GtOp, |a, b| a > b); +impl_f16_comparison_op!(GteOp, |a, b| a >= b); + // Comparison tensor expression types define_tensor_op!(@binary Eq, EqOp); define_tensor_op!(@binary Ne, NeOp); diff --git a/fusor-ml/cpu/src/conditional.rs b/fusor-ml/cpu/src/conditional.rs index 0e85c1277..7b6a9f0ee 100644 --- a/fusor-ml/cpu/src/conditional.rs +++ b/fusor-ml/cpu/src/conditional.rs @@ -22,7 +22,7 @@ macro_rules! impl_is_nonzero { } impl_is_nonzero!( - f32 => 0.0, f64 => 0.0, + f32 => 0.0, f64 => 0.0, half::f16 => half::f16::ZERO, i8 => 0, i16 => 0, i32 => 0, i64 => 0, u8 => 0, u16 => 0, u32 => 0, u64 => 0 ); diff --git a/fusor-ml/cpu/src/tensor.rs b/fusor-ml/cpu/src/tensor.rs index c91eccd57..be8b929f2 100644 --- a/fusor-ml/cpu/src/tensor.rs +++ b/fusor-ml/cpu/src/tensor.rs @@ -1082,6 +1082,7 @@ where pub trait Scalar: Copy {} impl Scalar for f32 {} +impl Scalar for half::f16 {} impl Scalar for f64 {} impl Scalar for i8 {} impl Scalar for i16 {} diff --git a/fusor-ml/fusor/Cargo.toml b/fusor-ml/fusor/Cargo.toml index 54702077b..7a90597e0 100644 --- a/fusor-ml/fusor/Cargo.toml +++ b/fusor-ml/fusor/Cargo.toml @@ -31,6 +31,11 @@ futures = "0.3.31" fusor-core = { path = "../core" } fusor-gguf = { path = "../gguf" } ndarray = "0.16.1" +ureq = { version = "3.2.0", default-features = false, features = ["rustls"] } +flate2 = "1.0.35" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +wgpu = { git = "https://github.com/ealmloff/wgpu", branch = "yield-now", default-features = false, features = ["std", "naga-ir", "wgsl"] } +futures-channel = "0.3" [lib] bench = false diff --git a/fusor-ml/fusor/benches/attention_comparison.rs b/fusor-ml/fusor/benches/attention_comparison.rs index 62300429e..4750701d1 100644 --- a/fusor-ml/fusor/benches/attention_comparison.rs +++ b/fusor-ml/fusor/benches/attention_comparison.rs @@ -123,18 +123,18 @@ fn bench_backend( }, ); - let device_flash = device.clone(); - let q_data_flash = q_data.to_vec(); - let k_data_flash = k_data.to_vec(); - let v_data_flash = v_data.to_vec(); + let device_attention = device.clone(); + let q_data_attention = q_data.to_vec(); + let k_data_attention = k_data.to_vec(); + let v_data_attention = v_data.to_vec(); group.bench_with_input( - BenchmarkId::new(format!("{backend}_flash"), size_str), + BenchmarkId::new(format!("{backend}_attention"), size_str), &shape, move |b, &shape| { - let device = device_flash.clone(); - let q_data = q_data_flash.clone(); - let k_data = k_data_flash.clone(); - let v_data = v_data_flash.clone(); + let device = device_attention.clone(); + let q_data = q_data_attention.clone(); + let k_data = k_data_attention.clone(); + let v_data = v_data_attention.clone(); b.to_async(FuturesExecutor).iter_custom(move |iters| { let device = device.clone(); @@ -149,7 +149,7 @@ fn bench_backend( while total.is_zero() { for _ in 0..iters { let start = std::time::Instant::now(); - let output = q.flash_attention(&k, &v, scale, None); + let output = q.attention(&k, &v, scale, None); resolve_tensor(&output).await; total += start.elapsed(); } diff --git a/fusor-ml/fusor/benches/bert_layer.rs b/fusor-ml/fusor/benches/bert_layer.rs index 41b8b09b7..c896dbed3 100644 --- a/fusor-ml/fusor/benches/bert_layer.rs +++ b/fusor-ml/fusor/benches/bert_layer.rs @@ -313,12 +313,8 @@ fn self_attention(c: &mut Criterion) { let k = k_reshaped.transpose(1, 2).to_concrete(); let v = v_reshaped.transpose(1, 2).to_concrete(); - let context = q.flash_attention( - &k, - &v, - 1.0 / (head_size as f32).sqrt(), - None, - ); + let context = + q.attention(&k, &v, 1.0 / (head_size as f32).sqrt(), None); let context = context.transpose(1, 2); let output = context.flatten_last_n::<1, _>(); diff --git a/fusor-ml/fusor/examples/bench_ce.rs b/fusor-ml/fusor/examples/bench_ce.rs new file mode 100644 index 000000000..8c846221b --- /dev/null +++ b/fusor-ml/fusor/examples/bench_ce.rs @@ -0,0 +1,75 @@ +//! Composed lm-head + softmax-cross-entropy cost at a real vocabulary size: +//! the baseline the streaming (logits-free) CE kernels must beat. The +//! example transformer's vocab of 65 makes this cluster ~1% of a step; +//! at 32k vocab the materialized logits and their gradient dominate. + +use fusor::autograd::{Graph, Tensor}; +use fusor::{Device, Tensor as RawTensor}; + +fn main() { + if std::env::var_os("RUST_LOG").is_some() { + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .init(); + } + pollster::block_on(async { + let device = Device::gpu().await.expect("gpu device"); + let rows: usize = std::env::var("ROWS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(16384); + let dim: usize = 384; + let vocab: usize = std::env::var("VOCAB") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(32768); + + let mut state = 0x5eed_u64; + let mut next = move || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5 + }; + let x_data: Vec = (0..rows * dim).map(|_| next() * 0.1).collect(); + let w_data: Vec = (0..vocab * dim).map(|_| next() * 0.05).collect(); + let targets: Vec = (0..rows).map(|i| (i * 2654435761 % vocab) as u32).collect(); + let targets = RawTensor::from_slice(&device, [rows], &targets); + + let step = || { + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [rows, dim], &x_data); + let w: Tensor<2> = Tensor::from_slice(&graph, &device, [vocab, dim], &w_data); + let logits = x.mat_mul_transposed_rhs(&w); + let loss = logits.softmax_cross_entropy(&targets); + let gradients = loss.backward().unwrap(); + let dw = gradients.get(&w).expect("dW"); + let dx = gradients.get(&x).expect("dX"); + (loss, dx, dw) + }; + + // Warm (compile kernels). + { + let (_loss, dx, dw) = step(); + pollster::block_on(dx.materialize()); + pollster::block_on(dw.materialize()); + } + let iters = 5; + let mut best = f64::MAX; + for _ in 0..3 { + let start = std::time::Instant::now(); + for _ in 0..iters { + let (_loss, dx, dw) = step(); + pollster::block_on(dx.materialize()); + pollster::block_on(dw.materialize()); + } + best = best.min(start.elapsed().as_secs_f64() / iters as f64); + } + let logits_bytes = (rows * vocab * 4) as f64; + println!( + "composed lm_head+CE fwd+bwd rows={rows} dim={dim} vocab={vocab}: {:.2} ms/iter (logits slab {:.2} GB)", + best * 1e3, + logits_bytes / 1e9 + ); + }); +} diff --git a/fusor-ml/fusor/examples/kernel_bench.rs b/fusor-ml/fusor/examples/kernel_bench.rs new file mode 100644 index 000000000..2899ddaa2 --- /dev/null +++ b/fusor-ml/fusor/examples/kernel_bench.rs @@ -0,0 +1,207 @@ +//! Microbenchmark for individual GPU kernel shapes — mirrors the transformer +//! training profile measurement mode. Runs batches of independent operations +//! per resolve; the resolver's GPU timestamp profile is read back +//! programmatically and printed as machine-diffable `kernel_profile` / +//! `kernel_profile_category` lines. Only the first two resolves per process +//! are profiled (later rounds replay the recorded materialization plan), so +//! run each case's process three times for six samples. +//! +//! Usage: +//! cargo run --release -p fusor --example kernel_bench -- [case] + +use fusor::{Device, FusorConfig, Tensor}; + +const REPEATS: usize = 16; +const ROUNDS: usize = 6; + +fn fill(seed: u32, len: usize) -> Vec { + let mut state = seed as u64 | 1; + (0..len) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5 + }) + .collect() +} + +fn print_profiles(case: &str, device: &Device) { + let Some(gpu) = device.as_gpu() else { + return; + }; + for profile in gpu.take_kernel_profiles() { + println!( + "kernel_profile case={case} mode={} kernels={} unmeasured={} accounted_ms={:.3} span_ms={}", + profile.timestamp_mode, + profile.kernels, + profile.unmeasured_kernels, + profile.accounted_ms, + profile + .span_ms + .map_or_else(|| "absent".to_string(), |span| format!("{span:.3}")) + ); + let mut categories = profile.categories; + categories.sort_by(|a, b| a.name.cmp(&b.name)); + for row in categories { + println!( + "kernel_profile_category case={case} category={} count={} total_ms={:.3} avg_us={:.1} max_us={:.1}", + row.name, row.count, row.total_ms, row.average_us, row.max_us + ); + } + } +} + +async fn bench_matmul(case: &str, device: &Device, m: usize, k: usize, n: usize) { + bench_matmul_repeats(case, device, m, k, n, REPEATS).await +} + +async fn bench_matmul_repeats( + case: &str, + device: &Device, + m: usize, + k: usize, + n: usize, + repeats: usize, +) { + println!("=== matmul {m}x{k} by {k}x{n} x{repeats} ==="); + for _ in 0..ROUNDS { + let outputs: Vec> = (0..repeats) + .map(|i| { + let a = Tensor::from_slice(device, [m, k], &fill(3 + i as u32, m * k)); + let b = Tensor::from_slice(device, [k, n], &fill(77 + i as u32, k * n)); + a.mat_mul(&b) + }) + .collect(); + device.flush(); + let mut total = 0.0f32; + for out in outputs { + let slice = out.as_slice().await.unwrap(); + total += slice[[0, 0]]; + } + println!("checksum {total}"); + print_profiles(case, device); + } +} + +async fn bench_batched_matmul( + case: &str, + device: &Device, + b0: usize, + b1: usize, + m: usize, + k: usize, + n: usize, +) { + println!("=== batched matmul [{b0},{b1}] {m}x{k} by {k}x{n} ==="); + for _ in 0..ROUNDS { + let outputs: Vec> = (0..REPEATS) + .map(|i| { + let a = Tensor::from_slice( + device, + [b0, b1, m, k], + &fill(3 + i as u32, b0 * b1 * m * k), + ); + let b = Tensor::from_slice( + device, + [b0, b1, k, n], + &fill(77 + i as u32, b0 * b1 * k * n), + ); + a.mat_mul(&b) + }) + .collect(); + device.flush(); + let mut total = 0.0f32; + for out in outputs { + let slice = out.as_slice().await.unwrap(); + total += slice[[0, 0, 0, 0]]; + } + println!("checksum {total}"); + print_profiles(case, device); + } +} + +async fn bench_softmax(case: &str, device: &Device, b0: usize, b1: usize, m: usize, k: usize) { + println!("=== softmax [{b0},{b1},{m},{k}] last dim ==="); + for _ in 0..ROUNDS { + let outputs: Vec> = (0..REPEATS) + .map(|i| { + let x = Tensor::from_slice( + device, + [b0, b1, m, k], + &fill(3 + i as u32, b0 * b1 * m * k), + ); + x.softmax_last_dim() + }) + .collect(); + device.flush(); + let mut total = 0.0f32; + for out in outputs { + let slice = out.as_slice().await.unwrap(); + total += slice[[0, 0, 0, 0]]; + } + println!("checksum {total}"); + print_profiles(case, device); + } +} + +/// Streaming elementwise over a working set far larger than cache: the +/// achievable-bandwidth roof for the roofline model. One add reads two +/// operands and writes one result, so the kernel moves `3 * n * 4` bytes and +/// does one flop per element. +async fn bench_stream(case: &str, device: &Device, n: usize) { + println!("=== stream add {n} elements ==="); + for _ in 0..ROUNDS { + let outputs: Vec> = (0..2) + .map(|i| { + let a = Tensor::from_slice(device, [n], &fill(3 + i as u32, n)); + let b = Tensor::from_slice(device, [n], &fill(77 + i as u32, n)); + (&a + &b).to_concrete() + }) + .collect(); + device.flush(); + let mut total = 0.0f32; + for out in outputs { + total += out.as_slice().await.unwrap()[[0]]; + } + println!("checksum {total}"); + print_profiles(case, device); + } +} + +#[tokio::main] +async fn main() { + if std::env::var_os("RUST_LOG").is_some() { + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .init(); + } + let mut config = FusorConfig::from_env(); + config.trace_gpu_kernels = true; + let device = Device::gpu_with_config(config).await.unwrap(); + let case = std::env::args().nth(1).unwrap_or_else(|| "all".to_string()); + match case.as_str() { + "wgrad" => bench_matmul("wgrad", &device, 64, 2048, 64).await, + "wgrad256" => bench_matmul("wgrad256", &device, 64, 2048, 256).await, + "wgrad256m" => bench_matmul("wgrad256m", &device, 256, 2048, 64).await, + "fwd" => bench_matmul("fwd", &device, 2048, 64, 64).await, + "fwd256" => bench_matmul("fwd256", &device, 2048, 256, 64).await, + "fwdup" => bench_matmul("fwdup", &device, 2048, 64, 256).await, + "attn" => bench_batched_matmul("attn", &device, 32, 4, 64, 64, 16).await, + // The shape the pre-unification split-K guard was calibrated on: its + // unsplit grid fills ~28% of the device, which the guard's flat + // threshold refused to split. Kept as a standing check that the + // scored split count does not resurrect that loss. + "wgradbig" => bench_matmul_repeats("wgradbig", &device, 384, 16384, 1536, 4).await, + "softmax" => bench_softmax("softmax", &device, 32, 4, 64, 64).await, + // Roofline anchors: the achievable bandwidth and compute roofs this + // machine sustains, measured through the same profile path as the + // shape cases so the comparison uses one clock. + "roof_bw" => bench_stream("roof_bw", &device, 32 << 20).await, + "roof_flops" => bench_matmul_repeats("roof_flops", &device, 2048, 2048, 2048, 4).await, + _ => { + bench_matmul("wgrad", &device, 64, 2048, 64).await; + bench_matmul("fwd", &device, 2048, 64, 64).await; + } + } +} diff --git a/fusor-ml/fusor/examples/mnist.rs b/fusor-ml/fusor/examples/mnist.rs new file mode 100644 index 000000000..52427910f --- /dev/null +++ b/fusor-ml/fusor/examples/mnist.rs @@ -0,0 +1,307 @@ +//! Train a small convolutional network on MNIST with the autograd API. +//! +//! Downloads the MNIST dataset on first run (cached under `examples/data`), +//! then trains conv(1->8) -> pool -> conv(8->16) -> pool -> linear(784->10) +//! with softmax cross-entropy and plain SGD, rebuilding the tape each step. +//! +//! Run with: +//! ```sh +//! cargo run --release --example mnist +//! ``` + +use std::io::Read; +use std::path::PathBuf; + +use fusor::autograd::layers::{ConvNd, ConvNdConfig, Linear}; +use fusor::autograd::{Gradients, Graph, Tensor}; +use fusor::{Device, Tensor as RawTensor, ToVec}; + +const BATCH_SIZE: usize = 64; +const EPOCHS: usize = 1; +const LEARNING_RATE: f32 = 0.05; +const IMAGE_SIZE: usize = 28; +const PIXELS: usize = IMAGE_SIZE * IMAGE_SIZE; +const CLASSES: usize = 10; + +const MNIST_MIRROR: &str = "https://ossci-datasets.s3.amazonaws.com/mnist"; + +fn data_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/data") +} + +/// Download and gunzip one MNIST idx file, caching the decompressed bytes. +fn fetch(name: &str) -> Vec { + let path = data_dir().join(name); + if let Ok(bytes) = std::fs::read(&path) { + return bytes; + } + println!("downloading {name}.gz"); + let mut response = ureq::get(format!("{MNIST_MIRROR}/{name}.gz")) + .call() + .unwrap_or_else(|err| panic!("failed to download {name}: {err}")); + let mut bytes = Vec::new(); + flate2::read::GzDecoder::new(response.body_mut().as_reader()) + .read_to_end(&mut bytes) + .unwrap_or_else(|err| panic!("failed to decompress {name}: {err}")); + std::fs::create_dir_all(data_dir()).unwrap(); + std::fs::write(&path, &bytes).unwrap(); + bytes +} + +fn read_be_u32(bytes: &[u8], offset: usize) -> usize { + u32::from_be_bytes(bytes[offset..offset + 4].try_into().unwrap()) as usize +} + +/// Parse an idx image file into pixels normalized to [0, 1]. +fn load_images(name: &str) -> Vec { + let bytes = fetch(name); + assert_eq!(read_be_u32(&bytes, 0), 0x803, "bad image file magic"); + assert_eq!(read_be_u32(&bytes, 8), IMAGE_SIZE); + assert_eq!(read_be_u32(&bytes, 12), IMAGE_SIZE); + bytes[16..] + .iter() + .map(|&pixel| pixel as f32 / 255.0) + .collect() +} + +fn load_labels(name: &str) -> Vec { + let bytes = fetch(name); + assert_eq!(read_be_u32(&bytes, 0), 0x801, "bad label file magic"); + bytes[8..].iter().map(|&label| label as u32).collect() +} + +/// Deterministic LCG so runs are reproducible without a rand dependency. +struct Lcg(u64); + +impl Lcg { + /// Uniform in (-bound, bound) with `bound = sqrt(1 / fan_in)`, the + /// PyTorch default init for conv and linear layers. + fn kaiming(&mut self, count: usize, fan_in: usize) -> Vec { + let bound = (1.0 / fan_in as f32).sqrt(); + (0..count) + .map(|_| { + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let uniform = (self.0 >> 32) as f32 / (1u64 << 32) as f32 - 0.5; + uniform * 2.0 * bound + }) + .collect() + } +} + +/// Raw parameter tensors that persist across training steps. +struct Params { + conv1_weight: RawTensor<4, f32>, + conv1_bias: RawTensor<1, f32>, + conv2_weight: RawTensor<4, f32>, + conv2_bias: RawTensor<1, f32>, + fc_weight: RawTensor<2, f32>, + fc_bias: RawTensor<1, f32>, +} + +impl Params { + fn new(device: &Device) -> Self { + let mut rng = Lcg(42); + let flat = 16 * (IMAGE_SIZE / 4) * (IMAGE_SIZE / 4); + Self { + conv1_weight: RawTensor::from_slice(device, [8, 1, 3, 3], &rng.kaiming(8 * 9, 9)), + conv1_bias: RawTensor::zeros(device, [8]), + conv2_weight: RawTensor::from_slice( + device, + [16, 8, 3, 3], + &rng.kaiming(16 * 8 * 9, 8 * 9), + ), + conv2_bias: RawTensor::zeros(device, [16]), + fc_weight: RawTensor::from_slice( + device, + [CLASSES, flat], + &rng.kaiming(CLASSES * flat, flat), + ), + fc_bias: RawTensor::zeros(device, [CLASSES]), + } + } +} + +/// conv(1->8) -> relu -> maxpool -> conv(8->16) -> relu -> maxpool -> linear. +struct Cnn { + conv1: ConvNd<2, 4>, + conv2: ConvNd<2, 4>, + fc: Linear, +} + +impl Cnn { + /// Build the model on `graph`, as trainable leaves (training) or + /// constants (evaluation). + fn new(graph: &Graph, params: &Params, trainable: bool) -> Self { + fn wrap( + graph: &Graph, + value: &RawTensor, + trainable: bool, + ) -> Tensor { + if trainable { + graph.leaf(value.clone()) + } else { + Tensor::constant_from_raw(graph, value.clone()) + } + } + let config = ConvNdConfig { + padding: [1, 1], + stride: [1, 1], + groups: 1, + }; + Self { + conv1: ConvNd::new( + wrap(graph, ¶ms.conv1_weight, trainable), + Some(wrap(graph, ¶ms.conv1_bias, trainable)), + config, + ), + conv2: ConvNd::new( + wrap(graph, ¶ms.conv2_weight, trainable), + Some(wrap(graph, ¶ms.conv2_bias, trainable)), + config, + ), + fc: Linear::new( + wrap(graph, ¶ms.fc_weight, trainable), + Some(wrap(graph, ¶ms.fc_bias, trainable)), + ), + } + } + + /// Input shape: (batch, 1, 28, 28). Output shape: (batch, 10). + fn forward(&self, images: &Tensor<4>) -> Tensor<2> { + let x = self + .conv1 + .forward(images) + .relu() + .pool_max::<2, 6, 7, 5>([2, 2]); + let x = self.conv2.forward(&x).relu().pool_max::<2, 6, 7, 5>([2, 2]); + self.fc.forward(&x.flatten_last_n::<2, 2>()) + } +} + +/// Numerically stable softmax cross-entropy averaged over the batch: +/// log softmax via log-sum-exp so a saturated class cannot underflow. +fn cross_entropy(logits: &Tensor<2>, targets: &RawTensor<1, u32>) -> Tensor<0> { + let batch = logits.shape()[0]; + let shifted = logits.sub_::<2, 2>(&logits.max_keepdim::<1>(1)); + let log_sum_exp = shifted.exp().sum_keepdim(1).log(); + let label_log_probs = shifted.sub_::<2, 2>(&log_sum_exp).gather_last(targets); + label_log_probs.sum().mul_scalar(-1.0 / batch as f32) +} + +fn sgd_step( + param: &mut RawTensor, + gradients: &Gradients, + leaf: &Tensor, +) { + let gradient = gradients.get(leaf).expect("missing gradient"); + *param = (param.clone() - gradient * LEARNING_RATE).to_concrete(); +} + +async fn to_scalar(value: RawTensor<0, f32>) -> f32 { + value.reshape([1]).as_slice().await.unwrap().to_vec()[0] +} + +fn argmax(row: &[f32]) -> u32 { + row.iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .unwrap() + .0 as u32 +} + +#[tokio::main] +async fn main() { + let train_images = load_images("train-images-idx3-ubyte"); + let train_labels = load_labels("train-labels-idx1-ubyte"); + let test_images = load_images("t10k-images-idx3-ubyte"); + let test_labels = load_labels("t10k-labels-idx1-ubyte"); + assert_eq!(train_images.len(), train_labels.len() * PIXELS); + + let device = match Device::gpu().await { + Ok(gpu) => gpu, + Err(_) => { + println!("GPU unavailable, training on CPU"); + Device::cpu() + } + }; + + let mut params = Params::new(&device); + + let steps_per_epoch = train_labels.len() / BATCH_SIZE; + for epoch in 0..EPOCHS { + for step in 0..steps_per_epoch { + let start = step * BATCH_SIZE; + let images = RawTensor::from_slice( + &device, + [BATCH_SIZE, 1, IMAGE_SIZE, IMAGE_SIZE], + &train_images[start * PIXELS..(start + BATCH_SIZE) * PIXELS], + ); + let targets = RawTensor::from_slice( + &device, + [BATCH_SIZE], + &train_labels[start..start + BATCH_SIZE], + ); + + let graph = Graph::new(); + let model = Cnn::new(&graph, ¶ms, true); + let logits = model.forward(&Tensor::constant_from_raw(&graph, images)); + let loss = cross_entropy(&logits, &targets); + + let loss_value = to_scalar(loss.raw().clone()).await; + let gradients = loss.backward().unwrap().into_detached(); + sgd_step(&mut params.conv1_weight, &gradients, model.conv1.weight()); + sgd_step( + &mut params.conv1_bias, + &gradients, + model.conv1.bias().unwrap(), + ); + sgd_step(&mut params.conv2_weight, &gradients, model.conv2.weight()); + sgd_step( + &mut params.conv2_bias, + &gradients, + model.conv2.bias().unwrap(), + ); + sgd_step(&mut params.fc_weight, &gradients, model.fc.weight()); + sgd_step(&mut params.fc_bias, &gradients, model.fc.bias().unwrap()); + + if step % 50 == 0 { + println!("epoch {epoch} step {step}/{steps_per_epoch}: loss {loss_value:.4}"); + } + } + } + + const EVAL_BATCH: usize = 500; + let mut correct = 0; + for (images, labels) in test_images + .chunks(EVAL_BATCH * PIXELS) + .zip(test_labels.chunks(EVAL_BATCH)) + { + let graph = Graph::new(); + let model = Cnn::new(&graph, ¶ms, false); + let x = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [labels.len(), 1, IMAGE_SIZE, IMAGE_SIZE], images), + ); + let logits = model + .forward(&x) + .raw() + .clone() + .as_slice() + .await + .unwrap() + .to_vec(); + correct += logits + .iter() + .zip(labels) + .filter(|(row, label)| argmax(row) == **label) + .count(); + } + let total = test_labels.len(); + println!( + "test accuracy: {correct}/{total} ({:.2}%)", + 100.0 * correct as f32 / total as f32 + ); +} diff --git a/fusor-ml/fusor/examples/probe_shapes.rs b/fusor-ml/fusor/examples/probe_shapes.rs new file mode 100644 index 000000000..fd0cf4890 --- /dev/null +++ b/fusor-ml/fusor/examples/probe_shapes.rs @@ -0,0 +1,711 @@ +//! Single-op probes at the exact shapes where the batch=64 transformer +//! example wedges the GPU. Each case is one op family in its own process: +//! run `probe_shapes `, one bounded submission, full CPU verification. +//! +//! Cases cross the 65535-workgroup / 16.78M-element dispatch boundary that +//! batch=32 stays under and batch=64 exceeds. + +use fusor::{Device, Tensor, ToVec}; + +const B: usize = 64; +const H: usize = 6; +const S: usize = 256; +const HD: usize = 64; +const ELEMS: usize = B * H * S * S; // 25_165_824 > 65535 * 256 + +fn fill(len: usize) -> Vec { + (0..len).map(|i| ((i % 251) as f32) * 0.01 - 1.25).collect() +} + +async fn read_flat(t: &Tensor, len: usize) -> Vec { + t.reshape([len]).as_slice().await.unwrap().to_vec() +} + +fn check(name: &str, got: &[f32], want: impl Fn(usize) -> f32, tol: f32) { + let mut bad = 0usize; + let mut first = None; + for (i, &g) in got.iter().enumerate() { + let w = want(i); + if (g - w).abs() > tol * w.abs().max(1.0) { + bad += 1; + if first.is_none() { + first = Some((i, g, w)); + } + } + } + if bad == 0 { + println!("{name}: PASS ({} elements)", got.len()); + } else { + println!( + "{name}: FAIL {bad}/{} mismatched, first at {:?}", + got.len(), + first + ); + std::process::exit(1); + } +} + +#[tokio::main] +async fn main() { + let case = std::env::args().nth(1).expect("usage: probe_shapes "); + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("error")), + ) + .init(); + let device = Device::gpu().await.unwrap(); + let start = std::time::Instant::now(); + println!("SUBMITTING {case}"); + + match case.as_str() { + // nary elementwise over 25.17M elements + "add25m" => { + let data = fill(ELEMS); + let x = Tensor::from_slice(&device, [B, H, S, S], &data); + let y = x.add_::<4, 4, _>(&x); + let got = read_flat(&y, ELEMS).await; + check("add25m", &got, |i| 2.0 * data[i], 1e-6); + } + // reduce over the last axis: 98304 rows of 256 + "sum25m" => { + let data = fill(ELEMS); + let x = Tensor::from_slice(&device, [B, H, S, S], &data); + let y: Tensor<3, f32> = x.sum(3); + let got = read_flat(&y, ELEMS / S).await; + let rows: Vec = (0..ELEMS / S) + .map(|r| data[r * S..(r + 1) * S].iter().sum()) + .collect(); + check("sum25m", &got, |r| rows[r], 1e-4); + } + // row-program softmax over 98304 rows (> 65535 workgroups) + "softmax98k" => { + let data = fill(ELEMS); + let x = Tensor::from_slice(&device, [B, H, S, S], &data); + let y = x.softmax_last_dim::<3>(); + let got = read_flat(&y, ELEMS).await; + let mut refs = vec![0.0f32; ELEMS]; + for r in 0..ELEMS / S { + let row = &data[r * S..(r + 1) * S]; + let max = row.iter().cloned().fold(f32::MIN, f32::max); + let exps: Vec = row.iter().map(|v| (v - max).exp()).collect(); + let denom: f32 = exps.iter().sum(); + for (c, e) in exps.iter().enumerate() { + refs[r * S + c] = e / denom; + } + } + check("softmax98k", &got, |i| refs[i], 1e-4); + } + // big-M coop matmul: [16384, 384] x [384, 384] + "matmul16k" => { + let (m, k, n) = (B * S, 384usize, 384usize); + let a = fill(m * k); + let b = fill(k * n); + let at = Tensor::from_slice(&device, [m, k], &a); + let bt = Tensor::from_slice(&device, [k, n], &b); + let y = at.mat_mul(&bt); + let got = read_flat(&y, m * n).await; + // verify rows around the tile-grid extremes + for row in [0usize, 127, 8191, 8192, 16256, m - 1] { + let mut want = vec![0.0f32; n]; + for kk in 0..k { + let av = a[row * k + kk]; + for j in 0..n { + want[j] += av * b[kk * n + j]; + } + } + for j in 0..n { + let g = got[row * n + j]; + let w = want[j]; + if (g - w).abs() > 1e-2 * w.abs().max(1.0) { + println!("matmul16k: FAIL at row {row} col {j}: got {g} want {w}"); + std::process::exit(1); + } + } + } + println!("matmul16k: PASS (6 rows exact)"); + } + // batched attention-shape matmuls + "qkt" => { + let a = fill(B * H * S * HD); + let b = fill(B * H * S * HD); + let at = Tensor::from_slice(&device, [B, H, S, HD], &a); + let bt = Tensor::from_slice(&device, [B, H, HD, S], &b); + let y = at.mat_mul(&bt); + let got = read_flat(&y, ELEMS).await; + // verify one full matrix at the start, middle, and last batch + for batch in [0usize, 192, B * H - 1] { + let a0 = &a[batch * S * HD..(batch + 1) * S * HD]; + let b0 = &b[batch * HD * S..(batch + 1) * HD * S]; + for i in [0usize, 255] { + for j in 0..S { + let mut w = 0.0f32; + for kk in 0..HD { + w += a0[i * HD + kk] * b0[kk * S + j]; + } + let g = got[batch * S * S + i * S + j]; + if (g - w).abs() > 1e-3 * w.abs().max(1.0) { + println!("qkt: FAIL batch {batch} [{i},{j}]: got {g} want {w}"); + std::process::exit(1); + } + } + } + } + println!("qkt: PASS (sampled)"); + } + "av" => { + let a = fill(ELEMS); + let v = fill(B * H * S * HD); + let at = Tensor::from_slice(&device, [B, H, S, S], &a); + let vt = Tensor::from_slice(&device, [B, H, S, HD], &v); + let y = at.mat_mul(&vt); + let got = read_flat(&y, B * H * S * HD).await; + for batch in [0usize, 192, B * H - 1] { + let a0 = &a[batch * S * S..(batch + 1) * S * S]; + let v0 = &v[batch * S * HD..(batch + 1) * S * HD]; + for i in [0usize, 255] { + for j in 0..HD { + let mut w = 0.0f32; + for kk in 0..S { + w += a0[i * S + kk] * v0[kk * HD + j]; + } + let g = got[batch * S * HD + i * HD + j]; + if (g - w).abs() > 1e-3 * w.abs().max(1.0) { + println!("av: FAIL batch {batch} [{i},{j}]: got {g} want {w}"); + std::process::exit(1); + } + } + } + } + println!("av: PASS (sampled)"); + } + // rank-2 adds bracketing the 65535-workgroup dispatch boundary + "add65535" | "add65536" => { + let rows: usize = if case == "add65535" { 65535 } else { 65536 }; + let len = rows * 256; + let data = fill(len); + let x = Tensor::from_slice(&device, [rows, 256], &data); + let y = x.add_::<2, 2, _>(&x); + let got = y.as_slice().await.unwrap().to_vec(); + let mut bad = 0usize; + let mut first = None; + for r in 0..rows { + for c in 0..256 { + let w = 2.0 * data[r * 256 + c]; + let g = got[r][c]; + if (g - w).abs() > 1e-6 * w.abs().max(1.0) { + bad += 1; + if first.is_none() { + first = Some((r, c, g, w)); + } + } + } + } + if bad == 0 { + println!("{case}: PASS ({len} elements)"); + } else { + println!("{case}: FAIL {bad}/{len} mismatched, first at {first:?}"); + std::process::exit(1); + } + } + // full failing size, rank-2, no reshape before readback + "add25m2d" => { + let rows = ELEMS / 256; + let data = fill(ELEMS); + let x = Tensor::from_slice(&device, [rows, 256], &data); + let y = x.add_::<2, 2, _>(&x); + let got = y.as_slice().await.unwrap().to_vec(); + let mut bad = 0usize; + let mut first = None; + for r in 0..rows { + for c in 0..256 { + let w = 2.0 * data[r * 256 + c]; + let g = got[r][c]; + if (g - w).abs() > 1e-6 * w.abs().max(1.0) { + bad += 1; + if first.is_none() { + first = Some((r, c, g, w)); + } + } + } + } + if bad == 0 { + println!("add25m2d: PASS ({ELEMS} elements)"); + } else { + println!("add25m2d: FAIL {bad}/{ELEMS} mismatched, first at {first:?}"); + std::process::exit(1); + } + } + // region map of the failing case: which flat ranges were written + "map25m" => { + let data = fill(ELEMS); + let x = Tensor::from_slice(&device, [B, H, S, S], &data); + let y = x.add_::<4, 4, _>(&x); + let got = read_flat(&y, ELEMS).await; + let state = |i: usize| -> u8 { + let w = 2.0 * data[i]; + let g = got[i]; + if (g - w).abs() <= 1e-6 * w.abs().max(1.0) { + 1 // correct + } else if g == 0.0 { + 0 // untouched + } else { + 2 // garbage + } + }; + let mut runs: Vec<(u8, usize, usize)> = Vec::new(); + let mut cur = state(0); + let mut start = 0usize; + let mut counts = [0usize; 3]; + for i in 0..ELEMS { + let s = state(i); + counts[s as usize] += 1; + if s != cur { + runs.push((cur, start, i)); + cur = s; + start = i; + } + } + runs.push((cur, start, ELEMS)); + println!( + "map25m: untouched={} correct={} garbage={} runs={}", + counts[0], + counts[1], + counts[2], + runs.len() + ); + for (s, a, b) in runs.iter().take(24) { + let label = ["ZERO", "OK", "GARBAGE"][*s as usize]; + println!( + " {label:8} [{a:>9}..{b:>9}) len {:>9} wg [{}..{}]", + b - a, + a / 256, + (b - 1) / 256 + ); + } + } + // rank-2 add read through a reshape view (same readback as add25m) + "add25m2d_reshaped" => { + let rows = ELEMS / 256; + let data = fill(ELEMS); + let x = Tensor::from_slice(&device, [rows, 256], &data); + let y = x.add_::<2, 2, _>(&x); + let got = read_flat(&y, ELEMS).await; + check("add25m2d_reshaped", &got, |i| 2.0 * data[i], 1e-6); + } + // rank-4 add read back directly, no reshape + "add25m4d_direct" => { + let data = fill(ELEMS); + let x = Tensor::from_slice(&device, [B, H, S, S], &data); + let y = x.add_::<4, 4, _>(&x); + let slice = y.as_slice().await.unwrap(); + let mut bad = 0usize; + let mut first = None; + for b in 0..B { + for h in 0..H { + for i in 0..S { + for j in 0..S { + let idx = ((b * H + h) * S + i) * S + j; + let w = 2.0 * data[idx]; + let g = slice[[b, h, i, j]]; + if (g - w).abs() > 1e-6 * w.abs().max(1.0) { + bad += 1; + if first.is_none() { + first = Some((idx, g, w)); + } + } + } + } + } + } + if bad == 0 { + println!("add25m4d_direct: PASS ({ELEMS} elements)"); + } else { + println!("add25m4d_direct: FAIL {bad}/{ELEMS}, first {first:?}"); + std::process::exit(1); + } + } + // generic rank-4 add over a shape given as 4 extra args + "addshape" => { + let dims: Vec = std::env::args() + .skip(2) + .map(|a| a.parse().unwrap()) + .collect(); + let [d0, d1, d2, d3] = dims[..] else { + panic!("addshape needs 4 dims") + }; + let len = d0 * d1 * d2 * d3; + let data = fill(len); + let x = Tensor::from_slice(&device, [d0, d1, d2, d3], &data); + let y = x.add_::<4, 4, _>(&x); + let slice = y.as_slice().await.unwrap(); + let mut bad = 0usize; + let mut first = None; + let mut idx = 0usize; + 'outer: for a in 0..d0 { + for b in 0..d1 { + for i in 0..d2 { + for j in 0..d3 { + let w = 2.0 * data[idx]; + let g = slice[[a, b, i, j]]; + if (g - w).abs() > 1e-6 * w.abs().max(1.0) { + bad += 1; + if first.is_none() { + first = Some((idx, g, w)); + } + if bad > 5_000_000 { + break 'outer; + } + } + idx += 1; + } + } + } + } + if bad == 0 { + println!("addshape {dims:?}: PASS ({len} elements)"); + } else { + println!("addshape {dims:?}: FAIL {bad}+/{len}, first {first:?}"); + std::process::exit(1); + } + } + // empirical write map: which source value lands at which target slot + "writemap" => { + let dims: Vec = std::env::args() + .skip(2) + .map(|a| a.parse().unwrap()) + .collect(); + let [d0, d1, d2, d3] = dims[..] else { + panic!("writemap needs 4 dims") + }; + let len = d0 * d1 * d2 * d3; + // encode source index as (quotient+1, remainder+1) across two runs + let dq: Vec = (0..len).map(|i| (i / 4096 + 1) as f32).collect(); + let dr: Vec = (0..len).map(|i| (i % 4096 + 1) as f32).collect(); + let read4 = |data: &[f32]| { + let x = Tensor::from_slice(&device, [d0, d1, d2, d3], data); + let y = x.add_::<4, 4, _>(&x); + async move { read_flat(&y, len).await } + }; + let gq = read4(&dq).await; + let gr = read4(&dr).await; + let mut untouched = 0usize; + let mut identity = 0usize; + let mut moved = 0usize; + let mut samples: Vec<(usize, usize)> = Vec::new(); + for t in 0..len { + if gq[t] == 0.0 && gr[t] == 0.0 { + untouched += 1; + continue; + } + let q = (gq[t] / 2.0 - 1.0) as usize; + let r = (gr[t] / 2.0 - 1.0) as usize; + let s = q * 4096 + r; + if s == t { + identity += 1; + } else { + moved += 1; + if samples.len() < 30 { + samples.push((t, s)); + } + } + } + println!("writemap {dims:?}: untouched={untouched} identity={identity} moved={moved}"); + for (t, s) in samples { + let tc = ( + t / (d1 * d2 * d3), + (t / (d2 * d3)) % d1, + (t / d3) % d2, + t % d3, + ); + let sc = ( + s / (d1 * d2 * d3), + (s / (d2 * d3)) % d1, + (s / d3) % d2, + s % d3, + ); + println!( + " t={t} {tc:?} <- s={s} {sc:?} (s-t={})", + s as i64 - t as i64 + ); + } + } + // same failing add, but materialize + sleep before reading back: + // complete data => submit/map race; still partial => work lost on GPU + "addsleep" => { + let data = fill(ELEMS); + let x = Tensor::from_slice(&device, [B, H, S, S], &data); + let y = x.add_::<4, 4, _>(&x); + y.materialize().await; + tokio::time::sleep(std::time::Duration::from_millis(1500)).await; + let got = read_flat(&y, ELEMS).await; + check("addsleep", &got, |i| 2.0 * data[i], 1e-6); + } + // raw wgpu micro-test: does the GPU compute the delinearize/relinearize + // identity exactly? Bounds checks left ON; independent of fusor kernels. + "divident" => { + let fusor::Device::Gpu(core_device) = &device else { + panic!("gpu device required") + }; + let wgpu_device = core_device.wgpu_device(); + let queue = core_device.wgpu_queue(); + let total: u32 = ELEMS as u32; // 25_165_824 + let shader = r#" +@group(0) @binding(0) var out: array; + +@compute @workgroup_size(256) +fn main(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_index) lane: u32) { + let group = wg.x + wg.y * 65535u; + let f = group * 256u + lane; + if (f < arrayLength(&out)) { + let a = (f / 393216u) % 64u; + let b = (f / 65536u) % 6u; + let i = (f / 256u) % 256u; + let j = f % 256u; + out[f] = a * 393216u + b * 65536u + i * 256u + j; + } +} +"#; + let module = wgpu_device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("divident"), + source: wgpu::ShaderSource::Wgsl(shader.into()), + }); + let out = wgpu_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: total as u64 * 4, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let pipeline = wgpu_device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: None, + module: &module, + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + let bind = wgpu_device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &pipeline.get_bind_group_layout(0), + entries: &[wgpu::BindGroupEntry { + binding: 0, + resource: out.as_entire_binding(), + }], + }); + let staging = wgpu_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: total as u64 * 4, + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut encoder = wgpu_device.create_command_encoder(&Default::default()); + { + let mut pass = encoder.begin_compute_pass(&Default::default()); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &bind, &[]); + pass.dispatch_workgroups(65535, 2, 1); + } + encoder.copy_buffer_to_buffer(&out, 0, &staging, 0, total as u64 * 4); + queue.submit(Some(encoder.finish())); + let (sender, receiver) = futures_channel::oneshot::channel(); + staging.slice(..).map_async(wgpu::MapMode::Read, move |r| { + let _ = sender.send(r); + }); + core_device.poll_wait(); + receiver.await.unwrap().unwrap(); + let view = staging.slice(..).get_mapped_range(); + let words: &[u32] = bytemuck::cast_slice(&view); + let mut bad = 0usize; + let mut samples = Vec::new(); + for (f, &got) in words.iter().enumerate() { + if got != f as u32 { + bad += 1; + if samples.len() < 20 { + samples.push((f, got)); + } + } + } + println!("divident: {bad}/{total} wrong"); + for (f, got) in samples { + println!( + " f={f} -> {got} (a={} b={} i={} j={})", + (f as u32 / 393216) % 64, + (f as u32 / 65536) % 6, + (f as u32 / 256) % 256, + f as u32 % 256 + ); + } + if bad > 0 { + std::process::exit(1); + } + } + // same expression, 1-D grid only (covers f < 16.78M) + "divident1d" | "divident2" => { + let fusor::Device::Gpu(core_device) = &device else { + panic!("gpu device required") + }; + let wgpu_device = core_device.wgpu_device(); + let queue = core_device.wgpu_queue(); + let one_d = case == "divident1d"; + let total: u32 = if one_d { 16_776_960 } else { ELEMS as u32 }; + // divident2 keeps the 2-D dispatch but derives the group id from a + // flat linearization that the compiler can't fold with wg.x alone + let group_expr = if one_d { + "let group = wg.x;" + } else { + "let group = wg.x + wg.y * 65535u;" + }; + let shader = format!( + r#" +@group(0) @binding(0) var out: array; + +@compute @workgroup_size(256) +fn main(@builtin(workgroup_id) wg: vec3, @builtin(local_invocation_index) lane: u32) {{ + {group_expr} + let f = group * 256u + lane; + if (f < arrayLength(&out)) {{ + let q = f / 65536u; + let b = q % 6u; + out[f] = b; + }} +}} +"# + ); + let module = wgpu_device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some(case.as_str()), + source: wgpu::ShaderSource::Wgsl(shader.into()), + }); + let out = wgpu_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: total as u64 * 4, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }); + let pipeline = wgpu_device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: None, + layout: None, + module: &module, + entry_point: Some("main"), + compilation_options: Default::default(), + cache: None, + }); + let bind = wgpu_device.create_bind_group(&wgpu::BindGroupDescriptor { + label: None, + layout: &pipeline.get_bind_group_layout(0), + entries: &[wgpu::BindGroupEntry { + binding: 0, + resource: out.as_entire_binding(), + }], + }); + let staging = wgpu_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + size: total as u64 * 4, + usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, + mapped_at_creation: false, + }); + let mut encoder = wgpu_device.create_command_encoder(&Default::default()); + { + let mut pass = encoder.begin_compute_pass(&Default::default()); + pass.set_pipeline(&pipeline); + pass.set_bind_group(0, &bind, &[]); + if one_d { + pass.dispatch_workgroups(65535, 1, 1); + } else { + pass.dispatch_workgroups(65535, 2, 1); + } + } + encoder.copy_buffer_to_buffer(&out, 0, &staging, 0, total as u64 * 4); + queue.submit(Some(encoder.finish())); + let (sender, receiver) = futures_channel::oneshot::channel(); + staging.slice(..).map_async(wgpu::MapMode::Read, move |r| { + let _ = sender.send(r); + }); + core_device.poll_wait(); + receiver.await.unwrap().unwrap(); + let view = staging.slice(..).get_mapped_range(); + let words: &[u32] = bytemuck::cast_slice(&view); + let mut bad = 0usize; + let mut samples = Vec::new(); + for (f, &got) in words.iter().enumerate() { + let want = (f as u32 / 65536) % 6; + if got != want { + bad += 1; + if samples.len() < 8 { + samples.push((f, got, want)); + } + } + } + println!("{case}: {bad}/{total} wrong"); + for (f, got, want) in samples { + println!(" f={f} got={got} want={want}"); + } + } + // timed matmul: probe_shapes matmul_bench M K N [iters] + "matmul_bench" => { + let dims: Vec = std::env::args() + .skip(2) + .map(|a| a.parse().unwrap()) + .collect(); + let (m, k, n) = (dims[0], dims[1], dims[2]); + let iters = dims.get(3).copied().unwrap_or(20); + let a = fill(m * k); + // Scale toward identity-ish magnitudes so a 40-deep chain + // neither overflows nor denormals out. + let b: Vec = fill(k * n).iter().map(|v| v * 0.05).collect(); + let b_back: Vec = fill(n * k).iter().map(|v| v * 0.05).collect(); + let at = Tensor::from_slice(&device, [m, k], &a); + let bt = Tensor::from_slice(&device, [k, n], &b); + // For k != n, alternate with an [n, k] partner so the chain's + // shapes ping-pong between m×n and m×k (both measured strides + // of the same geometry family). + let bt_back = Tensor::from_slice(&device, [n, k], &b_back); + // Warm the kernel cache and pipeline. + at.mat_mul(&bt).as_slice().await.unwrap(); + let mut best = f64::MAX; + for _ in 0..3 { + let start = std::time::Instant::now(); + // A dependency chain defeats semantic-identity dedup and + // amortizes the single readback sync across all iterations. + let mut y = at.mat_mul(&bt); + for i in 1..iters { + y = if i % 2 == 1 { + y.mat_mul(&bt_back) + } else { + y.mat_mul(&bt) + }; + } + let _ = y.as_slice().await.unwrap(); + best = best.min(start.elapsed().as_secs_f64() / iters as f64); + } + // Correctness tripwire: one un-chained product checked against a + // sampled f64 reference. A "fast" wrong kernel reads as an + // impossible TF/s without this. + let single = at.mat_mul(&bt); + let got = single.as_slice().await.unwrap(); + for &row in &[0usize, m - 1] { + for &col in &[0usize, n - 1] { + let mut want = 0.0f64; + for kk in 0..k { + want += a[row * k + kk] as f64 * b[kk * n + col] as f64; + } + let got = got[[row, col]]; + assert!( + (got as f64 - want).abs() < 1e-2 * want.abs().max(1.0), + "verification failed at [{row},{col}]: got {got}, want {want}" + ); + } + } + let flops = 2.0 * m as f64 * k as f64 * n as f64; + println!( + "matmul_bench {m}x{k}x{n}: min {:.3} ms/matmul, {:.2} TF/s", + best * 1e3, + flops / best / 1e12 + ); + } + other => { + eprintln!("unknown case {other}"); + std::process::exit(2); + } + } + println!("{case} elapsed: {:?}", start.elapsed()); +} diff --git a/fusor-ml/fusor/examples/transformer.rs b/fusor-ml/fusor/examples/transformer.rs new file mode 100644 index 000000000..34db899ce --- /dev/null +++ b/fusor-ml/fusor/examples/transformer.rs @@ -0,0 +1,1064 @@ +//! Train a small GPT-style transformer on tiny Shakespeare with the autograd API. +//! +//! Downloads the dataset on first run (cached under `examples/data`), builds a +//! character-level vocabulary, then trains a 2-layer pre-norm transformer +//! (learned positional embeddings, 4-head causal self-attention, gelu MLP) +//! with softmax cross-entropy and Adam, rebuilding the tape each step. +//! Reports train/test next-token accuracy while training, then samples text. +//! +//! The model body is written once over an element type alias and stamped out +//! per dtype: `--dtype f16` runs the standard mixed-precision recipe (f32 +//! master weights and optimizer, f16 model compute, final norm + lm head + +//! loss in f32), the default runs everything in f32. +//! +//! Run with: +//! ```sh +//! cargo run --release --example transformer +//! ``` + +use std::collections::{BTreeSet, VecDeque}; +use std::io::Read; +use std::path::PathBuf; + +use fusor::autograd::layers::{Embedding, LayerNorm, Linear}; +use fusor::autograd::{Graph, Tensor}; +use fusor::{ + Device, FloatDataType, MaskKind, StandardSamplerParams, Tensor as RawTensor, ToVec, cat, +}; + +const CONTEXT: usize = 256; +const BATCH_SIZE: usize = 64; +const DIM: usize = 512; +const HEADS: usize = 8; +const HEAD_DIM: usize = DIM / HEADS; +const MLP_DIM: usize = 4 * DIM; +const LAYERS: usize = 6; +const STEPS: usize = 300; +const LEARNING_RATE: f32 = 1e-3; +const BETA1: f32 = 0.9; +const BETA2: f32 = 0.999; +const EPSILON: f32 = 1e-8; +const LAYER_NORM_EPS: f32 = 1e-5; +const MASK_VALUE: f32 = -1e9; +const TEMPERATURE: f32 = 0.8; +const GENERATION_RUN_AHEAD: usize = 32; +const PROGRESS_EVERY: usize = 25; + +const DATA_URL: &str = + "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"; + +#[derive(Clone, Copy, PartialEq)] +enum Dtype { + F32, + F16, +} + +#[derive(Clone, Copy)] +struct RunConfig { + steps: usize, + min_steps_per_second: Option, + expect_loss: Option, + skip_eval: bool, + progress_every: usize, + trace_host: bool, + trace_resolve: bool, + trace_names: bool, + dtype: Dtype, +} + +impl RunConfig { + fn from_args() -> Self { + let mut config = Self { + steps: STEPS, + min_steps_per_second: None, + expect_loss: None, + skip_eval: false, + progress_every: PROGRESS_EVERY, + trace_host: false, + trace_resolve: false, + trace_names: false, + dtype: Dtype::F32, + }; + let mut args = std::env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--steps" => { + let value = args.next().expect("--steps requires a value"); + config.steps = value.parse().expect("--steps must be a positive integer"); + } + "--min-steps-per-sec" | "--min-steps-per-second" => { + let value = args.next().expect("--min-steps-per-sec requires a value"); + config.min_steps_per_second = + Some(value.parse().expect("--min-steps-per-sec must be a number")); + } + "--expect-loss" => { + let value = args.next().expect("--expect-loss requires a value"); + config.expect_loss = + Some(value.parse().expect("--expect-loss must be a number")); + } + "--skip-eval" => config.skip_eval = true, + "--progress-every" => { + let value = args.next().expect("--progress-every requires a value"); + config.progress_every = value + .parse() + .expect("--progress-every must be a non-negative integer"); + } + "--dtype" => { + let value = args.next().expect("--dtype requires a value"); + config.dtype = match value.as_str() { + "f32" => Dtype::F32, + "f16" => Dtype::F16, + _ => panic!("--dtype must be f32 or f16"), + }; + } + "--trace-host" => config.trace_host = true, + "--trace-resolve" => config.trace_resolve = true, + "--trace-names" => { + config.trace_resolve = true; + config.trace_names = true; + } + _ => panic!("unknown argument: {arg}"), + } + } + assert!(config.steps > 0, "--steps must be greater than zero"); + config + } +} + +/// Block until all submitted GPU work has completed (no-op on CPU). +fn wait_for_gpu(device: &Device) { + if let Device::Gpu(gpu) = device { + gpu.poll_wait(); + } +} + +fn data_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/data") +} + +/// Download the tiny Shakespeare corpus, caching it on disk. +fn fetch_text() -> String { + let path = data_dir().join("tinyshakespeare.txt"); + if let Ok(text) = std::fs::read_to_string(&path) { + return text; + } + println!("downloading tinyshakespeare.txt"); + let mut response = ureq::get(DATA_URL) + .call() + .unwrap_or_else(|err| panic!("failed to download tiny shakespeare: {err}")); + let mut text = String::new(); + response + .body_mut() + .as_reader() + .read_to_string(&mut text) + .unwrap(); + std::fs::create_dir_all(data_dir()).unwrap(); + std::fs::write(&path, &text).unwrap(); + text +} + +/// Deterministic LCG so runs are reproducible without a rand dependency. +struct Lcg(u64); + +impl Lcg { + /// Uniform in [0, 1). + fn next_f32(&mut self) -> f32 { + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (self.0 >> 32) as f32 / (1u64 << 32) as f32 + } + + /// Uniform in (-bound, bound) with `bound = sqrt(1 / fan_in)`, the + /// PyTorch default init for linear layers. + fn kaiming(&mut self, count: usize, fan_in: usize) -> Vec { + let bound = (1.0 / fan_in as f32).sqrt(); + (0..count) + .map(|_| (self.next_f32() - 0.5) * 2.0 * bound) + .collect() + } +} + +/// Raw parameter tensors that persist across training steps. Masters stay in +/// f32 regardless of the model compute dtype. +struct BlockParams { + ln1_weight: RawTensor<1, f32>, + ln1_bias: RawTensor<1, f32>, + query: RawTensor<2, f32>, + key: RawTensor<2, f32>, + value: RawTensor<2, f32>, + output: RawTensor<2, f32>, + ln2_weight: RawTensor<1, f32>, + ln2_bias: RawTensor<1, f32>, + mlp_up: RawTensor<2, f32>, + mlp_up_bias: RawTensor<1, f32>, + mlp_down: RawTensor<2, f32>, + mlp_down_bias: RawTensor<1, f32>, +} + +struct Params { + token_embedding: RawTensor<2, f32>, + position_embedding: RawTensor<2, f32>, + blocks: Vec, + ln_final_weight: RawTensor<1, f32>, + ln_final_bias: RawTensor<1, f32>, + lm_head: RawTensor<2, f32>, +} + +impl Params { + fn new(device: &Device, vocab_size: usize) -> Self { + let mut rng = Lcg(42); + let blocks = (0..LAYERS) + .map(|_| BlockParams { + ln1_weight: RawTensor::splat(device, 1.0, [DIM]), + ln1_bias: RawTensor::zeros(device, [DIM]), + query: RawTensor::from_slice(device, [DIM, DIM], &rng.kaiming(DIM * DIM, DIM)), + key: RawTensor::from_slice(device, [DIM, DIM], &rng.kaiming(DIM * DIM, DIM)), + value: RawTensor::from_slice(device, [DIM, DIM], &rng.kaiming(DIM * DIM, DIM)), + output: RawTensor::from_slice(device, [DIM, DIM], &rng.kaiming(DIM * DIM, DIM)), + ln2_weight: RawTensor::splat(device, 1.0, [DIM]), + ln2_bias: RawTensor::zeros(device, [DIM]), + mlp_up: RawTensor::from_slice( + device, + [MLP_DIM, DIM], + &rng.kaiming(MLP_DIM * DIM, DIM), + ), + mlp_up_bias: RawTensor::zeros(device, [MLP_DIM]), + mlp_down: RawTensor::from_slice( + device, + [DIM, MLP_DIM], + &rng.kaiming(DIM * MLP_DIM, MLP_DIM), + ), + mlp_down_bias: RawTensor::zeros(device, [DIM]), + }) + .collect(); + Self { + token_embedding: RawTensor::from_slice( + device, + [vocab_size, DIM], + &rng.kaiming(vocab_size * DIM, DIM), + ), + position_embedding: RawTensor::from_slice( + device, + [CONTEXT, DIM], + &rng.kaiming(CONTEXT * DIM, DIM), + ), + blocks, + ln_final_weight: RawTensor::splat(device, 1.0, [DIM]), + ln_final_bias: RawTensor::zeros(device, [DIM]), + lm_head: RawTensor::from_slice( + device, + [vocab_size, DIM], + &rng.kaiming(vocab_size * DIM, DIM), + ), + } + } + + /// Parameters in a fixed order; must match [`Gpt::leaves`]. + fn tensors_mut(&mut self) -> (Vec<&mut RawTensor<1, f32>>, Vec<&mut RawTensor<2, f32>>) { + let mut rank1 = vec![&mut self.ln_final_weight, &mut self.ln_final_bias]; + let mut rank2 = vec![ + &mut self.token_embedding, + &mut self.position_embedding, + &mut self.lm_head, + ]; + for block in &mut self.blocks { + rank1.extend([ + &mut block.ln1_weight, + &mut block.ln1_bias, + &mut block.ln2_weight, + &mut block.ln2_bias, + &mut block.mlp_up_bias, + &mut block.mlp_down_bias, + ]); + rank2.extend([ + &mut block.query, + &mut block.key, + &mut block.value, + &mut block.output, + &mut block.mlp_up, + &mut block.mlp_down, + ]); + } + (rank1, rank2) + } +} + +/// Numerically stable softmax cross-entropy averaged over all positions: +/// log softmax via log-sum-exp so a saturated class cannot underflow. +fn cross_entropy(logits: &Tensor<2>, targets: &RawTensor<1, u32>) -> Tensor<0> { + logits.softmax_cross_entropy(targets) +} + +fn correct_predictions(logits: &Tensor<2>, targets: &RawTensor<1, u32>) -> Tensor<0> { + assert_eq!(logits.shape()[0], targets.shape()[0]); + logits + .gather_last(targets) + .eq_tensor(&logits.max::<1>(1)) + .sum() +} + +/// First and second moment estimates for one parameter tensor, stored flat +/// so the update math is rank-independent. +struct AdamState { + momentum: RawTensor<1, f32>, + variance: RawTensor<1, f32>, +} + +impl AdamState { + fn zeros(device: &Device, elements: usize) -> Self { + let zeros = vec![0.0; elements]; + Self { + momentum: RawTensor::from_slice(device, [elements], &zeros), + variance: RawTensor::from_slice(device, [elements], &zeros), + } + } +} + +/// One Adam update. `lr` is the learning rate with warmup and the bias +/// correction for step `t` already folded in. It is passed as a [1] tensor +/// rather than an f32 on purpose: scalar constants are baked into the +/// generated kernel source, so a per-step scalar would force a shader +/// recompile every step, while tensor *contents* are runtime data and the +/// cached pipelines are reused. +fn adam_step( + param: &mut RawTensor, + state: &mut AdamState, + gradient: RawTensor, + lr: &RawTensor<1, f32>, +) { + let shape = param.shape(); + let elements = shape.iter().product(); + let gradient = gradient.reshape([elements]); + state.momentum = + (state.momentum.clone() * BETA1 + gradient.clone() * (1.0 - BETA1)).into_concrete(); + state.variance = (state.variance.clone() * BETA2 + + (gradient.clone() * gradient) * (1.0 - BETA2)) + .into_concrete(); + let update: RawTensor<1, f32> = state + .momentum + .mul_(lr) + .div_(&state.variance.sqrt().add_scalar(EPSILON).into_concrete()); + *param = (param.clone().reshape([elements]) - update) + .reshape(shape) + .into_concrete(); +} + +/// Pick `batch_size` random windows and return (inputs, next-char targets), +/// each flattened to batch_size * CONTEXT ids. +fn sample_batch(tokens: &[u32], rng: &mut Lcg, batch_size: usize) -> (Vec, Vec) { + let mut inputs = Vec::with_capacity(batch_size * CONTEXT); + let mut targets = Vec::with_capacity(batch_size * CONTEXT); + for _ in 0..batch_size { + let start = (rng.next_f32() * (tokens.len() - CONTEXT - 1) as f32) as usize; + inputs.extend_from_slice(&tokens[start..start + CONTEXT]); + targets.extend_from_slice(&tokens[start + 1..start + CONTEXT + 1]); + } + (inputs, targets) +} + +async fn read_metrics( + loss: RawTensor<0, f32>, + correct: RawTensor<0, f32>, + tokens: usize, +) -> (f32, f32) { + let metrics = cat([loss.reshape([1]), correct.reshape([1])], 0) + .as_slice() + .await + .unwrap() + .to_vec(); + (metrics[0], metrics[1] / tokens as f32) +} + +/// Sample a token id from logits with temperature. +fn sample(logits: &[f32], rng: &mut Lcg) -> u32 { + let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let weights: Vec = logits + .iter() + .map(|logit| ((logit - max) / TEMPERATURE).exp()) + .collect(); + let mut remaining = rng.next_f32() * weights.iter().sum::(); + for (index, weight) in weights.iter().enumerate() { + remaining -= weight; + if remaining <= 0.0 { + return index as u32; + } + } + (weights.len() - 1) as u32 +} + +/// The model and training loop, written once over the element alias `E` and +/// stamped out per dtype. Masters, the optimizer, the final norm, the lm +/// head, and the loss always run in f32; everything between the embedding +/// lookup and the final norm runs in `E`. +macro_rules! gpt_model { + ($mod_name:ident, $elem:ty) => { + mod $mod_name { + use super::*; + + type E = $elem; + + /// Model-dtype view of an f32 leaf: the identity for `E == f32` + /// (keeping that graph byte-identical to the pure-f32 model), a + /// differentiable cast otherwise (its backward casts the weight + /// gradient back to the f32 master). + fn to_model(leaf: &Tensor) -> Tensor { + let any: &dyn std::any::Any = leaf; + match any.downcast_ref::>() { + Some(same) => same.clone(), + None => leaf.cast::(), + } + } + + /// Back to f32 at norm and head boundaries; identity for `E == f32`. + fn to_f32(hidden: &Tensor) -> Tensor { + let any: &dyn std::any::Any = hidden; + match any.downcast_ref::>() { + Some(same) => same.clone(), + None => hidden.cast::(), + } + } + + /// Additive causal mask: 0 where position j <= i, a large negative + /// value elsewhere so softmax zeroes out future positions. The + /// mask value stays finite in the model dtype: -1e9 overflows f16 + /// to -inf, which turns masked-lane arithmetic into NaNs. + pub fn causal_mask(device: &Device, seq: usize) -> RawTensor<2, E> { + let mask_value = if std::mem::size_of::() == 2 { + -6.0e4 + } else { + MASK_VALUE + }; + let values: Vec = (0..seq * seq) + .map(|i| { + E::from_f32(if i % seq <= i / seq { 0.0 } else { mask_value }) + }) + .collect(); + RawTensor::from_slice(device, [seq, seq], &values) + } + + /// Pre-norm transformer block: causal self-attention then a gelu + /// MLP, each behind a residual connection. + pub struct Block { + ln1: LayerNorm<1>, + query: Linear, + key: Linear, + value: Linear, + output: Linear, + ln2: LayerNorm<1>, + mlp_up: Linear, + mlp_down: Linear, + } + + impl Block { + /// Multi-head causal self-attention for `x` of shape + /// (batch, seq, dim). The fused attention primitive keeps the + /// softmax probabilities out of the graph (its backward + /// replays a recompute), so the compiler's flash kernel claims + /// the forward cluster. + fn attention(&self, x: &Tensor<3, E>, mask: &RawTensor<2, E>) -> Tensor<3, E> { + let [batch, seq, dim] = x.shape(); + // (batch, seq, dim) -> (batch, heads, seq, head_dim) + let split = |projected: &Tensor<3, E>| { + projected + .reshape([batch, seq, HEADS, HEAD_DIM]) + .transpose(1, 2) + }; + let query = split(&self.query.forward(x)); + let key = split(&self.key.forward(x)); + let value = split(&self.value.forward(x)); + + let context = query + .attention( + &key, + &value, + 1.0 / (HEAD_DIM as f32).sqrt(), + Some((mask, MaskKind::Causal)), + ) + .transpose(1, 2) + .reshape([batch, seq, dim]); + self.output.forward(&context) + } + + fn forward(&self, x: &Tensor<3, E>, mask: &RawTensor<2, E>) -> Tensor<3, E> { + // Norm statistics stay in f32 (the standard mixed-precision + // recipe); everything else runs in the model dtype. + let normed = to_model(&self.ln1.forward(&to_f32(x))); + let x = x.add(&self.attention(&normed, mask)); + let normed = to_model(&self.ln2.forward(&to_f32(&x))); + let mlp = self.mlp_down.forward(&self.mlp_up.forward(&normed).gelu()); + x.add(&mlp) + } + } + + pub struct Gpt { + leaf_rank1: Vec>, + leaf_rank2: Vec>, + token_embedding: Embedding, + position_embedding: Embedding, + blocks: Vec, + ln_final: LayerNorm<1>, + lm_head: Linear, + } + + impl Gpt { + /// Build the model on `graph`, as trainable leaves (training) + /// or constants (evaluation). Leaves are always f32 masters; + /// the model weights are their `E` views. + pub fn new(graph: &Graph, params: &Params, trainable: bool) -> Self { + fn wrap( + graph: &Graph, + value: &RawTensor, + trainable: bool, + ) -> Tensor { + if trainable { + graph.leaf(value.clone()) + } else { + Tensor::constant_from_raw(graph, value.clone()) + } + } + // Leaf order must match [`Params::tensors_mut`]. + let mut leaf_rank1 = vec![ + wrap(graph, ¶ms.ln_final_weight, trainable), + wrap(graph, ¶ms.ln_final_bias, trainable), + ]; + let mut leaf_rank2 = vec![ + wrap(graph, ¶ms.token_embedding, trainable), + wrap(graph, ¶ms.position_embedding, trainable), + wrap(graph, ¶ms.lm_head, trainable), + ]; + let blocks = params + .blocks + .iter() + .map(|block| { + let ln1_weight = wrap(graph, &block.ln1_weight, trainable); + let ln1_bias = wrap(graph, &block.ln1_bias, trainable); + let ln2_weight = wrap(graph, &block.ln2_weight, trainable); + let ln2_bias = wrap(graph, &block.ln2_bias, trainable); + let mlp_up_bias = wrap(graph, &block.mlp_up_bias, trainable); + let mlp_down_bias = wrap(graph, &block.mlp_down_bias, trainable); + let query = wrap(graph, &block.query, trainable); + let key = wrap(graph, &block.key, trainable); + let value = wrap(graph, &block.value, trainable); + let output = wrap(graph, &block.output, trainable); + let mlp_up = wrap(graph, &block.mlp_up, trainable); + let mlp_down = wrap(graph, &block.mlp_down, trainable); + let built = Block { + ln1: LayerNorm::new( + ln1_weight.clone(), + Some(ln1_bias.clone()), + LAYER_NORM_EPS, + ), + query: Linear::new(to_model(&query), None), + key: Linear::new(to_model(&key), None), + value: Linear::new(to_model(&value), None), + output: Linear::new(to_model(&output), None), + ln2: LayerNorm::new( + ln2_weight.clone(), + Some(ln2_bias.clone()), + LAYER_NORM_EPS, + ), + mlp_up: Linear::new( + to_model(&mlp_up), + Some(to_model(&mlp_up_bias)), + ), + mlp_down: Linear::new( + to_model(&mlp_down), + Some(to_model(&mlp_down_bias)), + ), + }; + leaf_rank1.extend([ + ln1_weight, + ln1_bias, + ln2_weight, + ln2_bias, + mlp_up_bias, + mlp_down_bias, + ]); + leaf_rank2 + .extend([query, key, value, output, mlp_up, mlp_down]); + built + }) + .collect(); + Self { + token_embedding: Embedding::new_from_tensor(to_model(&leaf_rank2[0])), + position_embedding: Embedding::new_from_tensor(to_model( + &leaf_rank2[1], + )), + blocks, + ln_final: LayerNorm::new( + leaf_rank1[0].clone(), + Some(leaf_rank1[1].clone()), + LAYER_NORM_EPS, + ), + lm_head: Linear::new(leaf_rank2[2].clone(), None), + leaf_rank1, + leaf_rank2, + } + } + + /// Parameter leaves in a fixed order; must match + /// [`Params::tensors_mut`]. + pub fn leaves(&self) -> (Vec<&Tensor<1>>, Vec<&Tensor<2>>) { + ( + self.leaf_rank1.iter().collect(), + self.leaf_rank2.iter().collect(), + ) + } + + /// Input shape: (batch, seq) token ids. Output shape: + /// (batch, seq, vocab), always f32. + pub fn forward( + &self, + tokens: &RawTensor<2, u32>, + positions: &RawTensor<1, u32>, + mask: &RawTensor<2, E>, + ) -> Tensor<3> { + let [batch, seq] = tokens.shape(); + let position_embedded = self.position_embedding.forward(positions); + let mut x = self + .token_embedding + .forward(tokens) + .add(&position_embedded.broadcast_as([batch, seq, DIM])); + for block in &self.blocks { + x = block.forward(&x, mask); + } + self.lm_head.forward(&self.ln_final.forward(&to_f32(&x))) + } + } + + pub async fn evaluate_batch( + device: &Device, + params: &Params, + inputs: &[u32], + targets: &[u32], + positions: &RawTensor<1, u32>, + mask: &RawTensor<2, E>, + vocab_size: usize, + ) -> (f32, f32) { + assert_eq!(inputs.len(), targets.len()); + assert!(inputs.len().is_multiple_of(CONTEXT)); + let batch_size = inputs.len() / CONTEXT; + let inputs = RawTensor::from_slice(device, [batch_size, CONTEXT], inputs); + let targets = RawTensor::from_slice(device, [targets.len()], targets); + let graph = Graph::new(); + let model = Gpt::new(&graph, params, false); + let logits = model.forward(&inputs, positions, mask); + let flat_logits = logits.reshape([batch_size * CONTEXT, vocab_size]); + let loss = cross_entropy(&flat_logits, &targets); + let correct = correct_predictions(&flat_logits, &targets); + read_metrics( + loss.raw().clone(), + correct.raw().clone(), + targets.shape()[0], + ) + .await + } + + async fn generate_synchronized( + device: &Device, + params: &Params, + mut tokens: Vec, + length: usize, + rng: &mut Lcg, + ) -> Vec { + assert!(tokens.len() >= CONTEXT, "prompt shorter than CONTEXT"); + let positions: Vec = (0..CONTEXT as u32).collect(); + let positions = RawTensor::from_slice(device, [CONTEXT], &positions); + let mask = causal_mask(device, CONTEXT); + for _ in 0..length { + let window = &tokens[tokens.len() - CONTEXT..]; + let graph = Graph::new(); + let model = Gpt::new(&graph, params, false); + let input = RawTensor::from_slice(device, [1, CONTEXT], window); + let logits = model.forward(&input, &positions, &mask); + let vocab_size = logits.shape()[2]; + let last_logits = logits + .raw() + .clone() + .narrow(1, CONTEXT - 1, 1) + .reshape([vocab_size]) + .into_concrete(); + drop(logits); + let last_logits = last_logits.as_slice().await.unwrap().to_vec(); + tokens.push(sample(&last_logits, rng)); + } + tokens + } + + /// Keep a bounded window of autoregressive steps in flight using + /// GPU-resident sampled tokens. Queue ordering preserves the + /// sliding-window dependency; host readbacks are only needed to + /// assemble the final output text. + async fn generate_gpu_run_ahead( + device: &Device, + params: &Params, + mut tokens: Vec, + length: usize, + rng: &mut Lcg, + ) -> Option> { + assert!(tokens.len() >= CONTEXT, "prompt shorter than CONTEXT"); + let positions: Vec = (0..CONTEXT as u32).collect(); + let positions = RawTensor::from_slice(device, [CONTEXT], &positions); + let mask = causal_mask(device, CONTEXT); + let prompt = + RawTensor::from_slice(device, [1, CONTEXT], &tokens[tokens.len() - CONTEXT..]); + let mut input = prompt; + let mut pending_tokens = VecDeque::with_capacity(GENERATION_RUN_AHEAD + 1); + + for _ in 0..length { + let graph = Graph::new(); + let model = Gpt::new(&graph, params, false); + let logits = model.forward(&input, &positions, &mask); + let vocab_size = logits.shape()[2]; + let last_logits = logits + .raw() + .clone() + .narrow(1, CONTEXT - 1, 1) + .reshape([vocab_size]) + .into_concrete(); + drop(logits); + + let pending = last_logits + .sample_standard_token_pending( + &[], + None, + StandardSamplerParams { + top_k: vocab_size, + temperature: TEMPERATURE, + repetition_penalty: 1.0, + top_p: 1.0, + min_p: 0.0, + random: rng.next_f32(), + }, + ) + .ok()??; + let next = pending.token_tensor().reshape([1, 1]).into_concrete(); + let retained = input.narrow(1, 1, CONTEXT - 1).into_concrete(); + input = cat([retained, next], 1); + pending_tokens.push_back(pending); + if pending_tokens.len() > GENERATION_RUN_AHEAD { + let token = pending_tokens + .pop_front() + .expect("pending token queue cannot be empty") + .read_token() + .await + .ok()??; + tokens.push(token); + } + } + while let Some(pending) = pending_tokens.pop_front() { + let token = pending.read_token().await.ok()??; + tokens.push(token); + } + Some(tokens) + } + + /// Autoregressively extend `tokens` by `length` sampled + /// characters. The prompt must be at least CONTEXT tokens so every + /// forward pass sees the same shapes and reuses the same compiled + /// kernels. + pub async fn generate( + device: &Device, + params: &Params, + tokens: Vec, + length: usize, + rng: &mut Lcg, + ) -> Vec { + match device { + Device::Gpu(_) => { + let rng_state = rng.0; + if let Some(generated) = + generate_gpu_run_ahead(device, params, tokens.clone(), length, rng) + .await + { + generated + } else { + rng.0 = rng_state; + generate_synchronized(device, params, tokens, length, rng).await + } + } + Device::Cpu => { + generate_synchronized(device, params, tokens, length, rng).await + } + } + } + + pub async fn run( + config: RunConfig, + device: Device, + vocab: Vec, + train_tokens: &[u32], + test_tokens: &[u32], + all_tokens: &[u32], + ) { + let mut params = Params::new(&device, vocab.len()); + let (mut adam1, mut adam2) = { + let (rank1, rank2) = params.tensors_mut(); + ( + rank1 + .iter() + .map(|tensor| { + AdamState::zeros(&device, tensor.shape().iter().product()) + }) + .collect::>(), + rank2 + .iter() + .map(|tensor| { + AdamState::zeros(&device, tensor.shape().iter().product()) + }) + .collect::>(), + ) + }; + + let positions: Vec = (0..CONTEXT as u32).collect(); + let positions = RawTensor::from_slice(&device, [CONTEXT], &positions); + let mask = causal_mask(&device, CONTEXT); + let (progress_test_inputs, progress_test_targets) = { + let mut test_rng = Lcg(0x5eed); + sample_batch(test_tokens, &mut test_rng, BATCH_SIZE) + }; + + let mut rng = Lcg(7); + let mut final_loss = None; + let start = std::time::Instant::now(); + for step in 0..config.steps { + let report_progress = config.progress_every != 0 + && ((step + 1).is_multiple_of(config.progress_every) + || step + 1 == config.steps); + let progress = { + let (inputs, targets) = + sample_batch(train_tokens, &mut rng, BATCH_SIZE); + let inputs = + RawTensor::from_slice(&device, [BATCH_SIZE, CONTEXT], &inputs); + let targets = + RawTensor::from_slice(&device, [BATCH_SIZE * CONTEXT], &targets); + + let graph = Graph::new(); + let model = Gpt::new(&graph, ¶ms, true); + let logits = model.forward(&inputs, &positions, &mask); + let flat_logits = + logits.reshape([BATCH_SIZE * CONTEXT, vocab.len()]); + let loss = cross_entropy(&flat_logits, &targets); + let progress = report_progress.then(|| { + let correct = correct_predictions(&flat_logits, &targets); + (loss.raw().clone(), correct.raw().clone()) + }); + + // Gradients stay lazy: no per-step readback. The whole + // step (forward, backward, and the optimizer updates + // below) is submitted to the GPU by the flush after + // step-local temporaries are dropped. + // + // The f16 model seeds the backward with a loss scale: + // activation gradients entering the f16 blocks sit + // near 1e-6, deep in f16's subnormal range, and would + // quantize to noise. Adam's update is invariant to a + // constant gradient scale (m and sqrt(v) scale + // together), so no unscaling is needed. + let loss_scale = if std::mem::size_of::() == 2 { + 1024.0 + } else { + 1.0 + }; + let seed = RawTensor::splat(&device, loss_scale, []); + let gradients = loss.backward_with(seed).unwrap(); + + // Adam with warmup and bias correction folded into the + // learning rate. + let t = step as i32 + 1; + let warmup = (config.steps / 10).max(1); + let lr_value = LEARNING_RATE + * ((step + 1) as f32 / warmup as f32).min(1.0) + * (1.0 - BETA2.powi(t)).sqrt() + / (1.0 - BETA1.powi(t)); + let lr = RawTensor::from_slice(&device, [1], &[lr_value]); + let (leaves1, leaves2) = model.leaves(); + let (tensors1, tensors2) = params.tensors_mut(); + for ((param, state), leaf) in + tensors1.into_iter().zip(&mut adam1).zip(leaves1) + { + adam_step( + param, + state, + gradients.get(leaf).expect("missing gradient"), + &lr, + ); + } + for ((param, state), leaf) in + tensors2.into_iter().zip(&mut adam2).zip(leaves2) + { + adam_step( + param, + state, + gradients.get(leaf).expect("missing gradient"), + &lr, + ); + } + + progress + }; + + // Submit parameter updates after step-local + // forward/backward handles have dropped. Reporting steps + // intentionally retain two scalar metric targets; ordinary + // steps exclude loss, logits, and gradient views. + device.flush(); + + if let Some((loss, correct)) = progress { + let (loss, train_accuracy) = + read_metrics(loss, correct, BATCH_SIZE * CONTEXT).await; + final_loss = Some(loss); + let (_, test_accuracy) = evaluate_batch( + &device, + ¶ms, + &progress_test_inputs, + &progress_test_targets, + &positions, + &mask, + vocab.len(), + ) + .await; + println!( + "step {}/{}: loss {loss:.4} | train accuracy {:.2}% | test accuracy {:.2}%", + step + 1, + config.steps, + train_accuracy * 100.0, + test_accuracy * 100.0, + ); + } + } + // Drain the GPU before timing and before exit: without this, + // sparse progress reporting lets the loop report throughput + // (or terminate the process) with steps still queued on the + // GPU. + wait_for_gpu(&device); + let elapsed = start.elapsed(); + let steps_per_second = config.steps as f64 / elapsed.as_secs_f64(); + println!( + "trained {} steps ({} tokens) in {elapsed:.2?} ({steps_per_second:.1} steps/s)", + config.steps, + config.steps * BATCH_SIZE * CONTEXT, + ); + + if let Some(min_steps_per_second) = config.min_steps_per_second + && steps_per_second < min_steps_per_second + { + panic!( + "transformer throughput {steps_per_second:.1} steps/s below required {min_steps_per_second:.1}" + ); + } + + // Bit-exact loss gate for refactors that claim no numeric + // change: the run is deterministic, so the final reported + // training loss must reproduce exactly. + if let Some(expected) = config.expect_loss { + let final_loss = + final_loss.expect("--expect-loss requires the final step to report progress"); + assert!( + final_loss.to_bits() == expected.to_bits(), + "final loss {final_loss:?} does not match expected {expected:?}" + ); + } + + if config.skip_eval { + return; + } + + // Final metrics over a larger held-out sample. + const TEST_BATCHES: usize = 10; + const TEST_BATCH_SIZE: usize = TEST_BATCHES * BATCH_SIZE; + let (test_inputs, test_targets) = + sample_batch(test_tokens, &mut rng, TEST_BATCH_SIZE); + let (test_loss, test_accuracy) = evaluate_batch( + &device, + ¶ms, + &test_inputs, + &test_targets, + &positions, + &mask, + vocab.len(), + ) + .await; + println!( + "test loss: {test_loss:.4} | test accuracy {:.2}%", + test_accuracy * 100.0 + ); + + // Prompt with the tail of the corpus so the window is always + // full. + let prompt = all_tokens[all_tokens.len() - CONTEXT..].to_vec(); + let generated = generate(&device, ¶ms, prompt, 400, &mut rng).await; + let text: String = generated[CONTEXT..] + .iter() + .map(|&id| vocab[id as usize] as char) + .collect(); + println!("--- sample ---\n{text}"); + } + } + }; +} + +gpt_model!(gpt_f32, f32); +gpt_model!(gpt_f16, half::f16); + +#[tokio::main] +async fn main() { + let config = RunConfig::from_args(); + if std::env::var_os("RUST_LOG").is_some() || config.trace_host || config.trace_resolve { + let env_filter = if std::env::var_os("RUST_LOG").is_some() { + tracing_subscriber::EnvFilter::from_default_env() + } else { + tracing_subscriber::EnvFilter::new( + "fusor_core::compute_graph::resolve=info,fusor_tile_ir_runtime::plan_cache=warn", + ) + }; + tracing_subscriber::fmt().with_env_filter(env_filter).init(); + } + let text = fetch_text(); + // Character-level vocabulary over the corpus bytes (the text is ASCII). + let vocab: Vec = text.bytes().collect::>().into_iter().collect(); + let index_of: std::collections::HashMap = vocab + .iter() + .enumerate() + .map(|(index, &byte)| (byte, index as u32)) + .collect(); + let tokens: Vec = text.bytes().map(|byte| index_of[&byte]).collect(); + let split = tokens.len() * 9 / 10; + let (train_tokens, test_tokens) = tokens.split_at(split); + println!( + "corpus: {} chars, vocab {}, {} train / {} test", + tokens.len(), + vocab.len(), + train_tokens.len(), + test_tokens.len() + ); + + let device = match Device::gpu().await { + Ok(gpu) => gpu, + Err(_) => { + println!("GPU unavailable, training on CPU"); + Device::cpu() + } + }; + if config.trace_host || config.trace_resolve { + // SAFETY: this example sets resolver trace flags before building any + // training graph; the flags are read by Fusor during subsequent + // single-threaded graph resolution. + unsafe { + std::env::set_var("FUSOR_TRACE_RESOLVE_HOST", "1"); + if config.trace_resolve { + std::env::set_var("FUSOR_TRACE_RESOLVE", "1"); + } + if config.trace_names { + std::env::set_var("FUSOR_TRACE_DECODE_NAMES", "1"); + } + } + } + + match config.dtype { + Dtype::F32 => gpt_f32::run(config, device, vocab, train_tokens, test_tokens, &tokens).await, + Dtype::F16 => gpt_f16::run(config, device, vocab, train_tokens, test_tokens, &tokens).await, + } +} diff --git a/fusor-ml/fusor/src/autograd/composite.rs b/fusor-ml/fusor/src/autograd/composite.rs new file mode 100644 index 000000000..9df3a60de --- /dev/null +++ b/fusor-ml/fusor/src/autograd/composite.rs @@ -0,0 +1,1159 @@ +use crate::MaskKind; +use fusor_types::SlidingWindow; + +use super::*; + +impl Tensor +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + fn softmax_composite(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input_shape = self.shape(); + let max_values = self.max_keepdim_any::(axis); + let shifted = self.sub(&max_values.broadcast_as(input_shape)); + let exp_values = shifted.exp(); + let normalization = exp_values + .sum_keepdim_any::(axis) + .broadcast_as(input_shape); + exp_values.div(&normalization) + } + + fn rms_norm_composite( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.layer_norm_composite::(weight, bias, eps, false) + } + + fn layer_norm_composite( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + remove_mean: bool, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let centered = if remove_mean { + let mean = self.mean_keepdim_any::(R - 1); + self.sub(&mean.broadcast_as(self.shape())) + } else { + self.clone() + }; + let variance = centered.sqr().mean_keepdim_any::(R - 1); + let std = variance.add_scalar(eps).sqrt(); + let normalized = centered.div(&std.broadcast_as(self.shape())); + let scaled = normalized.mul(&weight.broadcast_as(self.shape())); + if let Some(bias) = bias { + scaled.add(&bias.broadcast_as(self.shape())) + } else { + scaled + } + } + + pub fn softmax(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + if axis == R - 1 { + // Fused forward with composite replay backward: fewer kernels, same math. + return self.softmax_last_dim_fused::(); + } + self.softmax_composite::(axis) + } + + pub fn softmax_last_dim(&self) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.softmax::(R - 1) + } + + pub fn softmax_slow(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.softmax::(axis) + } + + pub fn softmax_slow_last_dim(&self) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.softmax_slow::(R - 1) + } + + pub fn layer_norm( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + remove_mean: bool, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.layer_norm_composite::(weight, bias, eps, remove_mean) + } + + pub fn rms_norm(&self, weight: &Tensor, eps: f32) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.rms_norm_composite::(weight, None, eps) + } + + pub fn softmax_last_dim_fused(&self) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + { + let value = self.value.softmax_last_dim::().into_concrete(); + // Analytic softmax backward: dS = P * (dP - rowsum(dP * P)). + // The product inside the row sum is written inline so the reduce + // absorbs it, and the output is a single P * (dP - s) expression. + self.unary_from_value(value, move |grad, probs| { + let shape = probs.shape(); + let row_sum = (&probs * &grad) + .into_concrete() + .sum_keepdim::(R - 1); + let shifted = (&grad - &row_sum.broadcast_as(shape)).into_concrete(); + (&probs * &shifted).into_concrete() + }) + } + + pub fn rms_norm_fused( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::MulOp: crate::cpu::SimdBinaryOp, + crate::DivOp: crate::cpu::SimdBinaryOp, + crate::AddOp: crate::cpu::SimdBinaryOp, + crate::SqrtOp: crate::cpu::SimdUnaryOp, + (crate::gpu::Tensor, crate::gpu::Tensor): crate::gpu::MaxRank, + T: crate::CastTensor, + f32: crate::CastTensor, + crate::gpu::Tensor: crate::gpu::LastRank, + { + let value = self.value.rms_norm_fused::( + &weight.value, + bias.as_ref().map(|bias| &bias.value), + eps, + ); + if let Some(bias) = bias { + self.replay_ternary( + weight, + bias, + "rms_norm_fused", + value, + move |input, weight, bias| { + input.rms_norm_composite::(&weight, Some(&bias), eps) + }, + ) + } else { + self.replay_binary(weight, "rms_norm_fused", value, move |input, weight| { + input.rms_norm_composite::(&weight, None, eps) + }) + } + } + + pub fn rms_norm_fused_no_bias( + &self, + weight: &Tensor, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::MulOp: crate::cpu::SimdBinaryOp, + crate::DivOp: crate::cpu::SimdBinaryOp, + crate::AddOp: crate::cpu::SimdBinaryOp, + crate::SqrtOp: crate::cpu::SimdUnaryOp, + (crate::gpu::Tensor, crate::gpu::Tensor): crate::gpu::MaxRank, + T: crate::CastTensor, + f32: crate::CastTensor, + crate::gpu::Tensor: crate::gpu::LastRank, + { + self.rms_norm_fused::(weight, None, eps) + } + + pub fn rms_norm_residual_fused( + &self, + residual: &Self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::MulOp: crate::cpu::SimdBinaryOp, + crate::DivOp: crate::cpu::SimdBinaryOp, + crate::AddOp: crate::cpu::SimdBinaryOp, + crate::SqrtOp: crate::cpu::SimdUnaryOp, + (crate::gpu::Tensor, crate::gpu::Tensor): crate::gpu::MaxRank, + T: crate::CastTensor, + f32: crate::CastTensor, + crate::gpu::Tensor: crate::gpu::LastRank, + { + let value = self.value.rms_norm_residual_fused::( + &residual.value, + &weight.value, + bias.as_ref().map(|bias| &bias.value), + eps, + ); + match bias { + None => self.replay_ternary( + residual, + weight, + "rms_norm_residual_fused", + value, + move |input, residual, weight| { + input + .add(&residual) + .rms_norm_composite::(&weight, None, eps) + }, + ), + Some(bias) => self.replay_quaternary( + residual, + weight, + bias, + "rms_norm_residual_fused", + value, + move |input, residual, weight, bias| { + input.add(&residual).rms_norm_composite::( + &weight, + Some(&bias), + eps, + ) + }, + ), + } + } + + pub fn layer_norm_last_dim_fused( + &self, + weight: &Tensor, + bias: Option<&Tensor>, + eps: f32, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::AddOp: crate::cpu::SimdBinaryOp, + crate::SubOp: crate::cpu::SimdBinaryOp, + crate::MulOp: crate::cpu::SimdBinaryOp, + crate::DivOp: crate::cpu::SimdBinaryOp, + crate::SqrtOp: crate::cpu::SimdUnaryOp, + { + let mut param_shape = [1usize; R]; + param_shape[R - 1] = self.shape()[R - 1]; + let weight_row = weight.value.reshape(param_shape); + let weight_b = weight_row.broadcast_as(self.shape()); + let bias_row = bias.map(|bias| bias.value.reshape(param_shape)); + let bias_b = bias_row + .as_ref() + .map(|bias| bias.broadcast_as(self.shape())); + let value = self.value.layer_norm::( + &weight_b, + bias_b.as_ref(), + T::from_f32(eps), + true, + ); + + let input_id = self.handle.id; + let weight_id = weight.handle.id; + let bias_id = bias.map(|bias| bias.handle.id); + let input_value = self.value.clone(); + let weight_value = weight.value.clone(); + let weight_shape = weight.value.shape(); + // Analytic layer-norm backward. Row statistics are recomputed here + // instead of saved from the forward pass so the forward chain stays + // exclusively consumed (and therefore fusable into one row program). + let backward: BackwardRule = Arc::new(move |gradient| { + let dy = downcast_tensor::(&*gradient, "layer_norm_last_dim_fused")?; + let shape = input_value.shape(); + let n = shape[R - 1]; + let rows: usize = shape.iter().take(R - 1).product(); + + let x = input_value.to_concrete(); + let mean = x.mean_keepdim::(R - 1); + let centered = (&x - &mean.broadcast_as(shape)).into_concrete(); + let var = centered + .sqr() + .into_concrete() + .mean_keepdim::(R - 1); + let std = var.add_scalar(T::from_f32(eps)).sqrt().into_concrete(); + let xhat = (¢ered / &std.broadcast_as(shape)).into_concrete(); + + let weight_row = weight_value.reshape(param_shape).into_concrete(); + let dxhat = (&dy * &weight_row.broadcast_as(shape)).into_concrete(); + let m1 = dxhat.mean_keepdim::(R - 1); + let dxhat_xhat = (&dxhat * &xhat).into_concrete(); + let m2 = dxhat_xhat.mean_keepdim::(R - 1); + let recentered = (&dxhat - &m1.broadcast_as(shape)).into_concrete(); + let projected = (&xhat * &m2.broadcast_as(shape)).into_concrete(); + let dx_num = (recentered - projected).into_concrete(); + let dx = (&dx_num / &std.broadcast_as(shape)).into_concrete(); + + let dy_flat = dy.reshape([rows, n]).into_concrete(); + let xhat_flat = xhat.reshape([rows, n]).into_concrete(); + let dw_flat = (&dy_flat * &xhat_flat).into_concrete().sum::<1>(0); + let dw = dw_flat.reshape(weight_shape).into_concrete(); + + let mut targets = vec![ + BackwardTarget { + node: input_id, + gradient: Box::new(dx), + }, + BackwardTarget { + node: weight_id, + gradient: Box::new(dw), + }, + ]; + if let Some(bias_id) = bias_id { + let db = dy_flat.sum::<1>(0).reshape(weight_shape).into_concrete(); + targets.push(BackwardTarget { + node: bias_id, + gradient: Box::new(db), + }); + } + Ok(targets) + }); + let mut parents = vec![self.handle.clone(), weight.handle.clone()]; + if let Some(bias) = bias { + parents.push(bias.handle.clone()); + } + self.emit_op(value, parents, Some(backward)) + } + + fn pad_spatial(&self, padding: [usize; DIFF]) -> Self { + let mut padded = self.clone(); + for (i, padding) in padding.into_iter().enumerate() { + padded = padded.pad_axis(R - DIFF + i, padding); + } + padded + } + + fn conv_output_shape( + input_shape: [usize; R], + out_channels: usize, + kernel: [usize; DIFF], + padding: [usize; DIFF], + strides: [usize; DIFF], + ) -> [usize; R] { + let spatial_start = R - DIFF; + let mut output_shape = input_shape; + output_shape[1] = out_channels; + for i in 0..DIFF { + let padded_len = input_shape[spatial_start + i] + 2 * padding[i]; + output_shape[spatial_start + i] = (padded_len - kernel[i]) / strides[i] + 1; + } + output_shape + } + + /// Pad + sliding-window view + flatten to one matmul row per output + /// location: `(batch * out_spatial, in_channels * kernel_size)`. + fn conv_windows_flat( + &self, + kernel: [usize; DIFF], + padding: [usize; DIFF], + strides: [usize; DIFF], + ) -> Tensor<2, T> + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let input_shape = self.shape(); + let spatial_start = R - DIFF; + let output_shape = Self::conv_output_shape(input_shape, 0, kernel, padding, strides); + let windows: [SlidingWindow; DIFF] = + std::array::from_fn(|i| SlidingWindow::new(spatial_start + i, kernel[i], strides[i])); + let windows: Tensor = self.pad_spatial(padding).sliding_window_view(windows); + let permutation: [usize; R2] = std::array::from_fn(|index| { + if index == 0 { + 0 + } else if index <= DIFF { + index + 1 + } else if index == DIFF + 1 { + 1 + } else { + index + } + }); + let out_spatial_size: usize = output_shape[spatial_start..].iter().product(); + let kernel_size: usize = kernel.iter().product(); + windows.permute(permutation).reshape([ + input_shape[0] * out_spatial_size, + input_shape[1] * kernel_size, + ]) + } + + /// The same window operand as [`Self::conv_windows_flat`], built from one + /// shifted slice per kernel offset instead of a sliding-window view, with + /// columns ordered kernel-major (`offset * in_channels + channel`). + /// + /// A sliding-window view reads each input element from several output + /// locations, so its transpose is an overlap-add — which the generic view + /// backward expresses as a masked reduce over *every* (input position, + /// output position) pair. That is quadratic in the spatial extent: a + /// 768-long sequence spends 1.8 billion element visits per convolution to + /// scatter 5.7 million gradients. Concatenated slices carry the same + /// values with a linear backward — `cat` differentiates to slices and + /// `narrow` to a zero-fill assign — and still materialize exactly once, + /// because concatenating along the trailing (channel) axis of a + /// channels-last view already produces the contiguous matmul operand. + /// + /// Only valid for unit strides; strided windows keep the view path. + fn conv_windows_shifted( + &self, + kernel: [usize; DIFF], + padding: [usize; DIFF], + out_spatial: [usize; DIFF], + ) -> Tensor<2, T> { + let input_shape = self.shape(); + let in_channels = input_shape[1]; + let kernel_size: usize = kernel.iter().product(); + let out_spatial_size: usize = out_spatial.iter().product(); + // (batch, channels, ...spatial) -> (batch, ...spatial, channels) + let channels_last: [usize; R] = std::array::from_fn(|axis| { + if axis == 0 { + 0 + } else if axis < R - 1 { + axis + 1 + } else { + 1 + } + }); + let padded = self.pad_spatial(padding).permute(channels_last); + let slices: Vec = (0..kernel_size) + .map(|flat_offset| { + let mut window = padded.clone(); + let mut rest = flat_offset; + for axis in (0..DIFF).rev() { + let offset = rest % kernel[axis]; + rest /= kernel[axis]; + window = window.narrow(1 + axis, offset, out_spatial[axis]); + } + window + }) + .collect(); + Self::cat(slices, R - 1).reshape([ + input_shape[0] * out_spatial_size, + kernel_size * in_channels, + ]) + } + + /// Reshape the `(batch * out_spatial, out_channels)` matmul output back to + /// `(batch, out_channels, ...out_spatial)` and add the broadcast bias. + fn conv_reassemble( + output: Tensor<2, T>, + bias: Option<&Tensor<1, T>>, + output_shape: [usize; R], + ) -> Self { + let out_channels = output_shape[1]; + // Add the bias to the matmul's own `(rows, out_channels)` output, + // where it broadcasts along the trailing axis, rather than after the + // reassembly permute: an elementwise expression sitting directly on + // the matmul result can ride its epilogue, while one behind a permute + // costs a separate full pass over the activation. + let output = match bias { + Some(bias) => output.add_::<1, 2>(bias), + None => output, + }; + let output: Tensor = output.reshape(std::array::from_fn(|axis| { + if axis == 0 { + output_shape[0] + } else if axis <= DIFF { + output_shape[axis + 1] + } else { + out_channels + } + })); + let permutation: [usize; R] = std::array::from_fn(|index| { + if index == 0 { + 0 + } else if index == 1 { + DIFF + 1 + } else { + index - 1 + } + }); + output.permute(permutation) + } + + fn conv_composite( + &self, + weight: &Tensor, + bias: Option<&Tensor<1, T>>, + padding: [usize; DIFF], + strides: [usize; DIFF], + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + assert_eq!( + R, + 2 + DIFF, + "Conv expects (batch, channels, ...spatial) format where R = 2 + DIFF" + ); + let input_shape = self.shape(); + let weight_shape = weight.shape(); + let spatial_start = R - DIFF; + let in_channels = input_shape[1]; + let out_channels = weight_shape[0]; + assert_eq!( + weight_shape[1], in_channels, + "Weight in_channels must match input in_channels" + ); + + let kernel: [usize; DIFF] = std::array::from_fn(|i| weight_shape[spatial_start + i]); + let kernel_size: usize = kernel.iter().product(); + let output_shape = + Self::conv_output_shape(input_shape, out_channels, kernel, padding, strides); + + // Unit strides take the shifted-slice operand, whose backward is + // linear in the spatial extent; strided windows keep the view. + // Its columns are kernel-major, so the weight is laid out to match: + // (out, in, ...kernel) -> (...kernel, in, out). + let output = if strides.iter().all(|stride| *stride == 1) { + let out_spatial: [usize; DIFF] = + std::array::from_fn(|i| output_shape[spatial_start + i]); + let windows_flat = self.conv_windows_shifted::(kernel, padding, out_spatial); + let kernel_major: [usize; WEIGHT_RANK] = std::array::from_fn(|axis| { + if axis < DIFF { + axis + spatial_start + } else if axis == DIFF { + 1 + } else { + 0 + } + }); + let weight_rows = weight + .permute(kernel_major) + .reshape([kernel_size * in_channels, out_channels]); + windows_flat.mat_mul_internal(&weight_rows) + } else { + let windows_flat = self.conv_windows_flat::(kernel, padding, strides); + let weight_t = weight + .reshape([out_channels, in_channels * kernel_size]) + .transpose(0, 1); + windows_flat.mat_mul_internal(&weight_t) + }; + + Self::conv_reassemble::(output, bias, output_shape) + } + + pub fn conv( + &self, + weight: &Tensor, + bias: Option<&Tensor<1, T>>, + padding: [usize; DIFF], + strides: [usize; DIFF], + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let value = self.value.conv::( + &weight.value, + bias.map(|bias| &bias.value), + padding, + strides, + ); + match bias { + None => self.replay_binary(weight, "conv", value, move |input, weight| { + input.conv_composite::(&weight, None, padding, strides) + }), + Some(bias) => { + self.replay_ternary(weight, bias, "conv", value, move |input, weight, bias| { + input.conv_composite::( + &weight, + Some(&bias), + padding, + strides, + ) + }) + } + } + } + + fn grouped_conv_composite( + &self, + weight: &Tensor, + bias: Option<&Tensor<1, T>>, + padding: [usize; DIFF], + strides: [usize; DIFF], + groups: usize, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + assert_eq!(R, 2 + DIFF); + let input_shape = self.shape(); + let weight_shape = weight.shape(); + let spatial_start = R - DIFF; + let batch = input_shape[0]; + let in_channels = input_shape[1]; + let out_channels = weight_shape[0]; + assert_eq!(in_channels % groups, 0); + assert_eq!(out_channels % groups, 0); + let in_ch_per_group = in_channels / groups; + let out_ch_per_group = out_channels / groups; + assert_eq!(weight_shape[1], in_ch_per_group); + + let kernel: [usize; DIFF] = std::array::from_fn(|i| weight_shape[spatial_start + i]); + let kernel_size: usize = kernel.iter().product(); + let output_shape = + Self::conv_output_shape(input_shape, out_channels, kernel, padding, strides); + let out_spatial_size: usize = output_shape[spatial_start..].iter().product(); + + let windows_grouped = self + .conv_windows_flat::(kernel, padding, strides) + .reshape([ + batch * out_spatial_size, + groups, + in_ch_per_group * kernel_size, + ]) + .transpose(0, 1); + let weight_grouped_t = weight + .reshape([groups, out_ch_per_group, in_ch_per_group * kernel_size]) + .transpose(1, 2); + let output = windows_grouped + .mat_mul_internal(&weight_grouped_t) + .transpose(0, 1) + .reshape([batch * out_spatial_size, out_channels]); + Self::conv_reassemble::(output, bias, output_shape) + } + + pub fn grouped_conv( + &self, + weight: &Tensor, + bias: Option<&Tensor<1, T>>, + padding: [usize; DIFF], + strides: [usize; DIFF], + groups: usize, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let value = self.value.grouped_conv::( + &weight.value, + bias.map(|bias| &bias.value), + padding, + strides, + groups, + ); + match bias { + None => self.replay_binary(weight, "grouped_conv", value, move |input, weight| { + input.grouped_conv_composite::( + &weight, None, padding, strides, groups, + ) + }), + Some(bias) => self.replay_ternary( + weight, + bias, + "grouped_conv", + value, + move |input, weight, bias| { + input.grouped_conv_composite::( + &weight, + Some(&bias), + padding, + strides, + groups, + ) + }, + ), + } + } +} + +impl Tensor<4, T> +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + fn rotate_half(&self) -> Tensor<4, T> { + let [batch, heads, sequence_length, embed] = self.shape(); + let half = embed / 2; + let first_half = self.narrow(3, 0, half); + let second_half = self.narrow(3, half, embed - half).mul_scalar(-1.0); + let graph = self.graph(); + let device = self.device(); + let zeros = Tensor::zeros(&graph, &device, [batch, heads, sequence_length, embed]); + let combined = zeros.slice_assign( + [0..batch, 0..heads, 0..sequence_length, 0..half], + &second_half, + ); + combined.slice_assign( + [0..batch, 0..heads, 0..sequence_length, half..embed], + &first_half, + ) + } + + fn rope_interleaved_composite(&self, cos: &Tensor<2, T>, sin: &Tensor<2, T>) -> Tensor<4, T> { + assert_same_graph(self, cos); + assert_same_graph(self, sin); + + let [batch, heads, sequence_length, embed] = self.shape(); + let half = embed / 2; + let cos = cos + .narrow(0, 0, sequence_length) + .reshape([sequence_length, half, 1]) + .broadcast_as([batch, heads, sequence_length, half, 1]); + let sin = sin + .narrow(0, 0, sequence_length) + .reshape([sequence_length, half, 1]) + .broadcast_as([batch, heads, sequence_length, half, 1]); + let x = self.reshape([batch, heads, sequence_length, half, 2]); + let x0 = x.narrow(4, 0, 1); + let x1 = x.narrow(4, 1, 1); + let y0 = x0.mul(&cos).sub(&x1.mul(&sin)); + let y1 = x0.mul(&sin).add(&x1.mul(&cos)); + let graph = self.graph(); + let device = self.device(); + let zeros = Tensor::zeros(&graph, &device, [batch, heads, sequence_length, half, 2]); + let combined = + zeros.slice_assign([0..batch, 0..heads, 0..sequence_length, 0..half, 0..1], &y0); + combined + .slice_assign([0..batch, 0..heads, 0..sequence_length, 0..half, 1..2], &y1) + .flatten_last_n::<1, 4>() + } + + pub fn rope(&self, cos: &Tensor<2, T>, sin: &Tensor<2, T>) -> Tensor<4, T> { + assert_same_graph(self, cos); + assert_same_graph(self, sin); + + let [batch, heads, sequence_length, embed] = self.shape(); + let half = embed / 2; + let graph = self.graph(); + let device = self.device(); + let cos_base = cos.narrow(0, 0, sequence_length); + let sin_base = sin.narrow(0, 0, sequence_length); + let cos = Tensor::zeros(&graph, &device, [sequence_length, embed]) + .slice_assign([0..sequence_length, 0..half], &cos_base) + .slice_assign([0..sequence_length, half..embed], &cos_base) + .unsqueeze_dims::<2, 4>([0, 1]) + .broadcast_as([batch, heads, sequence_length, embed]); + let sin = Tensor::zeros(&graph, &device, [sequence_length, embed]) + .slice_assign([0..sequence_length, 0..half], &sin_base) + .slice_assign([0..sequence_length, half..embed], &sin_base) + .unsqueeze_dims::<2, 4>([0, 1]) + .broadcast_as([batch, heads, sequence_length, embed]); + let rotated = self.rotate_half(); + self.mul(&cos).add(&rotated.mul(&sin)) + } + + pub fn rope_interleaved(&self, cos: &Tensor<2, T>, sin: &Tensor<2, T>) -> Tensor<4, T> { + self.rope_interleaved_composite(cos, sin) + } + + pub fn attention( + &self, + k: &Tensor<4, T>, + v: &Tensor<4, T>, + scale: f32, + mask: Option<(&RawTensor<2, T>, MaskKind)>, + ) -> Tensor<4, T> { + let value = self.value.attention(&k.value, &v.value, scale, mask); + let mask_value = mask.map(|(mask, kind)| (mask.clone(), kind)); + // The explicit rule recomputes probabilities from the forward output + // and its row log-sum-exp, so no probability matrix survives into + // the graph and pattern recognition can stream every piece. + // Grouped-query, batch-key-masked, and CPU shapes replay the + // composite as before. + let explicit = self.shape()[1] == k.shape()[1] + && matches!(&self.value, RawTensor::Gpu(_)) + && !matches!(mask, Some((_, MaskKind::BatchKeyMask))); + if explicit { + let q_id = self.handle.id; + let k_id = k.handle.id; + let v_id = v.handle.id; + let q_value = self.value.clone(); + let k_value = k.value.clone(); + let v_value = v.value.clone(); + let o_value = value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let grad = downcast_tensor::<4, T>(&*gradient, "attention")?; + let ( + RawTensor::Gpu(q), + RawTensor::Gpu(k), + RawTensor::Gpu(v), + RawTensor::Gpu(o), + RawTensor::Gpu(grad), + ) = (&q_value, &k_value, &v_value, &o_value, &grad) + else { + return Err(Error::msg("attention gradient expects GPU tensors")); + }; + let causal = matches!(mask_value, Some((_, MaskKind::Causal))); + let mask_gpu = match &mask_value { + Some((RawTensor::Gpu(mask), MaskKind::QKMask)) => Some(mask), + Some((_, MaskKind::QKMask)) => { + return Err(Error::msg("attention mask must be a GPU tensor")); + } + _ => None, + }; + let lse = q.attention_lse(k, scale, mask_gpu, causal); + let (dq, dk, dv) = q.attention_grads(k, v, o, grad, &lse, scale, mask_gpu, causal); + Ok(vec![ + BackwardTarget { + node: q_id, + gradient: Box::new(RawTensor::Gpu(dq)), + }, + BackwardTarget { + node: k_id, + gradient: Box::new(RawTensor::Gpu(dk)), + }, + BackwardTarget { + node: v_id, + gradient: Box::new(RawTensor::Gpu(dv)), + }, + ]) + }); + return self.emit_op( + value, + vec![self.handle.clone(), k.handle.clone(), v.handle.clone()], + Some(backward), + ); + } + self.replay_ternary(k, v, "attention", value, move |q, k, v| { + q.attention_composite(&k, &v, scale, mask_value.as_ref()) + }) + } + + pub fn rope_fused(&self, cos: &Tensor<2, T>, sin: &Tensor<2, T>) -> Tensor<4, T> { + assert_same_graph(self, cos); + assert_same_graph(self, sin); + + let value = self + .value + .rope_fused(&cos.value, &sin.value) + .into_concrete(); + self.replay_ternary(cos, sin, "rope_fused", value, |input, cos, sin| { + input.rope_interleaved_composite(&cos, &sin) + }) + } + + pub fn rope_normal_fused(&self, cos: &Tensor<2, T>, sin: &Tensor<2, T>) -> Tensor<4, T> { + assert_same_graph(self, cos); + assert_same_graph(self, sin); + + let value = self + .value + .rope_normal_fused(&cos.value, &sin.value) + .into_concrete(); + self.replay_ternary(cos, sin, "rope_normal_fused", value, |input, cos, sin| { + input.rope(&cos, &sin) + }) + } + + pub fn rope_pair_fused( + &self, + k: &Self, + cos: &Tensor<2, T>, + sin: &Tensor<2, T>, + ) -> (Tensor<4, T>, Tensor<4, T>) { + let (q_value, k_value) = self.value.rope_pair_fused(&k.value, &cos.value, &sin.value); + ( + self.replay_ternary( + cos, + sin, + "rope_pair_fused", + q_value.into_concrete(), + |input, cos, sin| input.rope_interleaved_composite(&cos, &sin), + ), + k.replay_ternary( + cos, + sin, + "rope_pair_fused", + k_value.into_concrete(), + |input, cos, sin| input.rope_interleaved_composite(&cos, &sin), + ), + ) + } + + pub fn rope_normal_pair_fused( + &self, + k: &Self, + cos: &Tensor<2, T>, + sin: &Tensor<2, T>, + ) -> (Tensor<4, T>, Tensor<4, T>) { + let (q_value, k_value) = self + .value + .rope_normal_pair_fused(&k.value, &cos.value, &sin.value); + ( + self.replay_ternary( + cos, + sin, + "rope_normal_pair_fused", + q_value.into_concrete(), + |input, cos, sin| input.rope(&cos, &sin), + ), + k.replay_ternary( + cos, + sin, + "rope_normal_pair_fused", + k_value.into_concrete(), + |input, cos, sin| input.rope(&cos, &sin), + ), + ) + } + + pub fn upsample_nearest2d(&self, scale_h: usize, scale_w: usize) -> Tensor<4, T> { + let value = self.value.upsample_nearest2d(scale_h, scale_w); + self.replay_unary("upsample_nearest2d", value, move |input| { + let [b, c, h, w] = input.shape(); + input + .reshape([b, c, h, 1, w, 1]) + .broadcast_as([b, c, h, scale_h, w, scale_w]) + .reshape([b, c, h * scale_h, w * scale_w]) + }) + } + + pub(super) fn attention_composite( + &self, + k: &Tensor<4, T>, + v: &Tensor<4, T>, + scale: f32, + mask: Option<&(RawTensor<2, T>, MaskKind)>, + ) -> Tensor<4, T> { + let q_shape = self.shape(); + let k_shape = k.shape(); + let batch = q_shape[0]; + let num_heads = q_shape[1]; + let q_seq_len = q_shape[2]; + let head_dim = q_shape[3]; + let num_kv_heads = k_shape[1]; + let kv_seq_len = k_shape[2]; + assert!( + num_heads.is_multiple_of(num_kv_heads), + "Number of Q heads ({num_heads}) must be divisible by number of K/V heads ({num_kv_heads})" + ); + + let num_key_value_groups = num_heads / num_kv_heads; + let (k_expanded, v_expanded) = if num_key_value_groups > 1 { + let k_broadcast = k + .reshape([batch, num_kv_heads, 1, kv_seq_len, head_dim]) + .broadcast_as([ + batch, + num_kv_heads, + num_key_value_groups, + kv_seq_len, + head_dim, + ]); + let v_broadcast = v + .reshape([batch, num_kv_heads, 1, kv_seq_len, head_dim]) + .broadcast_as([ + batch, + num_kv_heads, + num_key_value_groups, + kv_seq_len, + head_dim, + ]); + ( + k_broadcast.reshape([batch, num_heads, kv_seq_len, head_dim]), + v_broadcast.reshape([batch, num_heads, kv_seq_len, head_dim]), + ) + } else { + (k.clone(), v.clone()) + }; + + let scores = self + .mat_mul_internal(&k_expanded.transpose(2, 3)) + .div_scalar(scale.recip()); + let masked_scores = if let Some((mask, kind)) = mask { + let mask_tensor = Tensor::constant_from_raw(&self.graph(), mask.clone()); + let mask_4d = match kind { + // Causal falls back to QKMask semantics here: the provided mask is the + // [seq_len, seq_len] (== [q_seq_len, kv_seq_len]) additive causal mask. + MaskKind::QKMask | MaskKind::Causal => { + assert_eq!(mask_tensor.shape(), [q_seq_len, kv_seq_len]); + mask_tensor.reshape([1, 1, q_seq_len, kv_seq_len]) + } + MaskKind::BatchKeyMask => { + assert_eq!(mask_tensor.shape(), [batch, kv_seq_len]); + mask_tensor.reshape([batch, 1, 1, kv_seq_len]) + } + }; + scores.add(&mask_4d.broadcast_as([batch, num_heads, q_seq_len, kv_seq_len])) + } else { + scores + }; + masked_scores + .softmax_last_dim::<3>() + .mat_mul_internal(&v_expanded) + } +} + +/// RoPE-cache companions stay on f32: [`crate::RopeCache`] holds f32 tables. +impl Tensor<4> { + fn rope_cache_tables( + &self, + cache: &crate::RopeCache, + start_pos: usize, + ) -> (Tensor<2>, Tensor<2>) { + let seq_len = self.shape()[2]; + let graph = self.graph(); + let table = |table: RawTensor<2, f32>| { + Tensor::constant_from_raw(&graph, table.narrow(0, start_pos, seq_len).into_concrete()) + }; + (table(cache.cos().clone()), table(cache.sin().clone())) + } + + /// Autograd companion of [`crate::RopeCache::forward`]: applies normal + /// RoPE to `self` (q) and `k` from the cache's tables at `start_pos`. + pub fn rope_cache_forward( + &self, + k: &Self, + cache: &crate::RopeCache, + start_pos: usize, + ) -> (Tensor<4>, Tensor<4>) { + let (cos, sin) = self.rope_cache_tables(cache, start_pos); + self.rope_normal_pair_fused(k, &cos, &sin) + } + + /// Autograd companion of [`crate::RopeCache::forward_interleaved`]. + pub fn rope_cache_forward_interleaved( + &self, + k: &Self, + cache: &crate::RopeCache, + start_pos: usize, + ) -> (Tensor<4>, Tensor<4>) { + let (cos, sin) = self.rope_cache_tables(cache, start_pos); + self.rope_pair_fused(k, &cos, &sin) + } +} diff --git a/fusor-ml/fusor/src/autograd/elementwise.rs b/fusor-ml/fusor/src/autograd/elementwise.rs new file mode 100644 index 000000000..094d8cd7b --- /dev/null +++ b/fusor-ml/fusor/src/autograd/elementwise.rs @@ -0,0 +1,1054 @@ +use super::*; + +impl Tensor +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + pub fn add(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + (self.value.clone() + rhs.value.clone()).into_concrete(), + |grad, _, _| vec![grad.clone().into_concrete(), grad.into_concrete()], + ) + } + + pub fn add_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.add(&rhs) + } + + pub fn sub(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + (self.value.clone() - rhs.value.clone()).into_concrete(), + |grad, _, _| vec![grad.clone().into_concrete(), (-grad).into_concrete()], + ) + } + + pub fn sub_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.sub(&rhs) + } + + pub fn mul(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + (self.value.clone() * rhs.value.clone()).into_concrete(), + |grad, lhs, rhs| { + vec![ + (grad.clone() * rhs).into_concrete(), + (grad * lhs).into_concrete(), + ] + }, + ) + } + + pub fn mul_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.mul(&rhs) + } + + pub fn div(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + (self.value.clone() / rhs.value.clone()).into_concrete(), + |grad, lhs, rhs| { + let lhs_grad = (grad.clone() / rhs.clone()).into_concrete(); + let rhs_grad = (-((grad * lhs) / rhs.sqr().into_concrete())).into_concrete(); + vec![lhs_grad, rhs_grad] + }, + ) + } + + pub fn div_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.div(&rhs) + } + + pub fn pow(&self, rhs: &Self) -> Self { + self.binary_op( + rhs, + self.value.pow(&rhs.value).into_concrete(), + |grad, lhs, rhs| { + let rhs_minus_one = rhs.sub_scalar(T::from_f32(1.0)).into_concrete(); + let lhs_power = lhs.pow(&rhs_minus_one).into_concrete(); + let lhs_grad = + ((grad.clone() * rhs.clone()).into_concrete() * lhs_power).into_concrete(); + let rhs_grad = ((grad * lhs.pow(&rhs).into_concrete()).into_concrete() + * lhs.log().into_concrete()) + .into_concrete(); + vec![lhs_grad, rhs_grad] + }, + ) + } + + pub fn pow_(&self, second: &Tensor) -> Tensor { + let out_shape: [usize; R3] = + crate::composite::broadcast_shapes(&self.shape(), &second.shape()); + let lhs = self.broadcast_as(out_shape); + let rhs = second.broadcast_as(out_shape); + lhs.pow(&rhs) + } + + pub fn pow_elementwise(&self, exponent: f32) -> Self { + let input = self.value.clone(); + self.unary_from_value( + self.value + .pow_elementwise(T::from_f32(exponent)) + .into_concrete(), + move |grad, _| { + let power = input + .pow_elementwise(T::from_f32(exponent - 1.0)) + .into_concrete(); + (grad * power) + .into_concrete() + .mul_scalar(T::from_f32(exponent)) + .into_concrete() + }, + ) + } + + pub fn pow_scalar(&self, exponent: f32) -> Self { + self.pow_elementwise(exponent) + } + + pub fn add_scalar(&self, scalar: f32) -> Self { + self.unary_from_value( + self.value.add_scalar(T::from_f32(scalar)), + move |grad, _| grad, + ) + } + + pub fn sub_scalar(&self, scalar: f32) -> Self { + self.unary_from_value( + self.value.sub_scalar(T::from_f32(scalar)), + move |grad, _| grad, + ) + } + + pub fn mul_scalar(&self, scalar: f32) -> Self { + self.unary_from_value( + self.value.mul_scalar(T::from_f32(scalar)).into_concrete(), + move |grad, _| grad.mul_scalar(T::from_f32(scalar)).into_concrete(), + ) + } + + pub fn div_scalar(&self, scalar: f32) -> Self { + self.unary_from_value( + self.value.div_scalar(T::from_f32(scalar)).into_concrete(), + move |grad, _| grad.div_scalar(T::from_f32(scalar)).into_concrete(), + ) + } + + pub fn neg(&self) -> Self { + self.unary_from_value((-self.value.clone()).into_concrete(), move |grad, _| { + (-grad).into_concrete() + }) + } + + pub fn sqr(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.sqr().into_concrete(), move |grad, _| { + ((grad * input.clone()) + .into_concrete() + .mul_scalar(T::from_f32(2.0))) + .into_concrete() + }) + } + + pub fn abs(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.abs().into_concrete(), move |grad, _| { + let positive = input.mt(T::from_f32(0.0)).into_concrete(); + let negative = input.lt(T::from_f32(0.0)).into_concrete(); + ((grad.clone() * positive).into_concrete() - (grad * negative).into_concrete()) + .into_concrete() + }) + } + + pub fn acos(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.acos().into_concrete(), move |grad, _| { + let denom = (RawTensor::splat(&input.device(), T::from_f32(1.0), input.shape()) + - input.sqr().into_concrete()) + .into_concrete() + .sqrt() + .into_concrete(); + (-(grad / denom).into_concrete()).into_concrete() + }) + } + + pub fn acosh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.acosh().into_concrete(), move |grad, _| { + let lower = input + .add_scalar(T::from_f32(-1.0)) + .into_concrete() + .sqrt() + .into_concrete(); + let upper = input + .add_scalar(T::from_f32(1.0)) + .into_concrete() + .sqrt() + .into_concrete(); + (grad / (lower * upper).into_concrete()).into_concrete() + }) + } + + pub fn approximate_exp(&self) -> Self { + self.unary_from_value( + self.value.approximate_exp().into_concrete(), + move |grad, out| (grad * out).into_concrete(), + ) + } + + pub fn asin(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.asin().into_concrete(), move |grad, _| { + let denom = (RawTensor::splat(&input.device(), T::from_f32(1.0), input.shape()) + - input.sqr().into_concrete()) + .into_concrete() + .sqrt() + .into_concrete(); + (grad / denom).into_concrete() + }) + } + + pub fn asinh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.asinh().into_concrete(), move |grad, _| { + let denom = input + .sqr() + .add_scalar(T::from_f32(1.0)) + .into_concrete() + .sqrt() + .into_concrete(); + (grad / denom).into_concrete() + }) + } + + pub fn atan(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.atan().into_concrete(), move |grad, _| { + let denom = input.sqr().add_scalar(T::from_f32(1.0)).into_concrete(); + (grad / denom).into_concrete() + }) + } + + pub fn atanh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.atanh().into_concrete(), move |grad, _| { + let denom = (RawTensor::splat(&input.device(), T::from_f32(1.0), input.shape()) + - input.sqr().into_concrete()) + .into_concrete(); + (grad / denom).into_concrete() + }) + } + + pub fn cos(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.cos().into_concrete(), move |grad, _| { + (-(grad * input.sin().into_concrete()).into_concrete()).into_concrete() + }) + } + + pub fn cosh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.cosh().into_concrete(), move |grad, _| { + (grad * input.sinh().into_concrete()).into_concrete() + }) + } + + pub fn exp2(&self) -> Self { + self.unary_from_value(self.value.exp2().into_concrete(), move |grad, out| { + (grad * out) + .into_concrete() + .mul_scalar(T::from_f32(std::f32::consts::LN_2)) + .into_concrete() + }) + } + + pub fn less_approximate_exp(&self) -> Self { + self.unary_from_value( + self.value.less_approximate_exp().into_concrete(), + move |grad, out| (grad * out).into_concrete(), + ) + } + + pub fn log2(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.log2().into_concrete(), move |grad, _| { + (grad / input.clone()) + .into_concrete() + .div_scalar(T::from_f32(std::f32::consts::LN_2)) + .into_concrete() + }) + } + + pub fn sin(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.sin().into_concrete(), move |grad, _| { + (grad * input.cos().into_concrete()).into_concrete() + }) + } + + pub fn sinh(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.sinh().into_concrete(), move |grad, _| { + (grad * input.cosh().into_concrete()).into_concrete() + }) + } + + pub fn tan(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.tan().into_concrete(), move |grad, _| { + let cos = input.cos().into_concrete(); + (grad / (cos.clone() * cos).into_concrete()).into_concrete() + }) + } + + pub fn tanh_exact(&self) -> Self { + self.unary_from_value(self.value.tanh_exact().into_concrete(), move |grad, out| { + let one_minus_sq = (RawTensor::splat(&out.device(), T::from_f32(1.0), out.shape()) + - out.sqr().into_concrete()) + .into_concrete(); + (grad * one_minus_sq).into_concrete() + }) + } + + /// Cast the raw value to another element type, dropping the tape. + pub fn cast_raw(&self) -> crate::Tensor + where + T: crate::CastTo + crate::CastTensor, + D2: crate::SimdElement + crate::DataType + Default, + { + self.value.cast() + } + + /// Differentiable cast between autograd element types: the forward casts + /// the value, the backward casts the gradient back. This is the bridge for + /// mixed-precision training (e.g. f32 master weights feeding an f16 model). + pub fn cast(&self) -> Tensor + where + T2: AutogradElement, + T: crate::CastElement, + T2: crate::CastElement, + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + { + let value = self.value.cast::(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "cast")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.cast::()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn to_concrete(&self) -> Self { + self.unary_from_value(self.value.to_concrete(), move |grad, _| grad) + } + + pub fn relu(&self) -> Self { + self.max_elementwise(0.0) + } + + pub fn clamp(&self, min: f32, max: f32) -> Self { + let input = self.value.clone(); + self.unary_from_value( + self.value + .clamp(T::from_f32(min), T::from_f32(max)) + .into_concrete(), + move |grad, _| { + let lower = input.mt(T::from_f32(min)).into_concrete(); + let upper = input.lt(T::from_f32(max)).into_concrete(); + ((grad * lower).into_concrete() * upper).into_concrete() + }, + ) + } + + pub fn eq(&self, rhs: f32) -> Self { + self.unary_from_value( + self.value.eq(T::from_f32(rhs)).into_concrete(), + move |_, out| RawTensor::zeros(&out.device(), out.shape()), + ) + } + + pub fn eq_scalar(&self, rhs: f32) -> Self { + self.eq(rhs) + } + + pub fn eq_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.eq_tensor(&rhs.value).into_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn gt_scalar(&self, rhs: f32) -> Self { + self.unary_from_value( + self.value.gt_scalar(T::from_f32(rhs)).into_concrete(), + move |_, out| RawTensor::zeros(&out.device(), out.shape()), + ) + } + + pub fn gt_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.gt_tensor(&rhs.value).into_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn gte_scalar(&self, rhs: f32) -> Self { + self.unary_from_value( + self.value.gte_scalar(T::from_f32(rhs)).into_concrete(), + move |_, out| RawTensor::zeros(&out.device(), out.shape()), + ) + } + + pub fn gte_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.gte_tensor(&rhs.value).into_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn lt(&self, rhs: f32) -> Self { + self.unary_from_value( + self.value.lt(T::from_f32(rhs)).into_concrete(), + move |_, out| RawTensor::zeros(&out.device(), out.shape()), + ) + } + + pub fn lt_scalar(&self, rhs: f32) -> Self { + self.lt(rhs) + } + + pub fn lt_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.lt_tensor(&rhs.value).into_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn lte(&self, rhs: f32) -> Self { + self.unary_from_value( + self.value.lte(T::from_f32(rhs)).into_concrete(), + move |_, out| RawTensor::zeros(&out.device(), out.shape()), + ) + } + + pub fn lte_scalar(&self, rhs: f32) -> Self { + self.lte(rhs) + } + + pub fn lte_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.lte_tensor(&rhs.value).into_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn max_elementwise(&self, rhs: f32) -> Self { + let input = self.value.clone(); + self.unary_from_value( + self.value.max_elementwise(T::from_f32(rhs)).into_concrete(), + move |grad, _| (grad * input.mt(T::from_f32(rhs)).into_concrete()).into_concrete(), + ) + } + + pub fn max_scalar(&self, rhs: f32) -> Self { + self.max_elementwise(rhs) + } + + pub fn min_elementwise(&self, rhs: f32) -> Self { + let input = self.value.clone(); + self.unary_from_value( + self.value.min_elementwise(T::from_f32(rhs)).into_concrete(), + move |grad, _| (grad * input.lt(T::from_f32(rhs)).into_concrete()).into_concrete(), + ) + } + + pub fn min_scalar(&self, rhs: f32) -> Self { + self.min_elementwise(rhs) + } + + pub fn mt(&self, rhs: f32) -> Self { + self.gt_scalar(rhs) + } + + pub fn mte(&self, rhs: f32) -> Self { + self.gte_scalar(rhs) + } + + pub fn ne(&self, rhs: f32) -> Self { + self.unary_from_value( + self.value.ne(T::from_f32(rhs)).into_concrete(), + move |_, out| RawTensor::zeros(&out.device(), out.shape()), + ) + } + + pub fn ne_scalar(&self, rhs: f32) -> Self { + self.ne(rhs) + } + + pub fn ne_tensor(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + self.binary_op( + rhs, + self.value.ne_tensor(&rhs.value).into_concrete(), + move |_, lhs, rhs| { + vec![ + RawTensor::zeros(&lhs.device(), lhs.shape()), + RawTensor::zeros(&rhs.device(), rhs.shape()), + ] + }, + ) + } + + pub fn sigmoid(&self) -> Self { + self.mul_scalar(-1.0).exp().add_scalar(1.0).pow_scalar(-1.0) + } + + pub fn silu(&self) -> Self { + let denom = self.mul_scalar(-1.0).exp().add_scalar(1.0); + self.div(&denom) + } + + pub fn gelu(&self) -> Self { + // Fused forward (single elementwise chain) with an analytic backward: + // gelu'(x) = 0.5 * (1 + t) + 0.5 * x * (1 - t^2) * c * (1 + 3 * 0.044715 * x^2) + // where t = tanh(c * (x + 0.044715 * x^3)) and c = sqrt(2 / pi). + let value = self.value.gelu(); + let input_value = self.value.clone(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let grad = downcast_tensor::(&*gradient, "gelu")?; + let coeff = T::from_f32((2.0f32 / std::f32::consts::PI).sqrt()); + let one = T::from_f32(1.0); + let x = input_value.to_concrete(); + let x_sq = x.sqr().into_concrete(); + let inner_factor = (&x_sq * T::from_f32(0.044_715) + one).into_concrete(); + let inner = ((&x * &inner_factor).into_concrete() * coeff).into_concrete(); + let t = inner.tanh().into_concrete(); + let sech_sq = (t.sqr() * T::from_f32(-1.0) + one).into_concrete(); + let du = ((&x_sq * T::from_f32(3.0 * 0.044_715) + one).into_concrete() * coeff) + .into_concrete(); + let tail = ((&x * &sech_sq).into_concrete() * du).into_concrete(); + let dgelu = (((t + one).into_concrete() + tail).into_concrete() * T::from_f32(0.5)) + .into_concrete(); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new((&grad * &dgelu).into_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn tanh(&self) -> Self { + self.unary_from_value(self.value.tanh().into_concrete(), move |grad, out| { + let one_minus_sq = (RawTensor::splat(&out.device(), T::from_f32(1.0), out.shape()) + - out.sqr().into_concrete()) + .into_concrete(); + (grad * one_minus_sq).into_concrete() + }) + } + + pub fn exp(&self) -> Self { + self.unary_from_value(self.value.exp().into_concrete(), move |grad, out| { + (grad * out).into_concrete() + }) + } + + pub fn where_cond(&self, on_true: &Self, on_false: &Self) -> Self { + assert_same_graph(self, on_true); + assert_same_graph(self, on_false); + + let value = self + .value + .where_cond(&on_true.value, &on_false.value) + .into_concrete(); + let condition_id = self.handle.id; + let true_id = on_true.handle.id; + let false_id = on_false.handle.id; + let condition = self.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "where_cond")?; + let zeros = RawTensor::zeros(&condition.device(), condition.shape()); + let ones = RawTensor::splat(&condition.device(), T::from_f32(1.0), condition.shape()); + let true_mask = condition.where_cond(&ones, &zeros).into_concrete(); + let false_mask = condition.where_cond(&zeros, &ones).into_concrete(); + Ok(vec![ + BackwardTarget { + node: condition_id, + gradient: Box::new(zeros), + }, + BackwardTarget { + node: true_id, + gradient: Box::new((gradient.clone() * true_mask).into_concrete()), + }, + BackwardTarget { + node: false_id, + gradient: Box::new((gradient * false_mask).into_concrete()), + }, + ]) + }); + self.emit_op( + value, + vec![ + self.handle.clone(), + on_true.handle.clone(), + on_false.handle.clone(), + ], + Some(backward), + ) + } + + pub fn log(&self) -> Self { + let input = self.value.clone(); + self.unary_from_value(self.value.log().into_concrete(), move |grad, _| { + (grad / input.clone()).into_concrete() + }) + } + + pub fn sqrt(&self) -> Self { + self.unary_from_value(self.value.sqrt().into_concrete(), move |grad, out| { + let denom = out.mul_scalar(T::from_f32(2.0)).into_concrete(); + (grad / denom).into_concrete() + }) + } +} + +macro_rules! impl_autograd_pairwise_op { + ($trait:ident, $method:ident) => { + impl std::ops::$trait> for Tensor + where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + { + type Output = Tensor; + + fn $method(self, rhs: Tensor) -> Tensor { + Tensor::$method(&self, &rhs) + } + } + + impl std::ops::$trait<&Tensor> for Tensor + where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + { + type Output = Tensor; + + fn $method(self, rhs: &Tensor) -> Tensor { + Tensor::$method(&self, rhs) + } + } + + impl std::ops::$trait> for &Tensor + where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + { + type Output = Tensor; + + fn $method(self, rhs: Tensor) -> Tensor { + Tensor::$method(self, &rhs) + } + } + + impl std::ops::$trait<&Tensor> for &Tensor + where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + { + type Output = Tensor; + + fn $method(self, rhs: &Tensor) -> Tensor { + Tensor::$method(self, rhs) + } + } + }; +} + +impl_autograd_pairwise_op!(Add, add); +impl_autograd_pairwise_op!(Sub, sub); +impl_autograd_pairwise_op!(Mul, mul); +impl_autograd_pairwise_op!(Div, div); + +macro_rules! impl_autograd_scalar_op { + ($trait:ident, $method:ident, $scalar_method:ident) => { + impl std::ops::$trait for Tensor + where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + { + type Output = Tensor; + + fn $method(self, rhs: f32) -> Tensor { + Tensor::$scalar_method(&self, rhs) + } + } + + impl std::ops::$trait for &Tensor + where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + { + type Output = Tensor; + + fn $method(self, rhs: f32) -> Tensor { + Tensor::$scalar_method(self, rhs) + } + } + }; +} + +impl_autograd_scalar_op!(Mul, mul, mul_scalar); +impl_autograd_scalar_op!(Add, add, add_scalar); +impl_autograd_scalar_op!(Sub, sub, sub_scalar); +impl_autograd_scalar_op!(Div, div, div_scalar); + +impl std::ops::Neg for Tensor +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + type Output = Tensor; + + fn neg(self) -> Tensor { + Tensor::neg(&self) + } +} + +impl std::ops::Neg for &Tensor +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + type Output = Tensor; + + fn neg(self) -> Tensor { + Tensor::neg(self) + } +} diff --git a/fusor-ml/fusor/src/autograd/indexing.rs b/fusor-ml/fusor/src/autograd/indexing.rs new file mode 100644 index 000000000..cbf53506e --- /dev/null +++ b/fusor-ml/fusor/src/autograd/indexing.rs @@ -0,0 +1,351 @@ +use std::ops::Range; + +use crate::composite::index::IndexOp; + +use super::*; + +impl Tensor +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + u32: crate::CastTo + crate::CastTensor, +{ + pub fn index_select(&self, dimension: usize, indices: &RawTensor<1, u32>) -> Self { + let input_shape = self.shape(); + assert!(dimension < R, "index_select dimension out of bounds"); + + let value = self.value.index_select(dimension, indices).into_concrete(); + let input_id = self.handle.id; + let indices = indices.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "index_select")?; + let one_hot = one_hot_matrix(&indices, input_shape[dimension]); + // transpose+reshape only commute through a copy, so the moved axis + // is materialized once on each side of the matmul; dimension 0 + // needs neither. + let moved = if dimension == 0 { + gradient + } else { + gradient.transpose(0, dimension).into_concrete() + }; + let moved_shape = moved.shape(); + let rest = moved_shape[1..].iter().product::(); + let flat = moved.reshape([moved_shape[0], rest]); + let scattered = one_hot.transpose(0, 1).mat_mul(&flat); + let mut unmoved_shape = moved_shape; + unmoved_shape[0] = input_shape[dimension]; + let scattered = scattered.reshape(unmoved_shape); + let input_gradient = if dimension == 0 { + scattered.into_concrete() + } else { + scattered + .into_concrete() + .transpose(0, dimension) + .into_concrete() + }; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(input_gradient), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn index_ops(&self, ops: [IndexOp; R]) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank<1, OUT, T>, + { + let shape = self.shape(); + let slices: [Range; R] = std::array::from_fn(|axis| ops[axis].to_range(shape[axis])); + let dim = crate::composite::index::removed_dim(ops.map(|op| op.removes_dim())); + self.slice(slices).squeeze_dims::<1, OUT>([dim]) + } +} + +impl Tensor<2, T> +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + u32: crate::CastTo + crate::CastTensor, +{ + pub fn i(&self, index: (I1, I2)) -> Tensor<1, T> + where + I1: Into, + I2: Into, + { + self.index_ops([index.0.into(), index.1.into()]) + } + + pub fn gather_last(&self, indices: &RawTensor<1, u32>) -> Tensor<1, T> { + let shape = self.shape(); + assert_eq!( + shape[0], + indices.shape()[0], + "gather_last expects one index per row" + ); + let width = shape[1]; + let device = self.device(); + let row_offsets = (0..shape[0]) + .map(|row| (row * width) as u32) + .collect::>(); + let row_offsets: RawTensor<1, u32> = + RawTensor::from_slice(&device, [shape[0]], &row_offsets); + let linear_indices = (row_offsets + indices.clone()).into_concrete(); + let flat = self.value.reshape([shape[0] * width]).into_concrete(); + let value = flat.index_select(0, &linear_indices).into_concrete(); + let input_id = self.handle.id; + let indices = indices.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::<1, T>(&*gradient, "gather_last")?; + let one_hot = one_hot_matrix(&indices, width); + let input_gradient: RawTensor<2, T> = one_hot.mul_(&gradient.reshape([shape[0], 1])); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(input_gradient), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + /// Numerically stable softmax cross-entropy against integer class + /// targets, averaged over rows: `mean_i(LSE(x_i) - x_i[t_i])` with the + /// log-sum-exp max-shifted. The backward is analytic — + /// `dlogits = (softmax(x) - onehot(t)) * grad / rows` — so the whole + /// loss runs in a handful of fused kernels instead of a taped chain. + pub fn softmax_cross_entropy(&self, targets: &RawTensor<1, u32>) -> Tensor<0, T> { + let [rows, width] = self.shape(); + assert_eq!( + targets.shape()[0], + rows, + "softmax_cross_entropy expects one target per row" + ); + let device = self.value.device(); + let logits = self.value.clone(); + + // Forward: the shifted exp-sum chain is exclusively consumed, so it + // fuses into one row program; the label logits gather reads the raw + // logits directly. + let max = logits.max_keepdim::<1>(1); + let shifted_exp = (&logits - &max.broadcast_as([rows, width])) + .into_concrete() + .exp() + .into_concrete(); + let lse_total = (shifted_exp + .sum_keepdim::<1>(1) + .into_concrete() + .log() + .into_concrete() + + max) + .into_concrete(); + let row_offsets: Vec = (0..rows).map(|row| (row * width) as u32).collect(); + let row_offsets: RawTensor<1, u32> = RawTensor::from_slice(&device, [rows], &row_offsets); + let linear = (row_offsets + targets.clone()).into_concrete(); + let flat = logits.reshape([rows * width]).into_concrete(); + let picked = flat.index_select(0, &linear).into_concrete(); + let per_row = (lse_total.reshape([rows]).into_concrete() - picked).into_concrete(); + let value = (per_row.sum::<0>(0) * T::from_f32(1.0 / rows as f32)).into_concrete(); + + let input_id = self.handle.id; + let logits_value = self.value.clone(); + let targets = targets.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let grad = downcast_tensor::<0, T>(&*gradient, "softmax_cross_entropy")?; + let probs = logits_value.softmax_last_dim::<1>(); + let one_hot = one_hot_matrix(&targets, width); + let scale = (grad.reshape([1, 1]).into_concrete() * T::from_f32(1.0 / rows as f32)) + .into_concrete(); + let diff = (probs - one_hot).into_concrete(); + let dlogits = (&diff * &scale.broadcast_as([rows, width])).into_concrete(); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(dlogits), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn embedding(&self, indices: &RawTensor<2, u32>) -> Tensor<3, T> { + let [rows, columns] = indices.shape(); + let width = self.shape()[1]; + let flat_indices = indices.clone().reshape([rows * columns]).into_concrete(); + self.index_select(0, &flat_indices) + .reshape([rows, columns, width]) + } +} + +impl Tensor<3, T> +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + u32: crate::CastTo + crate::CastTensor, +{ + pub fn i(&self, index: (I1, I2, I3)) -> Tensor<2, T> + where + I1: Into, + I2: Into, + I3: Into, + { + self.index_ops([index.0.into(), index.1.into(), index.2.into()]) + } +} + +impl Tensor<4, T> +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + u32: crate::CastTo + crate::CastTensor, +{ + pub fn i(&self, index: (I1, I2, I3, I4)) -> Tensor<3, T> + where + I1: Into, + I2: Into, + I3: Into, + I4: Into, + { + self.index_ops([ + index.0.into(), + index.1.into(), + index.2.into(), + index.3.into(), + ]) + } +} + +/// Build a `[indices.len(), size]` one-hot matrix with 1.0 at `[row, indices[row]]` +/// so scatter-adds stay on-device as matmuls/products against it; duplicate +/// indices accumulate through the contraction. +fn one_hot_matrix(indices: &RawTensor<1, u32>, size: usize) -> RawTensor<2, T> +where + u32: crate::CastTo + crate::CastTensor, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, +{ + let device = indices.device(); + let rows = indices.shape()[0]; + let positions = (0..size) + .map(|position| T::from_f32(position as f32)) + .collect::>(); + let positions: RawTensor<2, T> = RawTensor::from_slice(&device, [1, size], &positions); + let indices = indices.cast::().reshape([rows, 1]).into_concrete(); + indices.sub_(&positions).eq(T::from_f32(0.0)) +} diff --git a/fusor-ml/fusor/src/autograd/layers/mod.rs b/fusor-ml/fusor/src/autograd/layers/mod.rs new file mode 100644 index 000000000..280c79f11 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/layers/mod.rs @@ -0,0 +1,5 @@ +pub use crate::layers::conv::autograd::{ConvNd, ConvNdConfig}; +pub use crate::layers::embedding::autograd::Embedding; +pub use crate::layers::layer_norm::autograd::{LayerNorm, LayerNormNd}; +pub use crate::layers::linear::autograd::Linear; +pub use crate::layers::rms_norm::autograd::RmsNorm; diff --git a/fusor-ml/fusor/src/autograd/matmul.rs b/fusor-ml/fusor/src/autograd/matmul.rs new file mode 100644 index 000000000..3d9f80d77 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/matmul.rs @@ -0,0 +1,147 @@ +use super::*; + +impl Tensor +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + pub(super) fn mat_mul_internal(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + let value = self.value.mat_mul(&rhs.value); + let lhs_id = self.handle.id; + let rhs_id = rhs.handle.id; + let lhs_value = self.value.clone(); + let rhs_value = rhs.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "mat_mul")?; + Ok(vec![ + BackwardTarget { + node: lhs_id, + gradient: Box::new( + gradient.clone().mat_mul(&rhs_value.transpose(R - 2, R - 1)), + ), + }, + BackwardTarget { + node: rhs_id, + gradient: Box::new(lhs_value.transpose(R - 2, R - 1).mat_mul(&gradient)), + }, + ]) + }); + self.emit_op( + value, + vec![self.handle.clone(), rhs.handle.clone()], + Some(backward), + ) + } + + pub fn mat_mul(&self, rhs: &Self) -> Self { + self.mat_mul_internal(rhs) + } + + /// `self @ rhs^T` (transposing the last two dims of `rhs`), with a + /// backward that produces the `rhs` gradient directly in `rhs`'s own + /// layout: `d_rhs = grad^T @ self` is contiguous, so optimizer-side + /// flattens stay zero-cost views instead of gather kernels. + pub fn mat_mul_transposed_rhs(&self, rhs: &Self) -> Self { + assert_same_graph(self, rhs); + let value = self.value.mat_mul(&rhs.value.transpose(R - 2, R - 1)); + let lhs_id = self.handle.id; + let rhs_id = rhs.handle.id; + let lhs_value = self.value.clone(); + let rhs_value = rhs.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "mat_mul_transposed_rhs")?; + Ok(vec![ + BackwardTarget { + node: lhs_id, + gradient: Box::new(gradient.clone().mat_mul(&rhs_value)), + }, + BackwardTarget { + node: rhs_id, + gradient: Box::new(gradient.transpose(R - 2, R - 1).mat_mul(&lhs_value)), + }, + ]) + }); + self.emit_op( + value, + vec![self.handle.clone(), rhs.handle.clone()], + Some(backward), + ) + } + + pub fn matmul(&self, rhs: &Self) -> Self { + self.mat_mul_internal(rhs) + } + + pub fn t(&self) -> Self { + assert!(R >= 2, "t requires rank >= 2"); + self.transpose(R - 2, R - 1) + } +} + +impl Tensor { + pub fn q_mat_mul(&self, weights: &crate::QMatrix) -> Self { + if R == 1 { + let k = self.shape()[0]; + let n = weights.shape()[0]; + let out_shape: [usize; R] = std::array::from_fn(|_| n); + return self.reshape([1, k]).q_mat_mul(weights).reshape(out_shape); + } + let value = self.value.q_mat_mul(weights).into_concrete(); + if !self.requires_grad() { + return self.emit_op(value, vec![self.handle.clone()], None); + } + let weights = weights.clone(); + self.replay_unary("q_mat_mul", value, move |input| { + let n = weights.shape()[0]; + let k = weights.shape()[1]; + let batch_dims = R - 2; + let weight_shape: [usize; R] = std::array::from_fn(|i| { + if i < batch_dims { + 1 + } else if i == batch_dims { + k + } else { + n + } + }); + let dequantized: RawTensor<2, f32> = weights.dequantize(); + let weight = dequantized + .transpose(0, 1) + .reshape(weight_shape) + .into_concrete(); + let weight = Tensor::constant_from_raw(&input.graph(), weight); + input.mat_mul_internal(&weight) + }) + } +} diff --git a/fusor-ml/fusor/src/autograd/mod.rs b/fusor-ml/fusor/src/autograd/mod.rs new file mode 100644 index 000000000..92087d199 --- /dev/null +++ b/fusor-ml/fusor/src/autograd/mod.rs @@ -0,0 +1,954 @@ +use std::{ + any::Any, + collections::{HashMap, HashSet, VecDeque}, + sync::{Arc, Mutex}, +}; + +use crate::{Device, Error, Result, Tensor as RawTensor}; + +mod composite; +mod elementwise; +mod indexing; +pub mod layers; +mod matmul; +mod reduce; +mod view; + +#[cfg(test)] +mod tests; + +/// The element types the autograd layer can differentiate through: the +/// float types every backend implements end-to-end (`f32`, `half::f16`). +/// The SIMD op traits (`SimdUnaryOp`/`SimdBinaryOp`/`SimdReduceOp`) are +/// implemented per concrete element type in the CPU backend, so impl blocks +/// that use them still carry those bounds explicitly. +pub trait AutogradElement: + crate::FloatElement + + crate::MatmulElement + + crate::cpu::Scalar + + crate::IsNonZero + + Default + + PartialEq + + PartialOrd + + std::ops::Add + + std::ops::Sub + + std::ops::Mul + + std::ops::Div + + std::ops::Neg + + crate::WasmNotSend + + crate::WasmNotSync + + std::fmt::Debug + + 'static +{ +} + +impl AutogradElement for T where + T: crate::FloatElement + + crate::MatmulElement + + crate::cpu::Scalar + + crate::IsNonZero + + Default + + PartialEq + + PartialOrd + + std::ops::Add + + std::ops::Sub + + std::ops::Mul + + std::ops::Div + + std::ops::Neg + + crate::WasmNotSend + + crate::WasmNotSync + + std::fmt::Debug + + 'static +{ +} + +type NodeId = usize; +#[cfg(not(target_arch = "wasm32"))] +type BackwardRule = + Arc) -> Result> + Send + Sync>; +#[cfg(target_arch = "wasm32")] +type BackwardRule = Arc) -> Result>>; + +#[cfg(not(target_arch = "wasm32"))] +trait BackwardClosure: Send + Sync + 'static {} +#[cfg(not(target_arch = "wasm32"))] +impl BackwardClosure for T where T: Send + Sync + 'static {} + +#[cfg(target_arch = "wasm32")] +trait BackwardClosure: 'static {} +#[cfg(target_arch = "wasm32")] +impl BackwardClosure for T where T: 'static {} + +#[derive(Clone)] +pub struct Graph { + inner: Arc, +} + +#[derive(Clone)] +pub struct Tensor { + value: RawTensor, + handle: NodeHandle, +} + +pub struct Gradients { + gradients: HashMap>, +} + +pub struct BackwardTarget { + node: NodeId, + gradient: Box, +} + +#[derive(Clone)] +pub struct Parent { + handle: NodeHandle, +} + +#[derive(Clone)] +struct NodeHandle { + graph: Arc, + id: NodeId, +} + +#[derive(Clone)] +struct Node { + parents: Vec, + backward: Option, + requires_grad: bool, +} + +struct GraphInner { + state: Mutex, +} + +struct GraphState { + next_id: NodeId, + nodes: HashMap, +} + +#[cfg(not(target_arch = "wasm32"))] +trait AnyTensorValue: Send + Sync { + fn as_any(&self) -> &dyn Any; + fn clone_box(&self) -> Box; + fn into_detached(self: Box) -> Box; + fn add_box(&self, other: &dyn AnyTensorValue) -> Result>; +} + +#[cfg(target_arch = "wasm32")] +trait AnyTensorValue { + fn as_any(&self) -> &dyn Any; + fn clone_box(&self) -> Box; + fn into_detached(self: Box) -> Box; + fn add_box(&self, other: &dyn AnyTensorValue) -> Result>; +} + +impl Graph { + pub fn new() -> Self { + Self { + inner: Arc::new(GraphInner { + state: Mutex::new(GraphState { + next_id: 0, + nodes: HashMap::new(), + }), + }), + } + } + + pub fn leaf(&self, value: RawTensor) -> Tensor + where + crate::AddOp: crate::SimdBinaryOp, + { + self.tensor_with_grad(value, true) + } + + pub fn constant( + &self, + value: RawTensor, + ) -> Tensor + where + crate::AddOp: crate::SimdBinaryOp, + { + self.tensor_with_grad(value, false) + } + + pub fn tensor( + &self, + device: &Device, + data: A, + ) -> Tensor + where + RawTensor: fusor_types::FromArray, + crate::AddOp: crate::SimdBinaryOp, + { + self.leaf(RawTensor::new(device, data)) + } + + pub fn constant_from_data( + &self, + device: &Device, + data: A, + ) -> Tensor + where + RawTensor: fusor_types::FromArray, + crate::AddOp: crate::SimdBinaryOp, + { + self.constant(RawTensor::new(device, data)) + } + + fn tensor_with_grad( + &self, + value: RawTensor, + requires_grad: bool, + ) -> Tensor + where + crate::AddOp: crate::SimdBinaryOp, + { + let id = self.inner.add_node(Vec::new(), None, requires_grad); + Tensor { + value, + handle: NodeHandle { + graph: self.inner.clone(), + id, + }, + } + } +} + +impl Default for Graph { + fn default() -> Self { + Self::new() + } +} + +impl Tensor +where + crate::AddOp: crate::SimdBinaryOp, +{ + pub fn from_raw(graph: &Graph, value: RawTensor) -> Self { + graph.leaf(value) + } + + pub fn constant_from_raw(graph: &Graph, value: RawTensor) -> Self { + graph.constant(value) + } + + pub fn new(graph: &Graph, device: &Device, data: A) -> Self + where + RawTensor: fusor_types::FromArray, + { + graph.tensor(device, data) + } + + pub fn from_array(graph: &Graph, device: &Device, data: A) -> Self + where + RawTensor: fusor_types::FromArray, + { + Self::new(graph, device, data) + } + + pub fn from_slice(graph: &Graph, device: &Device, shape: [usize; R], data: &[T]) -> Self { + graph.leaf(RawTensor::from_slice(device, shape, data)) + } + + pub fn zeros(graph: &Graph, device: &Device, shape: [usize; R]) -> Self { + graph.leaf(RawTensor::zeros(device, shape)) + } + + pub fn ones(graph: &Graph, device: &Device, shape: [usize; R]) -> Self { + Self::splat(graph, device, 1.0, shape) + } + + pub fn splat(graph: &Graph, device: &Device, value: f32, shape: [usize; R]) -> Self { + graph.leaf(RawTensor::splat(device, T::from_f32(value), shape)) + } + + pub fn full(graph: &Graph, device: &Device, shape: [usize; R], value: f32) -> Self { + Self::splat(graph, device, value, shape) + } + + pub fn zeros_like(&self) -> Self { + Self::zeros(&self.graph(), &self.device(), self.shape()) + } + + pub fn ones_like(&self) -> Self { + Self::ones(&self.graph(), &self.device(), self.shape()) + } + + pub fn raw(&self) -> &RawTensor { + &self.value + } + + pub fn into_raw(self) -> RawTensor { + self.value + } + + pub fn shape(&self) -> [usize; R] { + self.value.shape() + } + + pub fn device(&self) -> Device { + self.value.device() + } + + pub fn graph(&self) -> Graph { + Graph { + inner: self.handle.graph.clone(), + } + } + + pub fn requires_grad(&self) -> bool { + self.handle.graph.requires_grad(self.handle.id) + } + + pub fn parent(&self) -> Parent { + Parent { + handle: self.handle.clone(), + } + } + + /// This tensor's gradient destination, without a handle on the graph. + /// See [`GradientSlot`]: custom backward closures must capture this + /// rather than a `Tensor`. + pub fn slot(&self) -> GradientSlot { + GradientSlot(self.handle.id) + } + + pub fn detach(&self) -> Self { + let requires_grad = self.requires_grad(); + let id = self.handle.graph.add_node(Vec::new(), None, requires_grad); + Self { + value: self.value.to_concrete(), + handle: NodeHandle { + graph: self.handle.graph.clone(), + id, + }, + } + } + + #[cfg(not(target_arch = "wasm32"))] + pub fn with_backwards(self, parents: I, backwards: F) -> Self + where + I: IntoIterator, + F: Fn(RawTensor) -> Result> + Send + Sync + 'static, + { + self.with_backwards_impl(parents, backwards) + } + + #[cfg(target_arch = "wasm32")] + pub fn with_backwards(self, parents: I, backwards: F) -> Self + where + I: IntoIterator, + F: Fn(RawTensor) -> Result> + 'static, + { + self.with_backwards_impl(parents, backwards) + } + + fn with_backwards_impl(self, parents: I, backwards: F) -> Self + where + I: IntoIterator, + F: Fn(RawTensor) -> Result> + BackwardClosure, + { + let parent_handles = parents + .into_iter() + .map(|parent| parent.handle) + .collect::>(); + let requires_grad = parent_handles + .iter() + .any(|parent| parent.graph.requires_grad(parent.id)); + let parent_ids = parent_handles + .iter() + .map(|parent| parent.id) + .collect::>(); + // The graph owns this closure, so the closure may only hold a *weak* + // reference back: parent handles would close a reference cycle that + // keeps the whole tape — and every activation buffer it cached — + // alive for the rest of the process. The closure never runs outside + // `backward()`, where the graph is alive by construction. + let owner = Arc::downgrade(&self.handle.graph); + let checked_ids = parent_ids.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = gradient + .as_any() + .downcast_ref::>() + .ok_or_else(|| Error::msg("gradient rank mismatch in custom backward"))? + .clone(); + let targets = backwards(gradient)?; + // The scheduler only unlocks a parent once every child sends it a + // gradient, so a missing target would silently starve that + // parent's whole subgraph. + if let Some(graph) = owner.upgrade() { + for &parent in &checked_ids { + if graph.requires_grad(parent) + && !targets.iter().any(|target| target.node == parent) + { + return Err(Error::msg( + "custom backward omitted a gradient for a parent that requires grad", + )); + } + } + } + Ok(targets) + }); + self.handle.graph.replace_node( + self.handle.id, + Node { + parents: parent_ids, + backward: Some(backward), + requires_grad, + }, + ); + self + } + + pub fn backward(&self) -> Result { + let elements = self.shape().iter().product::(); + if elements != 1 { + return Err(Error::msg( + "backward() requires a single-element tensor; use backward_with() for non-scalars", + )); + } + let seed = RawTensor::splat(&self.device(), T::from_f32(1.0), self.shape()); + self.backward_with(seed) + } + + pub fn backward_with(&self, seed: RawTensor) -> Result { + self.handle.graph.backward(self.handle.id, Box::new(seed)) + } + + fn emit_op( + &self, + value: RawTensor, + parents: Vec, + backward: Option, + ) -> Tensor { + for parent in &parents { + assert!( + Arc::ptr_eq(&self.handle.graph, &parent.graph), + "cannot mix autograd tensors from different graphs" + ); + } + let requires_grad = parents + .iter() + .any(|parent| parent.graph.requires_grad(parent.id)); + let parent_ids = parents.into_iter().map(|parent| parent.id).collect(); + let id = self + .handle + .graph + .add_node(parent_ids, backward, requires_grad); + Tensor { + value, + handle: NodeHandle { + graph: self.handle.graph.clone(), + id, + }, + } + } + + fn unary_from_value( + &self, + value: RawTensor, + backward: impl Fn(RawTensor, RawTensor) -> RawTensor + BackwardClosure, + ) -> Self { + let input_id = self.handle.id; + let output = value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "unary")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(backward(gradient, output.clone()).into_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn binary_op( + &self, + rhs: &Self, + value: RawTensor, + backward: impl Fn(RawTensor, RawTensor, RawTensor) -> Vec> + + BackwardClosure, + ) -> Self { + assert!( + Arc::ptr_eq(&self.handle.graph, &rhs.handle.graph), + "cannot mix autograd tensors from different graphs" + ); + let lhs_id = self.handle.id; + let rhs_id = rhs.handle.id; + let lhs_value = self.value.clone(); + let rhs_value = rhs.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "binary")?; + let gradients = backward(gradient, lhs_value.clone(), rhs_value.clone()); + Ok(vec![ + BackwardTarget { + node: lhs_id, + gradient: Box::new(gradients[0].clone().into_concrete()), + }, + BackwardTarget { + node: rhs_id, + gradient: Box::new(gradients[1].clone().into_concrete()), + }, + ]) + }); + self.emit_op( + value, + vec![self.handle.clone(), rhs.handle.clone()], + Some(backward), + ) + } + + fn replay_unary( + &self, + context: &'static str, + value: RawTensor, + replay: impl Fn(Tensor) -> Tensor + BackwardClosure, + ) -> Tensor { + let input_id = self.handle.id; + let input_value = self.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, context)?; + let graph = Graph::new(); + let replay_input = Tensor::from_raw(&graph, input_value.clone()); + let replay_output = replay(replay_input.clone()); + let gradients = replay_output.backward_with(gradient)?; + let input_gradient = gradients + .get(&replay_input) + .ok_or_else(|| Error::msg(format!("missing replay gradient in {context}")))?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(input_gradient), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn replay_binary( + &self, + rhs: &Tensor, + context: &'static str, + value: RawTensor, + replay: impl Fn(Tensor, Tensor) -> Tensor + BackwardClosure, + ) -> Tensor { + assert_same_graph(self, rhs); + let lhs_id = self.handle.id; + let rhs_id = rhs.handle.id; + let lhs_value = self.value.clone(); + let rhs_value = rhs.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, context)?; + let graph = Graph::new(); + let replay_lhs = Tensor::from_raw(&graph, lhs_value.clone()); + let replay_rhs = Tensor::from_raw(&graph, rhs_value.clone()); + let replay_output = replay(replay_lhs.clone(), replay_rhs.clone()); + let gradients = replay_output.backward_with(gradient)?; + let lhs_gradient = gradients + .get(&replay_lhs) + .ok_or_else(|| Error::msg(format!("missing lhs replay gradient in {context}")))?; + let rhs_gradient = gradients + .get(&replay_rhs) + .ok_or_else(|| Error::msg(format!("missing rhs replay gradient in {context}")))?; + Ok(vec![ + BackwardTarget { + node: lhs_id, + gradient: Box::new(lhs_gradient), + }, + BackwardTarget { + node: rhs_id, + gradient: Box::new(rhs_gradient), + }, + ]) + }); + self.emit_op( + value, + vec![self.handle.clone(), rhs.handle.clone()], + Some(backward), + ) + } + + fn replay_quaternary( + &self, + second: &Tensor, + third: &Tensor, + fourth: &Tensor, + context: &'static str, + value: RawTensor, + replay: impl Fn(Tensor, Tensor, Tensor, Tensor) -> Tensor + + BackwardClosure, + ) -> Tensor { + assert_same_graph(self, second); + assert_same_graph(self, third); + assert_same_graph(self, fourth); + let ids = [ + self.handle.id, + second.handle.id, + third.handle.id, + fourth.handle.id, + ]; + let first_value = self.value.clone(); + let second_value = second.value.clone(); + let third_value = third.value.clone(); + let fourth_value = fourth.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, context)?; + let graph = Graph::new(); + let replay_first = Tensor::from_raw(&graph, first_value.clone()); + let replay_second = Tensor::from_raw(&graph, second_value.clone()); + let replay_third = Tensor::from_raw(&graph, third_value.clone()); + let replay_fourth = Tensor::from_raw(&graph, fourth_value.clone()); + let replay_output = replay( + replay_first.clone(), + replay_second.clone(), + replay_third.clone(), + replay_fourth.clone(), + ); + let gradients = replay_output.backward_with(gradient)?; + let missing = || Error::msg(format!("missing replay gradient in {context}")); + Ok(vec![ + BackwardTarget { + node: ids[0], + gradient: Box::new(gradients.get(&replay_first).ok_or_else(missing)?), + }, + BackwardTarget { + node: ids[1], + gradient: Box::new(gradients.get(&replay_second).ok_or_else(missing)?), + }, + BackwardTarget { + node: ids[2], + gradient: Box::new(gradients.get(&replay_third).ok_or_else(missing)?), + }, + BackwardTarget { + node: ids[3], + gradient: Box::new(gradients.get(&replay_fourth).ok_or_else(missing)?), + }, + ]) + }); + self.emit_op( + value, + vec![ + self.handle.clone(), + second.handle.clone(), + third.handle.clone(), + fourth.handle.clone(), + ], + Some(backward), + ) + } + + fn replay_ternary( + &self, + second: &Tensor, + third: &Tensor, + context: &'static str, + value: RawTensor, + replay: impl Fn(Tensor, Tensor, Tensor) -> Tensor + BackwardClosure, + ) -> Tensor { + assert_same_graph(self, second); + assert_same_graph(self, third); + let first_id = self.handle.id; + let second_id = second.handle.id; + let third_id = third.handle.id; + let first_value = self.value.clone(); + let second_value = second.value.clone(); + let third_value = third.value.clone(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, context)?; + let graph = Graph::new(); + let replay_first = Tensor::from_raw(&graph, first_value.clone()); + let replay_second = Tensor::from_raw(&graph, second_value.clone()); + let replay_third = Tensor::from_raw(&graph, third_value.clone()); + let replay_output = replay( + replay_first.clone(), + replay_second.clone(), + replay_third.clone(), + ); + let gradients = replay_output.backward_with(gradient)?; + let first_gradient = gradients + .get(&replay_first) + .ok_or_else(|| Error::msg(format!("missing first replay gradient in {context}")))?; + let second_gradient = gradients.get(&replay_second).ok_or_else(|| { + Error::msg(format!("missing second replay gradient in {context}")) + })?; + let third_gradient = gradients + .get(&replay_third) + .ok_or_else(|| Error::msg(format!("missing third replay gradient in {context}")))?; + Ok(vec![ + BackwardTarget { + node: first_id, + gradient: Box::new(first_gradient), + }, + BackwardTarget { + node: second_id, + gradient: Box::new(second_gradient), + }, + BackwardTarget { + node: third_id, + gradient: Box::new(third_gradient), + }, + ]) + }); + self.emit_op( + value, + vec![ + self.handle.clone(), + second.handle.clone(), + third.handle.clone(), + ], + Some(backward), + ) + } +} + +impl Tensor<1> { + pub fn arange(graph: &Graph, device: &Device, start: f32, end: f32) -> Tensor<1> { + graph.leaf(crate::arange(device, start, end)) + } + + pub fn arange_step( + graph: &Graph, + device: &Device, + start: f32, + end: f32, + step: f32, + ) -> Tensor<1> { + graph.leaf(crate::arange_step(device, start, end, step)) + } +} + +impl Gradients { + pub fn get( + &self, + tensor: &Tensor, + ) -> Option> { + self.gradients + .get(&tensor.handle.id) + .and_then(|gradient| gradient.as_any().downcast_ref::>()) + .cloned() + } + + pub fn into_detached(self) -> Self { + Self { + gradients: self + .gradients + .into_iter() + .map(|(id, gradient)| (id, gradient.into_detached())) + .collect(), + } + } +} + +impl BackwardTarget { + pub fn wrt( + tensor: &Tensor, + gradient: RawTensor, + ) -> Self + where + crate::AddOp: crate::SimdBinaryOp, + { + Self::to(tensor.slot(), gradient) + } + + /// Route a gradient to a slot captured with [`Tensor::slot`]. + pub fn to( + slot: GradientSlot, + gradient: RawTensor, + ) -> Self + where + crate::AddOp: crate::SimdBinaryOp, + { + Self { + node: slot.0, + gradient: Box::new(gradient), + } + } +} + +/// Where a custom backward sends a gradient: a node identity detached from +/// the graph that owns it. +/// +/// A [`Tensor::with_backwards`] closure is stored *inside* the graph, so +/// capturing a `Tensor` to name the destination would give the graph an +/// `Arc` to itself — a cycle that keeps every buffer the tape cached alive +/// for the rest of the process. Capture a slot instead. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct GradientSlot(NodeId); + +impl GraphInner { + fn add_node( + &self, + mut parents: Vec, + mut backward: Option, + requires_grad: bool, + ) -> NodeId { + if !requires_grad { + parents.clear(); + backward = None; + } + let mut state = self.state.lock().unwrap(); + let id = state.next_id; + state.next_id += 1; + state.nodes.insert( + id, + Node { + parents, + backward, + requires_grad, + }, + ); + id + } + + fn replace_node(&self, id: NodeId, mut node: Node) { + if !node.requires_grad { + node.parents.clear(); + node.backward = None; + } + self.state.lock().unwrap().nodes.insert(id, node); + } + + fn requires_grad(&self, id: NodeId) -> bool { + self.state + .lock() + .unwrap() + .nodes + .get(&id) + .map(|node| node.requires_grad) + .unwrap_or(false) + } + + fn backward(&self, root: NodeId, seed: Box) -> Result { + let nodes = self.reachable_nodes(root); + let mut pending_children = HashMap::::new(); + for (id, node) in &nodes { + pending_children.entry(*id).or_insert(0); + for parent in &node.parents { + *pending_children.entry(*parent).or_insert(0) += 1; + } + } + + let mut gradients = HashMap::>::new(); + gradients.insert(root, seed); + + let mut queue = VecDeque::new(); + queue.push_back(root); + + while let Some(node_id) = queue.pop_front() { + let Some(node) = nodes.get(&node_id) else { + continue; + }; + let Some(backward) = node.backward.as_ref() else { + continue; + }; + let gradient = gradients + .get(&node_id) + .ok_or_else(|| Error::msg(format!("missing gradient for node {node_id}")))? + .clone_box(); + + for target in backward(gradient)? { + let Some(parent_node) = nodes.get(&target.node) else { + continue; + }; + if !parent_node.requires_grad { + continue; + } + accumulate_gradient(&mut gradients, target.node, target.gradient)?; + let remaining = pending_children.get_mut(&target.node).ok_or_else(|| { + Error::msg(format!("missing child count for node {}", target.node)) + })?; + *remaining = remaining.saturating_sub(1); + if *remaining == 0 { + queue.push_back(target.node); + } + } + } + + Ok(Gradients { gradients }) + } + + fn reachable_nodes(&self, root: NodeId) -> HashMap { + let snapshot = self.state.lock().unwrap().nodes.clone(); + let mut reachable = HashMap::new(); + let mut stack = vec![root]; + let mut visited = HashSet::new(); + while let Some(node_id) = stack.pop() { + if !visited.insert(node_id) { + continue; + } + if let Some(node) = snapshot.get(&node_id) { + reachable.insert(node_id, node.clone()); + stack.extend(node.parents.iter().copied()); + } + } + reachable + } +} + +impl AnyTensorValue for RawTensor +where + crate::AddOp: crate::SimdBinaryOp, +{ + fn as_any(&self) -> &dyn Any { + self + } + + fn clone_box(&self) -> Box { + Box::new(self.clone()) + } + + fn into_detached(self: Box) -> Box { + match *self { + RawTensor::Cpu(tensor) => Box::new(RawTensor::Cpu(tensor.to_concrete())), + RawTensor::Gpu(tensor) => Box::new(RawTensor::Gpu(tensor.detach())), + } + } + + fn add_box(&self, other: &dyn AnyTensorValue) -> Result> { + let other = other + .as_any() + .downcast_ref::>() + .ok_or_else(|| Error::msg("gradient rank mismatch while accumulating"))?; + Ok(Box::new((self.clone() + other.clone()).into_concrete())) + } +} + +fn accumulate_gradient( + gradients: &mut HashMap>, + node: NodeId, + gradient: Box, +) -> Result<()> { + match gradients.get(&node) { + Some(existing) => { + let accumulated = existing.add_box(&*gradient)?; + gradients.insert(node, accumulated); + } + None => { + gradients.insert(node, gradient); + } + } + Ok(()) +} + +fn downcast_tensor( + value: &dyn AnyTensorValue, + context: &str, +) -> Result> { + value + .as_any() + .downcast_ref::>() + .cloned() + .ok_or_else(|| Error::msg(format!("gradient rank mismatch in {context}"))) +} + +fn assert_same_graph< + const R: usize, + const R2: usize, + T: crate::cpu::SimdElement, + T2: crate::cpu::SimdElement, +>( + lhs: &Tensor, + rhs: &Tensor, +) { + assert!( + Arc::ptr_eq(&lhs.handle.graph, &rhs.handle.graph), + "cannot mix autograd tensors from different graphs" + ); +} diff --git a/fusor-ml/fusor/src/autograd/reduce.rs b/fusor-ml/fusor/src/autograd/reduce.rs new file mode 100644 index 000000000..5f696507e --- /dev/null +++ b/fusor-ml/fusor/src/autograd/reduce.rs @@ -0,0 +1,535 @@ +use fusor_types::{SlidingWindow, StrideSpec}; + +use super::*; + +impl Tensor +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + pub(super) fn sum_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input_shape = self.shape(); + let value = self.value.sum_keepdim::(axis).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "sum_keepdim")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.broadcast_as(input_shape).into_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub(super) fn sum_any(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input_shape = self.shape(); + let value = self.value.sum::(axis).into_concrete(); + let input_id = self.handle.id; + let mut keepdim_shape = input_shape; + keepdim_shape[axis] = 1; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "sum")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new( + gradient + .reshape(keepdim_shape) + .into_concrete() + .broadcast_as(input_shape) + .into_concrete(), + ), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub(super) fn max_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input = self.value.clone(); + let value = input.max_keepdim::(axis).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "max_keepdim")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(reduction_extrema_keepdim_grad::( + input.clone(), + axis, + gradient, + true, + )), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn min_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let input = self.value.clone(); + let value = input.min_keepdim::(axis).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "min_keepdim")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(reduction_extrema_keepdim_grad::( + input.clone(), + axis, + gradient, + false, + )), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub(super) fn mean_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.sum_keepdim_any::(axis) + .div_scalar(self.shape()[axis] as f32) + } + + fn product_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::ProdOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + { + let input = self.value.clone(); + let input_shape = self.shape(); + let value = input.product_keepdim::(axis).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "product_keepdim")?; + let upstream = gradient.broadcast_as(input_shape).into_concrete(); + let zeros = RawTensor::zeros(&input.device(), input_shape); + let ones = RawTensor::splat(&input.device(), T::from_f32(1.0), input_shape); + let zero_mask = input.eq(T::from_f32(0.0)).into_concrete(); + let safe_input = zero_mask.where_cond(&ones, &input).into_concrete(); + let zero_count = zero_mask.sum_keepdim::(axis).into_concrete(); + let zero_count_broadcast = zero_count.broadcast_as(input_shape).into_concrete(); + let product_non_zero = safe_input + .product_keepdim::(axis) + .broadcast_as(input_shape) + .into_concrete(); + let no_zero_grad = (upstream.clone() + * (product_non_zero.clone() / safe_input).into_concrete()) + .into_concrete(); + let single_zero_grad = zero_mask + .where_cond(&(upstream * product_non_zero).into_concrete(), &zeros) + .into_concrete(); + let gradient = ((no_zero_grad + * zero_count_broadcast.eq(T::from_f32(0.0)).into_concrete()) + .into_concrete() + + (single_zero_grad * zero_count_broadcast.eq(T::from_f32(1.0)).into_concrete()) + .into_concrete()) + .into_concrete(); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + fn var_keepdim_any(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + let mean = self.mean_keepdim_any::(axis); + let centered = self.sub(&mean.broadcast_as(self.shape())); + centered.sqr().mean_keepdim_any::(axis) + } + + pub fn pool( + &self, + pools: [impl Into; DIFF], + with: impl Fn(&Tensor, usize) -> Self + Copy, + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank<1, R3, T>, + crate::gpu::Tensor: crate::gpu::SmallerRank, + { + let pools: [crate::composite::pool::PoolSize; DIFF] = pools.map(|pool| pool.into()); + let axis_start = R - DIFF; + let windows: [SlidingWindow; DIFF] = std::array::from_fn(|i| { + let pool = pools[i]; + SlidingWindow::new(axis_start + i, pool.size, pool.stride) + }); + let shape = self.shape(); + let mut sorted_windows = windows; + sorted_windows.sort_by_key(|window| window.axis); + let specs: [StrideSpec; R2] = std::array::from_fn(|out_i| { + if out_i < R { + if let Some(window) = sorted_windows.iter().find(|window| window.axis == out_i) { + let positions = (shape[out_i] - window.window_size) / window.step + 1; + StrideSpec::dim_with(out_i, positions, window.step) + } else { + StrideSpec::dim(out_i, shape[out_i]) + } + } else { + let window = &sorted_windows[out_i - R]; + StrideSpec::dim(window.axis, window.window_size) + } + }); + + let tiled: Tensor = self.restride(specs); + let unsqueezed: Tensor = tiled.unsqueeze_dims::<1, R3>([R2]); + let flattened: Tensor = unsqueezed.flatten_last_n::(); + with(&flattened, O - 1) + } + + pub fn pool_max( + &self, + pools: [impl Into; DIFF], + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank<1, R3, T>, + crate::gpu::Tensor: crate::gpu::SmallerRank, + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank + crate::gpu::SmallerRank<1, R, T>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.pool::(pools, |windowed, axis| windowed.max::(axis)) + } + + pub fn pool_min( + &self, + pools: [impl Into; DIFF], + ) -> Self + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank<1, R3, T>, + crate::gpu::Tensor: crate::gpu::SmallerRank, + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank + crate::gpu::SmallerRank<1, R, T>, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.pool::(pools, |windowed, axis| windowed.min::(axis)) + } + + pub fn max_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.max_keepdim_any::(axis) + } + + pub fn max(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, T>, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.max_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn min_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.min_keepdim_any::(axis) + } + + pub fn min(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, T>, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.min_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn mean_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.mean_keepdim_any::(axis) + } + + pub fn mean(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, T>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.mean_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn product(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, T>, + crate::cpu::ProdOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + { + self.product_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn product_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::ProdOp: crate::cpu::SimdReduceOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + { + self.product_keepdim_any::(axis) + } + + pub fn var(&self, axis: usize) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: + crate::gpu::LastRank + crate::gpu::SmallerRank<1, OUT_RANK, T>, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.var_keepdim_any::(axis) + .squeeze_dims::<1, OUT_RANK>([axis]) + } + + pub fn var_keepdim(&self, axis: usize) -> Self + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + { + self.var_keepdim_any::(axis) + } +} + +impl Tensor<1, T> +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + pub fn sum(&self) -> Tensor<0, T> { + let input_shape = self.shape(); + let value = self.value.sum::<0>(0); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::<0, T>(&*gradient, "sum")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.broadcast_as(input_shape).into_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn sum_keepdim(&self, axis: usize) -> Tensor<1, T> { + self.sum_keepdim_any::<0>(axis) + } +} + +macro_rules! sum_wrappers { + ($($rank:literal => $out:literal),* $(,)?) => {$( + impl Tensor<$rank, T> +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + { + pub fn sum(&self, axis: usize) -> Tensor<$out, T> + where + crate::ConcreteTensor: crate::cpu::LastRank<$out, T>, + crate::gpu::Tensor<$rank, T>: crate::gpu::LastRank<$out, T>, + { + self.sum_any::<$out>(axis) + } + + pub fn sum_keepdim(&self, axis: usize) -> Tensor<$rank, T> + where + crate::ConcreteTensor: crate::cpu::LastRank<$out, T>, + crate::gpu::Tensor<$rank, T>: crate::gpu::LastRank<$out, T>, + { + self.sum_keepdim_any::<$out>(axis) + } + } + )*}; +} + +sum_wrappers!(2 => 1, 3 => 2, 4 => 3, 5 => 4, 6 => 5, 7 => 6, 8 => 7, 9 => 8, 10 => 9); + +fn reduction_extrema_keepdim_grad( + input: RawTensor, + axis: usize, + gradient: RawTensor, + is_max: bool, +) -> RawTensor +where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, +{ + let input_shape = input.shape(); + let extrema = if is_max { + input.max_keepdim::(axis) + } else { + input.min_keepdim::(axis) + } + .into_concrete(); + let extrema_broadcast = extrema.broadcast_as(input_shape).into_concrete(); + let mask = (input - extrema_broadcast) + .into_concrete() + .eq(T::from_f32(0.0)) + .into_concrete(); + let tie_count = mask + .sum_keepdim::(axis) + .broadcast_as(input_shape) + .into_concrete(); + ((mask * gradient.broadcast_as(input_shape)).into_concrete() / tie_count).into_concrete() +} diff --git a/fusor-ml/fusor/src/autograd/tests.rs b/fusor-ml/fusor/src/autograd/tests.rs new file mode 100644 index 000000000..549d8a09e --- /dev/null +++ b/fusor-ml/fusor/src/autograd/tests.rs @@ -0,0 +1,6663 @@ +use super::*; +use crate::{Layout, ToVec}; +use fusor_types::StrideSpec; + +fn assert_close(left: f32, right: f32) { + assert!((left - right).abs() < 1e-3, "expected {right}, got {left}"); +} + +fn assert_slice_close(left: &[f32], right: &[f32]) { + assert_eq!(left.len(), right.len(), "slice lengths differ"); + for (index, (left, right)) in left.iter().zip(right.iter()).enumerate() { + assert!( + (*left - *right).abs() < 1e-3, + "mismatch at index {index}: expected {right}, got {left}", + ); + } +} + +async fn test_devices() -> Vec { + let mut devices = vec![Device::cpu()]; + match Device::gpu().await { + Ok(gpu) => devices.push(gpu), + Err(_) => eprintln!("skipping GPU coverage: GPU unavailable"), + } + devices +} + +async fn flatten(tensor: RawTensor) -> Vec { + let elements = tensor.shape().into_iter().product(); + tensor + .reshape([elements]) + .as_slice() + .await + .unwrap() + .to_vec() +} + +#[test] +fn non_grad_nodes_do_not_retain_a_backward_tape() { + let graph = Graph::new(); + let input: Tensor<1> = graph.constant_from_data(&Device::cpu(), &[1.0f32, 2.0, 3.0]); + let output = input.sqr().sum(); + let state = graph.inner.state.lock().unwrap(); + + for node in state.nodes.values() { + assert!(!node.requires_grad); + assert!( + node.parents.is_empty(), + "non-grad nodes must not retain parent graph structure" + ); + assert!( + node.backward.is_none(), + "non-grad nodes must not retain tensor-capturing backward closures" + ); + } + assert!(state.nodes.contains_key(&output.handle.id)); +} + +async fn finite_difference_gradient( + device: &Device, + shape: [usize; R], + data: &[f32], + loss: &F, +) -> Vec +where + F: Fn(&Graph, Tensor) -> Tensor<0>, +{ + let epsilon = 1e-2f32; + let mut numeric = Vec::with_capacity(data.len()); + for index in 0..data.len() { + let mut perturbed = data.to_vec(); + perturbed[index] = data[index] + epsilon; + let graph = Graph::new(); + let plus = loss( + &graph, + Tensor::from_slice(&graph, device, shape, &perturbed), + ); + let plus = plus.raw().to_scalar().await.unwrap(); + perturbed[index] = data[index] - epsilon; + let graph = Graph::new(); + let minus = loss( + &graph, + Tensor::from_slice(&graph, device, shape, &perturbed), + ); + let minus = minus.raw().to_scalar().await.unwrap(); + numeric.push((plus - minus) / (2.0 * epsilon)); + } + numeric +} + +async fn assert_gradient_matches_finite_difference( + device: &Device, + shape: [usize; R], + data: &[f32], + loss: F, +) where + F: Fn(&Graph, Tensor) -> Tensor<0>, +{ + let graph = Graph::new(); + let input = Tensor::from_slice(&graph, device, shape, data); + let gradients = loss(&graph, input.clone()).backward().unwrap(); + let analytic = flatten(gradients.get(&input).unwrap()).await; + let numeric = finite_difference_gradient(device, shape, data, &loss).await; + assert_eq!(analytic.len(), numeric.len(), "gradient lengths differ"); + for (index, (analytic, numeric)) in analytic.iter().zip(numeric.iter()).enumerate() { + assert!( + (analytic - numeric).abs() < 1e-2 + 1e-2 * numeric.abs(), + "gradient mismatch at index {index}: analytic {analytic}, finite difference {numeric}", + ); + } +} + +#[tokio::test] +async fn test_backward_squared_sum() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let loss = x.sqr().sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(dx[0], 2.0); + assert_close(dx[1], 4.0); + assert_close(dx[2], 6.0); + } +} + +#[tokio::test] +async fn test_autograd_silu() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 0.5]); + + let output = x.silu(); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + + let expected = [1.0f32, -2.0, 0.5].map(|v| v / (1.0 + (-v).exp())); + for (value, expected) in values.iter().zip(expected) { + assert_close(*value, expected); + } + + let gradients = output.sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + let expected_grads = [1.0f32, -2.0, 0.5].map(|v| { + let sigmoid = 1.0 / (1.0 + (-v).exp()); + sigmoid * (1.0 + v * (1.0 - sigmoid)) + }); + for (value, expected) in dx.iter().zip(expected_grads) { + assert_close(*value, expected); + } + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, -2.0, 0.5], |_, x| { + x.silu().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_gelu() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 0.5]); + + let output = x.gelu(); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + + let expected = [1.0f32, -2.0, 0.5].map(|v| { + 0.5 * v + * (1.0 + ((2.0 / std::f32::consts::PI).sqrt() * (v + 0.044_715 * v.powi(3))).tanh()) + }); + for (value, expected) in values.iter().zip(expected) { + assert_close(*value, expected); + } + + let gradients = output.sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + let expected_grads = [1.0f32, -2.0, 0.5].map(|v| { + let scale = (2.0 / std::f32::consts::PI).sqrt(); + let inner = scale * (v + 0.044_715 * v.powi(3)); + let tanh = inner.tanh(); + let dinner = scale * (1.0 + 3.0 * 0.044_715 * v * v); + 0.5 * (1.0 + tanh) + 0.5 * v * (1.0 - tanh * tanh) * dinner + }); + for (value, expected) in dx.iter().zip(expected_grads) { + assert_close(*value, expected); + } + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, -2.0, 0.5], |_, x| { + x.gelu().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_where_cond() { + for device in test_devices().await { + let graph = Graph::new(); + let condition: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 0.0, -2.0]); + let on_true: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 3.0, 4.0]); + let on_false: Tensor<1> = Tensor::new(&graph, &device, &[10.0f32, 20.0, 30.0]); + + let output = condition.where_cond(&on_true, &on_false); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.flatten_all().sum().backward().unwrap(); + + let dcondition = gradients + .get(&condition) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let dtrue = gradients + .get(&on_true) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let dfalse = gradients + .get(&on_false) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![2.0, 20.0, 4.0]); + assert_eq!(dcondition, vec![0.0, 0.0, 0.0]); + assert_eq!(dtrue, vec![1.0, 0.0, 1.0]); + assert_eq!(dfalse, vec![0.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_index_select() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let indices = RawTensor::from_slice(&device, [3], &[2u32, 0, 2]); + + let output = input.index_select(1, &indices); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![vec![3.0, 1.0, 3.0], vec![6.0, 4.0, 6.0]] + ); + assert_eq!(dinput, vec![vec![1.0, 0.0, 2.0], vec![1.0, 0.0, 2.0]]); + } +} + +#[tokio::test] +async fn test_backward_slice_assign() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], + ); + let value: Tensor<2> = Tensor::new(&graph, &device, &[[10.0f32, 11.0], [12.0, 13.0]]); + + let output = input.slice_assign([0..2, 1..3], &value); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let dvalue = gradients + .get(&value) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 10.0, 11.0], + vec![4.0, 12.0, 13.0], + vec![7.0, 8.0, 9.0] + ] + ); + assert_eq!( + dinput, + vec![ + vec![1.0, 0.0, 0.0], + vec![1.0, 0.0, 0.0], + vec![1.0, 1.0, 1.0] + ] + ); + assert_eq!(dvalue, vec![vec![1.0, 1.0], vec![1.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_expand() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[2.0f32, 3.0, 4.0]]); + + let output = input.expand([2, 3]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![vec![2.0, 3.0, 4.0], vec![2.0, 3.0, 4.0]] + ); + assert_eq!(dinput, vec![vec![2.0, 2.0, 2.0]]); + } +} + +#[tokio::test] +async fn test_backward_flatten_all() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + + let output = input.flatten_all(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 2.0, 3.0, 4.0]); + assert_eq!(dinput, vec![vec![1.0, 1.0], vec![1.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_flatten_last_n() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new( + &graph, + &device, + &[ + [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + ); + + let output = input.flatten_last_n::<1, 2>(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .reshape([2, 6]) + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0] + ] + ); + assert_eq!( + dinput, + vec![ + vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_flatten_first_n() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new( + &graph, + &device, + &[ + [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]], + [[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]], + ], + ); + + let output = input.flatten_first_n::<1, 2>(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .reshape([4, 3]) + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 2.0, 3.0], + vec![4.0, 5.0, 6.0], + vec![7.0, 8.0, 9.0], + vec![10.0, 11.0, 12.0] + ] + ); + assert_eq!( + dinput, + vec![ + vec![1.0, 1.0, 1.0], + vec![1.0, 1.0, 1.0], + vec![1.0, 1.0, 1.0], + vec![1.0, 1.0, 1.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_narrow() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], + ); + + let output = input.narrow(1usize, 1, 2); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![vec![2.0, 3.0], vec![5.0, 6.0], vec![8.0, 9.0]] + ); + assert_eq!( + dinput, + vec![ + vec![0.0, 1.0, 1.0], + vec![0.0, 1.0, 1.0], + vec![0.0, 1.0, 1.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_repeat() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + + let output = input.repeat([2, 3]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0], + vec![3.0, 4.0, 3.0, 4.0, 3.0, 4.0], + vec![1.0, 2.0, 1.0, 2.0, 1.0, 2.0], + vec![3.0, 4.0, 3.0, 4.0, 3.0, 4.0] + ] + ); + assert_eq!(dinput, vec![vec![6.0, 6.0], vec![6.0, 6.0]]); + } +} + +#[tokio::test] +async fn test_backward_resize() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]], + ); + + let output = input.resize([2, 2]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![vec![1.0, 2.0], vec![4.0, 5.0]]); + assert_eq!( + dinput, + vec![ + vec![1.0, 1.0, 0.0], + vec![1.0, 1.0, 0.0], + vec![0.0, 0.0, 0.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_restride() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0]); + + let output = input.restride([StrideSpec::dim(0, 2), StrideSpec::dim(0, 3)]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![vec![1.0, 2.0, 3.0], vec![2.0, 3.0, 4.0]] + ); + assert_eq!(dinput, vec![1.0, 2.0, 2.0, 1.0]); + } +} + +#[tokio::test] +async fn test_backward_restride_strided_overlap() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new( + &graph, + &device, + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + ); + + let output = input.restride([StrideSpec::dim_with(0, 3, 2), StrideSpec::dim(0, 3)]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![ + vec![1.0, 2.0, 3.0], + vec![3.0, 4.0, 5.0], + vec![5.0, 6.0, 7.0] + ] + ); + assert_eq!(dinput, vec![1.0, 1.0, 2.0, 1.0, 2.0, 1.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_restride_layout() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0, 5.0]); + let layout = Layout::contiguous(&[5]) + .restride(&[StrideSpec::dim(0, 2).with_offset(1), StrideSpec::dim(0, 2)]); + + let output: Tensor<2> = input.restride_layout(layout); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![vec![2.0, 3.0], vec![3.0, 4.0]]); + assert_eq!(dinput, vec![0.0, 1.0, 2.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_squeeze_dims() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32], [2.0], [3.0]]], [[[4.0], [5.0], [6.0]]]], + ); + + let output = input.squeeze_dims::<2, 2>([1, 3]); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .reshape([2, 3]) + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!( + output_values, + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]] + ); + assert_eq!(dinput, vec![vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_unsqueeze_dims() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = input.unsqueeze_dims::<2, 4>([0, 2]); + let output_values = output + .raw() + .clone() + .reshape([2, 3]) + .as_slice() + .await + .unwrap() + .to_vec(); + let gradients = output.sum(3).sum(2).sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output.shape(), [1, 2, 1, 3]); + assert_eq!( + output_values, + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]] + ); + assert_eq!(dinput, vec![vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_max() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 5.0, 5.0], [4.0, 2.0, 0.0]]); + + let output = input.max::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![5.0, 4.0]); + assert_eq!(dinput, vec![vec![0.0, 0.5, 0.5], vec![1.0, 0.0, 0.0]]); + } +} + +#[tokio::test] +async fn test_backward_min() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0, 5.0], [4.0, 2.0, 0.0]]); + + let output = input.min::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 0.0]); + assert_eq!(dinput, vec![vec![0.5, 0.5, 0.0], vec![0.0, 0.0, 1.0]]); + } +} + +#[tokio::test] +async fn test_backward_mean() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = input.mean::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![2.0, 5.0]); + assert_eq!( + dinput, + vec![ + vec![1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0], + vec![1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_product() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[2.0f32, 3.0, 4.0], [5.0, 0.0, 7.0], [0.0, 0.0, 9.0]], + ); + + let output = input.product::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![24.0, 0.0, 0.0]); + assert_eq!( + dinput, + vec![ + vec![12.0, 8.0, 6.0], + vec![0.0, 35.0, 0.0], + vec![0.0, 0.0, 0.0] + ] + ); + } +} + +#[tokio::test] +async fn test_backward_product_keepdim() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[2.0f32, 3.0, 4.0], [5.0, 0.0, 7.0]]); + + let output = input.product_keepdim::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![vec![24.0], vec![0.0]]); + assert_eq!(dinput, vec![vec![12.0, 8.0, 6.0], vec![0.0, 35.0, 0.0]]); + } +} + +#[tokio::test] +async fn test_backward_var() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = input.var::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_slice_close(&output_values, &[2.0 / 3.0, 2.0 / 3.0]); + assert_slice_close( + &dinput.into_iter().flatten().collect::>(), + &[-2.0 / 3.0, 0.0, 2.0 / 3.0, -2.0 / 3.0, 0.0, 2.0 / 3.0], + ); + } +} + +#[tokio::test] +async fn test_backward_var_keepdim() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + + let output = input.var_keepdim::<1>(1); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_slice_close( + &output_values.into_iter().flatten().collect::>(), + &[2.0 / 3.0, 2.0 / 3.0], + ); + assert_slice_close( + &dinput.into_iter().flatten().collect::>(), + &[-2.0 / 3.0, 0.0, 2.0 / 3.0, -2.0 / 3.0, 0.0, 2.0 / 3.0], + ); + } +} + +#[tokio::test] +async fn test_backward_clamp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[-1.0f32, 0.0, 2.0, 5.0]); + + let output = input.clamp(0.0, 3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 0.0, 2.0, 3.0]); + assert_eq!(dinput, vec![0.0, 0.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_eq() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 1.0]); + + let output = input.eq(1.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 0.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_eq_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32, 2.0, 3.0]); + + let output = input.eq_scalar(3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 0.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_eq_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 0.0, 3.0]); + + let output = lhs.eq_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 0.0, 1.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_gt_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.gt_scalar(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 0.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_gt_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 1.0, 3.0]); + + let output = lhs.gt_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 1.0, 0.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_gte_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.gte_scalar(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 1.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_gte_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 4.0, 2.0]); + + let output = lhs.gte_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 1.0, 1.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lt() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.lt(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 0.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lt_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.lt_scalar(3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 1.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lt_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 1.0, 3.0]); + + let output = lhs.lt_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 0.0, 0.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lte() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.lte(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 1.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lte_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.lte_scalar(1.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 0.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_lte_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 2.0, 1.0]); + + let output = lhs.lte_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 1.0, 0.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_max_elementwise() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[-1.0f32, 0.0, 2.0]); + + let output = input.max_elementwise(0.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 0.0, 2.0]); + assert_eq!(dinput, vec![0.0, 0.0, 1.0]); + } +} + +#[tokio::test] +async fn test_backward_max_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.max_scalar(3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![3.0, 4.0, 3.0]); + assert_eq!(dinput, vec![0.0, 1.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_min_elementwise() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.min_elementwise(3.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 3.0, 2.0]); + assert_eq!(dinput, vec![1.0, 0.0, 1.0]); + } +} + +#[tokio::test] +async fn test_backward_min_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.min_scalar(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 2.0, 2.0]); + assert_eq!(dinput, vec![1.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_mt() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.mt(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 1.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_mte() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.mte(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 1.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_ne() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.ne(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 1.0, 0.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_ne_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + + let output = input.ne_scalar(4.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![1.0, 0.0, 1.0]); + assert_eq!(dinput, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_ne_tensor() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 4.0, 2.0]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 0.0, 3.0]); + + let output = lhs.ne_tensor(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![0.0, 1.0, 1.0]); + assert_eq!(dlhs, vec![0.0, 0.0, 0.0]); + assert_eq!(drhs, vec![0.0, 0.0, 0.0]); + } +} + +#[tokio::test] +async fn test_backward_abs() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[-2.0f32, 0.0, 3.0]); + + let output = input.abs(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![2.0, 0.0, 3.0]); + assert_eq!(dinput, vec![-1.0, 0.0, 1.0]); + } +} + +#[tokio::test] +async fn test_backward_acos() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.acos(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.acos()); + assert_close(dinput[0], -1.0f32 / (1.0f32 - 0.25f32).sqrt()); + } +} + +#[tokio::test] +async fn test_backward_acosh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32]); + + let output = input.acosh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 2.0f32.acosh()); + assert_close( + dinput[0], + 1.0f32 / ((2.0f32 - 1.0f32).sqrt() * (2.0f32 + 1.0f32).sqrt()), + ); + } +} + +#[tokio::test] +async fn test_backward_approximate_exp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32]); + + let output = input.approximate_exp(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 1.0f32.exp()); + assert_close(dinput[0], 1.0f32.exp()); + } +} + +#[tokio::test] +async fn test_backward_asin() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.asin(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.asin()); + assert_close(dinput[0], 1.0f32 / (1.0f32 - 0.25f32).sqrt()); + } +} + +#[tokio::test] +async fn test_backward_asinh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.5f32]); + + let output = input.asinh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 1.5f32.asinh()); + assert_close(dinput[0], 1.0f32 / (1.5f32 * 1.5f32 + 1.0f32).sqrt()); + } +} + +#[tokio::test] +async fn test_backward_atan() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.atan(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.atan()); + assert_close(dinput[0], 1.0f32 / (1.0f32 + 0.25f32)); + } +} + +#[tokio::test] +async fn test_backward_atanh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.atanh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.atanh()); + assert_close(dinput[0], 1.0f32 / (1.0f32 - 0.25f32)); + } +} + +#[tokio::test] +async fn test_backward_cos() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.cos(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.cos()); + assert_close(dinput[0], -0.5f32.sin()); + } +} + +#[tokio::test] +async fn test_backward_cosh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.cosh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.cosh()); + assert_close(dinput[0], 0.5f32.sinh()); + } +} + +#[tokio::test] +async fn test_backward_exp2() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32]); + + let output = input.exp2(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 2.0f32.exp2()); + assert_close(dinput[0], std::f32::consts::LN_2 * 2.0f32.exp2()); + } +} + +#[tokio::test] +async fn test_backward_less_approximate_exp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32]); + + let output = input.less_approximate_exp(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 1.0f32.exp()); + assert_close(dinput[0], 1.0f32.exp()); + } +} + +#[tokio::test] +async fn test_backward_log2() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32]); + + let output = input.log2(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 4.0f32.log2()); + assert_close(dinput[0], 1.0f32 / (4.0f32 * std::f32::consts::LN_2)); + } +} + +#[tokio::test] +async fn test_backward_sin() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.sin(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.sin()); + assert_close(dinput[0], 0.5f32.cos()); + } +} + +#[tokio::test] +async fn test_backward_sinh() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.sinh(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.sinh()); + assert_close(dinput[0], 0.5f32.cosh()); + } +} + +#[tokio::test] +async fn test_backward_tan() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.tan(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.tan()); + assert_close(dinput[0], 1.0f32 / (0.5f32.cos() * 0.5f32.cos())); + } +} + +#[tokio::test] +async fn test_backward_tanh_exact() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32]); + + let output = input.tanh_exact(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 0.5f32.tanh()); + assert_close(dinput[0], 1.0f32 - 0.5f32.tanh().powi(2)); + } +} + +#[tokio::test] +async fn test_cast() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + let output = input.cast::(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + + assert_close(f32::from(output_values[0]), 1.0); + assert_close(f32::from(output_values[1]), 2.0); + assert_close(f32::from(output_values[2]), 3.0); + } +} + +#[tokio::test] +async fn test_arange() { + for device in test_devices().await { + let graph = Graph::new(); + + let output = Tensor::<1>::arange(&graph, &device, 1.0, 5.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + + assert_eq!(output_values, vec![1.0, 2.0, 3.0, 4.0]); + } +} + +#[tokio::test] +async fn test_arange_step() { + for device in test_devices().await { + let graph = Graph::new(); + + let output = Tensor::<1>::arange_step(&graph, &device, 1.0, 6.0, 2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + + assert_eq!(output_values, vec![1.0, 3.0, 5.0]); + } +} + +#[tokio::test] +async fn test_full() { + for device in test_devices().await { + let graph = Graph::new(); + + let output: Tensor<2> = Tensor::full(&graph, &device, [2, 3], 1.5); + let output_values = output.raw().clone().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 3]); + for row in 0..2 { + for col in 0..3 { + assert_close(output_values[[row, col]], 1.5); + } + } + } +} + +#[tokio::test] +async fn test_zeros_like() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + + let output = input.zeros_like(); + let output_values = output.raw().clone().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 0.0); + assert_close(output_values[[0, 1]], 0.0); + assert_close(output_values[[1, 0]], 0.0); + assert_close(output_values[[1, 1]], 0.0); + } +} + +#[tokio::test] +async fn test_from_array() { + for device in test_devices().await { + let graph = Graph::new(); + + let output: Tensor<2> = Tensor::from_array(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let output_values = output.raw().clone().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 1.0); + assert_close(output_values[[0, 1]], 2.0); + assert_close(output_values[[1, 0]], 3.0); + assert_close(output_values[[1, 1]], 4.0); + } +} + +#[tokio::test] +async fn test_backward_add_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0]); + let rhs: Tensor<2> = Tensor::new(&graph, &device, &[[10.0f32], [20.0]]); + + let output: Tensor<2> = lhs.add_::<2, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0][0], 11.0); + assert_close(output_values[0][1], 12.0); + assert_close(output_values[1][0], 21.0); + assert_close(output_values[1][1], 22.0); + assert_close(dlhs[0], 2.0); + assert_close(dlhs[1], 2.0); + assert_close(drhs[0][0], 2.0); + assert_close(drhs[1][0], 2.0); + } +} + +#[tokio::test] +async fn test_backward_sub_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<2> = Tensor::new(&graph, &device, &[[3.0f32], [4.0]]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0]); + + let output: Tensor<2> = lhs.sub_::<1, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0][0], 2.0); + assert_close(output_values[0][1], 1.0); + assert_close(output_values[1][0], 3.0); + assert_close(output_values[1][1], 2.0); + assert_close(dlhs[0][0], 2.0); + assert_close(dlhs[1][0], 2.0); + assert_close(drhs[0], -2.0); + assert_close(drhs[1], -2.0); + } +} + +#[tokio::test] +async fn test_backward_mul_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 3.0]); + let rhs: Tensor<2> = Tensor::new(&graph, &device, &[[10.0f32], [20.0]]); + + let output: Tensor<2> = lhs.mul_::<2, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0][0], 20.0); + assert_close(output_values[0][1], 30.0); + assert_close(output_values[1][0], 40.0); + assert_close(output_values[1][1], 60.0); + assert_close(dlhs[0], 30.0); + assert_close(dlhs[1], 30.0); + assert_close(drhs[0][0], 5.0); + assert_close(drhs[1][0], 5.0); + } +} + +#[tokio::test] +async fn test_backward_div_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<2> = Tensor::new(&graph, &device, &[[10.0f32], [20.0]]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 4.0]); + + let output: Tensor<2> = lhs.div_::<1, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0][0], 5.0); + assert_close(output_values[0][1], 2.5); + assert_close(output_values[1][0], 10.0); + assert_close(output_values[1][1], 5.0); + assert_close(dlhs[0][0], 0.75); + assert_close(dlhs[1][0], 0.75); + assert_close(drhs[0], -7.5); + assert_close(drhs[1], -1.875); + } +} + +#[tokio::test] +async fn test_backward_pow_broadcast_api() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, 3.0]); + let rhs: Tensor<2> = Tensor::new(&graph, &device, &[[2.0f32], [1.0]]); + + let output: Tensor<2> = lhs.pow_::<2, 2>(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0][0], 4.0); + assert_close(output_values[0][1], 9.0); + assert_close(output_values[1][0], 2.0); + assert_close(output_values[1][1], 3.0); + assert_close(dlhs[0], 5.0); + assert_close(dlhs[1], 7.0); + assert_close(drhs[0][0], 12.660099); + assert_close(drhs[1][0], 4.6821313); + } +} + +#[tokio::test] +async fn test_backward_chunk() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new( + &graph, + &device, + &[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]], + ); + + let chunks = input.chunk(2, 1); + assert_eq!(chunks.len(), 2); + let first = chunks[0].raw().clone().as_slice().await.unwrap().to_vec(); + let second = chunks[1].raw().clone().as_slice().await.unwrap().to_vec(); + let loss = chunks[0] + .flatten_all() + .sum() + .add(&chunks[1].flatten_all().sum().mul_scalar(2.0)); + let gradients = loss.backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(first[0][0], 1.0); + assert_close(first[0][1], 2.0); + assert_close(first[1][0], 5.0); + assert_close(first[1][1], 6.0); + assert_close(second[0][0], 3.0); + assert_close(second[0][1], 4.0); + assert_close(second[1][0], 7.0); + assert_close(second[1][1], 8.0); + assert_close(dinput[0][0], 1.0); + assert_close(dinput[0][1], 1.0); + assert_close(dinput[0][2], 2.0); + assert_close(dinput[0][3], 2.0); + assert_close(dinput[1][0], 1.0); + assert_close(dinput[1][1], 1.0); + assert_close(dinput[1][2], 2.0); + assert_close(dinput[1][3], 2.0); + } +} + +#[tokio::test] +async fn test_backward_matmul() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let rhs: Tensor<2> = Tensor::new(&graph, &device, &[[5.0f32, 6.0], [7.0, 8.0]]); + + let output = lhs.matmul(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = gradients.get(&lhs).unwrap().as_slice().await.unwrap(); + let drhs = gradients.get(&rhs).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 19.0); + assert_close(output_values[[0, 1]], 22.0); + assert_close(output_values[[1, 0]], 43.0); + assert_close(output_values[[1, 1]], 50.0); + + assert_close(dlhs[[0, 0]], 11.0); + assert_close(dlhs[[0, 1]], 15.0); + assert_close(dlhs[[1, 0]], 11.0); + assert_close(dlhs[[1, 1]], 15.0); + + assert_close(drhs[[0, 0]], 4.0); + assert_close(drhs[[0, 1]], 4.0); + assert_close(drhs[[1, 0]], 6.0); + assert_close(drhs[[1, 1]], 6.0); + } +} + +#[tokio::test] +async fn test_backward_t() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + + let output = input.t(); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 1.0); + assert_close(output_values[[0, 1]], 3.0); + assert_close(output_values[[1, 0]], 2.0); + assert_close(output_values[[1, 1]], 4.0); + + assert_close(dinput[[0, 0]], 1.0); + assert_close(dinput[[0, 1]], 1.0); + assert_close(dinput[[1, 0]], 1.0); + assert_close(dinput[[1, 1]], 1.0); + } +} + +#[tokio::test] +async fn test_backward_pool() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new(&graph, &device, &[[[1.0f32, 2.0, 3.0, 4.0]]]); + + let output = input.pool::<1, 4, 5, 4>([(2, 1)], |windowed, axis| windowed.mean::<3>(axis)); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 3]); + assert_close(output_values[[0, 0, 0]], 1.5); + assert_close(output_values[[0, 0, 1]], 2.5); + assert_close(output_values[[0, 0, 2]], 3.5); + + assert_close(dinput[[0, 0, 0]], 0.5); + assert_close(dinput[[0, 0, 1]], 1.0); + assert_close(dinput[[0, 0, 2]], 1.0); + assert_close(dinput[[0, 0, 3]], 0.5); + } +} + +#[tokio::test] +async fn test_backward_pool_max() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new(&graph, &device, &[[[1.0f32, 4.0, 2.0, 3.0]]]); + + let output = input.pool_max::<1, 4, 5, 4>([(2, 1)]); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 3]); + assert_close(output_values[[0, 0, 0]], 4.0); + assert_close(output_values[[0, 0, 1]], 4.0); + assert_close(output_values[[0, 0, 2]], 3.0); + + assert_close(dinput[[0, 0, 0]], 0.0); + assert_close(dinput[[0, 0, 1]], 2.0); + assert_close(dinput[[0, 0, 2]], 0.0); + assert_close(dinput[[0, 0, 3]], 1.0); + } +} + +#[tokio::test] +async fn test_backward_pool_min() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = Tensor::new(&graph, &device, &[[[1.0f32, 4.0, 2.0, 3.0]]]); + + let output = input.pool_min::<1, 4, 5, 4>([(2, 1)]); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 3]); + assert_close(output_values[[0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 2]], 2.0); + + assert_close(dinput[[0, 0, 0]], 1.0); + assert_close(dinput[[0, 0, 1]], 0.0); + assert_close(dinput[[0, 0, 2]], 2.0); + assert_close(dinput[[0, 0, 3]], 0.0); + } +} + +#[tokio::test] +async fn test_backward_q_mat_mul() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0, 1.0, 1.0]]); + let weight_bytes: Vec = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] + .into_iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let weights = crate::QMatrix::from_raw_bytes( + &device, + [2, 4], + &weight_bytes, + fusor_gguf::GgmlType::F32, + ) + .unwrap(); + + let output = input.q_mat_mul(&weights); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 2]); + assert_close(output_values[[0, 0]], 10.0); + assert_close(output_values[[0, 1]], 26.0); + + assert_close(dinput[[0, 0]], 6.0); + assert_close(dinput[[0, 1]], 8.0); + assert_close(dinput[[0, 2]], 10.0); + assert_close(dinput[[0, 3]], 12.0); + } +} + +#[tokio::test] +async fn test_backward_stack() { + for device in test_devices().await { + let graph = Graph::new(); + let first: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0]); + let second: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32, 4.0]); + + let output = Tensor::stack::<2>(vec![first.clone(), second.clone()], 0); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dfirst = gradients + .get(&first) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let dsecond = gradients + .get(&second) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values.shape(), &[2, 2]); + assert_close(output_values[[0, 0]], 1.0); + assert_close(output_values[[0, 1]], 2.0); + assert_close(output_values[[1, 0]], 3.0); + assert_close(output_values[[1, 1]], 4.0); + + assert_close(dfirst[0], 1.0); + assert_close(dfirst[1], 1.0); + assert_close(dsecond[0], 1.0); + assert_close(dsecond[1], 1.0); + } +} + +#[tokio::test] +async fn test_backward_rope() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]], + ); + let cos: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0], [1.0, 1.0]]); + let sin: Tensor<2> = Tensor::new(&graph, &device, &[[0.0f32, 0.0], [0.0, 0.0]]); + + let output = input.rope(&cos, &sin); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + let dcos = gradients.get(&cos).unwrap().as_slice().await.unwrap(); + let dsin = gradients.get(&sin).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 2, 4]); + assert_close(output_values[[0, 0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 0, 2]], 3.0); + assert_close(output_values[[0, 0, 0, 3]], 4.0); + assert_close(output_values[[0, 0, 1, 0]], 5.0); + assert_close(output_values[[0, 0, 1, 1]], 6.0); + assert_close(output_values[[0, 0, 1, 2]], 7.0); + assert_close(output_values[[0, 0, 1, 3]], 8.0); + + for index in [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 0, 2], + [0, 0, 0, 3], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 0, 1, 2], + [0, 0, 1, 3], + ] { + assert_close(dinput[index], 1.0); + } + + assert_close(dcos[[0, 0]], 4.0); + assert_close(dcos[[0, 1]], 6.0); + assert_close(dcos[[1, 0]], 12.0); + assert_close(dcos[[1, 1]], 14.0); + + assert_close(dsin[[0, 0]], -2.0); + assert_close(dsin[[0, 1]], -2.0); + assert_close(dsin[[1, 0]], -2.0); + assert_close(dsin[[1, 1]], -2.0); + } +} + +#[tokio::test] +async fn test_backward_rope_fused() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]], + ); + let cos: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0], [1.0, 1.0]]); + let sin: Tensor<2> = Tensor::new(&graph, &device, &[[0.0f32, 0.0], [0.0, 0.0]]); + + let output = input.rope_fused(&cos, &sin); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + let dcos = gradients.get(&cos).unwrap().as_slice().await.unwrap(); + let dsin = gradients.get(&sin).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 2, 4]); + assert_close(output_values[[0, 0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 0, 2]], 3.0); + assert_close(output_values[[0, 0, 0, 3]], 4.0); + assert_close(output_values[[0, 0, 1, 0]], 5.0); + assert_close(output_values[[0, 0, 1, 1]], 6.0); + assert_close(output_values[[0, 0, 1, 2]], 7.0); + assert_close(output_values[[0, 0, 1, 3]], 8.0); + + for index in [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 0, 2], + [0, 0, 0, 3], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 0, 1, 2], + [0, 0, 1, 3], + ] { + assert_close(dinput[index], 1.0); + } + + assert_close(dcos[[0, 0]], 3.0); + assert_close(dcos[[0, 1]], 7.0); + assert_close(dcos[[1, 0]], 11.0); + assert_close(dcos[[1, 1]], 15.0); + + assert_close(dsin[[0, 0]], -1.0); + assert_close(dsin[[0, 1]], -1.0); + assert_close(dsin[[1, 0]], -1.0); + assert_close(dsin[[1, 1]], -1.0); + } +} + +#[tokio::test] +async fn test_backward_rope_interleaved() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]], + ); + let cos: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0], [1.0, 1.0]]); + let sin: Tensor<2> = Tensor::new(&graph, &device, &[[0.0f32, 0.0], [0.0, 0.0]]); + + let output = input.rope_interleaved(&cos, &sin); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + let dcos = gradients.get(&cos).unwrap().as_slice().await.unwrap(); + let dsin = gradients.get(&sin).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 2, 4]); + assert_close(output_values[[0, 0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 0, 2]], 3.0); + assert_close(output_values[[0, 0, 0, 3]], 4.0); + assert_close(output_values[[0, 0, 1, 0]], 5.0); + assert_close(output_values[[0, 0, 1, 1]], 6.0); + assert_close(output_values[[0, 0, 1, 2]], 7.0); + assert_close(output_values[[0, 0, 1, 3]], 8.0); + + for index in [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 0, 2], + [0, 0, 0, 3], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 0, 1, 2], + [0, 0, 1, 3], + ] { + assert_close(dinput[index], 1.0); + } + + assert_close(dcos[[0, 0]], 3.0); + assert_close(dcos[[0, 1]], 7.0); + assert_close(dcos[[1, 0]], 11.0); + assert_close(dcos[[1, 1]], 15.0); + + assert_close(dsin[[0, 0]], -1.0); + assert_close(dsin[[0, 1]], -1.0); + assert_close(dsin[[1, 0]], -1.0); + assert_close(dsin[[1, 1]], -1.0); + } +} + +#[tokio::test] +async fn test_backward_rope_normal_fused() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<4> = Tensor::new( + &graph, + &device, + &[[[[1.0f32, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]]], + ); + let cos: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 1.0], [1.0, 1.0]]); + let sin: Tensor<2> = Tensor::new(&graph, &device, &[[0.0f32, 0.0], [0.0, 0.0]]); + + let output = input.rope_normal_fused(&cos, &sin); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dinput = gradients.get(&input).unwrap().as_slice().await.unwrap(); + let dcos = gradients.get(&cos).unwrap().as_slice().await.unwrap(); + let dsin = gradients.get(&sin).unwrap().as_slice().await.unwrap(); + + assert_eq!(output_values.shape(), &[1, 1, 2, 4]); + assert_close(output_values[[0, 0, 0, 0]], 1.0); + assert_close(output_values[[0, 0, 0, 1]], 2.0); + assert_close(output_values[[0, 0, 0, 2]], 3.0); + assert_close(output_values[[0, 0, 0, 3]], 4.0); + assert_close(output_values[[0, 0, 1, 0]], 5.0); + assert_close(output_values[[0, 0, 1, 1]], 6.0); + assert_close(output_values[[0, 0, 1, 2]], 7.0); + assert_close(output_values[[0, 0, 1, 3]], 8.0); + + for index in [ + [0, 0, 0, 0], + [0, 0, 0, 1], + [0, 0, 0, 2], + [0, 0, 0, 3], + [0, 0, 1, 0], + [0, 0, 1, 1], + [0, 0, 1, 2], + [0, 0, 1, 3], + ] { + assert_close(dinput[index], 1.0); + } + + assert_close(dcos[[0, 0]], 4.0); + assert_close(dcos[[0, 1]], 6.0); + assert_close(dcos[[1, 0]], 12.0); + assert_close(dcos[[1, 1]], 14.0); + + assert_close(dsin[[0, 0]], -2.0); + assert_close(dsin[[0, 1]], -2.0); + assert_close(dsin[[1, 0]], -2.0); + assert_close(dsin[[1, 1]], -2.0); + } +} + +#[tokio::test] +async fn test_backward_pow() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32]); + let rhs: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32]); + + let output = lhs.pow(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dlhs = gradients + .get(&lhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let drhs = gradients + .get(&rhs) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 8.0); + assert_close(dlhs[0], 12.0); + assert_close(drhs[0], 8.0 * 2.0f32.ln()); + } +} + +#[tokio::test] +async fn test_backward_pow_elementwise() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32]); + + let output = input.pow_elementwise(2.0); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 9.0); + assert_close(dinput[0], 6.0); + } +} + +#[tokio::test] +async fn test_backward_pow_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32]); + + let output = input.pow_scalar(0.5); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(output_values[0], 2.0); + assert_close(dinput[0], 0.25); + } +} + +#[tokio::test] +async fn test_autograd_rms_norm() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let weight: Tensor<2> = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [1, 3], &[1.0f32, 1.0, 1.0]), + ); + + let output = input.rms_norm(&weight, 1e-5); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + + let expected = [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]].map(|row| { + let mean_sq = row.iter().map(|value| value * value).sum::() / row.len() as f32; + let scale = 1.0 / (mean_sq + 1e-5).sqrt(); + row.map(|value| value * scale) + }); + + for (actual_row, expected_row) in output_values.iter().zip(expected.iter()) { + for (actual, expected) in actual_row.iter().zip(expected_row.iter()) { + assert_close(*actual, *expected); + } + } + + let gradients = output.sum(1).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + // d/dx_j sum_k x_k * (mean(x^2) + eps)^-1/2 + // = 1/rms - x_j * sum(x) / (n * rms^3) + let expected_grads = [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]].map(|row| { + let n = row.len() as f32; + let mean_sq = row.iter().map(|value| value * value).sum::() / n; + let rms = (mean_sq + 1e-5).sqrt(); + let sum = row.iter().sum::(); + row.map(|value| 1.0 / rms - value * sum / (n * rms.powi(3))) + }); + for (actual_row, expected_row) in dinput.iter().zip(expected_grads.iter()) { + for (actual, expected) in actual_row.iter().zip(expected_row.iter()) { + assert_close(*actual, *expected); + } + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + move |graph, x| { + let weight = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &[1.0f32, 1.0, 1.0]), + ); + x.rms_norm(&weight, 1e-5).sum(1).sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_matmul_with_broadcast_bias() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let w: Tensor<2> = Tensor::new(&graph, &device, &[[0.5f32], [1.0], [1.5]]); + let b: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32]); + + let y = x.mat_mul(&w).add(&b.broadcast_as([2, 1])); + let loss = y.sum(1).sum(); + + let gradients = loss.backward().unwrap(); + let dw = gradients + .get(&w) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let db = gradients + .get(&b) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(dw[0][0], 5.0); + assert_close(dw[1][0], 7.0); + assert_close(dw[2][0], 9.0); + assert_close(db[0], 2.0); + } +} + +#[tokio::test] +async fn test_backward_embedding() { + for device in test_devices().await { + let graph = Graph::new(); + + let table: Tensor<2> = + Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]]); + let indices: RawTensor<2, u32> = RawTensor::new(&device, &[[0u32, 2u32]]); + let embedded = table.embedding(&indices); + let loss = embedded.sum(2).sum(1).sum(); + + let gradients = loss.backward().unwrap(); + let dtable = gradients + .get(&table) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(dtable[0][0], 1.0); + assert_close(dtable[0][1], 1.0); + assert_close(dtable[1][0], 0.0); + assert_close(dtable[1][1], 0.0); + assert_close(dtable[2][0], 1.0); + assert_close(dtable[2][1], 1.0); + } +} + +#[tokio::test] +async fn test_backward_gather_last() { + for device in test_devices().await { + let graph = Graph::new(); + + let values: Tensor<2> = + Tensor::new(&graph, &device, &[[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]); + let indices: RawTensor<1, u32> = RawTensor::new(&device, &[2u32, 0u32]); + let gathered = values.gather_last(&indices); + let loss = gathered.sum(); + + let gradients = loss.backward().unwrap(); + let dvalues = gradients + .get(&values) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_close(dvalues[0][0], 0.0); + assert_close(dvalues[0][1], 0.0); + assert_close(dvalues[0][2], 1.0); + assert_close(dvalues[1][0], 1.0); + assert_close(dvalues[1][1], 0.0); + assert_close(dvalues[1][2], 0.0); + } +} + +#[tokio::test] +async fn test_backward_softmax_last_dim_fused_matches_composite() { + for device in test_devices().await { + let input_data = &[ + [[0.2f32, -0.4, 1.1], [0.5, 0.3, -0.7]], + [[-1.0, 0.8, 0.6], [0.9, -0.2, 0.1]], + ]; + + let fused_graph = Graph::new(); + let fused_input: Tensor<3> = Tensor::new(&fused_graph, &device, input_data); + let fused_output = fused_input.softmax_last_dim_fused::<2>(); + let fused_loss = fused_output.sqr().reshape([12]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_input: Tensor<3> = Tensor::new(&composite_graph, &device, input_data); + let composite_output = composite_input.softmax_last_dim::<2>(); + let composite_loss = composite_output.sqr().reshape([12]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dx = flatten(fused_gradients.get(&fused_input).unwrap()).await; + let composite_dx = flatten(composite_gradients.get(&composite_input).unwrap()).await; + + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dx, &composite_dx); + } +} + +#[tokio::test] +async fn test_backward_rms_norm_fused_matches_composite() { + for device in test_devices().await { + let input_data = &[ + [[0.3f32, -1.2, 0.7], [1.5, 0.1, -0.8]], + [[-0.4, 0.9, 1.3], [0.2, -0.6, 0.5]], + ]; + let weight_data = &[1.0f32, 0.75, 1.25]; + let eps = 1e-5; + + let fused_graph = Graph::new(); + let fused_input: Tensor<3> = Tensor::new(&fused_graph, &device, input_data); + let fused_weight: Tensor<1> = Tensor::new(&fused_graph, &device, weight_data); + let fused_output = fused_input.rms_norm_fused_no_bias::<1, 2>(&fused_weight, eps); + let fused_loss = fused_output.sqr().reshape([12]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_input: Tensor<3> = Tensor::new(&composite_graph, &device, input_data); + let composite_weight: Tensor<3> = + Tensor::from_slice(&composite_graph, &device, [1, 1, 3], weight_data); + let composite_output = composite_input.rms_norm(&composite_weight, eps); + let composite_loss = composite_output.sqr().reshape([12]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dx = flatten(fused_gradients.get(&fused_input).unwrap()).await; + let composite_dx = flatten(composite_gradients.get(&composite_input).unwrap()).await; + let fused_dw = flatten(fused_gradients.get(&fused_weight).unwrap()).await; + let composite_dw = flatten(composite_gradients.get(&composite_weight).unwrap()).await; + + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dx, &composite_dx); + assert_slice_close(&fused_dw, &composite_dw); + } +} + +#[tokio::test] +async fn test_backward_layer_norm_last_dim_fused_matches_composite() { + for device in test_devices().await { + let input_data = &[ + [[0.25f32, -0.5, 1.0], [1.25, -1.5, 0.75]], + [[-0.8, 0.4, 1.2], [0.6, -0.1, -0.9]], + ]; + let weight_data = &[1.0f32, 0.9, 1.1]; + let bias_data = &[0.1f32, -0.2, 0.05]; + let eps = 1e-5; + + let fused_graph = Graph::new(); + let fused_input: Tensor<3> = Tensor::new(&fused_graph, &device, input_data); + let fused_weight: Tensor<1> = Tensor::new(&fused_graph, &device, weight_data); + let fused_bias: Tensor<1> = Tensor::new(&fused_graph, &device, bias_data); + let fused_output = + fused_input.layer_norm_last_dim_fused::<2, 1>(&fused_weight, Some(&fused_bias), eps); + let fused_loss = fused_output.sqr().reshape([12]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_input: Tensor<3> = Tensor::new(&composite_graph, &device, input_data); + let composite_weight: Tensor<3> = + Tensor::from_slice(&composite_graph, &device, [1, 1, 3], weight_data); + let composite_bias: Tensor<3> = + Tensor::from_slice(&composite_graph, &device, [1, 1, 3], bias_data); + let composite_output = + composite_input.layer_norm(&composite_weight, Some(&composite_bias), eps, true); + let composite_loss = composite_output.sqr().reshape([12]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dx = flatten(fused_gradients.get(&fused_input).unwrap()).await; + let composite_dx = flatten(composite_gradients.get(&composite_input).unwrap()).await; + let fused_dw = flatten(fused_gradients.get(&fused_weight).unwrap()).await; + let composite_dw = flatten(composite_gradients.get(&composite_weight).unwrap()).await; + let fused_db = flatten(fused_gradients.get(&fused_bias).unwrap()).await; + let composite_db = flatten(composite_gradients.get(&composite_bias).unwrap()).await; + + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dx, &composite_dx); + assert_slice_close(&fused_dw, &composite_dw); + assert_slice_close(&fused_db, &composite_db); + } +} + +#[tokio::test] +async fn test_backward_attention_matches_composite() { + for device in test_devices().await { + let q_data = &[[[[0.2f32, 0.6], [1.0, -0.3]]]]; + let k_data = &[[[[0.4f32, -0.7], [0.9, 0.1]]]]; + let v_data = &[[[[1.1f32, -0.5], [0.3, 0.8]]]]; + let scale = (2.0f32).sqrt(); + + let fused_graph = Graph::new(); + let fused_q: Tensor<4> = Tensor::new(&fused_graph, &device, q_data); + let fused_k: Tensor<4> = Tensor::new(&fused_graph, &device, k_data); + let fused_v: Tensor<4> = Tensor::new(&fused_graph, &device, v_data); + let fused_output = fused_q.attention(&fused_k, &fused_v, scale, None); + let fused_loss = fused_output.sqr().reshape([4]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_q: Tensor<4> = Tensor::new(&composite_graph, &device, q_data); + let composite_k: Tensor<4> = Tensor::new(&composite_graph, &device, k_data); + let composite_v: Tensor<4> = Tensor::new(&composite_graph, &device, v_data); + let composite_output = + composite_q.attention_composite(&composite_k, &composite_v, scale, None); + let composite_loss = composite_output.sqr().reshape([4]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dq = flatten(fused_gradients.get(&fused_q).unwrap()).await; + let composite_dq = flatten(composite_gradients.get(&composite_q).unwrap()).await; + let fused_dk = flatten(fused_gradients.get(&fused_k).unwrap()).await; + let composite_dk = flatten(composite_gradients.get(&composite_k).unwrap()).await; + let fused_dv = flatten(fused_gradients.get(&fused_v).unwrap()).await; + let composite_dv = flatten(composite_gradients.get(&composite_v).unwrap()).await; + + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dq, &composite_dq); + assert_slice_close(&fused_dk, &composite_dk); + assert_slice_close(&fused_dv, &composite_dv); + } +} + +#[tokio::test] +async fn test_backward_attention_flash_ineligible_shape() { + // A causal shape every flash gate rejects (q and kv not multiples of the + // tile sizes, head_dim not a multiple of 32): the forward lowers through + // the attention row program and the explicit backward's lse/grad + // clusters stay composed. The whole pipeline must still produce finite + // gradients that match the composite replay. + const BATCH: usize = 1; + const HEADS: usize = 2; + const SEQ: usize = 50; + const DIM: usize = 48; + let elements = BATCH * HEADS * SEQ * DIM; + let mut state = 0x2458_71a3_u64; + let mut next = move || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5 + }; + let q_data: Vec = (0..elements).map(|_| next()).collect(); + let k_data: Vec = (0..elements).map(|_| next()).collect(); + let v_data: Vec = (0..elements).map(|_| next()).collect(); + let mask_data: Vec = (0..SEQ * SEQ) + .map(|i| if i % SEQ <= i / SEQ { 0.0 } else { -1.0e9 }) + .collect(); + let scale = 1.0 / (DIM as f32).sqrt(); + + for device in test_devices().await { + let mask = crate::Tensor::from_slice(&device, [SEQ, SEQ], &mask_data); + + let fused_graph = Graph::new(); + let fused_q: Tensor<4> = + Tensor::from_slice(&fused_graph, &device, [BATCH, HEADS, SEQ, DIM], &q_data); + let fused_k: Tensor<4> = + Tensor::from_slice(&fused_graph, &device, [BATCH, HEADS, SEQ, DIM], &k_data); + let fused_v: Tensor<4> = + Tensor::from_slice(&fused_graph, &device, [BATCH, HEADS, SEQ, DIM], &v_data); + let fused_output = fused_q.attention( + &fused_k, + &fused_v, + scale, + Some((&mask, crate::MaskKind::Causal)), + ); + let fused_loss = fused_output.sqr().reshape([elements]).sum(); + let fused_gradients = fused_loss.backward().unwrap(); + + let composite_graph = Graph::new(); + let composite_q: Tensor<4> = + Tensor::from_slice(&composite_graph, &device, [BATCH, HEADS, SEQ, DIM], &q_data); + let composite_k: Tensor<4> = + Tensor::from_slice(&composite_graph, &device, [BATCH, HEADS, SEQ, DIM], &k_data); + let composite_v: Tensor<4> = + Tensor::from_slice(&composite_graph, &device, [BATCH, HEADS, SEQ, DIM], &v_data); + let composite_output = composite_q.attention_composite( + &composite_k, + &composite_v, + scale, + Some(&(mask.clone(), crate::MaskKind::Causal)), + ); + let composite_loss = composite_output.sqr().reshape([elements]).sum(); + let composite_gradients = composite_loss.backward().unwrap(); + + let fused_output = flatten(fused_output.raw().clone()).await; + let composite_output = flatten(composite_output.raw().clone()).await; + let fused_dq = flatten(fused_gradients.get(&fused_q).unwrap()).await; + let composite_dq = flatten(composite_gradients.get(&composite_q).unwrap()).await; + let fused_dk = flatten(fused_gradients.get(&fused_k).unwrap()).await; + let composite_dk = flatten(composite_gradients.get(&composite_k).unwrap()).await; + let fused_dv = flatten(fused_gradients.get(&fused_v).unwrap()).await; + let composite_dv = flatten(composite_gradients.get(&composite_v).unwrap()).await; + + for (name, values) in [ + ("output", &fused_output), + ("dq", &fused_dq), + ("dk", &fused_dk), + ("dv", &fused_dv), + ] { + assert!( + values.iter().all(|value| value.is_finite()), + "{name} contains non-finite values" + ); + } + assert_slice_close(&fused_output, &composite_output); + assert_slice_close(&fused_dq, &composite_dq); + assert_slice_close(&fused_dk, &composite_dk); + assert_slice_close(&fused_dv, &composite_dv); + } +} + +#[tokio::test] +async fn test_backward_mat_mul_rank3() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs_data = (1..=24).map(|n| n as f32).collect::>(); + let rhs_data = (1..=40).map(|n| n as f32).collect::>(); + let lhs: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 3, 4], &lhs_data); + let rhs: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 4, 5], &rhs_data); + + let output = lhs.mat_mul(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = flatten(gradients.get(&lhs).unwrap()).await; + let drhs = flatten(gradients.get(&rhs).unwrap()).await; + + assert_eq!(output_values.shape(), &[2, 3, 5]); + assert_close(output_values[[0, 0, 0]], 110.0); + assert_close(output_values[[1, 2, 4]], 2950.0); + + // with an all-ones seed, dlhs[b, i, k] = sum_j rhs[b, k, j] and + // drhs[b, k, j] = sum_i lhs[b, i, k] + for b in 0..2 { + for i in 0..3 { + for k in 0..4 { + let expected = (0..5).map(|j| rhs_data[b * 20 + k * 5 + j]).sum::(); + assert_close(dlhs[b * 12 + i * 4 + k], expected); + } + } + for k in 0..4 { + for j in 0..5 { + let expected = (0..3).map(|i| lhs_data[b * 12 + i * 4 + k]).sum::(); + assert_close(drhs[b * 20 + k * 5 + j], expected); + } + } + } + + let lhs_small = lhs_data + .iter() + .map(|value| value * 0.05) + .collect::>(); + let rhs_small = rhs_data + .iter() + .map(|value| value * 0.03) + .collect::>(); + let fd_device = device.clone(); + let fd_rhs = rhs_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 3, 4], + &lhs_small, + move |graph, lhs| { + let rhs = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 4, 5], &fd_rhs), + ); + lhs.mat_mul(&rhs).sqr().flatten_all().sum() + }, + ) + .await; + let fd_device = device.clone(); + let fd_lhs = lhs_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 4, 5], + &rhs_small, + move |graph, rhs| { + let lhs = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3, 4], &fd_lhs), + ); + lhs.mat_mul(&rhs).sqr().flatten_all().sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_cat_dim0() { + for device in test_devices().await { + let graph = Graph::new(); + let first_data = (1..=6).map(|n| n as f32).collect::>(); + let second_data = (7..=18).map(|n| n as f32).collect::>(); + let first: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 2, 3], &first_data); + let second: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &second_data); + + let output = Tensor::cat(vec![first.clone(), second.clone()], 0); + let output_values = flatten(output.raw().clone()).await; + let seed_data = (0..18).map(|n| n as f32 + 10.0).collect::>(); + let seed = RawTensor::from_slice(&device, [3, 2, 3], &seed_data); + let gradients = output.backward_with(seed).unwrap(); + let dfirst = flatten(gradients.get(&first).unwrap()).await; + let dsecond = flatten(gradients.get(&second).unwrap()).await; + + assert_eq!(output.shape(), [3, 2, 3]); + assert_eq!( + output_values, + (1..=18).map(|n| n as f32).collect::>() + ); + assert_eq!(dfirst, seed_data[..6].to_vec()); + assert_eq!(dsecond, seed_data[6..].to_vec()); + } +} + +#[tokio::test] +async fn test_backward_cat_dim1() { + for device in test_devices().await { + let graph = Graph::new(); + let first_data = (1..=6).map(|n| n as f32).collect::>(); + let second_data = (10..=21).map(|n| n as f32).collect::>(); + let first: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 1, 3], &first_data); + let second: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &second_data); + + let output = Tensor::cat(vec![first.clone(), second.clone()], 1); + let output_values = flatten(output.raw().clone()).await; + let seed_data = (0..18).map(|n| n as f32 + 10.0).collect::>(); + let seed = RawTensor::from_slice(&device, [2, 3, 3], &seed_data); + let gradients = output.backward_with(seed).unwrap(); + let dfirst = flatten(gradients.get(&first).unwrap()).await; + let dsecond = flatten(gradients.get(&second).unwrap()).await; + + assert_eq!(output.shape(), [2, 3, 3]); + assert_eq!( + output_values, + vec![ + 1.0, 2.0, 3.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 4.0, 5.0, 6.0, 16.0, 17.0, 18.0, + 19.0, 20.0, 21.0 + ] + ); + + let mut expected_dfirst = Vec::new(); + let mut expected_dsecond = Vec::new(); + for i in 0..2 { + for j in 0..3 { + for k in 0..3 { + let value = seed_data[i * 9 + j * 3 + k]; + if j < 1 { + expected_dfirst.push(value); + } else { + expected_dsecond.push(value); + } + } + } + } + assert_eq!(dfirst, expected_dfirst); + assert_eq!(dsecond, expected_dsecond); + + let first_small = first_data + .iter() + .map(|value| value * 0.1) + .collect::>(); + let second_small = second_data + .iter() + .map(|value| value * 0.1) + .collect::>(); + let fd_device = device.clone(); + let fd_second = second_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 1, 3], + &first_small, + move |graph, first| { + let second = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &fd_second), + ); + Tensor::cat(vec![first, second], 1) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + let fd_device = device.clone(); + let fd_first = first_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 3], + &second_small, + move |graph, second| { + let first = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 1, 3], &fd_first), + ); + Tensor::cat(vec![first, second], 1) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_cat_dim2() { + for device in test_devices().await { + let graph = Graph::new(); + let first_data = (1..=4).map(|n| n as f32).collect::>(); + let second_data = (5..=12).map(|n| n as f32).collect::>(); + let first: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 1], &first_data); + let second: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 2], &second_data); + + let output = Tensor::cat(vec![first.clone(), second.clone()], 2); + let output_values = flatten(output.raw().clone()).await; + let seed_data = (0..12).map(|n| n as f32 + 10.0).collect::>(); + let seed = RawTensor::from_slice(&device, [2, 2, 3], &seed_data); + let gradients = output.backward_with(seed).unwrap(); + let dfirst = flatten(gradients.get(&first).unwrap()).await; + let dsecond = flatten(gradients.get(&second).unwrap()).await; + + assert_eq!(output.shape(), [2, 2, 3]); + assert_eq!( + output_values, + vec![ + 1.0, 5.0, 6.0, 2.0, 7.0, 8.0, 3.0, 9.0, 10.0, 4.0, 11.0, 12.0 + ] + ); + + let mut expected_dfirst = Vec::new(); + let mut expected_dsecond = Vec::new(); + for i in 0..2 { + for j in 0..2 { + for k in 0..3 { + let value = seed_data[i * 6 + j * 3 + k]; + if k < 1 { + expected_dfirst.push(value); + } else { + expected_dsecond.push(value); + } + } + } + } + assert_eq!(dfirst, expected_dfirst); + assert_eq!(dsecond, expected_dsecond); + + let first_small = first_data + .iter() + .map(|value| value * 0.1) + .collect::>(); + let second_small = second_data + .iter() + .map(|value| value * 0.1) + .collect::>(); + let fd_device = device.clone(); + let fd_second = second_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 1], + &first_small, + move |graph, first| { + let second = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2], &fd_second), + ); + Tensor::cat(vec![first, second], 2) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + let fd_device = device.clone(); + let fd_first = first_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 2], + &second_small, + move |graph, second| { + let first = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 1], &fd_first), + ); + Tensor::cat(vec![first, second], 2) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_log() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.5f32, 1.5, 2.5]); + + let output = input.log(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + for (value, input) in output_values.iter().zip([0.5f32, 1.5, 2.5]) { + assert_close(*value, input.ln()); + } + for (value, input) in dinput.iter().zip([0.5f32, 1.5, 2.5]) { + assert_close(*value, 1.0 / input); + } + + assert_gradient_matches_finite_difference(&device, [3], &[0.5, 1.5, 2.5], |_, x| { + x.log().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_neg() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.5f32, -2.0, 0.5]); + + let output = input.neg(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![-1.5, 2.0, -0.5]); + assert_eq!(dinput, vec![-1.0, -1.0, -1.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.5, -2.0, 0.5], |_, x| { + x.neg().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_exp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[0.0f32, 0.5, -1.0]); + + let output = input.exp(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + for (value, input) in output_values.iter().zip([0.0f32, 0.5, -1.0]) { + assert_close(*value, input.exp()); + } + for (value, input) in dinput.iter().zip([0.0f32, 0.5, -1.0]) { + assert_close(*value, input.exp()); + } + + assert_gradient_matches_finite_difference(&device, [3], &[0.0, 0.5, -1.0], |_, x| { + x.exp().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_log_sum_exp() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = + Tensor::new(&graph, &device, &[[0.0f32, 0.5, 1.0], [1.0, -1.0, 0.0]]); + + let output = input.exp().sum_keepdim(1).log(); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.reshape([2]).sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + // d/dx_j log(sum_k exp(x_k)) = softmax(x)_j + let rows = [[0.0f32, 0.5, 1.0], [1.0, -1.0, 0.0]]; + for (row_index, row) in rows.iter().enumerate() { + let sum_exp = row.iter().map(|value| value.exp()).sum::(); + assert_close(output_values[row_index][0], sum_exp.ln()); + for (column, value) in row.iter().enumerate() { + assert_close(dinput[row_index][column], value.exp() / sum_exp); + } + } + + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[0.0, 0.5, 1.0, 1.0, -1.0, 0.0], + |_, x| x.exp().sum_keepdim(1).log().reshape([2]).sum(), + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_with_backwards() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let y: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32, 5.0, 6.0]); + + let x_target = x.clone(); + let y_target = y.clone(); + let output = x + .add(&y) + .with_backwards([x.parent(), y.parent()], move |grad| { + Ok(vec![ + BackwardTarget::wrt(&x_target, grad.clone().mul_scalar(2.0).to_concrete()), + BackwardTarget::wrt(&y_target, grad.mul_scalar(-3.0).to_concrete()), + ]) + }); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let dy = gradients + .get(&y) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_eq!(output_values, vec![5.0, 7.0, 9.0]); + // the custom rule replaces add's backward, so the gradients are the + // custom 2x/-3x rather than add's 1/1 + assert_eq!(dx, vec![2.0, 2.0, 2.0]); + assert_eq!(dy, vec![-3.0, -3.0, -3.0]); + } +} + +#[tokio::test] +async fn test_backward_with_backwards_missing_parent_errors() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let y: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32, 5.0, 6.0]); + + let x_target = x.clone(); + let output = x + .add(&y) + .with_backwards([x.parent(), y.parent()], move |grad| { + Ok(vec![BackwardTarget::wrt(&x_target, grad)]) + }); + + // The scheduler waits on a gradient from every child edge, so a custom + // rule that skips a live parent must fail loudly instead of silently + // dropping every gradient upstream of it. + let Err(error) = output.sum().backward() else { + panic!("backward succeeded despite an omitted parent gradient"); + }; + assert!( + error.to_string().contains("omitted a gradient"), + "expected missing-parent error, got: {error}", + ); + } +} + +#[tokio::test] +async fn test_graph_drops_after_backward() { + for device in test_devices().await { + let graph = Graph::new(); + let weak = Arc::downgrade(&graph.inner); + + let x: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let w: Tensor<2> = Tensor::new(&graph, &device, &[[0.5f32, -1.0], [1.5, 2.0]]); + let loss = x.mat_mul(&w).sum(1).sum(); + let gradients = loss.backward().unwrap(); + assert!(gradients.get(&x).is_some()); + assert!(gradients.get(&w).is_some()); + + drop(gradients); + drop(loss); + drop(x); + drop(w); + drop(graph); + + assert!( + weak.upgrade().is_none(), + "autograd graph stayed alive after all tensors were dropped", + ); + } +} + +#[test] +fn test_gpu_gradients_can_detach() { + let Ok(device) = Device::gpu_blocking() else { + eprintln!("skipping GPU gradient detach regression test: GPU unavailable"); + return; + }; + + let graph = Graph::new(); + let x: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let w: Tensor<2> = Tensor::new(&graph, &device, &[[0.5f32, -1.0], [1.5, 2.0]]); + let gradients = x + .mat_mul(&w) + .sum(1) + .sum() + .backward() + .unwrap() + .into_detached(); + let dx = gradients.get(&x).expect("missing x gradient"); + let dw = gradients.get(&w).expect("missing w gradient"); + + assert!( + dx.as_gpu() + .expect("expected GPU x gradient") + .resolves_in::<0>(), + "detached x gradient should not retain backward compute graph", + ); + assert!( + dw.as_gpu() + .expect("expected GPU w gradient") + .resolves_in::<0>(), + "detached w gradient should not retain backward compute graph", + ); +} + +/// 8x8 XOR grid plus deterministic LCG weight init for a 2-`hidden`-2 MLP, +/// shared by the XOR training tests. +fn xor_training_data(hidden: usize) -> (Vec, Vec, Vec, Vec) { + let mut features = Vec::with_capacity(128); + let mut labels = Vec::with_capacity(64); + for row in 0..8 { + for column in 0..8 { + let x = -0.875 + 0.25 * row as f32; + let y = -0.875 + 0.25 * column as f32; + features.extend([x, y]); + labels.push(u32::from((x > 0.0) != (y > 0.0))); + } + } + + let mut state = 42u64; + let mut next_uniform = move || { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (state >> 33) as f32 / (1u64 << 32) as f32 - 0.5 + }; + let w1_init: Vec = (0..2 * hidden).map(|_| next_uniform()).collect(); + let w2_init: Vec = (0..hidden * 2).map(|_| next_uniform() * 0.5).collect(); + (features, labels, w1_init, w2_init) +} + +/// End-to-end training: a 2-16-2 MLP learns XOR over an 8x8 grid of 2D +/// points with softmax cross-entropy and full-batch SGD. Exercises the +/// whole tape (matmul, broadcast bias, relu, softmax, gather, log, +/// reduce) plus the detach/update loop across many resolves per device. +#[tokio::test] +async fn test_train_xor_classifier() { + const SAMPLES: usize = 64; + const HIDDEN: usize = 16; + const STEPS: usize = 500; + const LEARNING_RATE: f32 = 1.0; + + let (features, labels, w1_init, w2_init) = xor_training_data(HIDDEN); + + for (device, name) in test_devices().await.into_iter().zip(["cpu", "gpu"]) { + let inputs = RawTensor::from_slice(&device, [SAMPLES, 2], &features); + let targets = RawTensor::from_slice(&device, [SAMPLES], &labels); + + let mut w1 = RawTensor::from_slice(&device, [2, HIDDEN], &w1_init); + let mut b1 = RawTensor::zeros(&device, [HIDDEN]); + let mut w2 = RawTensor::from_slice(&device, [HIDDEN, 2], &w2_init); + let mut b2 = RawTensor::zeros(&device, [2]); + + let mut final_loss = f32::INFINITY; + for step in 0..STEPS { + let graph = Graph::new(); + let x = Tensor::constant_from_raw(&graph, inputs.clone()); + let w1_t = Tensor::from_raw(&graph, w1.clone()); + let b1_t = Tensor::from_raw(&graph, b1.clone()); + let w2_t = Tensor::from_raw(&graph, w2.clone()); + let b2_t = Tensor::from_raw(&graph, b2.clone()); + + let hidden = b1_t.add_::<2, 2>(&x.mat_mul(&w1_t)).relu(); + let logits = b2_t.add_::<2, 2>(&hidden.mat_mul(&w2_t)); + // Numerically stable cross-entropy: log softmax via log-sum-exp + // so a saturated class cannot underflow to log(0). + let shifted = logits.sub_::<2, 2>(&logits.max_keepdim::<1>(1)); + let log_sum_exp = shifted.exp().sum_keepdim(1).log(); + let label_log_probs = shifted.sub_::<2, 2>(&log_sum_exp).gather_last(&targets); + let loss: Tensor<0> = label_log_probs.sum().mul_scalar(-1.0 / SAMPLES as f32); + + let loss_value = flatten(loss.raw().clone()).await[0]; + let gradients = loss.backward().unwrap().into_detached(); + let dw1 = gradients.get(&w1_t).unwrap(); + let db1 = gradients.get(&b1_t).unwrap(); + let dw2 = gradients.get(&w2_t).unwrap(); + let db2 = gradients.get(&b2_t).unwrap(); + + w1 = (w1 - dw1 * LEARNING_RATE).to_concrete(); + b1 = (b1 - db1 * LEARNING_RATE).to_concrete(); + w2 = (w2 - dw2 * LEARNING_RATE).to_concrete(); + b2 = (b2 - db2 * LEARNING_RATE).to_concrete(); + + final_loss = loss_value; + if step % 100 == 0 { + eprintln!("[{name}] step {step}: loss {loss_value:.4}"); + } + } + eprintln!("[{name}] final loss {final_loss:.4}"); + + let graph = Graph::new(); + let x = Tensor::constant_from_raw(&graph, inputs.clone()); + let w1_t = Tensor::constant_from_raw(&graph, w1.clone()); + let b1_t = Tensor::constant_from_raw(&graph, b1.clone()); + let w2_t = Tensor::constant_from_raw(&graph, w2.clone()); + let b2_t = Tensor::constant_from_raw(&graph, b2.clone()); + let hidden = b1_t.add_::<2, 2>(&x.mat_mul(&w1_t)).relu(); + let logits = b2_t.add_::<2, 2>(&hidden.mat_mul(&w2_t)); + let logits = logits.raw().clone().as_slice().await.unwrap().to_vec(); + let correct = logits + .iter() + .zip(&labels) + .filter(|(row, label)| u32::from(row[1] > row[0]) == **label) + .count(); + eprintln!("[{name}] accuracy {correct}/{SAMPLES}"); + + assert!( + final_loss < 0.1, + "training did not converge: final loss {final_loss}", + ); + assert_eq!(correct, SAMPLES, "classifier misclassified training points"); + } +} + +#[tokio::test] +async fn test_autograd_sigmoid() { + for device in test_devices().await { + let graph = Graph::new(); + let inputs = [-2.0f32, -0.5, 0.0, 1.0, 3.0]; + let x: Tensor<1> = Tensor::new(&graph, &device, &inputs); + + let output = x.sigmoid(); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + + let expected = inputs.map(|v| 1.0 / (1.0 + (-v).exp())); + for (value, expected) in values.iter().zip(expected) { + assert_close(*value, expected); + } + + let gradients = output.sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + let expected_grads = inputs.map(|v| { + let sigmoid = 1.0 / (1.0 + (-v).exp()); + sigmoid * (1.0 - sigmoid) + }); + for (value, expected) in dx.iter().zip(expected_grads) { + assert_close(*value, expected); + } + + assert_gradient_matches_finite_difference(&device, [5], &inputs, |_, x| x.sigmoid().sum()) + .await; + } +} + +#[tokio::test] +async fn test_autograd_to_concrete() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 3.0]); + + let output = x.mul_scalar(2.0).to_concrete(); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + assert_eq!(values, vec![2.0, -4.0, 6.0]); + + let gradients = output.sqr().sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dx, &[8.0, -16.0, 24.0]); + + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 3.0]); + let gradients = x.to_concrete().sqr().sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dx, &[2.0, -4.0, 6.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0f32, -2.0, 3.0], |_, x| { + x.to_concrete().sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_index_select_rank_generic() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0]); + let indices = RawTensor::from_slice(&device, [3], &[2u32, 0, 2]); + let selected = x.index_select(0, &indices); + assert_eq!( + selected.raw().clone().as_slice().await.unwrap().to_vec(), + vec![3.0, 1.0, 3.0] + ); + let gradients = selected.sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!(dx, vec![1.0, 0.0, 2.0, 0.0]); + + let graph = Graph::new(); + let data: Vec = (0..12).map(|v| v as f32).collect(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 3, 2], &data); + let indices = RawTensor::from_slice(&device, [4], &[2u32, 0, 1, 0]); + let selected = x.index_select(1, &indices); + assert_eq!(selected.shape(), [2, 4, 2]); + assert_eq!( + flatten(selected.raw().clone()).await, + vec![ + 4.0, 5.0, 0.0, 1.0, 2.0, 3.0, 0.0, 1.0, 10.0, 11.0, 6.0, 7.0, 8.0, 9.0, 6.0, 7.0 + ] + ); + let gradients = selected.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!( + dx, + vec![2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0] + ); + + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[0.5f32, -1.0, 2.0, 3.0, -0.5, 1.5], + |graph, x| { + let indices = RawTensor::from_slice(&x.device(), [2], &[2u32, 0]); + let weights = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [2, 2], &[1.0f32, 2.0, 3.0, 4.0]), + ); + x.index_select(1, &indices) + .mul(&weights) + .flatten_all() + .sum() + }, + ) + .await; + + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[0.5f32, -1.0, 2.0, 3.0, -0.5, 1.5], + |graph, x| { + let indices = RawTensor::from_slice(&x.device(), [3], &[1u32, 1, 0]); + let weights = Tensor::constant_from_raw( + graph, + RawTensor::from_slice( + &x.device(), + [3, 3], + &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], + ), + ); + x.index_select(0, &indices) + .mul(&weights) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_i_indexing() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let row = x.i((1, ..)); + assert_eq!( + row.raw().clone().as_slice().await.unwrap().to_vec(), + vec![4.0, 5.0, 6.0] + ); + let gradients = row.sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!(dx, vec![0.0, 0.0, 0.0, 1.0, 1.0, 1.0]); + + let graph = Graph::new(); + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let column = x.i((0..2, 1)); + assert_eq!( + column.raw().clone().as_slice().await.unwrap().to_vec(), + vec![2.0, 5.0] + ); + let gradients = column.sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!(dx, vec![0.0, 1.0, 0.0, 0.0, 1.0, 0.0]); + + let graph = Graph::new(); + let data: Vec = (0..12).map(|v| v as f32).collect(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &data); + let plane = x.i((.., 1, ..)); + assert_eq!(plane.shape(), [2, 3]); + assert_eq!( + flatten(plane.raw().clone()).await, + vec![3.0, 4.0, 5.0, 9.0, 10.0, 11.0] + ); + let gradients = plane.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_eq!( + dx, + vec![0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0] + ); + + let graph = Graph::new(); + let data: Vec = (0..16).map(|v| v as f32).collect(); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [2, 2, 2, 2], &data); + let cube = x.i((1, .., .., ..)); + assert_eq!(cube.shape(), [2, 2, 2]); + assert_eq!( + flatten(cube.raw().clone()).await, + (8..16).map(|v| v as f32).collect::>() + ); + let gradients = cube.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + let expected: Vec = (0..16).map(|v| if v < 8 { 0.0 } else { 1.0 }).collect(); + assert_eq!(dx, expected); + + assert_gradient_matches_finite_difference( + &device, + [2, 3], + &[0.5f32, -1.0, 2.0, 3.0, -0.5, 1.5], + |graph, x| { + let weights = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [3], &[1.0f32, 2.0, 3.0]), + ); + x.i((1, ..)).mul(&weights).sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_squeeze() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<3> = + Tensor::from_slice(&graph, &device, [2, 1, 2], &[1.0, 2.0, 3.0, 4.0]); + let output: Tensor<2> = input.squeeze::<2>(1); + assert_eq!(output.shape(), [2, 2]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + assert_eq!(values, vec![vec![1.0, 2.0], vec![3.0, 4.0]]); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&input).unwrap()).await; + assert_slice_close(&dx, &[2.0, 4.0, 6.0, 8.0]); + + assert_gradient_matches_finite_difference( + &device, + [2, 1, 2], + &[1.0, 2.0, 3.0, 4.0], + |_, x| x.squeeze::<2>(1).sqr().flatten_all().sum(), + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_unsqueeze() { + for device in test_devices().await { + let graph = Graph::new(); + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; + let input: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 2], &data); + let output: Tensor<4> = input.unsqueeze::<4>(1); + assert_eq!(output.shape(), [2, 1, 2, 2]); + assert_eq!(flatten(output.raw().clone()).await, data.to_vec()); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&input).unwrap()).await; + assert_slice_close(&dx, &data.map(|value| 2.0 * value)); + + assert_gradient_matches_finite_difference(&device, [2, 2, 2], &data, |_, x| { + x.unsqueeze::<4>(3).sqr().flatten_all().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_cat_rank1() { + for device in test_devices().await { + let graph = Graph::new(); + let first: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0]); + let second: Tensor<1> = Tensor::new(&graph, &device, &[3.0f32, 4.0, 5.0]); + let output = Tensor::cat(vec![first.clone(), second.clone()], 0); + assert_eq!(output.shape(), [5]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0, 5.0]); + + let gradients = output.sqr().sum().backward().unwrap(); + assert_slice_close(&flatten(gradients.get(&first).unwrap()).await, &[2.0, 4.0]); + assert_slice_close( + &flatten(gradients.get(&second).unwrap()).await, + &[6.0, 8.0, 10.0], + ); + + assert_gradient_matches_finite_difference(&device, [2], &[1.0, 2.0], |graph, x| { + let other = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [3], &[3.0, 4.0, 5.0]), + ); + Tensor::cat(vec![x, other], 0).sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_cat_rank2() { + for device in test_devices().await { + let graph = Graph::new(); + let first: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 2], &[1.0, 2.0, 3.0, 4.0]); + let second: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 1], &[5.0, 6.0]); + let output = Tensor::cat(vec![first.clone(), second.clone()], 1); + assert_eq!(output.shape(), [2, 3]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + assert_eq!(values, vec![vec![1.0, 2.0, 5.0], vec![3.0, 4.0, 6.0]]); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&first).unwrap()).await, + &[2.0, 4.0, 6.0, 8.0], + ); + assert_slice_close( + &flatten(gradients.get(&second).unwrap()).await, + &[10.0, 12.0], + ); + + assert_gradient_matches_finite_difference(&device, [2, 1], &[5.0, 6.0], |graph, x| { + let other = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [2, 2], &[1.0, 2.0, 3.0, 4.0]), + ); + Tensor::cat(vec![other, x], 1).sqr().flatten_all().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_pad_with_zeros() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let output = input.pad_with_zeros(0, 1, 2); + assert_eq!(output.shape(), [6]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + assert_eq!(values, vec![0.0, 1.0, 2.0, 3.0, 0.0, 0.0]); + + let gradients = output.sqr().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[2.0, 4.0, 6.0], + ); + + let passthrough = input.pad_with_zeros(0, 0, 0); + let gradients = passthrough.sqr().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[2.0, 4.0, 6.0], + ); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |_, x| { + x.pad_with_zeros(0, 1, 2).sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_pad_axis() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 2], &[1.0, 2.0, 3.0, 4.0]); + let output = input.pad_axis(1, 1); + assert_eq!(output.shape(), [2, 4]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + assert_eq!( + values, + vec![vec![0.0, 1.0, 2.0, 0.0], vec![0.0, 3.0, 4.0, 0.0]] + ); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[2.0, 4.0, 6.0, 8.0], + ); + + assert_gradient_matches_finite_difference( + &device, + [2, 2], + &[1.0, 2.0, 3.0, 4.0], + |_, x| x.pad_axis(0, 2).sqr().flatten_all().sum(), + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_sliding_window_view() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0, 5.0]); + let output: Tensor<2> = + input.sliding_window_view::<1, 2>([fusor_types::SlidingWindow::new(0, 3, 1)]); + assert_eq!(output.shape(), [3, 3]); + let values = output.raw().clone().as_slice().await.unwrap().to_vec(); + assert_eq!( + values, + vec![ + vec![1.0, 2.0, 3.0], + vec![2.0, 3.0, 4.0], + vec![3.0, 4.0, 5.0] + ] + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[1.0, 2.0, 3.0, 2.0, 1.0], + ); + + assert_gradient_matches_finite_difference( + &device, + [5], + &[1.0, 2.0, 3.0, 4.0, 5.0], + |_, x| { + x.sliding_window_view::<1, 2>([fusor_types::SlidingWindow::new(0, 3, 1)]) + .sqr() + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_backward_sliding_window_view_strided() { + for device in test_devices().await { + let graph = Graph::new(); + let data: Vec = (1..=10).map(|value| value as f32).collect(); + let input: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 5], &data); + let output: Tensor<3> = + input.sliding_window_view::<1, 3>([fusor_types::SlidingWindow::new(1, 3, 2)]); + assert_eq!(output.shape(), [2, 2, 3]); + assert_eq!( + flatten(output.raw().clone()).await, + vec![1.0, 2.0, 3.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 8.0, 9.0, 10.0] + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert_slice_close( + &flatten(gradients.get(&input).unwrap()).await, + &[1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 2.0, 1.0, 1.0], + ); + + assert_gradient_matches_finite_difference(&device, [2, 5], &data, |_, x| { + x.sliding_window_view::<1, 3>([fusor_types::SlidingWindow::new(1, 3, 2)]) + .sqr() + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_sum_axis() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let summed = x.sum(1); + let forward = summed.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&forward, &[6.0, 15.0]); + + let weight = + Tensor::constant_from_raw(&graph, RawTensor::from_slice(&device, [2], &[1.0, 2.0])); + let loss = summed.mul(&weight).sum(); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0]); + + let graph = Graph::new(); + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let summed = x.sum(0); + let forward = summed.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&forward, &[5.0, 7.0, 9.0]); + + let weight = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [3], &[1.0, 2.0, 3.0]), + ); + let loss = summed.mul(&weight).sum(); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0, 2.0, 3.0, 1.0, 2.0, 3.0]); + + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + assert_gradient_matches_finite_difference(&device, [2, 3], &data, |_, x| { + x.sum(1).sqr().sum() + }) + .await; + assert_gradient_matches_finite_difference(&device, [2, 3], &data, |_, x| { + x.sum(0).sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_sum_high_rank() { + for device in test_devices().await { + let graph = Graph::new(); + + let data: Vec = (1..=8).map(|v| v as f32).collect(); + let x: Tensor<5> = Tensor::from_slice(&graph, &device, [2, 1, 2, 1, 2], &data); + let summed = x.sum(2); + assert_eq!(summed.shape(), [2, 1, 1, 2]); + let forward = flatten(summed.raw().clone()).await; + assert_slice_close(&forward, &[4.0, 6.0, 12.0, 14.0]); + + let loss = summed.sqr().sum(3).sum(2).sum(1).sum(); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[8.0, 12.0, 8.0, 12.0, 24.0, 28.0, 24.0, 28.0]); + + assert_gradient_matches_finite_difference(&device, [2, 1, 2, 1, 2], &data, |_, x| { + x.sum(2).sqr().sum(3).sum(2).sum(1).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_sum_keepdim() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<2> = + Tensor::from_slice(&graph, &device, [2, 3], &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + let summed = x.sum_keepdim(1); + assert_eq!(summed.shape(), [2, 1]); + let forward = flatten(summed.raw().clone()).await; + assert_slice_close(&forward, &[6.0, 15.0]); + + let weight = + Tensor::constant_from_raw(&graph, RawTensor::from_slice(&device, [2, 1], &[1.0, 2.0])); + let loss = summed.mul(&weight).flatten_all().sum(); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0, 1.0, 1.0, 2.0, 2.0, 2.0]); + + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let summed = x.sum_keepdim(0); + assert_eq!(summed.shape(), [1]); + let forward = summed.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&forward, &[6.0]); + let gradients = summed.sqr().sum().backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dx, &[12.0, 12.0, 12.0]); + + let graph = Graph::new(); + let data: Vec = (1..=8).map(|v| v as f32).collect(); + let x: Tensor<5> = Tensor::from_slice(&graph, &device, [2, 1, 2, 1, 2], &data); + let summed = x.sum_keepdim(2); + assert_eq!(summed.shape(), [2, 1, 1, 1, 2]); + let forward = flatten(summed.raw().clone()).await; + assert_slice_close(&forward, &[4.0, 6.0, 12.0, 14.0]); + let gradients = summed.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0; 8]); + + let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + assert_gradient_matches_finite_difference(&device, [2, 3], &data, |_, x| { + x.sum_keepdim(0).sqr().flatten_all().sum() + }) + .await; + assert_gradient_matches_finite_difference(&device, [3], &data[..3], |_, x| { + x.sum_keepdim(0).sqr().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_backward_q_mat_mul_rank1() { + for device in test_devices().await { + let graph = Graph::new(); + let input: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0, 4.0]); + let weight_bytes: Vec = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] + .into_iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let weights = crate::QMatrix::from_raw_bytes( + &device, + [2, 4], + &weight_bytes, + fusor_gguf::GgmlType::F32, + ) + .unwrap(); + + let output = input.q_mat_mul(&weights); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + let gradients = output.sum().backward().unwrap(); + let dinput = gradients + .get(&input) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + + assert_slice_close(&output_values, &[30.0, 70.0]); + assert_slice_close(&dinput, &[6.0, 8.0, 10.0, 12.0]); + + assert_gradient_matches_finite_difference(&device, [4], &[0.5, -1.0, 2.0, 0.25], |_, x| { + x.q_mat_mul(&weights).sum() + }) + .await; + } +} + +fn composite_ramp(graph: &Graph, device: &Device, shape: [usize; R]) -> Tensor { + let elements: usize = shape.iter().product(); + Tensor::constant_from_raw( + graph, + crate::arange(device, 1.0, elements as f32 + 1.0) + .reshape(shape) + .to_concrete(), + ) +} + +#[tokio::test] +async fn test_autograd_conv1d() { + for device in test_devices().await { + let x_data: Vec = (0..8).map(|i| (i as f32 * 0.7).sin()).collect(); + let w_data: Vec = (0..8).map(|i| (i as f32 * 0.3).cos()).collect(); + let b_data = [0.5f32, -1.0]; + + let graph = Graph::new(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 2, 4], &x_data); + let w: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 2], &w_data); + let b: Tensor<1> = Tensor::from_slice(&graph, &device, [2], &b_data); + let output = x.conv(&w, Some(&b), [1], [1]); + + let raw_x = RawTensor::from_slice(&device, [1, 2, 4], &x_data); + let raw_w = RawTensor::from_slice(&device, [2, 2, 2], &w_data); + let raw_b = RawTensor::from_slice(&device, [2], &b_data); + let expected = raw_x.conv(&raw_w, Some(&raw_b), [1], [1]); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let fd_device = device.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference(&device, [1, 2, 4], &x_data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2], &w_fd), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2], &[0.5f32, -1.0]), + ); + let out = x.conv(&w, Some(&b), [1], [1]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference(&device, [2, 2, 2], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 4], &x_fd), + ); + let out = x.conv(&w, None, [1], [1]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference(&device, [2], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 4], &x_fd), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2], &w_fd), + ); + let out = x.conv(&w, Some(&b), [1], [1]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_conv2d_strided() { + for device in test_devices().await { + let x_data: Vec = (0..18).map(|i| (i as f32 * 0.41).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.23).cos()).collect(); + + let graph = Graph::new(); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 3], &x_data); + let w: Tensor<4> = Tensor::from_slice(&graph, &device, [2, 2, 2, 2], &w_data); + let output = x.conv(&w, None, [1, 1], [2, 2]); + + let raw_x = RawTensor::from_slice(&device, [1, 2, 3, 3], &x_data); + let raw_w = RawTensor::from_slice(&device, [2, 2, 2, 2], &w_data); + let expected = raw_x.conv(&raw_w, None, [1, 1], [2, 2]); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let fd_device = device.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 3, 3], + &x_data, + move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2, 2], &w_fd), + ); + let out = x.conv(&w, None, [1, 1], [2, 2]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 2, 2], + &w_data, + move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 3, 3], &x_fd), + ); + let out = x.conv(&w, None, [1, 1], [2, 2]); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_grouped_conv() { + for device in test_devices().await { + let x_data: Vec = (0..16).map(|i| (i as f32 * 0.57).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.31).cos()).collect(); + let b_data = [0.2f32, -0.4, 0.6, -0.8]; + + let graph = Graph::new(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 4, 4], &x_data); + let w: Tensor<3> = Tensor::from_slice(&graph, &device, [4, 2, 2], &w_data); + let b: Tensor<1> = Tensor::from_slice(&graph, &device, [4], &b_data); + let output = x.grouped_conv(&w, Some(&b), [1], [2], 2); + + let raw_x = RawTensor::from_slice(&device, [1, 4, 4], &x_data); + let raw_w = RawTensor::from_slice(&device, [4, 2, 2], &w_data); + let raw_b = RawTensor::from_slice(&device, [4], &b_data); + let expected = raw_x.grouped_conv(&raw_w, Some(&raw_b), [1], [2], 2); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let fd_device = device.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference(&device, [1, 4, 4], &x_data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [4, 2, 2], &w_fd), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [4], &[0.2f32, -0.4, 0.6, -0.8]), + ); + let out = x.grouped_conv(&w, Some(&b), [1], [2], 2); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference(&device, [4, 2, 2], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 4, 4], &x_fd), + ); + let out = x.grouped_conv(&w, None, [1], [2], 2); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_upsample_nearest2d() { + for device in test_devices().await { + let data: Vec = (0..12).map(|i| i as f32 * 0.5).collect(); + let graph = Graph::new(); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 2, 3], &data); + let output = x.upsample_nearest2d(2, 3); + assert_eq!(output.shape(), [1, 2, 4, 9]); + + let raw_x = RawTensor::from_slice(&device, [1, 2, 2, 3], &data); + let expected = raw_x.upsample_nearest2d(2, 3); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + for value in dx { + assert_close(value, 6.0); + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 2, 2, 3], &data, move |graph, x| { + let out = x.upsample_nearest2d(2, 3); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_softmax_slow() { + for device in test_devices().await { + let data = [1.0f32, 2.0, 3.0, -1.0, 0.5, 0.0]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &data); + let output = x.softmax_slow(1); + + let raw_x = RawTensor::from_slice(&device, [2, 3], &data); + let expected = raw_x.softmax_slow::<1>(1); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let last = x.softmax_slow_last_dim(); + assert_slice_close( + &flatten(last.raw().clone()).await, + &flatten(output.raw().clone()).await, + ); + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &data, move |graph, x| { + x.softmax_slow(1) + .mul(&composite_ramp(graph, &fd_device, [2, 3])) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_norm_rms_norm_rank4() { + for device in test_devices().await { + let data: Vec = (0..12).map(|i| (i as f32 * 0.83).sin() * 2.0).collect(); + let w_data = [0.5f32, 1.5, 2.0]; + let b_data = [0.1f32, -0.2, 0.3]; + let eps = 1e-5; + + let graph = Graph::new(); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 2, 3], &data); + let w: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 1, 1, 3], &w_data); + let b: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 1, 1, 3], &b_data); + let output = x.layer_norm(&w, Some(&b), eps, true); + + let raw_x = RawTensor::from_slice(&device, [1, 2, 2, 3], &data); + let raw_w = RawTensor::from_slice(&device, [3], &w_data); + let raw_b = RawTensor::from_slice(&device, [3], &b_data); + let expected = raw_x.layer_norm_last_dim_fused::<3, 1, _, _>(&raw_w, Some(&raw_b), eps); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 2, 2, 3], &data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 1, 1, 3], &[0.5f32, 1.5, 2.0]), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 1, 1, 3], &[0.1f32, -0.2, 0.3]), + ); + let out = x.layer_norm(&w, Some(&b), eps, true); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let rms = x.rms_norm(&w, eps); + let expected_rms = RawTensor::from_slice(&device, [1, 2, 2, 3], &data) + .rms_norm_fused::<1, 3>(&raw_w, None, eps); + assert_slice_close( + &flatten(rms.raw().clone()).await, + &flatten(expected_rms).await, + ); + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 2, 2, 3], &data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 1, 1, 3], &[0.5f32, 1.5, 2.0]), + ); + let out = x.rms_norm(&w, eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rms_norm_residual_fused() { + for device in test_devices().await { + let x_data: Vec = (0..12).map(|i| (i as f32 * 0.29).sin()).collect(); + let r_data: Vec = (0..12).map(|i| (i as f32 * 0.61).cos()).collect(); + let w_data = [0.5f32, 1.5, 2.0]; + let b_data = [0.1f32, -0.2, 0.3]; + let eps = 1e-5; + + let graph = Graph::new(); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &x_data); + let r: Tensor<3> = Tensor::from_slice(&graph, &device, [2, 2, 3], &r_data); + let w: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &w_data); + let b: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &b_data); + + let raw_x = RawTensor::from_slice(&device, [2, 2, 3], &x_data); + let raw_r = RawTensor::from_slice(&device, [2, 2, 3], &r_data); + let raw_w = RawTensor::from_slice(&device, [3], &w_data); + let raw_b = RawTensor::from_slice(&device, [3], &b_data); + + let output = x.rms_norm_residual_fused(&r, &w, Some(&b), eps); + let expected = raw_x.rms_norm_residual_fused::<1, 2, _>(&raw_r, &raw_w, Some(&raw_b), eps); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let no_bias = x.rms_norm_residual_fused(&r, &w, None, eps); + let expected_no_bias = RawTensor::from_slice(&device, [2, 2, 3], &x_data) + .rms_norm_residual_fused::<1, 2, _>(&raw_r, &raw_w, None, eps); + assert_slice_close( + &flatten(no_bias.raw().clone()).await, + &flatten(expected_no_bias).await, + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert!(gradients.get(&x).is_some()); + assert!(gradients.get(&r).is_some()); + assert!(gradients.get(&w).is_some()); + assert!(gradients.get(&b).is_some()); + + let fd_device = device.clone(); + let r_fd = r_data.clone(); + assert_gradient_matches_finite_difference(&device, [2, 2, 3], &x_data, move |graph, x| { + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &r_fd), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.5f32, 1.5, 2.0]), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.1f32, -0.2, 0.3]), + ); + let out = x.rms_norm_residual_fused(&r, &w, Some(&b), eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference(&device, [2, 2, 3], &r_data, move |graph, r| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &x_fd), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.5f32, 1.5, 2.0]), + ); + let out = x.rms_norm_residual_fused(&r, &w, None, eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let r_fd = r_data.clone(); + assert_gradient_matches_finite_difference(&device, [3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &x_fd), + ); + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &r_fd), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.1f32, -0.2, 0.3]), + ); + let out = x.rms_norm_residual_fused(&r, &w, Some(&b), eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let r_fd = r_data.clone(); + assert_gradient_matches_finite_difference(&device, [3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &x_fd), + ); + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3], &r_fd), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3], &[0.5f32, 1.5, 2.0]), + ); + let out = x.rms_norm_residual_fused(&r, &w, Some(&b), eps); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rope_pair_fused() { + for device in test_devices().await { + let q_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin()).collect(); + let k_data: Vec = (0..12).map(|i| (i as f32 * 0.53).cos()).collect(); + let cos_data: Vec = (0..6).map(|i| (i as f32 * 0.7).cos()).collect(); + let sin_data: Vec = (0..6).map(|i| (i as f32 * 0.7).sin()).collect(); + + let graph = Graph::new(); + let q: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 4], &q_data); + let k: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 1, 3, 4], &k_data); + let cos: Tensor<2> = + Tensor::constant_from_raw(&graph, RawTensor::from_slice(&device, [3, 2], &cos_data)); + let sin: Tensor<2> = + Tensor::constant_from_raw(&graph, RawTensor::from_slice(&device, [3, 2], &sin_data)); + let (q_out, k_out) = q.rope_pair_fused(&k, &cos, &sin); + + let raw_q = RawTensor::from_slice(&device, [1, 2, 3, 4], &q_data); + let raw_k = RawTensor::from_slice(&device, [1, 1, 3, 4], &k_data); + let raw_cos = RawTensor::from_slice(&device, [3, 2], &cos_data); + let raw_sin = RawTensor::from_slice(&device, [3, 2], &sin_data); + let (expected_q, expected_k) = raw_q.rope_pair_fused(&raw_k, &raw_cos, &raw_sin); + assert_slice_close( + &flatten(q_out.raw().clone()).await, + &flatten(expected_q).await, + ); + assert_slice_close( + &flatten(k_out.raw().clone()).await, + &flatten(expected_k).await, + ); + + let (normal_q, normal_k) = q.rope_normal_pair_fused(&k, &cos, &sin); + let (expected_nq, expected_nk) = RawTensor::from_slice(&device, [1, 2, 3, 4], &q_data) + .rope_normal_pair_fused(&raw_k, &raw_cos, &raw_sin); + assert_slice_close( + &flatten(normal_q.raw().clone()).await, + &flatten(expected_nq).await, + ); + assert_slice_close( + &flatten(normal_k.raw().clone()).await, + &flatten(expected_nk).await, + ); + + let fd_device = device.clone(); + let k_fd = k_data.clone(); + let cos_fd = cos_data.clone(); + let sin_fd = sin_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 3, 4], + &q_data, + move |graph, q| { + let k = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 1, 3, 4], &k_fd), + ); + let cos = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3, 2], &cos_fd), + ); + let sin = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3, 2], &sin_fd), + ); + let (q_out, k_out) = q.rope_pair_fused(&k, &cos, &sin); + let q_loss = q_out + .mul(&composite_ramp(graph, &fd_device, q_out.shape())) + .flatten_all() + .sum(); + q_loss.add(&k_out.flatten_all().sum()) + }, + ) + .await; + + let fd_device = device.clone(); + let q_fd = q_data.clone(); + let cos_fd = cos_data.clone(); + let sin_fd = sin_data.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 1, 3, 4], + &k_data, + move |graph, k| { + let q = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 3, 4], &q_fd), + ); + let cos = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3, 2], &cos_fd), + ); + let sin = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [3, 2], &sin_fd), + ); + let (_, k_out) = q.rope_normal_pair_fused(&k, &cos, &sin); + k_out + .mul(&composite_ramp(graph, &fd_device, k_out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rope_cache_forward() { + for device in test_devices().await { + let q_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin()).collect(); + let k_data: Vec = (0..24).map(|i| (i as f32 * 0.53).cos()).collect(); + let cache = crate::RopeCache::new(4, 8, 10000.0, &device).unwrap(); + + let graph = Graph::new(); + let q: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 4], &q_data); + let k: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 4], &k_data); + + let raw_q = RawTensor::from_slice(&device, [1, 2, 3, 4], &q_data); + let raw_k = RawTensor::from_slice(&device, [1, 2, 3, 4], &k_data); + + let (q_out, k_out) = q.rope_cache_forward(&k, &cache, 2); + let (expected_q, expected_k) = cache.forward(&raw_q, &raw_k, 2); + assert_slice_close( + &flatten(q_out.raw().clone()).await, + &flatten(expected_q).await, + ); + assert_slice_close( + &flatten(k_out.raw().clone()).await, + &flatten(expected_k).await, + ); + + let (qi_out, ki_out) = q.rope_cache_forward_interleaved(&k, &cache, 2); + let (expected_qi, expected_ki) = cache.forward_interleaved( + &RawTensor::from_slice(&device, [1, 2, 3, 4], &q_data), + &RawTensor::from_slice(&device, [1, 2, 3, 4], &k_data), + 2, + ); + assert_slice_close( + &flatten(qi_out.raw().clone()).await, + &flatten(expected_qi).await, + ); + assert_slice_close( + &flatten(ki_out.raw().clone()).await, + &flatten(expected_ki).await, + ); + + let fd_device = device.clone(); + let k_fd = k_data.clone(); + let fd_cache = cache.clone(); + assert_gradient_matches_finite_difference( + &device, + [1, 2, 3, 4], + &q_data, + move |graph, q| { + let k = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 2, 3, 4], &k_fd), + ); + let (q_out, _) = q.rope_cache_forward(&k, &fd_cache, 2); + q_out + .mul(&composite_ramp(graph, &fd_device, q_out.shape())) + .flatten_all() + .sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_ones() { + for device in test_devices().await { + let graph = Graph::new(); + + let ones: Tensor<1> = Tensor::ones(&graph, &device, [3]); + let forward = ones.raw().clone().as_slice().await.unwrap().to_vec(); + assert_eq!(forward, vec![1.0, 1.0, 1.0]); + + let x: Tensor<1> = Tensor::new(&graph, &device, &[2.0f32, -3.0, 4.0]); + let loss = x.mul(&ones).sum(); + assert_close(loss.raw().to_scalar().await.unwrap(), 3.0); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dx, &[1.0, 1.0, 1.0]); + let dones = gradients + .get(&ones) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dones, &[2.0, -3.0, 4.0]); + } +} + +#[tokio::test] +async fn test_autograd_ones_like() { + for device in test_devices().await { + let graph = Graph::new(); + + let x: Tensor<2> = Tensor::new(&graph, &device, &[[1.0f32, 2.0], [3.0, 4.0]]); + let ones = x.ones_like(); + assert_eq!(ones.shape(), [2, 2]); + assert_eq!(flatten(ones.raw().clone()).await, vec![1.0, 1.0, 1.0, 1.0]); + + let loss = x.mul(&ones).flatten_all().sum(); + assert_close(loss.raw().to_scalar().await.unwrap(), 10.0); + let gradients = loss.backward().unwrap(); + let dx = flatten(gradients.get(&x).unwrap()).await; + assert_slice_close(&dx, &[1.0, 1.0, 1.0, 1.0]); + let dones = flatten(gradients.get(&ones).unwrap()).await; + assert_slice_close(&dones, &[1.0, 2.0, 3.0, 4.0]); + } +} + +#[tokio::test] +async fn test_backward_mat_mul_rank4() { + for device in test_devices().await { + let graph = Graph::new(); + let lhs_data = (1..=24).map(|n| n as f32).collect::>(); + let rhs_data = (1..=24).map(|n| n as f32).collect::>(); + let lhs: Tensor<4> = Tensor::from_slice(&graph, &device, [2, 2, 2, 3], &lhs_data); + let rhs: Tensor<4> = Tensor::from_slice(&graph, &device, [2, 2, 3, 2], &rhs_data); + + let output = lhs.mat_mul(&rhs); + let output_values = output.raw().clone().as_slice().await.unwrap(); + let gradients = output.flatten_all().sum().backward().unwrap(); + let dlhs = flatten(gradients.get(&lhs).unwrap()).await; + let drhs = flatten(gradients.get(&rhs).unwrap()).await; + + assert_eq!(output_values.shape(), &[2, 2, 2, 2]); + assert_close(output_values[[0, 0, 0, 0]], 22.0); + assert_close(output_values[[1, 1, 1, 1]], 1522.0); + + // with an all-ones seed, dlhs[b, i, k] = sum_j rhs[b, k, j] and + // drhs[b, k, j] = sum_i lhs[b, i, k] + for batch in 0..4 { + for i in 0..2 { + for k in 0..3 { + let expected = (0..2).map(|j| rhs_data[batch * 6 + k * 2 + j]).sum::(); + assert_close(dlhs[batch * 6 + i * 3 + k], expected); + } + } + for k in 0..3 { + for j in 0..2 { + let expected = (0..2).map(|i| lhs_data[batch * 6 + i * 3 + k]).sum::(); + assert_close(drhs[batch * 6 + k * 2 + j], expected); + } + } + } + + let lhs_small = lhs_data + .iter() + .map(|value| value * 0.05) + .collect::>(); + let rhs_small = rhs_data + .iter() + .map(|value| value * 0.03) + .collect::>(); + let fd_device = device.clone(); + let fd_rhs = rhs_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 2, 3], + &lhs_small, + move |graph, lhs| { + let rhs = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 3, 2], &fd_rhs), + ); + lhs.mat_mul(&rhs).sqr().flatten_all().sum() + }, + ) + .await; + let fd_device = device.clone(); + let fd_lhs = lhs_small.clone(); + assert_gradient_matches_finite_difference( + &device, + [2, 2, 3, 2], + &rhs_small, + move |graph, rhs| { + let lhs = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 2, 2, 3], &fd_lhs), + ); + lhs.mat_mul(&rhs).sqr().flatten_all().sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_std_ops_add_sub() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let y: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32, -5.0, 8.0]); + + for add in [ + &x + &y, + x.clone() + y.clone(), + &x + y.clone(), + x.clone() + &y, + ] { + let values = add.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[5.0, -3.0, 11.0]); + } + for sub in [ + &x - &y, + x.clone() - y.clone(), + &x - y.clone(), + x.clone() - &y, + ] { + let values = sub.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[-3.0, 7.0, -5.0]); + } + + let loss = ((&x + &y) * (&x - &y)).sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let dy = gradients + .get(&y) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dx, &[2.0, 4.0, 6.0]); + assert_slice_close(&dy, &[-8.0, 10.0, -16.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |graph, x| { + let y = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [3], &[4.0, -5.0, 8.0]), + ); + ((&x + &y) * (&x - &y)).sum() + }) + .await; + assert_gradient_matches_finite_difference(&device, [3], &[4.0, -5.0, 8.0], |graph, y| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&y.device(), [3], &[1.0, 2.0, 3.0]), + ); + ((&x + &y) * (&x - &y)).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_std_ops_mul_div() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + let y: Tensor<1> = Tensor::new(&graph, &device, &[4.0f32, -5.0, 8.0]); + + for mul in [ + &x * &y, + x.clone() * y.clone(), + &x * y.clone(), + x.clone() * &y, + ] { + let values = mul.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[4.0, -10.0, 24.0]); + } + for div in [ + &x / &y, + x.clone() / y.clone(), + &x / y.clone(), + x.clone() / &y, + ] { + let values = div.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[0.25, -0.4, 0.375]); + } + + let loss = ((&x * &y) + (&x / &y)).sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + let dy = gradients + .get(&y) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dx, &[4.25, -5.2, 8.125]); + assert_slice_close(&dy, &[0.9375, 1.92, 2.953125]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |graph, x| { + let y = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&x.device(), [3], &[4.0, -5.0, 8.0]), + ); + ((&x * &y) + (&x / &y)).sum() + }) + .await; + assert_gradient_matches_finite_difference(&device, [3], &[4.0, -5.0, 8.0], |graph, y| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&y.device(), [3], &[1.0, 2.0, 3.0]), + ); + ((&x * &y) + (&x / &y)).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_std_ops_neg() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, -2.0, 3.0]); + + for neg in [-&x, -x.clone()] { + let values = neg.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[-1.0, 2.0, -3.0]); + } + + let loss = ((-&x) * &x).sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dx, &[-2.0, 4.0, -6.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, -2.0, 3.0], |_, x| { + ((-&x) * &x).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_std_ops_scalar() { + for device in test_devices().await { + let graph = Graph::new(); + let x: Tensor<1> = Tensor::new(&graph, &device, &[1.0f32, 2.0, 3.0]); + + for mul in [&x * 2.5, x.clone() * 2.5] { + let values = mul.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[2.5, 5.0, 7.5]); + } + for add in [&x + 1.5, x.clone() + 1.5] { + let values = add.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[2.5, 3.5, 4.5]); + } + for sub in [&x - 0.5, x.clone() - 0.5] { + let values = sub.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[0.5, 1.5, 2.5]); + } + for div in [&x / 2.0, x.clone() / 2.0] { + let values = div.raw().clone().as_slice().await.unwrap().to_vec(); + assert_slice_close(&values, &[0.5, 1.0, 1.5]); + } + + let loss = ((&x * 3.0) + 2.0).sum(); + let gradients = loss.backward().unwrap(); + let dx = gradients + .get(&x) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + assert_slice_close(&dx, &[3.0, 3.0, 3.0]); + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |_, x| { + ((&x * 3.0) + 2.0).sum() + }) + .await; + + assert_gradient_matches_finite_difference(&device, [3], &[1.0, 2.0, 3.0], |_, x| { + ((&x - 1.5) / 4.0).sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_norm_rank_generic() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_rows = [[1.0f32, 2.0, 4.0], [-1.0, 0.5, 3.0]]; + let x_data = [1.0f32, 2.0, 4.0, -1.0, 0.5, 3.0]; + let w_data = [0.5f32, 1.0, 1.5]; + let b_data = [0.1f32, -0.2, 0.3]; + for remove_mean in [true, false] { + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let w: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + let b: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &b_data); + + let output = x.layer_norm(&w, Some(&b), eps, remove_mean); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + for (row_index, row) in x_rows.iter().enumerate() { + let mean = if remove_mean { + row.iter().sum::() / 3.0 + } else { + 0.0 + }; + let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::() / 3.0; + let std = (var + eps).sqrt(); + for column in 0..3 { + let expected = (row[column] - mean) / std * w_data[column] + b_data[column]; + assert_close(output_values[row_index][column], expected); + } + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &x_data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &b_data), + ); + x.layer_norm(&w, Some(&b), eps, remove_mean) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &b_data), + ); + x.layer_norm(&w, Some(&b), eps, remove_mean) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + x.layer_norm(&w, Some(&b), eps, remove_mean) + .flatten_all() + .sum() + }) + .await; + } + } +} + +#[tokio::test] +async fn test_autograd_rms_norm_rank_generic_weight() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_rows = [[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]]; + let x_data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; + let w_data = [0.5f32, 1.0, 1.5]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let w: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + + let output = x.rms_norm(&w, eps); + let output_values = output.raw().clone().as_slice().await.unwrap().to_vec(); + for (row_index, row) in x_rows.iter().enumerate() { + let mean_sq = row.iter().map(|v| v * v).sum::() / 3.0; + let rms = (mean_sq + eps).sqrt(); + for column in 0..3 { + assert_close( + output_values[row_index][column], + row[column] / rms * w_data[column], + ); + } + } + + let gradients = output.flatten_all().sum().backward().unwrap(); + let dw = gradients + .get(&w) + .unwrap() + .as_slice() + .await + .unwrap() + .to_vec(); + for column in 0..3 { + let mut expected = 0.0f32; + for row in x_rows.iter() { + let mean_sq = row.iter().map(|v| v * v).sum::() / 3.0; + expected += row[column] / (mean_sq + eps).sqrt(); + } + assert_close(dw[0][column], expected); + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &x_data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + x.rms_norm(&w, eps).flatten_all().sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + x.rms_norm(&w, eps).flatten_all().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rms_norm_fused_weight_rank() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_rows = [[1.0f32, -2.0, 3.0], [0.5, 4.0, -1.5]]; + let x_data = [1.0f32, -2.0, 3.0, 0.5, 4.0, -1.5]; + let w_data = [0.5f32, 1.0, 1.5]; + let b_data = [0.1f32, -0.2, 0.3]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let w1: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &w_data); + let b1: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &b_data); + let w2: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + + let biased = x.rms_norm_fused::<1, 1>(&w1, Some(&b1), eps); + let biased_values = biased.raw().clone().as_slice().await.unwrap().to_vec(); + let no_bias = x.rms_norm_fused_no_bias::<2, 1>(&w2, eps); + let no_bias_values = no_bias.raw().clone().as_slice().await.unwrap().to_vec(); + for (row_index, row) in x_rows.iter().enumerate() { + let mean_sq = row.iter().map(|v| v * v).sum::() / 3.0; + let rms = (mean_sq + eps).sqrt(); + for column in 0..3 { + let scaled = row[column] / rms * w_data[column]; + assert_close(biased_values[row_index][column], scaled + b_data[column]); + assert_close(no_bias_values[row_index][column], scaled); + } + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &x_data, move |graph, x| { + let w = + Tensor::constant_from_raw(graph, RawTensor::from_slice(&fd_device, [3], &w_data)); + let b = + Tensor::constant_from_raw(graph, RawTensor::from_slice(&fd_device, [3], &b_data)); + x.rms_norm_fused::<1, 1>(&w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + x.rms_norm_fused_no_bias::<2, 1>(&w, eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let w = + Tensor::constant_from_raw(graph, RawTensor::from_slice(&fd_device, [3], &w_data)); + x.rms_norm_fused::<1, 1>(&w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_rms_norm_residual_fused_weight_rank() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_data = [1.0f32, -2.0, 3.0, 0.5, 4.0, -1.5]; + let r_data = [0.5f32, 1.5, -1.0, 2.0, -0.5, 1.0]; + let w_data = [0.5f32, 1.0, 1.5]; + let b_data = [0.1f32, -0.2, 0.3]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let r: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &r_data); + let w: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + let b: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &b_data); + + let biased = x.rms_norm_residual_fused::<2, 1>(&r, &w, Some(&b), eps); + let biased_values = biased.raw().clone().as_slice().await.unwrap().to_vec(); + let no_bias = x.rms_norm_residual_fused::<2, 1>(&r, &w, None, eps); + let no_bias_values = no_bias.raw().clone().as_slice().await.unwrap().to_vec(); + for row_index in 0..2 { + let combined: Vec = (0..3) + .map(|column| x_data[row_index * 3 + column] + r_data[row_index * 3 + column]) + .collect(); + let mean_sq = combined.iter().map(|v| v * v).sum::() / 3.0; + let rms = (mean_sq + eps).sqrt(); + for column in 0..3 { + let scaled = combined[column] / rms * w_data[column]; + assert_close(biased_values[row_index][column], scaled + b_data[column]); + assert_close(no_bias_values[row_index][column], scaled); + } + } + + // Input/residual gradients are covered by + // test_autograd_rms_norm_residual_fused; only the rank-2 weight/bias + // path is new here. + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &r_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &b_data), + ); + x.rms_norm_residual_fused::<2, 1>(&r, &w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let r = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &r_data), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + x.rms_norm_residual_fused::<2, 1>(&r, &w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_norm_last_dim_fused_weight_rank() { + for device in test_devices().await { + let eps = 1e-5f32; + let x_rows = [[1.0f32, 2.0, 4.0], [-1.0, 0.5, 3.0]]; + let x_data = [1.0f32, 2.0, 4.0, -1.0, 0.5, 3.0]; + let w_data = [0.5f32, 1.0, 1.5]; + let b_data = [0.1f32, -0.2, 0.3]; + let graph = Graph::new(); + let x: Tensor<2> = Tensor::from_slice(&graph, &device, [2, 3], &x_data); + let w1: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &w_data); + let b1: Tensor<1> = Tensor::from_slice(&graph, &device, [3], &b_data); + let w2: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &w_data); + let b2: Tensor<2> = Tensor::from_slice(&graph, &device, [1, 3], &b_data); + + let rank1 = x.layer_norm_last_dim_fused::<1, 1>(&w1, Some(&b1), eps); + let rank1_values = rank1.raw().clone().as_slice().await.unwrap().to_vec(); + let rank2 = x.layer_norm_last_dim_fused::<1, 2>(&w2, Some(&b2), eps); + let rank2_values = rank2.raw().clone().as_slice().await.unwrap().to_vec(); + for (row_index, row) in x_rows.iter().enumerate() { + let mean = row.iter().sum::() / 3.0; + let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::() / 3.0; + let std = (var + eps).sqrt(); + for column in 0..3 { + let expected = (row[column] - mean) / std * w_data[column] + b_data[column]; + assert_close(rank1_values[row_index][column], expected); + assert_close(rank2_values[row_index][column], expected); + } + } + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [2, 3], &x_data, move |graph, x| { + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + let b = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &b_data), + ); + x.layer_norm_last_dim_fused::<1, 2>(&w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &w_data, move |graph, w| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + x.layer_norm_last_dim_fused::<1, 2>(&w, None, eps) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + assert_gradient_matches_finite_difference(&device, [1, 3], &b_data, move |graph, b| { + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3], &x_data), + ); + let w = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3], &w_data), + ); + x.layer_norm_last_dim_fused::<1, 2>(&w, Some(&b), eps) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_linear_forward_matches_inference() { + for device in test_devices().await { + let weight_data = [ + 0.5f32, -1.0, 0.25, 2.0, 1.5, -0.75, 0.1, 0.4, -0.2, 0.9, -1.3, 0.6, + ]; + let bias_data = [0.3f32, -0.6, 1.1]; + let input_data = [ + 1.0f32, -2.0, 0.5, 0.25, 0.75, 1.5, -0.5, 2.0, -1.25, 0.4, 0.8, -0.3, 0.15, -0.9, 1.2, + 0.7, + ]; + let weight_bytes: Vec = weight_data + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let qweight = |device: &Device| { + crate::QMatrix::from_raw_bytes(device, [3, 4], &weight_bytes, fusor_gguf::GgmlType::F32) + .unwrap() + }; + let raw_bias = RawTensor::from_slice(&device, [3], &bias_data); + let inference = crate::layers::Linear::new(qweight(&device), Some(raw_bias)); + let inference_no_bias = crate::layers::Linear::::new(qweight(&device), None); + + let graph = Graph::new(); + let layer = layers::Linear::new( + Tensor::from_slice(&graph, &device, [3, 4], &weight_data), + Some(Tensor::from_slice(&graph, &device, [3], &bias_data)), + ); + assert_eq!(layer.in_features(), 4); + assert_eq!(layer.out_features(), 3); + + let raw_input = RawTensor::from_slice(&device, [2, 2, 4], &input_data); + let input = Tensor::constant_from_raw(&graph, raw_input.clone()); + let output = layer.forward(&input); + assert_eq!(output.shape(), [2, 2, 3]); + let expected = flatten(inference.forward(&raw_input)).await; + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let raw_input_2d = RawTensor::from_slice(&device, [4, 4], &input_data); + let input_2d = Tensor::constant_from_raw(&graph, raw_input_2d.clone()); + let output_2d = layer.forward(&input_2d); + assert_eq!(output_2d.shape(), [4, 3]); + let expected_2d = flatten(inference.forward(&raw_input_2d)).await; + assert_slice_close(&flatten(output_2d.raw().clone()).await, &expected_2d); + + let layer_no_bias = layers::Linear::new( + Tensor::from_slice(&graph, &device, [3, 4], &weight_data), + None, + ); + let output_no_bias = layer_no_bias.forward(&input); + let expected_no_bias = flatten(inference_no_bias.forward(&raw_input)).await; + assert_slice_close( + &flatten(output_no_bias.raw().clone()).await, + &expected_no_bias, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_linear_weight_gradient() { + for device in test_devices().await { + let weight_data = [ + 0.5f32, -1.0, 0.25, 2.0, 1.5, -0.75, 0.1, 0.4, -0.2, 0.9, -1.3, 0.6, + ]; + let bias_data = [0.3f32, -0.6, 1.1]; + let input_data = [ + 1.0f32, -2.0, 0.5, 0.25, 0.75, 1.5, -0.5, 2.0, -1.25, 0.4, 0.8, -0.3, 0.15, -0.9, 1.2, + 0.7, + ]; + assert_gradient_matches_finite_difference( + &device, + [3, 4], + &weight_data, + |graph, weight| { + let bias = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&device, [3], &bias_data), + ); + let layer = layers::Linear::new(weight, Some(bias)); + let input = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&device, [2, 2, 4], &input_data), + ); + layer.forward(&input).sqr().flatten_all().sum() + }, + ) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_linear_bias_gradient() { + for device in test_devices().await { + let weight_data = [ + 0.5f32, -1.0, 0.25, 2.0, 1.5, -0.75, 0.1, 0.4, -0.2, 0.9, -1.3, 0.6, + ]; + let bias_data = [0.3f32, -0.6, 1.1]; + let input_data = [ + 1.0f32, -2.0, 0.5, 0.25, 0.75, 1.5, -0.5, 2.0, -1.25, 0.4, 0.8, -0.3, 0.15, -0.9, 1.2, + 0.7, + ]; + assert_gradient_matches_finite_difference(&device, [3], &bias_data, |graph, bias| { + let weight = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&device, [3, 4], &weight_data), + ); + let layer = layers::Linear::new(weight, Some(bias)); + let input = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&device, [2, 2, 4], &input_data), + ); + layer.forward(&input).sqr().flatten_all().sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_linear_from_inference() { + for device in test_devices().await { + let weight_data = [ + 0.5f32, -1.0, 0.25, 2.0, 1.5, -0.75, 0.1, 0.4, -0.2, 0.9, -1.3, 0.6, + ]; + let bias_data = [0.3f32, -0.6, 1.1]; + let input_data = [ + 1.0f32, -2.0, 0.5, 0.25, 0.75, 1.5, -0.5, 2.0, -1.25, 0.4, 0.8, -0.3, 0.15, -0.9, 1.2, + 0.7, + ]; + let weight_bytes: Vec = weight_data + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + let qweight = crate::QMatrix::from_raw_bytes( + &device, + [3, 4], + &weight_bytes, + fusor_gguf::GgmlType::F32, + ) + .unwrap(); + let raw_bias = RawTensor::from_slice(&device, [3], &bias_data); + let inference = crate::layers::Linear::new(qweight, Some(raw_bias)); + + let graph = Graph::new(); + let layer = layers::Linear::from_inference(&graph, &inference); + assert_eq!(layer.in_features(), 4); + assert_eq!(layer.out_features(), 3); + + let raw_input = RawTensor::from_slice(&device, [2, 2, 4], &input_data); + let input = Tensor::constant_from_raw(&graph, raw_input.clone()); + let output = layer.forward(&input); + let expected = flatten(inference.forward(&raw_input)).await; + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let gradients = output.sqr().flatten_all().sum().backward().unwrap(); + let dweight = gradients.get(layer.weight()).unwrap(); + assert_eq!(dweight.shape(), [3, 4]); + let dbias = gradients.get(layer.bias().unwrap()).unwrap(); + let mut expected_dbias = [0.0f32; 3]; + for row in expected.chunks(3) { + for (acc, value) in expected_dbias.iter_mut().zip(row) { + *acc += 2.0 * value; + } + } + assert_slice_close(&flatten(dbias).await, &expected_dbias); + } +} + +#[tokio::test] +async fn test_autograd_layer_embedding_forward_parity() { + let weights: Vec = (0..12).map(|index| (index as f32 * 0.7).sin()).collect(); + for device in test_devices().await { + let table = RawTensor::from_slice(&device, [4, 3], &weights); + let raw_layer = crate::layers::Embedding::new_from_tensor(table.clone()); + let graph = Graph::new(); + let layer = layers::Embedding::new_from_tensor(graph.leaf(table)); + assert_eq!(layer.num_embeddings(), 4); + assert_eq!(layer.embedding_dim(), 3); + + let indices: RawTensor<2, u32> = RawTensor::from_slice(&device, [2, 2], &[0, 2, 1, 3]); + let expected: RawTensor<3, f32> = raw_layer.forward(&indices); + let output = layer.forward(&indices); + assert_eq!(output.shape(), [2, 2, 3]); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let flat_indices: RawTensor<1, u32> = RawTensor::from_slice(&device, [3], &[2, 0, 2]); + let expected_flat: RawTensor<2, f32> = raw_layer.forward(&flat_indices); + let output_flat = layer.forward(&flat_indices); + assert_eq!(output_flat.shape(), [3, 3]); + assert_slice_close( + &flatten(output_flat.raw().clone()).await, + &flatten(expected_flat).await, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_embedding_weight_gradient() { + let weights: Vec = (0..12).map(|index| (index as f32 * 0.7).sin()).collect(); + for device in test_devices().await { + let indices: RawTensor<2, u32> = RawTensor::from_slice(&device, [2, 2], &[0, 2, 1, 2]); + assert_gradient_matches_finite_difference(&device, [4, 3], &weights, |_graph, table| { + layers::Embedding::new_from_tensor(table) + .forward(&indices) + .sqr() + .flatten_all() + .sum() + }) + .await; + + let flat_indices: RawTensor<1, u32> = RawTensor::from_slice(&device, [3], &[3, 1, 3]); + assert_gradient_matches_finite_difference(&device, [4, 3], &weights, |_graph, table| { + layers::Embedding::new_from_tensor(table) + .forward(&flat_indices) + .sqr() + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_embedding_from_inference() { + let weights: Vec = (0..12).map(|index| index as f32 * 0.25 - 1.0).collect(); + let bytes: Vec = weights + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect(); + for device in test_devices().await { + let quantized = + crate::QMatrix::from_raw_bytes(&device, [4usize, 3], &bytes, fusor_gguf::GgmlType::F32) + .unwrap(); + let quantized_layer = crate::layers::Embedding::::new(quantized); + let dense_layer = crate::layers::Embedding::::new_from_tensor(RawTensor::from_slice( + &device, + [4, 3], + &weights, + )); + for raw_layer in [&quantized_layer, &dense_layer] { + let indices: RawTensor<2, u32> = RawTensor::from_slice(&device, [2, 2], &[3, 0, 2, 2]); + let expected: RawTensor<3, f32> = raw_layer.forward(&indices); + + let graph = Graph::new(); + let layer = layers::Embedding::from_inference(&graph, raw_layer); + assert_eq!(layer.num_embeddings(), 4); + assert_eq!(layer.embedding_dim(), 3); + assert_slice_close(&flatten(layer.embeddings().raw().clone()).await, &weights); + let output = layer.forward(&indices); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let loss = layer.forward(&indices).sqr().flatten_all().sum(); + let gradients = loss.backward().unwrap(); + assert!(gradients.get(layer.embeddings()).is_some()); + } + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_forward_matches_inference() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, 0.4]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + + let inference = crate::layers::LayerNorm::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let raw_input = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let raw_input_2d = RawTensor::from_slice(&device, [6, 4], &input_data); + let expected = flatten(inference.forward(&raw_input)).await; + let expected_2d = flatten(inference.forward(&raw_input_2d)).await; + + let graph = Graph::new(); + let layer = layers::LayerNorm::new( + Tensor::from_slice(&graph, &device, [4], &weight_data), + Some(Tensor::from_slice(&graph, &device, [4], &bias_data)), + 1e-5, + ); + let input = Tensor::from_slice(&graph, &device, [2, 3, 4], &input_data); + let input_2d = Tensor::from_slice(&graph, &device, [6, 4], &input_data); + assert_slice_close(&flatten(layer.forward(&input).into_raw()).await, &expected); + assert_slice_close(&flatten(layer.forward(&input).into_raw()).await, &expected); + assert_slice_close( + &flatten(layer.forward(&input_2d).into_raw()).await, + &expected_2d, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_nd_forward_matches_inference() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, 0.4]; + let axis_weight_data = [0.5f32, 1.5, -1.0]; + let axis_bias_data = [0.1f32, -0.2, 0.3]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + let raw_input = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + + let inference = crate::layers::LayerNormNd::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let expected = flatten(inference.forward::<3, 2, _>(&raw_input)).await; + let raw_input_2d = RawTensor::from_slice(&device, [6, 4], &input_data); + let expected_2d = flatten(inference.forward(&raw_input_2d)).await; + + let inference_axis = crate::layers::LayerNormNd::new_over_axis( + RawTensor::from_slice(&device, [3], &axis_weight_data), + Some(RawTensor::from_slice(&device, [3], &axis_bias_data)), + 1, + 1e-5, + ); + let expected_axis = flatten(inference_axis.forward::<3, 2, _>(&raw_input)).await; + + let graph = Graph::new(); + let input = Tensor::from_slice(&graph, &device, [2, 3, 4], &input_data); + let input_2d = Tensor::from_slice(&graph, &device, [6, 4], &input_data); + let layer = layers::LayerNormNd::new( + Tensor::from_slice(&graph, &device, [4], &weight_data), + Some(Tensor::from_slice(&graph, &device, [4], &bias_data)), + 1e-5, + ); + assert_slice_close( + &flatten(layer.forward::<3, 2>(&input).into_raw()).await, + &expected, + ); + assert_slice_close( + &flatten(layer.forward_fused(&input).into_raw()).await, + &expected, + ); + assert_slice_close( + &flatten(layer.forward(&input_2d).into_raw()).await, + &expected_2d, + ); + + let layer_axis = layers::LayerNormNd::new_over_axis( + Tensor::from_slice(&graph, &device, [3], &axis_weight_data), + Some(Tensor::from_slice(&graph, &device, [3], &axis_bias_data)), + 1, + 1e-5, + ); + assert_slice_close( + &flatten(layer_axis.forward::<3, 2>(&input).into_raw()).await, + &expected_axis, + ); + assert_slice_close( + &flatten(layer_axis.forward_fused(&input).into_raw()).await, + &expected_axis, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_parameter_gradients() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, 0.4]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let bias = Tensor::from_slice(graph, &device, [4], &bias_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNorm::new(weight, Some(bias), 1e-5) + .forward(&input) + .flatten_all() + .sum() + }) + .await; + + assert_gradient_matches_finite_difference(&device, [4], &bias_data, |graph, bias| { + let weight = Tensor::from_slice(graph, &device, [4], &weight_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNorm::new(weight, Some(bias), 1e-5) + .forward(&input) + .flatten_all() + .sum() + }) + .await; + + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let bias = Tensor::from_slice(graph, &device, [4], &bias_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNorm::new(weight, Some(bias), 1e-5) + .forward(&input) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_nd_parameter_gradients() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0]; + let bias_data = [0.1f32, -0.2, 0.3]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + + assert_gradient_matches_finite_difference(&device, [3], &weight_data, |graph, weight| { + let bias = Tensor::from_slice(graph, &device, [3], &bias_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNormNd::new_over_axis(weight, Some(bias), 1, 1e-5) + .forward::<3, 2>(&input) + .flatten_all() + .sum() + }) + .await; + + assert_gradient_matches_finite_difference(&device, [3], &bias_data, |graph, bias| { + let weight = Tensor::from_slice(graph, &device, [3], &weight_data); + let input = Tensor::from_slice(graph, &device, [2, 3, 4], &input_data); + layers::LayerNormNd::new_over_axis(weight, Some(bias), 1, 1e-5) + .forward::<3, 2>(&input) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_layer_norm_from_inference_roundtrip() { + for device in test_devices().await { + let weight_data = [0.5f32, 1.5, -1.0, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, 0.4]; + let axis_weight_data = [0.5f32, 1.5, -1.0]; + let input_data: Vec = (0..24).map(|index| (index as f32 * 0.7).sin()).collect(); + let raw_input = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + + let inference = crate::layers::LayerNorm::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let expected = flatten(inference.forward(&raw_input)).await; + + let graph = Graph::new(); + let input = Tensor::constant_from_raw(&graph, raw_input.clone()); + let layer = layers::LayerNorm::from_inference(&graph, &inference); + let output = layer.forward(&input); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert_eq!(gradients.get(layer.weight()).unwrap().shape(), [4]); + assert_eq!(gradients.get(layer.bias().unwrap()).unwrap().shape(), [4]); + assert!(gradients.get(&input).is_none()); + + let inference_nd = crate::layers::LayerNormNd::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let expected_nd = flatten(inference_nd.forward::<3, 2, _>(&raw_input)).await; + let layer_nd = layers::LayerNormNd::from_inference(&graph, &inference_nd); + assert_slice_close( + &flatten(layer_nd.forward::<3, 2>(&input).into_raw()).await, + &expected_nd, + ); + + let inference_axis = crate::layers::LayerNormNd::new_over_axis( + RawTensor::from_slice(&device, [3], &axis_weight_data), + None, + 1, + 1e-5, + ); + let expected_axis = flatten(inference_axis.forward::<3, 2, _>(&raw_input)).await; + let layer_axis = layers::LayerNormNd::from_inference_over_axis(&graph, &inference_axis, 1); + assert_slice_close( + &flatten(layer_axis.forward::<3, 2>(&input).into_raw()).await, + &expected_axis, + ); + } +} + +#[tokio::test] +async fn test_autograd_layer_rms_norm_forward_parity() { + let weight_data = [0.5f32, 1.5, -0.75, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, -0.4]; + let input_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin()).collect(); + let residual_data: Vec = (0..24).map(|i| (i as f32 * 0.61).cos()).collect(); + for device in test_devices().await { + let raw_weight = RawTensor::from_slice(&device, [4], &weight_data); + let raw_bias = RawTensor::from_slice(&device, [4], &bias_data); + let raw_layer = + crate::layers::RmsNorm::new(raw_weight.clone(), Some(raw_bias.clone()), 1e-5); + + let graph = Graph::new(); + let layer = layers::RmsNorm::new(graph.leaf(raw_weight), Some(graph.leaf(raw_bias)), 1e-5); + + let input_2d = RawTensor::from_slice(&device, [6, 4], &input_data); + let expected = flatten(raw_layer.forward(&input_2d)).await; + let output = layer.forward(&Tensor::constant_from_raw(&graph, input_2d)); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let input_3d = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let expected = flatten(raw_layer.forward(&input_3d)).await; + let output = layer.forward(&Tensor::constant_from_raw(&graph, input_3d.clone())); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let input_4d = RawTensor::from_slice(&device, [2, 1, 3, 4], &input_data); + let expected = flatten(raw_layer.forward(&input_4d)).await; + let output = layer.forward(&Tensor::constant_from_raw(&graph, input_4d)); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let residual_3d = RawTensor::from_slice(&device, [2, 3, 4], &residual_data); + let expected = flatten(raw_layer.forward_residual_f32(&input_3d, &residual_3d)).await; + let output = layer.forward_residual( + &Tensor::constant_from_raw(&graph, input_3d), + &Tensor::constant_from_raw(&graph, residual_3d), + ); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + } +} + +#[tokio::test] +async fn test_autograd_layer_rms_norm_weight_gradient() { + let weight_data = [0.5f32, 1.5, -0.75, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, -0.4]; + let input_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin() + 0.25).collect(); + let residual_data: Vec = (0..24).map(|i| (i as f32 * 0.61).cos()).collect(); + for device in test_devices().await { + let input_2d = RawTensor::from_slice(&device, [6, 4], &input_data); + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let layer = layers::RmsNorm::new(weight, None, 1e-5); + layer + .forward(&Tensor::constant_from_raw(graph, input_2d.clone())) + .sqr() + .flatten_all() + .sum() + }) + .await; + + let input_3d = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let bias = RawTensor::from_slice(&device, [4], &bias_data); + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let layer = layers::RmsNorm::new( + weight, + Some(Tensor::constant_from_raw(graph, bias.clone())), + 1e-5, + ); + layer + .forward(&Tensor::constant_from_raw(graph, input_3d.clone())) + .sqr() + .flatten_all() + .sum() + }) + .await; + + let residual_3d = RawTensor::from_slice(&device, [2, 3, 4], &residual_data); + assert_gradient_matches_finite_difference(&device, [4], &weight_data, |graph, weight| { + let layer = layers::RmsNorm::new( + weight, + Some(Tensor::constant_from_raw(graph, bias.clone())), + 1e-5, + ); + layer + .forward_residual( + &Tensor::constant_from_raw(graph, input_3d.clone()), + &Tensor::constant_from_raw(graph, residual_3d.clone()), + ) + .sqr() + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_rms_norm_bias_gradient() { + let weight_data = [0.5f32, 1.5, -0.75, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, -0.4]; + let input_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin() + 0.25).collect(); + let residual_data: Vec = (0..24).map(|i| (i as f32 * 0.61).cos()).collect(); + for device in test_devices().await { + let input_3d = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let weight = RawTensor::from_slice(&device, [4], &weight_data); + assert_gradient_matches_finite_difference(&device, [4], &bias_data, |graph, bias| { + let layer = layers::RmsNorm::new( + Tensor::constant_from_raw(graph, weight.clone()), + Some(bias), + 1e-5, + ); + layer + .forward(&Tensor::constant_from_raw(graph, input_3d.clone())) + .sqr() + .flatten_all() + .sum() + }) + .await; + + let residual_3d = RawTensor::from_slice(&device, [2, 3, 4], &residual_data); + assert_gradient_matches_finite_difference(&device, [4], &bias_data, |graph, bias| { + let layer = layers::RmsNorm::new( + Tensor::constant_from_raw(graph, weight.clone()), + Some(bias), + 1e-5, + ); + layer + .forward_residual( + &Tensor::constant_from_raw(graph, input_3d.clone()), + &Tensor::constant_from_raw(graph, residual_3d.clone()), + ) + .sqr() + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_rms_norm_from_inference() { + let weight_data = [0.5f32, 1.5, -0.75, 2.0]; + let bias_data = [0.1f32, -0.2, 0.3, -0.4]; + let input_data: Vec = (0..24).map(|i| (i as f32 * 0.37).sin()).collect(); + for device in test_devices().await { + let raw_layer = crate::layers::RmsNorm::new( + RawTensor::from_slice(&device, [4], &weight_data), + Some(RawTensor::from_slice(&device, [4], &bias_data)), + 1e-5, + ); + let graph = Graph::new(); + let layer = layers::RmsNorm::from_inference(&graph, &raw_layer); + assert_eq!(layer.eps(), raw_layer.eps()); + + let input = RawTensor::from_slice(&device, [2, 3, 4], &input_data); + let expected = flatten(raw_layer.forward(&input)).await; + let output = layer.forward(&Tensor::constant_from_raw(&graph, input)); + assert_slice_close(&flatten(output.raw().clone()).await, &expected); + + let gradients = output.flatten_all().sum().backward().unwrap(); + let weight_gradient = gradients.get(layer.weight()).unwrap(); + let bias_gradient = gradients.get(layer.bias().unwrap()).unwrap(); + assert_eq!(weight_gradient.shape(), [4]); + assert_eq!(bias_gradient.shape(), [4]); + } +} + +#[tokio::test] +async fn test_autograd_layer_conv_nd_1d_forward_parity() { + for device in test_devices().await { + let w_data: Vec = (0..12).map(|i| (i as f32 * 0.37).cos()).collect(); + let b_data = [0.3f32, -0.7]; + let x_data: Vec = (0..15).map(|i| (i as f32 * 0.61).sin()).collect(); + let config = crate::layers::ConvNdConfig { + padding: [1], + stride: [2], + groups: 1, + }; + + let graph = Graph::new(); + let layer = layers::ConvNd::<1, 3>::new( + Tensor::from_slice(&graph, &device, [2, 3, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [2], &b_data)), + config, + ); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 3, 5], &x_data); + let output = layer.forward(&x); + + let raw_layer = crate::layers::ConvNd::<1, 3, f32>::new( + RawTensor::from_slice(&device, [2, 3, 2], &w_data), + Some(RawTensor::from_slice(&device, [2], &b_data)), + config, + ); + let expected = raw_layer.forward(&RawTensor::from_slice(&device, [1, 3, 5], &x_data)); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + assert_eq!(layer.in_channels(), 3); + assert_eq!(layer.out_channels(), 2); + assert_eq!(layer.config().stride, [2]); + } +} + +#[tokio::test] +async fn test_autograd_layer_conv_nd_1d_parameter_gradients() { + for device in test_devices().await { + let w_data: Vec = (0..12).map(|i| (i as f32 * 0.37).cos()).collect(); + let b_data = [0.3f32, -0.7]; + let x_data: Vec = (0..15).map(|i| (i as f32 * 0.61).sin()).collect(); + let config = crate::layers::ConvNdConfig { + padding: [1], + stride: [2], + groups: 1, + }; + + let graph = Graph::new(); + let layer = layers::ConvNd::<1, 3>::new( + Tensor::from_slice(&graph, &device, [2, 3, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [2], &b_data)), + config, + ); + let x = + Tensor::constant_from_raw(&graph, RawTensor::from_slice(&device, [1, 3, 5], &x_data)); + let gradients = layer.forward(&x).flatten_all().sum().backward().unwrap(); + assert!(gradients.get(layer.weight()).is_some()); + assert!(gradients.get(layer.bias().unwrap()).is_some()); + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let b_fd = b_data; + assert_gradient_matches_finite_difference(&device, [2, 3, 2], &w_data, move |graph, w| { + let layer = layers::ConvNd::<1, 3>::new( + w, + Some(Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2], &b_fd), + )), + config, + ); + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3, 5], &x_fd), + ); + let out = layer.forward(&x); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference(&device, [2], &b_data, move |graph, b| { + let layer = layers::ConvNd::<1, 3>::new( + Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [2, 3, 2], &w_fd), + ), + Some(b), + config, + ); + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 3, 5], &x_fd), + ); + let out = layer.forward(&x); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_conv_nd_forward_parity() { + for device in test_devices().await { + { + let x_data: Vec = (0..18).map(|i| (i as f32 * 0.41).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.23).cos()).collect(); + let b_data = [0.5f32, -1.0]; + let config = crate::layers::ConvNdConfig { + padding: [1, 1], + stride: [2, 2], + groups: 1, + }; + + let raw_layer = crate::layers::ConvNd::<2, 4, f32>::new( + RawTensor::from_slice(&device, [2, 2, 2, 2], &w_data), + Some(RawTensor::from_slice(&device, [2], &b_data)), + config, + ); + let expected = + raw_layer.forward(&RawTensor::from_slice(&device, [1, 2, 3, 3], &x_data)); + + let graph = Graph::new(); + let layer = layers::ConvNd::<2, 4>::new( + Tensor::from_slice(&graph, &device, [2, 2, 2, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [2], &b_data)), + config, + ); + let x: Tensor<4> = Tensor::from_slice(&graph, &device, [1, 2, 3, 3], &x_data); + let output = layer.forward(&x); + + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + } + + { + let x_data: Vec = (0..16).map(|i| (i as f32 * 0.57).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.31).cos()).collect(); + let b_data = [0.2f32, -0.4, 0.6, -0.8]; + let config = crate::layers::ConvNdConfig { + padding: [1], + stride: [2], + groups: 2, + }; + + let raw_layer = crate::layers::ConvNd::<1, 3, f32>::new( + RawTensor::from_slice(&device, [4, 2, 2], &w_data), + Some(RawTensor::from_slice(&device, [4], &b_data)), + config, + ); + let expected = raw_layer.forward(&RawTensor::from_slice(&device, [1, 4, 4], &x_data)); + + let graph = Graph::new(); + let layer = layers::ConvNd::<1, 3>::new( + Tensor::from_slice(&graph, &device, [4, 2, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [4], &b_data)), + config, + ); + let x: Tensor<3> = Tensor::from_slice(&graph, &device, [1, 4, 4], &x_data); + let output = layer.forward(&x); + + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + } + } +} + +#[tokio::test] +async fn test_autograd_layer_conv_nd_parameter_gradients() { + for device in test_devices().await { + let x_data: Vec = (0..16).map(|i| (i as f32 * 0.57).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.31).cos()).collect(); + let b_data = [0.2f32, -0.4, 0.6, -0.8]; + let config = crate::layers::ConvNdConfig { + padding: [1], + stride: [2], + groups: 2, + }; + + { + let graph = Graph::new(); + let layer = layers::ConvNd::<1, 3>::new( + Tensor::from_slice(&graph, &device, [4, 2, 2], &w_data), + Some(Tensor::from_slice(&graph, &device, [4], &b_data)), + config, + ); + let x = Tensor::constant_from_raw( + &graph, + RawTensor::from_slice(&device, [1, 4, 4], &x_data), + ); + let gradients = layer.forward(&x).flatten_all().sum().backward().unwrap(); + assert_eq!(gradients.get(layer.weight()).unwrap().shape(), [4, 2, 2]); + assert_eq!(gradients.get(layer.bias().unwrap()).unwrap().shape(), [4]); + } + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + assert_gradient_matches_finite_difference(&device, [4, 2, 2], &w_data, move |graph, w| { + let layer = layers::ConvNd::<1, 3>::new( + w, + Some(Tensor::from_slice(graph, &fd_device, [4], &b_data)), + config, + ); + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 4, 4], &x_fd), + ); + let out = layer.forward(&x); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + + let fd_device = device.clone(); + let x_fd = x_data.clone(); + let w_fd = w_data.clone(); + assert_gradient_matches_finite_difference(&device, [4], &b_data, move |graph, b| { + let layer = layers::ConvNd::<1, 3>::new( + Tensor::from_slice(graph, &fd_device, [4, 2, 2], &w_fd), + Some(b), + config, + ); + let x = Tensor::constant_from_raw( + graph, + RawTensor::from_slice(&fd_device, [1, 4, 4], &x_fd), + ); + let out = layer.forward(&x); + out.mul(&composite_ramp(graph, &fd_device, out.shape())) + .flatten_all() + .sum() + }) + .await; + } +} + +#[tokio::test] +async fn test_autograd_layer_conv_nd_from_inference() { + for device in test_devices().await { + let x_data: Vec = (0..16).map(|i| (i as f32 * 0.57).sin()).collect(); + let w_data: Vec = (0..16).map(|i| (i as f32 * 0.31).cos()).collect(); + let b_data = [0.2f32, -0.4, 0.6, -0.8]; + let config = crate::layers::ConvNdConfig { + padding: [1], + stride: [2], + groups: 2, + }; + let raw_layer = crate::layers::ConvNd::<1, 3, f32>::new( + RawTensor::from_slice(&device, [4, 2, 2], &w_data), + Some(RawTensor::from_slice(&device, [4], &b_data)), + config, + ); + + let graph = Graph::new(); + let layer = layers::ConvNd::<1, 3>::from_inference(&graph, &raw_layer); + assert_slice_close(&flatten(layer.weight().raw().clone()).await, &w_data); + assert_slice_close(&flatten(layer.bias().unwrap().raw().clone()).await, &b_data); + + let x_raw = RawTensor::from_slice(&device, [1, 4, 4], &x_data); + let expected = raw_layer.forward(&x_raw); + let output = layer.forward(&Tensor::constant_from_raw(&graph, x_raw)); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + + let gradients = output.flatten_all().sum().backward().unwrap(); + assert!(gradients.get(layer.weight()).is_some()); + assert!(gradients.get(layer.bias().unwrap()).is_some()); + + let raw_no_bias = crate::layers::ConvNd::<1, 3, f32>::new( + RawTensor::from_slice(&device, [4, 2, 2], &w_data), + None, + config, + ); + let imported = layers::ConvNd::<1, 3>::from_inference(&graph, &raw_no_bias); + assert!(imported.bias().is_none()); + let x_raw = RawTensor::from_slice(&device, [1, 4, 4], &x_data); + let expected = raw_no_bias.forward(&x_raw); + let output = imported.forward(&Tensor::constant_from_raw(&graph, x_raw)); + assert_slice_close( + &flatten(output.raw().clone()).await, + &flatten(expected).await, + ); + } +} + +/// End-to-end training with the trainable autograd layers: the same 2-16-2 +/// XOR MLP as [`test_train_xor_classifier`], but built from two +/// `layers::Linear` layers with gradients fetched through the +/// layer's `weight()`/`bias()` parameter handles. +#[tokio::test] +async fn test_train_xor_with_layers() { + const SAMPLES: usize = 64; + const HIDDEN: usize = 16; + const STEPS: usize = 500; + const LEARNING_RATE: f32 = 1.0; + + let (features, labels, w1_init, w2_init) = xor_training_data(HIDDEN); + + for (device, name) in test_devices().await.into_iter().zip(["cpu", "gpu"]) { + let inputs = RawTensor::from_slice(&device, [SAMPLES, 2], &features); + let targets = RawTensor::from_slice(&device, [SAMPLES], &labels); + + // Linear stores weights as (out_features, in_features). + let mut w1 = RawTensor::from_slice(&device, [HIDDEN, 2], &w1_init); + let mut b1 = RawTensor::zeros(&device, [HIDDEN]); + let mut w2 = RawTensor::from_slice(&device, [2, HIDDEN], &w2_init); + let mut b2 = RawTensor::zeros(&device, [2]); + + let mut final_loss = f32::INFINITY; + for step in 0..STEPS { + let graph = Graph::new(); + let x = Tensor::constant_from_raw(&graph, inputs.clone()); + let layer1 = layers::Linear::new(graph.leaf(w1.clone()), Some(graph.leaf(b1.clone()))); + let layer2 = layers::Linear::new(graph.leaf(w2.clone()), Some(graph.leaf(b2.clone()))); + + let hidden = layer1.forward(&x).relu(); + let logits = layer2.forward(&hidden); + // Numerically stable cross-entropy: log softmax via log-sum-exp + // so a saturated class cannot underflow to log(0). + let shifted = logits.sub_::<2, 2>(&logits.max_keepdim::<1>(1)); + let log_sum_exp = shifted.exp().sum_keepdim(1).log(); + let label_log_probs = shifted.sub_::<2, 2>(&log_sum_exp).gather_last(&targets); + let loss: Tensor<0> = label_log_probs.sum().mul_scalar(-1.0 / SAMPLES as f32); + + let loss_value = flatten(loss.raw().clone()).await[0]; + let gradients = loss.backward().unwrap().into_detached(); + let dw1 = gradients.get(layer1.weight()).unwrap(); + let db1 = gradients.get(layer1.bias().unwrap()).unwrap(); + let dw2 = gradients.get(layer2.weight()).unwrap(); + let db2 = gradients.get(layer2.bias().unwrap()).unwrap(); + + w1 = (w1 - dw1 * LEARNING_RATE).to_concrete(); + b1 = (b1 - db1 * LEARNING_RATE).to_concrete(); + w2 = (w2 - dw2 * LEARNING_RATE).to_concrete(); + b2 = (b2 - db2 * LEARNING_RATE).to_concrete(); + + final_loss = loss_value; + if step % 100 == 0 { + eprintln!("[{name}] step {step}: loss {loss_value:.4}"); + } + } + eprintln!("[{name}] final loss {final_loss:.4}"); + + let graph = Graph::new(); + let x = Tensor::constant_from_raw(&graph, inputs.clone()); + let layer1 = layers::Linear::new( + Tensor::constant_from_raw(&graph, w1.clone()), + Some(Tensor::constant_from_raw(&graph, b1.clone())), + ); + let layer2 = layers::Linear::new( + Tensor::constant_from_raw(&graph, w2.clone()), + Some(Tensor::constant_from_raw(&graph, b2.clone())), + ); + let hidden = layer1.forward(&x).relu(); + let logits = layer2.forward(&hidden); + let logits = logits.raw().clone().as_slice().await.unwrap().to_vec(); + let correct = logits + .iter() + .zip(&labels) + .filter(|(row, label)| u32::from(row[1] > row[0]) == **label) + .count(); + eprintln!("[{name}] accuracy {correct}/{SAMPLES}"); + + assert!( + final_loss < 0.1, + "training did not converge: final loss {final_loss}", + ); + assert_eq!(correct, SAMPLES, "classifier misclassified training points"); + } +} diff --git a/fusor-ml/fusor/src/autograd/view.rs b/fusor-ml/fusor/src/autograd/view.rs new file mode 100644 index 000000000..590d9667c --- /dev/null +++ b/fusor-ml/fusor/src/autograd/view.rs @@ -0,0 +1,954 @@ +use std::ops::Range; + +use crate::{Dim, Layout}; +use fusor_types::{SlidingWindow, StrideSpec}; + +use super::*; + +impl Tensor +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + pub fn reshape(&self, shape: [usize; OUT]) -> Tensor { + let input_shape = self.shape(); + let value = self.value.reshape(shape).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "reshape")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.reshape(input_shape).into_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn transpose(&self, dim0: usize, dim1: usize) -> Self { + let value = self.value.transpose(dim0, dim1).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "transpose")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.transpose(dim0, dim1).into_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn permute(&self, axes: [usize; R]) -> Self { + let value = self.value.permute(axes).into_concrete(); + let input_id = self.handle.id; + let mut inverse = [0usize; R]; + for (index, axis) in axes.iter().copied().enumerate() { + inverse[axis] = index; + } + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "permute")?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(gradient.permute(inverse).into_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn slice(&self, slices: [Range; R]) -> Self { + let input_shape = self.shape(); + let value = self.value.slice(slices.clone()).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "slice")?; + let zeros = RawTensor::zeros(&gradient.device(), input_shape); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new( + zeros + .slice_assign(slices.clone(), &gradient) + .into_concrete(), + ), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn broadcast_as(&self, shape: [usize; OUT]) -> Tensor { + let input_shape = self.shape(); + let value = self.value.broadcast_as(shape).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "broadcast_as")?; + let reduced = reduce_broadcast_gradient(gradient, input_shape)?; + Ok(vec![BackwardTarget { + node: input_id, + gradient: reduced, + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn expand(&self, shape: [usize; OUT]) -> Tensor { + self.broadcast_as(shape) + } + + pub fn flatten_all(&self) -> Tensor<1, T> { + self.reshape([self.shape().iter().product()]) + } + + pub fn flatten_last_n(&self) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank, + { + let shape = self.shape(); + let new_shape: [usize; OUT] = std::array::from_fn(|i| { + if i < R - 1 - FROM_END { + shape[i] + } else if i == R - 1 - FROM_END { + shape[R - 1 - FROM_END..].iter().product() + } else { + 1 + } + }); + self.reshape(new_shape) + } + + pub fn flatten_first_n(&self) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank, + { + let shape = self.shape(); + let new_shape: [usize; OUT] = std::array::from_fn(|i| { + if i == 0 { + shape[..=FROM_START].iter().product() + } else { + shape[i + FROM_START] + } + }); + self.reshape(new_shape) + } + + pub fn narrow(&self, dim: impl Dim, start: usize, length: usize) -> Self { + let dim = dim.resolve(); + let shape = self.shape(); + let slices: [Range; R] = std::array::from_fn(|axis| { + if axis == dim { + start..start + length + } else { + 0..shape[axis] + } + }); + self.slice(slices) + } + + pub fn chunk(&self, chunks: usize, dim: impl Dim) -> Vec { + let dim = dim.resolve(); + let shape = self.shape(); + let dim_size = shape[dim]; + let chunk_size = dim_size.div_ceil(chunks); + + let mut result = Vec::with_capacity(chunks); + let mut start = 0; + while start < dim_size { + let length = chunk_size.min(dim_size - start); + result.push(self.narrow(dim, start, length)); + start += length; + } + result + } + + pub fn repeat(&self, repeats: [usize; R]) -> Self { + let input_shape = self.shape(); + let value = self.value.repeat(repeats).into_concrete(); + let input_id = self.handle.id; + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "repeat")?; + let total: usize = gradient.shape().iter().product(); + let mut flat = gradient.reshape([total]).into_concrete(); + for axis in (0..R).rev() { + if repeats[axis] == 1 { + continue; + } + let before: usize = (0..axis) + .map(|dim| repeats[dim] * input_shape[dim]) + .product(); + let after: usize = input_shape[axis + 1..].iter().product(); + flat = flat + .reshape([before, repeats[axis], input_shape[axis], after]) + .into_concrete() + .sum::<3>(1) + .reshape([before * input_shape[axis] * after]) + .into_concrete(); + } + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(flat.reshape(input_shape).into_concrete()), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn resize(&self, new_shape: [usize; R]) -> Self { + let input_shape = self.shape(); + let value = self.value.resize(new_shape).into_concrete(); + let input_id = self.handle.id; + let copy_shape = std::array::from_fn(|axis| input_shape[axis].min(new_shape[axis])); + let copy_slices = copy_shape.map(|size| 0..size); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "resize")?; + let patch = gradient.slice(copy_slices.clone()).into_concrete(); + let zeros = RawTensor::zeros(&gradient.device(), input_shape); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new( + zeros + .slice_assign(copy_slices.clone(), &patch) + .into_concrete(), + ), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn restride(&self, specs: [StrideSpec; OUT]) -> Tensor { + let input_shape = self.shape(); + let value = self.value.restride(specs).into_concrete(); + let input_id = self.handle.id; + let output_shape: [usize; OUT] = specs.map(|spec| spec.size); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "restride")?; + let reduced = reduce_restride_gradient(&gradient, &specs, [0; R], input_shape) + .unwrap_or_else(|| { + scatter_restride_gradient( + &gradient, + output_shape, + input_shape, + |output_index| restride_input_index(specs, output_index), + ) + }); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(reduced), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn restride_layout(&self, new_layout: Layout) -> Tensor { + assert_eq!(new_layout.rank(), OUT, "restride_layout rank mismatch"); + let input_shape = self.shape(); + let value = self + .value + .restride_layout(new_layout.clone()) + .into_concrete(); + let input_id = self.handle.id; + let output_shape: [usize; OUT] = std::array::from_fn(|axis| new_layout.shape()[axis]); + let input_strides = Layout::continuous_strides(&input_shape); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "restride_layout")?; + let reduced = layout_restride_specs(&new_layout, input_shape) + .and_then(|(specs, offsets)| { + reduce_restride_gradient(&gradient, &specs, offsets, input_shape) + }) + .unwrap_or_else(|| { + scatter_restride_gradient( + &gradient, + output_shape, + input_shape, + |output_index| { + let linear = new_layout.linear_index(&output_index); + contiguous_index_from_linear::(linear, &input_strides) + }, + ) + }); + Ok(vec![BackwardTarget { + node: input_id, + gradient: Box::new(reduced), + }]) + }); + self.emit_op(value, vec![self.handle.clone()], Some(backward)) + } + + pub fn squeeze_dims( + &self, + axes: [usize; DIFF], + ) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank, + { + let shape = self.shape(); + for &axis in &axes { + assert_eq!( + shape[axis], 1, + "Squeeze dimension {} must have size 1", + axis + ); + } + let mut sorted_axes = axes; + sorted_axes.sort_unstable(); + let mut input_axis = 0; + let mut axis_index = 0; + let specs: [StrideSpec; OUT] = std::array::from_fn(|_| { + while axis_index < DIFF && input_axis == sorted_axes[axis_index] { + input_axis += 1; + axis_index += 1; + } + let spec = StrideSpec::dim(input_axis, shape[input_axis]); + input_axis += 1; + spec + }); + self.restride(specs) + } + + pub fn unsqueeze_dims( + &self, + axes: [usize; DIFF], + ) -> Tensor + where + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let shape = self.shape(); + let mut sorted_axes = axes; + sorted_axes.sort_unstable(); + let mut input_axis = 0; + let mut axis_index = 0; + let specs: [StrideSpec; OUT] = std::array::from_fn(|output_axis| { + if axis_index < DIFF && output_axis == sorted_axes[axis_index] { + axis_index += 1; + StrideSpec::dim_with(0, 1, 0) + } else { + let spec = StrideSpec::dim(input_axis, shape[input_axis]); + input_axis += 1; + spec + } + }); + self.restride(specs) + } + + pub fn squeeze(&self, dim: usize) -> Tensor + where + crate::gpu::Tensor: crate::gpu::SmallerRank<1, OUT, T>, + { + self.squeeze_dims::<1, OUT>([dim]) + } + + pub fn unsqueeze(&self, dim: usize) -> Tensor + where + crate::gpu::Tensor: crate::gpu::LargerRank<1, OUT, T>, + { + self.unsqueeze_dims::<1, OUT>([dim]) + } + + pub fn sliding_window_view( + &self, + windows: [SlidingWindow; DIFF], + ) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + { + let shape = self.shape(); + let mut sorted_windows = windows; + sorted_windows.sort_by_key(|window| window.axis); + let specs: [StrideSpec; R2] = std::array::from_fn(|out_i| { + if out_i < R { + if let Some(window) = sorted_windows.iter().find(|window| window.axis == out_i) { + let positions = (shape[out_i] - window.window_size) / window.step + 1; + StrideSpec::dim_with(out_i, positions, window.step) + } else { + StrideSpec::dim(out_i, shape[out_i]) + } + } else { + let window = &sorted_windows[out_i - R]; + StrideSpec::dim(window.axis, window.window_size) + } + }); + self.restride(specs) + } + + pub fn pad_axis(&self, axis: usize, padding: usize) -> Self { + self.pad_with_zeros(axis, padding, padding) + } + + pub fn pad_with_zeros(&self, axis: usize, left: usize, right: usize) -> Self { + if left == 0 && right == 0 { + return self.clone(); + } + let shape = self.shape(); + let mut padded_shape = shape; + padded_shape[axis] += left + right; + let padded = Self::constant_from_raw( + &self.graph(), + RawTensor::zeros(&self.device(), padded_shape), + ); + let slices: [Range; R] = std::array::from_fn(|dim| { + if dim == axis { + left..left + shape[axis] + } else { + 0..shape[dim] + } + }); + padded.slice_assign(slices, self) + } + + pub fn slice_assign(&self, slices: [Range; R], value: &Self) -> Self { + assert_same_graph(self, value); + + let output = self + .value + .slice_assign(slices.clone(), &value.value) + .into_concrete(); + let input_id = self.handle.id; + let value_id = value.handle.id; + let slice_shape = slices + .clone() + .map(|range| range.end.saturating_sub(range.start)); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "slice_assign")?; + let zeros = RawTensor::zeros(&gradient.device(), slice_shape); + Ok(vec![ + BackwardTarget { + node: input_id, + gradient: Box::new( + gradient + .slice_assign(slices.clone(), &zeros) + .into_concrete(), + ), + }, + BackwardTarget { + node: value_id, + gradient: Box::new(gradient.slice(slices.clone()).into_concrete()), + }, + ]) + }); + self.emit_op( + output, + vec![self.handle.clone(), value.handle.clone()], + Some(backward), + ) + } + + pub fn stack( + tensors: impl IntoIterator, + dim: usize, + ) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank<1, OUT, T>, + { + let tensors: Vec = tensors.into_iter().collect(); + assert!(!tensors.is_empty(), "stack requires at least one tensor"); + + let graph = tensors[0].handle.graph.clone(); + let input_shape = tensors[0].shape(); + let raw = tensors + .iter() + .map(|tensor| { + assert!( + Arc::ptr_eq(&graph, &tensor.handle.graph), + "cannot mix autograd tensors from different graphs" + ); + assert_eq!( + tensor.shape(), + input_shape, + "stack requires matching shapes" + ); + tensor.value.unsqueeze_dims::<1, OUT>([dim]).into_concrete() + }) + .collect::>(); + let value = RawTensor::cat(raw, dim); + let parents = tensors + .iter() + .map(|tensor| tensor.handle.clone()) + .collect::>(); + let parent_ids = parents.iter().map(|parent| parent.id).collect::>(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "stack")?; + let mut targets = Vec::with_capacity(parent_ids.len()); + for (index, &parent_id) in parent_ids.iter().enumerate() { + let slices: [Range; OUT] = std::array::from_fn(|axis| { + if axis == dim { + index..index + 1 + } else { + 0..gradient.shape()[axis] + } + }); + let grad = gradient.slice(slices).reshape(input_shape).into_concrete(); + targets.push(BackwardTarget { + node: parent_id, + gradient: Box::new(grad), + }); + } + Ok(targets) + }); + let id = graph.add_node( + parents.iter().map(|parent| parent.id).collect(), + Some(backward), + parents + .iter() + .any(|parent| parent.graph.requires_grad(parent.id)), + ); + Tensor { + value, + handle: NodeHandle { graph, id }, + } + } + pub fn cat(tensors: Vec, dim: usize) -> Self { + assert!(!tensors.is_empty(), "cat requires at least one tensor"); + let graph = tensors[0].handle.graph.clone(); + let raw = tensors + .iter() + .map(|tensor| tensor.value.clone()) + .collect::>(); + let value = RawTensor::cat(raw, dim); + let parents = tensors + .iter() + .map(|tensor| tensor.handle.clone()) + .collect::>(); + let parent_ids = parents.iter().map(|parent| parent.id).collect::>(); + let slices = tensors + .iter() + .scan(0usize, |offset, tensor| { + let start = *offset; + let length = tensor.shape()[dim]; + *offset += length; + Some(start..start + length) + }) + .collect::>(); + let backward: BackwardRule = Arc::new(move |gradient| { + let gradient = downcast_tensor::(&*gradient, "cat")?; + let mut targets = Vec::with_capacity(parent_ids.len()); + for (&parent_id, slice) in parent_ids.iter().zip(slices.iter()) { + let ranges: [Range; R] = std::array::from_fn(|axis| { + if axis == dim { + slice.clone() + } else { + 0..gradient.shape()[axis] + } + }); + targets.push(BackwardTarget { + node: parent_id, + gradient: Box::new(gradient.slice(ranges).into_concrete()), + }); + } + Ok(targets) + }); + let id = graph.add_node( + parents.iter().map(|parent| parent.id).collect(), + Some(backward), + parents + .iter() + .any(|parent| parent.graph.requires_grad(parent.id)), + ); + Tensor { + value, + handle: NodeHandle { graph, id }, + } + } +} + +fn for_each_index(limits: [usize; R], mut visitor: impl FnMut([usize; R])) { + if limits.contains(&0) { + return; + } + + let mut index = [0; R]; + loop { + visitor(index); + + let mut axis = R; + loop { + if axis == 0 { + return; + } + axis -= 1; + index[axis] += 1; + if index[axis] < limits[axis] { + break; + } + index[axis] = 0; + } + } +} + +fn restride_input_index( + specs: [StrideSpec; OUT], + output_index: [usize; OUT], +) -> [usize; R] { + let mut input_index = [0; R]; + for axis in 0..OUT { + let spec = specs[axis]; + input_index[spec.input_dim] += spec.offset + output_index[axis] * spec.multiplier; + } + input_index +} + +fn contiguous_index_from_linear( + mut linear: usize, + strides: &[usize], +) -> [usize; R] { + let mut input_index = [0; R]; + for axis in 0..R { + input_index[axis] = linear / strides[axis]; + linear %= strides[axis]; + } + input_index +} + +/// Grouped view of a restride's specs for one input axis: at most one strided +/// "position" run and one unit-stride "window" run plus a constant offset, so +/// the forward map factors per input axis as +/// `input = offset + position * step + window`. +#[derive(Clone, Copy)] +struct RestrideRuns { + offset: usize, + /// `(output_axis, step, count)` of the strided run. + position: Option<(usize, usize, usize)>, + /// `(output_axis, size)` of the unit-stride run. + window: Option<(usize, usize)>, +} + +impl RestrideRuns { + fn counts(&self) -> (usize, usize, usize) { + let (_, step, positions) = self.position.unwrap_or((0, 1, 1)); + let (_, window) = self.window.unwrap_or((0, 1)); + (positions, step, window) + } + + fn output_len(&self) -> usize { + let (positions, _, window) = self.counts(); + positions * window + } + + fn fold_len(&self) -> usize { + let (positions, step, window) = self.counts(); + (positions - 1) * step + window + } + + /// True when the runs enumerate the input axis exactly once in row-major + /// order, so the gradient maps back with a plain reshape. + fn is_reshape(&self, size: usize) -> bool { + let (positions, step, window) = self.counts(); + self.offset == 0 && self.fold_len() == size && (positions == 1 || step == window) + } +} + +fn group_restride_runs( + specs: &[StrideSpec], + base_offsets: [usize; IN], + input_shape: [usize; IN], +) -> Option<[RestrideRuns; IN]> { + let mut runs: [RestrideRuns; IN] = std::array::from_fn(|axis| RestrideRuns { + offset: base_offsets[axis], + position: None, + window: None, + }); + for (output_axis, spec) in specs.iter().enumerate() { + let axis = &mut runs[spec.input_dim]; + axis.offset += spec.offset; + if spec.size == 1 { + continue; + } + if spec.multiplier == 1 && axis.window.is_none() { + axis.window = Some((output_axis, spec.size)); + } else if spec.multiplier >= 1 && axis.position.is_none() { + axis.position = Some((output_axis, spec.multiplier, spec.size)); + } else { + return None; + } + } + runs.iter() + .zip(&input_shape) + .all(|(axis, &size)| axis.offset + axis.fold_len() <= size) + .then_some(runs) +} + +/// Backward of a position/window restride as compiled graph ops: group the +/// output axes per input axis, canonicalize their order with one permute, +/// then fold each axis with [`fold_restride_axis`]. Returns `None` for +/// restrides that do not factor into per-axis runs; those go through the +/// host-loop fallback. +fn reduce_restride_gradient( + gradient: &RawTensor, + specs: &[StrideSpec], + base_offsets: [usize; IN], + input_shape: [usize; IN], +) -> Option> +where + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, +{ + let output_shape = gradient.shape(); + if input_shape.contains(&0) || output_shape.contains(&0) { + return None; + } + let runs = group_restride_runs(specs, base_offsets, input_shape)?; + + let mut order = Vec::with_capacity(OUT); + for axis in &runs { + if let Some((output_axis, _, _)) = axis.position { + order.push(output_axis); + } + if let Some((output_axis, _)) = axis.window { + order.push(output_axis); + } + } + // Size-1 output axes reshape away wherever they sit, so only the run axes + // decide whether the gradient needs a permute into canonical order. + let canonical = if order.windows(2).all(|pair| pair[0] < pair[1]) { + gradient.clone() + } else { + let mut in_runs = [false; OUT]; + for &axis in &order { + in_runs[axis] = true; + } + order.extend((0..OUT).filter(|&axis| !in_runs[axis])); + let permutation: [usize; OUT] = order + .as_slice() + .try_into() + .expect("every output axis appears in the permutation once"); + gradient.permute(permutation).into_concrete() + }; + + let mut flat = canonical + .reshape([output_shape.iter().product()]) + .into_concrete(); + let mut after = 1usize; + for dim in (0..IN).rev() { + let size = input_shape[dim]; + if !runs[dim].is_reshape(size) { + let before = runs[..dim].iter().map(RestrideRuns::output_len).product(); + let (positions, step, window) = runs[dim].counts(); + flat = fold_restride_axis( + flat, + before, + positions, + step, + window, + runs[dim].offset, + size, + after, + ); + } + after *= size; + } + Some(flat.reshape(input_shape).into_concrete()) +} + +/// Scatter-add one folded axis of the gradient with padded views and a +/// reduce: `out[offset + p*step + w] += g[.., p, w, ..]` for every position +/// `p` and window element `w`, with the surrounding axes flattened into +/// `before` and `after` batch extents. +#[allow(clippy::too_many_arguments)] +fn fold_restride_axis( + flat: RawTensor<1, T>, + before: usize, + positions: usize, + step: usize, + window: usize, + offset: usize, + size: usize, + after: usize, +) -> RawTensor<1, T> +where + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, +{ + let fold_len = (positions - 1) * step + window; + let block = flat + .reshape([before, positions, window, after]) + .into_concrete(); + let folded: RawTensor<3, T> = if positions == 1 { + block.reshape([before, window, after]).into_concrete() + } else if step >= window { + // Injective: interleave the windows with zeros and trim the overhang. + block + .pad_with_zeros(2, 0, step - window) + .reshape([before, positions * step, after]) + .into_concrete() + .narrow(1usize, 0, fold_len) + .into_concrete() + } else { + // Overlapping: reverse the window axis (`u = window - 1 - w`), + // right-pad each window row to `step * window` elements and left-pad + // `(window - 1) * window` zeros; the affine view + // `f(v, u) = v*window + u*(window + 1)` then reads `g[p, w]` exactly + // when `p*step + w == v` and a zero cell otherwise, so one reduce + // over `u` folds every overlapping window. + let reversed: Vec = (0..window as u32).rev().collect(); + let indices = RawTensor::from_slice(&block.device(), [window], &reversed); + block + .index_select(2, &indices) + .pad_with_zeros(2, 0, (step - 1) * window) + .reshape([before, positions * step * window, after]) + .into_concrete() + .pad_with_zeros(1, (window - 1) * window, window * (window - step)) + .restride([ + StrideSpec::dim(0, before), + StrideSpec::dim_with(1, fold_len, window), + StrideSpec::dim_with(1, window, window + 1), + StrideSpec::dim(2, after), + ]) + .into_concrete() + .sum::<3>(2) + }; + folded + .pad_with_zeros(1, offset, size - offset - fold_len) + .reshape([before * size * after]) + .into_concrete() +} + +/// Express a layout over a contiguous input as per-output-axis stride specs +/// plus per-input-axis base offsets. Returns `None` when a stride does not +/// decompose into a single input dimension or when an axis' reach could carry +/// into the next dimension, where per-axis factoring would diverge from the +/// layout's linear indexing. +fn layout_restride_specs( + layout: &Layout, + input_shape: [usize; IN], +) -> Option<(Vec, [usize; IN])> { + if input_shape.contains(&0) { + return None; + } + let input_strides = Layout::continuous_strides(&input_shape); + let mut reach = [0usize; IN]; + let mut specs = Vec::with_capacity(layout.rank()); + for (&size, &stride) in layout.shape().iter().zip(layout.strides()) { + if size == 1 || stride == 0 { + specs.push(StrideSpec::dim_with(0, size, 0)); + continue; + } + let dim = (0..IN).find(|&dim| input_strides[dim] <= stride)?; + if stride % input_strides[dim] != 0 { + return None; + } + let multiplier = stride / input_strides[dim]; + reach[dim] += (size - 1) * multiplier; + specs.push(StrideSpec::dim_with(dim, size, multiplier)); + } + let mut offsets = [0usize; IN]; + let mut offset = layout.offset(); + for dim in 0..IN { + offsets[dim] = offset / input_strides[dim]; + offset %= input_strides[dim]; + reach[dim] += offsets[dim]; + } + reach + .iter() + .zip(&input_shape) + .all(|(&reach, &size)| reach < size) + .then_some((specs, offsets)) +} + +/// Host-loop fallback for the restride patterns [`reduce_restride_gradient`] +/// cannot factor into per-axis runs. +fn scatter_restride_gradient( + gradient: &RawTensor, + output_shape: [usize; OUT], + input_shape: [usize; IN], + input_index: impl Fn([usize; OUT]) -> [usize; IN], +) -> RawTensor +where + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, +{ + let mut input_gradient = RawTensor::zeros(&gradient.device(), input_shape); + for_each_index(output_shape, |output_index| { + let input_index = input_index(output_index); + let output_slices: [Range; OUT] = + std::array::from_fn(|axis| output_index[axis]..output_index[axis] + 1); + let patch = gradient + .slice(output_slices) + .reshape([1; IN]) + .into_concrete(); + let target: [Range; IN] = + std::array::from_fn(|axis| input_index[axis]..input_index[axis] + 1); + let current = input_gradient.slice(target.clone()).into_concrete(); + let updated = (current + patch).into_concrete(); + input_gradient = input_gradient + .slice_assign(target, &updated) + .into_concrete(); + }); + input_gradient +} + +fn reduce_broadcast_gradient( + gradient: RawTensor, + input_shape: [usize; IN], +) -> Result> +where + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, +{ + let output_shape = gradient.shape(); + let mut aligned_input_shape = [1usize; OUT]; + for axis in 0..IN { + aligned_input_shape[OUT - IN + axis] = input_shape[axis]; + } + + for axis in 0..OUT { + let output_dim = output_shape[axis]; + let input_dim = aligned_input_shape[axis]; + if input_dim != 1 && input_dim != output_dim { + return Err(Error::msg("incompatible broadcast gradient shape")); + } + } + + if aligned_input_shape == output_shape { + if IN == OUT { + return Ok(Box::new(gradient)); + } + return Ok(Box::new(gradient.reshape(input_shape).into_concrete())); + } + + // Sum the axes the forward broadcast expanded, one compiled reduce per axis. + let mut remaining = output_shape; + let mut flat = gradient + .reshape([output_shape.iter().product()]) + .into_concrete(); + for axis in 0..OUT { + if aligned_input_shape[axis] != 1 || remaining[axis] == 1 { + continue; + } + let before: usize = remaining[..axis].iter().product(); + let after: usize = remaining[axis + 1..].iter().product(); + flat = flat + .reshape([before, remaining[axis], after]) + .into_concrete() + .sum::<2>(1) + .reshape([before * after]) + .into_concrete(); + remaining[axis] = 1; + } + Ok(Box::new(flat.reshape(input_shape).into_concrete())) +} diff --git a/fusor-ml/fusor/src/cache/attention_mask.rs b/fusor-ml/fusor/src/cache/attention_mask.rs index 945b7f0d7..f4fe89ee1 100644 --- a/fusor-ml/fusor/src/cache/attention_mask.rs +++ b/fusor-ml/fusor/src/cache/attention_mask.rs @@ -10,8 +10,8 @@ use crate::{ConcreteTensor, Device, SimdElement, Tensor}; pub struct AttentionMask { mask: Tensor<2, D, ConcreteTensor>, /// `true` when the mask is exactly the strict lower-triangular causal - /// mask of shape `[n, n]`. Hint to GPU flash attention so it can skip the - /// mask tensor entirely and prune upper-triangle work. + /// mask of shape `[n, n]`. This lets the compiler omit the mask tensor + /// and prune upper-triangle work. is_strict_causal: bool, } @@ -28,7 +28,7 @@ where } /// Returns true if this is a strict lower-triangular causal mask. The - /// GPU flash attention kernel can then skip masking work entirely. + /// compiler can then skip masking work entirely. pub fn is_strict_causal(&self) -> bool { self.is_strict_causal } diff --git a/fusor-ml/fusor/src/composite/flash_attention.rs b/fusor-ml/fusor/src/composite/attention.rs similarity index 61% rename from fusor-ml/fusor/src/composite/flash_attention.rs rename to fusor-ml/fusor/src/composite/attention.rs index 32568638c..7b26f398b 100644 --- a/fusor-ml/fusor/src/composite/flash_attention.rs +++ b/fusor-ml/fusor/src/composite/attention.rs @@ -1,4 +1,4 @@ -//! Flash attention operations that work on both CPU and GPU backends. +//! Attention operations that work on both CPU and GPU backends. use crate::cpu::{MatmulImpl, MaxOp, SimdReduceOp, SumOp}; use crate::gpu::{DataType, FloatDataType}; @@ -17,9 +17,9 @@ pub enum MaskKind { /// Used for padding masks in encoder/embedding models. BatchKeyMask, /// Mask is a strict lower-triangular causal mask of shape [seq_len, seq_len]. - /// The GPU flash-attention kernel can skip the upper-triangle Q·K work - /// entirely and does not load the mask tensor at all. Falls back to - /// `QKMask` semantics on backends that don't support the optimisation. + /// GPU graphs encode causality directly so the compiler can skip the + /// upper-triangle Q·K work without loading the mask tensor. Other + /// backends apply the provided mask with `QKMask` semantics. Causal, } @@ -44,7 +44,7 @@ where SumOp: SimdReduceOp, ExpOp: SimdUnaryOp, { - /// Computes flash attention with optional masking. + /// Computes scaled dot-product attention with optional masking. /// /// Supports grouped-query attention (GQA) and multi-query attention (MQA) where /// K and V may have fewer heads than Q. The number of Q heads must be divisible @@ -55,7 +55,7 @@ where /// - v: Value tensor with shape [batch, num_kv_heads, kv_seq_len, head_dim] /// - scale: Scale factor (typically 1/sqrt(head_dim)) /// - mask: Optional attention mask with a [`MaskKind`] describing its layout - pub fn flash_attention( + pub fn attention( &self, k: &Self, v: &Self, @@ -63,54 +63,29 @@ where mask: Option<(&Tensor<2, D, ConcreteTensor>, MaskKind)>, ) -> Self { match (self, k, v) { - // GPU path - use the optimized fused kernel (QKMask/Causal only) + // Preserve the canonical GPU graph shape for the compiler's + // attention recognizer. Batch-key masks use the backend-neutral + // composition below because the recognizer accepts QK masks. #[cfg(feature = "gpu")] (Tensor::Gpu(q), Tensor::Gpu(k_gpu), Tensor::Gpu(v_gpu)) if !matches!(mask, Some((_, MaskKind::BatchKeyMask))) => { - // Decode (q_seq_len == 1) runs the DecodeSmall attention kernel, - // which uses workgroup reductions and needs no subgroups, so it - // works on browser adapters that report no subgroup support. Only - // the prefill/streaming kernels require subgroups — keep the - // fallback for those (q_seq_len > 1). - if !q.device().subgroups_supported() && self.shape()[2] != 1 { - #[cfg(target_arch = "wasm32")] - { - return self.flash_attention_composite_impl(k, v, scale, mask); - } - #[cfg(not(target_arch = "wasm32"))] - { - let cpu_q = tensor4_to_cpu(q); - let cpu_k = tensor4_to_cpu(k_gpu); - let cpu_v = tensor4_to_cpu(v_gpu); - let cpu_mask = mask.map(|(m, kind)| { - let Tensor::Gpu(mask) = m else { - panic!("Mask must be on the same device as other tensors"); - }; - (tensor2_to_cpu(mask), kind) - }); - let cpu_mask_ref = cpu_mask.as_ref().map(|(mask, kind)| (mask, *kind)); - let cpu_output = cpu_q.flash_attention(&cpu_k, &cpu_v, scale, cpu_mask_ref); - return tensor4_to_gpu(cpu_output, q.device()); - } - } if matches!(mask, Some((_, MaskKind::Causal))) { - return Tensor::Gpu(q.flash_attention_causal(k_gpu, v_gpu, scale)); + return Tensor::Gpu(q.attention_causal(k_gpu, v_gpu, scale)); } let gpu_mask = mask.map(|(m, _kind)| match m { Tensor::Gpu(mask) => mask, _ => panic!("Mask must be on the same device as other tensors"), }); - Tensor::Gpu(q.flash_attention(k_gpu, v_gpu, scale, gpu_mask)) + Tensor::Gpu(q.attention(k_gpu, v_gpu, scale, gpu_mask)) } - // CPU path and GPU+BatchKeyMask fallback - use composite operations via Tensor methods - _ => self.flash_attention_composite_impl(k, v, scale, mask), + // CPU path and GPU+BatchKeyMask path. + _ => self.attention_composite_impl(k, v, scale, mask), } } - /// Implementation of flash attention using Tensor composite operations. - /// Works on both CPU and GPU tensors (GPU uses individual ops instead of fused kernel). - fn flash_attention_composite_impl( + /// Backend-neutral attention composition. + fn attention_composite_impl( &self, k: &Self, v: &Self, @@ -235,62 +210,3 @@ where attn_weights.mat_mul(&v_expanded) } } - -#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] -fn tensor4_to_cpu(tensor: &crate::gpu::Tensor<4, D>) -> Tensor<4, D> -where - D: SimdElement + DataType + Copy, -{ - let shape = *tensor.shape(); - let slice = pollster::block_on(tensor.as_slice()).expect("failed to read tensor"); - let mut values = Vec::with_capacity(shape.iter().product()); - for b in 0..shape[0] { - for h in 0..shape[1] { - for s in 0..shape[2] { - for d in 0..shape[3] { - values.push(slice[[b, h, s, d]]); - } - } - } - } - Tensor::Cpu(crate::cpu::TypedTensor::from_slice(shape, &values)) -} - -#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] -fn tensor4_to_gpu(tensor: Tensor<4, D>, device: &crate::gpu::Device) -> Tensor<4, D> -where - D: SimdElement + DataType + Copy, -{ - let Tensor::Cpu(tensor) = tensor else { - unreachable!("subgroup fallback should produce a CPU tensor"); - }; - let shape = tensor.shape(); - let slice = tensor.as_slice(); - let mut values = Vec::with_capacity(shape.iter().product()); - for b in 0..shape[0] { - for h in 0..shape[1] { - for s in 0..shape[2] { - for d in 0..shape[3] { - values.push(slice[[b, h, s, d]]); - } - } - } - } - Tensor::Gpu(crate::gpu::Tensor::from_slice(device, shape, &values)) -} - -#[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] -fn tensor2_to_cpu(tensor: &crate::gpu::Tensor<2, D>) -> Tensor<2, D> -where - D: SimdElement + DataType + Copy, -{ - let shape = *tensor.shape(); - let slice = pollster::block_on(tensor.as_slice()).expect("failed to read tensor"); - let mut values = Vec::with_capacity(shape.iter().product()); - for row in 0..shape[0] { - for col in 0..shape[1] { - values.push(slice[[row, col]]); - } - } - Tensor::Cpu(crate::cpu::TypedTensor::from_slice(shape, &values)) -} diff --git a/fusor-ml/fusor/src/composite/comparison.rs b/fusor-ml/fusor/src/composite/comparison.rs index 95ce32db3..f6232f324 100644 --- a/fusor-ml/fusor/src/composite/comparison.rs +++ b/fusor-ml/fusor/src/composite/comparison.rs @@ -27,7 +27,7 @@ macro_rules! scalar_cmp { }; } -/// Emit a tensor-tensor comparison method that runs on CPU only. +/// Emit a tensor-tensor comparison method that dispatches CPU/GPU. macro_rules! tensor_cmp { ($(#[$meta:meta])* $method:ident, $op:ident, $cpu_method:ident) => { $(#[$meta])* @@ -36,7 +36,11 @@ macro_rules! tensor_cmp { B2: Fusion, $op: SimdBinaryOp, { - self.dispatch_cpu_only_pair(rhs, |a, b| a.as_ref().$cpu_method(b.as_ref()).to_concrete()) + self.dispatch_pair_concrete( + rhs, + |a, b| a.as_ref().$cpu_method(b.as_ref()).to_concrete(), + |a, b| a.$method(b), + ) } }; } @@ -63,7 +67,6 @@ where /// Element-wise equality comparison between two tensors. /// /// Returns 1.0 where elements are equal, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. eq_tensor, EqOp, eq ); @@ -71,7 +74,6 @@ where /// Element-wise inequality comparison between two tensors. /// /// Returns 1.0 where elements are not equal, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. ne_tensor, NeOp, ne ); @@ -79,7 +81,6 @@ where /// Element-wise less-than comparison between two tensors. /// /// Returns 1.0 where self < rhs, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. lt_tensor, LtOp, lt ); @@ -87,7 +88,6 @@ where /// Element-wise less-than-or-equal comparison between two tensors. /// /// Returns 1.0 where self <= rhs, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. lte_tensor, LteOp, lte ); @@ -95,7 +95,6 @@ where /// Element-wise greater-than comparison between two tensors. /// /// Returns 1.0 where self > rhs, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. gt_tensor, GtOp, gt ); @@ -103,7 +102,6 @@ where /// Element-wise greater-than-or-equal comparison between two tensors. /// /// Returns 1.0 where self >= rhs, 0.0 otherwise. - /// Note: GPU comparison is only available for CPU tensors at this time. gte_tensor, GteOp, gte ); diff --git a/fusor-ml/fusor/src/composite/index.rs b/fusor-ml/fusor/src/composite/index.rs index 2018ca101..1f2d7a82f 100644 --- a/fusor-ml/fusor/src/composite/index.rs +++ b/fusor-ml/fusor/src/composite/index.rs @@ -48,7 +48,7 @@ impl From for IndexOp { } impl IndexOp { - fn to_range(&self, dim_size: usize) -> Range { + pub(crate) fn to_range(&self, dim_size: usize) -> Range { match self { IndexOp::Full => 0..dim_size, IndexOp::Range(r) => r.clone(), @@ -58,7 +58,7 @@ impl IndexOp { } } - fn removes_dim(&self) -> bool { + pub(crate) fn removes_dim(&self) -> bool { matches!(self, IndexOp::Index(_)) } } @@ -88,7 +88,7 @@ where } } -fn removed_dim(removes: [bool; R]) -> usize { +pub(crate) fn removed_dim(removes: [bool; R]) -> usize { let num_removes = removes.iter().filter(|&&removed| removed).count(); assert!( num_removes == 1, diff --git a/fusor-ml/fusor/src/composite/mod.rs b/fusor-ml/fusor/src/composite/mod.rs index 36cb1dcba..292a9eb46 100644 --- a/fusor-ml/fusor/src/composite/mod.rs +++ b/fusor-ml/fusor/src/composite/mod.rs @@ -4,10 +4,10 @@ //! across CPU and GPU tensors via the Tensor abstraction. mod activations; +mod attention; mod comparison; mod construction; mod conv; -mod flash_attention; pub mod index; mod index_select; mod math; @@ -20,8 +20,8 @@ mod to_vec; mod upsample; mod where_cond; -pub use flash_attention::MaskKind; +pub use attention::MaskKind; pub use rope::{RopeCache, base_inverse_frequency}; pub(crate) use shape::broadcast_shapes; pub use shape::{arange, arange_step, cat, stack}; -pub use to_vec::{ToVec, ToVec1, ToVec2, ToVec3}; +pub use to_vec::ToVec; diff --git a/fusor-ml/fusor/src/composite/to_vec.rs b/fusor-ml/fusor/src/composite/to_vec.rs index 811812c58..455ce92b8 100644 --- a/fusor-ml/fusor/src/composite/to_vec.rs +++ b/fusor-ml/fusor/src/composite/to_vec.rs @@ -6,32 +6,18 @@ use bytemuck::{AnyBitPattern, NoUninit}; use crate::gpu::TensorSlice; -/// Extension trait for TensorSlice to convert to Vec types +/// Converts a tensor slice into nested `Vec`s matching its rank. pub trait ToVec { type Output; fn to_vec(&self) -> Self::Output; } -/// Extension trait for TensorSlice to convert to Vec types -pub trait ToVec1 { - fn to_vec1(&self) -> Vec; -} - -/// Extension trait for TensorSlice to convert to Vec types -pub trait ToVec2 { - fn to_vec2(&self) -> Vec>; -} - -/// Extension trait for TensorSlice to convert to Vec types -pub trait ToVec3 { - fn to_vec3(&self) -> Vec>>; -} - -impl> ToVec1 +impl> ToVec for TensorSlice<1, D, Bytes> { - /// Convert a 1D tensor slice to a `Vec` - fn to_vec1(&self) -> Vec { + type Output = Vec; + + fn to_vec(&self) -> Self::Output { let shape = self.shape(); let len = shape[0]; @@ -44,21 +30,11 @@ impl> ToVec1 } impl> ToVec - for TensorSlice<1, D, Bytes> + for TensorSlice<2, D, Bytes> { - type Output = Vec; + type Output = Vec>; - /// Convert a 1D tensor slice to a `Vec` fn to_vec(&self) -> Self::Output { - self.to_vec1() - } -} - -impl> ToVec2 - for TensorSlice<2, D, Bytes> -{ - /// Convert a 2D tensor slice to a `Vec>` - fn to_vec2(&self) -> Vec> { let shape = self.shape(); let rows = shape[0]; let cols = shape[1]; @@ -76,21 +52,11 @@ impl> ToVec2 } impl> ToVec - for TensorSlice<2, D, Bytes> + for TensorSlice<3, D, Bytes> { - type Output = Vec>; + type Output = Vec>>; - /// Convert a 2D tensor slice to a `Vec>` fn to_vec(&self) -> Self::Output { - self.to_vec2() - } -} - -impl> ToVec3 - for TensorSlice<3, D, Bytes> -{ - /// Convert a 3D tensor slice to a `Vec>>` - fn to_vec3(&self) -> Vec>> { let shape = self.shape(); let dim0 = shape[0]; let dim1 = shape[1]; @@ -111,14 +77,3 @@ impl> ToVec3 result } } - -impl> ToVec - for TensorSlice<3, D, Bytes> -{ - type Output = Vec>>; - - /// Convert a 3D tensor slice to a `Vec>>` - fn to_vec(&self) -> Self::Output { - self.to_vec3() - } -} diff --git a/fusor-ml/fusor/src/device.rs b/fusor-ml/fusor/src/device.rs index eb8648883..81655d4c8 100644 --- a/fusor-ml/fusor/src/device.rs +++ b/fusor-ml/fusor/src/device.rs @@ -38,6 +38,14 @@ impl Device { } } + /// Create a new GPU device with an explicit [`crate::FusorConfig`] + /// instead of reading the process environment. + #[cfg(feature = "gpu")] + pub async fn gpu_with_config(config: crate::FusorConfig) -> Result { + let device = crate::gpu::Device::new_with_config(config).await?; + Ok(Device::Gpu(device)) + } + /// Create a new GPU device, blocking until ready. pub fn gpu_blocking() -> Result { #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))] @@ -56,6 +64,18 @@ impl Device { } } + /// Resolve every pending lazy tensor now, submitting the work to the GPU + /// without waiting for it or downloading anything. Call at iteration + /// boundaries in training-style loops to keep the pending graph bounded + /// while the GPU runs ahead of the host. No-op on CPU, where lazy + /// expressions are evaluated when they are consumed. + pub fn flush(&self) { + match self { + Device::Cpu => {} + Device::Gpu(device) => device.flush(), + } + } + /// Create a device, preferring GPU if available, otherwise falling back to CPU. pub async fn auto() -> Self { #[cfg(not(feature = "gpu"))] @@ -66,10 +86,10 @@ impl Device { match Self::gpu().await { Ok(gpu) => gpu, Err(err) => { - if std::env::var_os("KALOSM_TRACE_DECODE_TIMING").is_some() - || std::env::var_os("FUSOR_TRACE_DECODE").is_some() - || std::env::var_os("FUSOR_TRACE_RESOLVE").is_some() - { + // No device exists to carry a config when creation fails, so + // this bootstrap diagnostic reads the environment directly. + let config = crate::FusorConfig::from_env(); + if config.trace_decode_timing || config.trace_decode || config.trace_resolve { tracing::warn!("fusor_device_auto_gpu_error={err}"); } Device::Cpu diff --git a/fusor-ml/fusor/src/gpu.rs b/fusor-ml/fusor/src/gpu.rs index 615368c84..e36611f5c 100644 --- a/fusor-ml/fusor/src/gpu.rs +++ b/fusor-ml/fusor/src/gpu.rs @@ -7,9 +7,10 @@ use std::{ pub use fusor_core::{ CastTensor, DataType, DataTypeEnum, Device, Dim, Error, FloatDataType, GgufReadError, - GpuMirostat2Sampler, GpuMirostat2SamplerParams, GpuStandardSamplerParams, Layout, MappedBuffer, - MatMulParams, NodeIndex, PendingGpuSampledToken as CorePendingGpuSampledToken, QMatrix, Result, - ShapeWithOneHole, StrideSpec, TensorSlice, WasmNotSend, WasmNotSync, + GpuMirostat2Sampler, GpuMirostat2SamplerParams, GpuStandardSamplerParams, KernelProfile, + KernelProfileRow, Layout, MappedBuffer, NodeIndex, + PendingGpuSampledToken as CorePendingGpuSampledToken, QMatrix, Result, ShapeWithOneHole, + StrideSpec, TensorSlice, WasmNotSend, WasmNotSync, }; type CoreTensor = fusor_core::Tensor; @@ -156,6 +157,11 @@ impl Tensor { self.inner.count_kernels_to_resolve() } + #[inline] + pub fn resolves_in(&self) -> bool { + self.inner.resolves_in::() + } + #[inline] pub async fn as_slice(&self) -> Result> { self.inner.as_slice::().await.map_err(Error::from) @@ -261,14 +267,6 @@ impl Tensor { Self::from_core(self.inner.mat_mul(other.as_core())) } - #[inline] - pub fn mat_mul_with_parameters(&self, other: &Self, parameters: MatMulParams) -> Self { - Self::from_core( - self.inner - .mat_mul_with_parameters(other.as_core(), parameters), - ) - } - #[inline] pub fn sum(&self, dim: impl Dim) -> Tensor { Tensor::from_core(self.inner.sum(dim.resolve())) @@ -334,6 +332,36 @@ impl Tensor { Tensor::from_core(self.inner.mte::(rhs)) } + #[inline] + pub fn eq_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.eq_tensor::(other.as_core())) + } + + #[inline] + pub fn ne_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.ne_tensor::(other.as_core())) + } + + #[inline] + pub fn lt_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.lt_tensor::(other.as_core())) + } + + #[inline] + pub fn lte_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.lte_tensor::(other.as_core())) + } + + #[inline] + pub fn gt_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.gt_tensor::(other.as_core())) + } + + #[inline] + pub fn gte_tensor(&self, other: &Self) -> Tensor { + Tensor::from_core(self.inner.gte_tensor::(other.as_core())) + } + #[inline] pub fn cast(&self) -> Tensor where @@ -633,22 +661,62 @@ impl Tensor<1, f32> { impl Tensor<4, T> { #[inline] - pub fn flash_attention_causal(&self, k: &Self, v: &Self, scale: f32) -> Self { - Self::from_core( - self.inner - .flash_attention_causal(k.as_core(), v.as_core(), scale), - ) + pub fn attention_causal(&self, k: &Self, v: &Self, scale: f32) -> Self { + Self::from_core(self.inner.attention_causal(k.as_core(), v.as_core(), scale)) } + /// Row log-sum-exp of the attention scores over the KV axis. #[inline] - pub fn flash_attention( + pub fn attention_lse( + &self, + k: &Self, + scale: f32, + mask: Option<&Tensor<2, T>>, + causal: bool, + ) -> Tensor<3, T> { + Tensor::from_core(self.inner.attention_lse( + k.as_core(), + scale, + mask.map(|mask| mask.as_core()), + causal, + )) + } + + /// Gradients of [`Self::attention`] with respect to q, k, and v, + /// recomputed from the forward output and its row log-sum-exp. + #[inline] + #[allow(clippy::too_many_arguments)] + pub fn attention_grads( &self, k: &Self, v: &Self, + o: &Self, + grad_o: &Self, + lse: &Tensor<3, T>, scale: f32, mask: Option<&Tensor<2, T>>, - ) -> Self { - Self::from_core(self.inner.flash_attention( + causal: bool, + ) -> (Self, Self, Self) { + let (dq, dk, dv) = self.inner.attention_grads( + k.as_core(), + v.as_core(), + o.as_core(), + grad_o.as_core(), + lse.as_core(), + scale, + mask.map(|mask| mask.as_core()), + causal, + ); + ( + Self::from_core(dq), + Self::from_core(dk), + Self::from_core(dv), + ) + } + + #[inline] + pub fn attention(&self, k: &Self, v: &Self, scale: f32, mask: Option<&Tensor<2, T>>) -> Self { + Self::from_core(self.inner.attention( k.as_core(), v.as_core(), scale, diff --git a/fusor-ml/fusor/src/layers/conv/autograd.rs b/fusor-ml/fusor/src/layers/conv/autograd.rs new file mode 100644 index 000000000..03c3eee88 --- /dev/null +++ b/fusor-ml/fusor/src/layers/conv/autograd.rs @@ -0,0 +1,161 @@ +//! Trainable N-dimensional convolution layer. + +use crate::autograd::{AutogradElement, Graph, Tensor}; + +pub use crate::layers::ConvNdConfig; + +/// N-dimensional convolution layer with trainable parameters. +/// +/// Input / output tensors have rank `RANK = SPATIAL + 2`: +/// `(batch, channels, ...spatial)`. +/// Weight has shape `(out_channels, in_channels / groups, ...kernel)`. +pub struct ConvNd { + weight: Tensor, + bias: Option>, + config: ConvNdConfig, + in_channels: usize, + out_channels: usize, +} + +impl ConvNd +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, +{ + /// Create a new convolution layer. + /// + /// `weight` shape: `(out_channels, in_channels / groups, ...kernel)`. + /// `bias` shape: `(out_channels,)`. + pub fn new( + weight: Tensor, + bias: Option>, + config: ConvNdConfig, + ) -> Self { + const { + assert!(RANK == SPATIAL + 2); + } + let shape = weight.shape(); + let out_channels = shape[0]; + let in_channels = shape[1] * config.groups; + + assert!(config.groups >= 1, "groups must be >= 1"); + assert_eq!( + out_channels % config.groups, + 0, + "out_channels ({out_channels}) must be divisible by groups ({})", + config.groups + ); + + if let Some(ref b) = bias { + assert_eq!( + b.shape()[0], + out_channels, + "Bias shape must match out_channels" + ); + } + + Self { + weight, + bias, + config, + in_channels, + out_channels, + } + } + + /// Get the configuration. + pub fn config(&self) -> &ConvNdConfig { + &self.config + } + + /// Number of input channels. + pub fn in_channels(&self) -> usize { + self.in_channels + } + + /// Number of output channels. + pub fn out_channels(&self) -> usize { + self.out_channels + } + + /// The weight leaf: `(out_channels, in_channels / groups, ...kernel)`. + pub fn weight(&self) -> &Tensor { + &self.weight + } + + /// The bias leaf: `(out_channels,)`. + pub fn bias(&self) -> Option<&Tensor<1, T>> { + self.bias.as_ref() + } + + /// Forward pass for any spatial rank. The free const generic `R2` equals + /// `RANK + SPATIAL` and is determined by the `LargerRank` bound, exactly + /// the same way the underlying `conv` operation infers it. + pub fn forward(&self, input: &Tensor) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LargerRank, + crate::gpu::Tensor: crate::gpu::LargerRank, + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + { + if self.config.groups == 1 { + input.conv( + &self.weight, + self.bias.as_ref(), + self.config.padding, + self.config.stride, + ) + } else { + input.grouped_conv( + &self.weight, + self.bias.as_ref(), + self.config.padding, + self.config.stride, + self.config.groups, + ) + } + } +} + +impl ConvNd { + /// Import an inference layer's parameters as trainable leaves on `graph`. + pub fn from_inference( + graph: &Graph, + layer: &crate::layers::ConvNd, + ) -> Self { + Self::new( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + *layer.config(), + ) + } +} diff --git a/fusor-ml/fusor/src/layers/conv.rs b/fusor-ml/fusor/src/layers/conv/mod.rs similarity index 94% rename from fusor-ml/fusor/src/layers/conv.rs rename to fusor-ml/fusor/src/layers/conv/mod.rs index 9786fd554..065000c5f 100644 --- a/fusor-ml/fusor/src/layers/conv.rs +++ b/fusor-ml/fusor/src/layers/conv/mod.rs @@ -1,5 +1,7 @@ //! N-dimensional convolution layer. +pub(crate) mod autograd; + use crate::fusion::Concrete; use crate::{ DataType, Device, FloatDataType, FloatOps, Fusion, MatmulImpl, SimdElement, Tensor, VarBuilder, @@ -98,6 +100,16 @@ where pub fn out_channels(&self) -> usize { self.out_channels } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor> { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor<1, D, Concrete>> { + self.bias.as_ref() + } } impl ConvNd diff --git a/fusor-ml/fusor/src/layers/conv1d.rs b/fusor-ml/fusor/src/layers/conv1d.rs deleted file mode 100644 index 6e4d90cc9..000000000 --- a/fusor-ml/fusor/src/layers/conv1d.rs +++ /dev/null @@ -1,125 +0,0 @@ -//! Conv1d layer implementation. - -use crate::fusion::Concrete; -use crate::{DataType, FloatDataType, FloatOps, MatmulImpl, SimdElement, Tensor}; - -/// Configuration for Conv1d layer. -#[derive(Debug, Clone, Copy)] -pub struct Conv1dConfig { - pub padding: usize, - pub stride: usize, - pub groups: usize, - pub dilation: usize, -} - -impl Default for Conv1dConfig { - fn default() -> Self { - Self { - padding: 0, - stride: 1, - groups: 1, - dilation: 1, - } - } -} - -/// 1D Convolution layer. -/// -/// Applies a 1D convolution over an input signal. -/// Input shape: (batch, in_channels, length) -/// Output shape: (batch, out_channels, out_length) -/// where out_length = (length + 2*padding - kernel_size) / stride + 1 -pub struct Conv1d { - weight: Tensor<3, D, Concrete>, // (out_channels, in_channels, kernel_size) - bias: Option>>, // (out_channels,) - config: Conv1dConfig, - in_channels: usize, - out_channels: usize, - kernel_size: usize, -} - -impl Conv1d -where - D: SimdElement - + DataType - + FloatDataType - + FloatOps - + Default - + MatmulImpl - + std::ops::Mul - + std::ops::Add, -{ - /// Create a new Conv1d layer with given weights and configuration. - /// - /// Weight shape: (out_channels, in_channels, kernel_size) - /// Bias shape: (out_channels,) - pub fn new( - weight: Tensor<3, D, Concrete>, - bias: Option>>, - config: Conv1dConfig, - ) -> Self { - let shape = weight.shape(); - let out_channels = shape[0]; - let in_channels = shape[1]; - let kernel_size = shape[2]; - - // Validate configuration - assert_eq!(config.groups, 1, "Only groups=1 is currently supported"); - assert_eq!(config.dilation, 1, "Only dilation=1 is currently supported"); - - if let Some(ref b) = bias { - assert_eq!( - b.shape()[0], - out_channels, - "Bias shape must match out_channels" - ); - } - - Self { - weight, - bias, - config, - in_channels, - out_channels, - kernel_size, - } - } - - /// Forward pass. - /// - /// Input shape: (batch, in_channels, length) - /// Output shape: (batch, out_channels, out_length) - pub fn forward(&self, input: &Tensor<3, D, Concrete>) -> Tensor<3, D, Concrete> - where - crate::MulOp: crate::SimdBinaryOp, - crate::AddOp: crate::SimdBinaryOp, - crate::SumOp: crate::SimdReduceOp, - { - input.conv( - &self.weight, - self.bias.as_ref(), - [self.config.padding], - [self.config.stride], - ) - } - - /// Get the configuration. - pub fn config(&self) -> &Conv1dConfig { - &self.config - } - - /// Get the number of input channels. - pub fn in_channels(&self) -> usize { - self.in_channels - } - - /// Get the number of output channels. - pub fn out_channels(&self) -> usize { - self.out_channels - } - - /// Get the kernel size. - pub fn kernel_size(&self) -> usize { - self.kernel_size - } -} diff --git a/fusor-ml/fusor/src/layers/embedding/autograd.rs b/fusor-ml/fusor/src/layers/embedding/autograd.rs new file mode 100644 index 000000000..24898d791 --- /dev/null +++ b/fusor-ml/fusor/src/layers/embedding/autograd.rs @@ -0,0 +1,120 @@ +//! Trainable embedding layer implementation. + +use crate::Tensor as RawTensor; +use crate::autograd::{AutogradElement, Graph, Tensor}; + +/// Embedding layer for token/position embeddings. +/// +/// Maps integer indices to dense vectors. +/// Embedding table shape: (num_embeddings, embedding_dim) +#[derive(Clone)] +pub struct Embedding { + embeddings: Tensor<2, T>, + num_embeddings: usize, + embedding_dim: usize, +} + +impl Embedding +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, +{ + /// Create a new embedding layer with the given embedding table. + pub fn new_from_tensor(embeddings: Tensor<2, T>) -> Self { + let shape = embeddings.shape(); + let num_embeddings = shape[0]; + let embedding_dim = shape[1]; + + Self { + embeddings, + num_embeddings, + embedding_dim, + } + } + + /// Get the embedding table. + pub fn embeddings(&self) -> &Tensor<2, T> { + &self.embeddings + } + + /// Get the number of embeddings. + pub fn num_embeddings(&self) -> usize { + self.num_embeddings + } + + /// Get the embedding dimension. + pub fn embedding_dim(&self) -> usize { + self.embedding_dim + } +} + +impl Embedding { + /// Import an inference embedding layer as a trainable layer whose + /// embedding table is a gradient leaf on `graph`. + pub fn from_inference(graph: &Graph, layer: &crate::layers::Embedding) -> Self { + let table = match layer.dense_embeddings() { + Some(dense) => dense.clone(), + None => layer.embeddings_quantized().dequantize().into_concrete(), + }; + Self::new_from_tensor(graph.leaf(table)) + } +} + +impl Embedding +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + u32: crate::CastTo + crate::CastTensor, +{ + /// Looks up embeddings for an index tensor, appending the embedding dimension. + pub fn forward( + &self, + indices: &RawTensor, + ) -> Tensor + where + crate::gpu::Tensor: crate::gpu::NextRank, + { + assert_eq!(M, N + 1, "embedding output rank must be input rank + 1"); + let input_shape = indices.shape(); + let output_shape = std::array::from_fn(|axis| { + if axis < N { + input_shape[axis] + } else { + self.embedding_dim + } + }); + let indices = indices.flatten_all().to_concrete(); + self.embeddings + .index_select(0, &indices) + .reshape(output_shape) + } +} diff --git a/fusor-ml/fusor/src/layers/embedding.rs b/fusor-ml/fusor/src/layers/embedding/mod.rs similarity index 97% rename from fusor-ml/fusor/src/layers/embedding.rs rename to fusor-ml/fusor/src/layers/embedding/mod.rs index 73c2f3bf0..d2b1553f7 100644 --- a/fusor-ml/fusor/src/layers/embedding.rs +++ b/fusor-ml/fusor/src/layers/embedding/mod.rs @@ -1,5 +1,7 @@ //! Embedding layer implementation. +pub(crate) mod autograd; + use crate::{ CastTensor, CastTo, DataType, Device, Fusion, QMatrix, SimdElement, Tensor, VarBuilder, }; @@ -115,6 +117,11 @@ impl Embedding { .expect("dense embeddings unavailable for this embedding table") } + /// Get the dense embedding table, if this layer holds one. + pub fn dense_embeddings(&self) -> Option<&Tensor<2, T>> { + self.embeddings.as_ref() + } + /// Get the number of embeddings. pub fn num_embeddings(&self) -> usize { self.num_embeddings diff --git a/fusor-ml/fusor/src/layers/layer_norm/autograd.rs b/fusor-ml/fusor/src/layers/layer_norm/autograd.rs new file mode 100644 index 000000000..ee8ab6e21 --- /dev/null +++ b/fusor-ml/fusor/src/layers/layer_norm/autograd.rs @@ -0,0 +1,251 @@ +//! Trainable layer normalization. + +use crate::autograd::{AutogradElement, Graph, Tensor}; + +/// Layer Normalization. +/// +/// Normalizes the input over the last dimension. +/// Formula: output = (input - mean) / sqrt(variance + eps) * weight + bias +pub struct LayerNorm { + weight: Tensor, + bias: Option>, + eps: f32, +} + +impl LayerNorm { + /// Create a new LayerNorm layer. + /// + /// Weight and bias should have shape (normalized_dim,). + pub fn new(weight: Tensor, bias: Option>, eps: f32) -> Self { + Self { weight, bias, eps } + } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor> { + self.bias.as_ref() + } + + /// Get the epsilon value. + pub fn eps(&self) -> f32 { + self.eps + } +} + +impl LayerNorm { + /// Import an inference [`crate::layers::LayerNorm`]'s weights as trainable leaves. + pub fn from_inference(graph: &Graph, layer: &crate::layers::LayerNorm) -> Self { + Self::new( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + layer.eps(), + ) + } +} + +impl LayerNorm<1, T> +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + /// Normalizes the last dimension of an input tensor. + pub fn forward( + &self, + input: &Tensor, + ) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + { + input.layer_norm_last_dim_fused::(&self.weight, self.bias.as_ref(), self.eps) + } +} + +/// Layer normalization with a selectable reduction axis. +/// +/// `axis == None` normalizes the last dimension. `axis == Some(a)` normalizes +/// dimension `a` by transposing that axis to the end, applying the common +/// last-dimension path, then transposing back. +pub struct LayerNormNd { + weight: Tensor<1, T>, + bias: Option>, + axis: Option, + eps: f32, +} + +impl LayerNormNd { + /// Create a LayerNorm that normalizes the last dimension. + pub fn new(weight: Tensor<1, T>, bias: Option>, eps: f32) -> Self { + Self { + weight, + bias, + axis: None, + eps, + } + } + + /// Create a LayerNorm that normalizes the given axis. + pub fn new_over_axis( + weight: Tensor<1, T>, + bias: Option>, + axis: usize, + eps: f32, + ) -> Self { + Self { + weight, + bias, + axis: Some(axis), + eps, + } + } + + pub fn weight(&self) -> &Tensor<1, T> { + &self.weight + } + + pub fn bias(&self) -> Option<&Tensor<1, T>> { + self.bias.as_ref() + } + + pub fn eps(&self) -> f32 { + self.eps + } +} + +impl LayerNormNd { + /// Import an inference [`crate::layers::LayerNormNd`]'s weights as trainable + /// leaves, normalizing the last dimension. + pub fn from_inference(graph: &Graph, layer: &crate::layers::LayerNormNd) -> Self { + Self::new( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + layer.eps(), + ) + } + + /// Import an inference [`crate::layers::LayerNormNd`]'s weights as trainable + /// leaves, normalizing `axis`. + pub fn from_inference_over_axis( + graph: &Graph, + layer: &crate::layers::LayerNormNd, + axis: usize, + ) -> Self { + Self::new_over_axis( + graph.leaf(layer.weight().clone()), + layer.bias().map(|bias| graph.leaf(bias.clone())), + axis, + layer.eps(), + ) + } +} + +impl LayerNormNd +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + /// Forward pass for any input rank. `OUT_RANK` equals `N - 1`. + pub fn forward( + &self, + input: &Tensor, + ) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + { + let shape = input.shape(); + let axis = self.axis.unwrap_or(N - 1); + + if axis == N - 1 { + let weight_b = self.weight.broadcast_as(shape); + let bias_b = self.bias.as_ref().map(|bias| bias.broadcast_as(shape)); + return input.layer_norm::(&weight_b, bias_b.as_ref(), self.eps, true); + } + + let mut permuted_shape = shape; + permuted_shape.swap(axis, N - 1); + let permuted = input.transpose(axis, N - 1); + let weight_b = self.weight.broadcast_as(permuted_shape); + let bias_b = self + .bias + .as_ref() + .map(|bias| bias.broadcast_as(permuted_shape)); + let normed = permuted.layer_norm::(&weight_b, bias_b.as_ref(), self.eps, true); + normed.transpose(axis, N - 1) + } + + /// Fused fast path for normalizing the last dim of a rank-3 tensor. + pub fn forward_fused(&self, input: &Tensor<3, T>) -> Tensor<3, T> { + if matches!(self.axis, Some(axis) if axis != 2) { + return self.forward::<3, 2>(input); + } + + input.layer_norm_last_dim_fused::<2, 1>(&self.weight, self.bias.as_ref(), self.eps) + } +} diff --git a/fusor-ml/fusor/src/layers/layer_norm.rs b/fusor-ml/fusor/src/layers/layer_norm/mod.rs similarity index 81% rename from fusor-ml/fusor/src/layers/layer_norm.rs rename to fusor-ml/fusor/src/layers/layer_norm/mod.rs index 65a0aa55d..7a743c2b9 100644 --- a/fusor-ml/fusor/src/layers/layer_norm.rs +++ b/fusor-ml/fusor/src/layers/layer_norm/mod.rs @@ -1,5 +1,7 @@ //! Layer normalization implementation. +pub(crate) mod autograd; + use crate::fusion::Concrete; use crate::{ DataType, Device, DivOp, FloatDataType, FloatOps, Fusion, MulOp, SimdBinaryOp, SimdElement, @@ -99,39 +101,11 @@ impl LayerNorm<1, D> where D: SimdElement + DataType + FloatDataType + FloatOps + Default, { - /// Forward pass for 2D input (batch, features). - /// - /// Normalizes over the last dimension (features). - pub fn forward_2d(&self, input: &Tensor<2, D, B>) -> Tensor<2, D, Concrete> - where - D: std::ops::Add - + std::ops::Sub - + std::ops::Mul - + std::ops::Div, - crate::AddOp: SimdBinaryOp, - SubOp: SimdBinaryOp, - MulOp: SimdBinaryOp, - DivOp: SimdBinaryOp, - SumOp: SimdReduceOp, - SqrtOp: SimdUnaryOp, - B: Fusion<2, D>, - { - // Broadcast weight to input shape - let weight_broadcast: Tensor<2, D, _> = self.weight.broadcast_as(input.shape()); - let bias_broadcast: Option> = - self.bias.as_ref().map(|b| b.broadcast_as(input.shape())); - input.layer_norm( - &weight_broadcast, - bias_broadcast.as_ref(), - D::from_f32(self.eps), - true, - ) - } - - /// Forward pass for 3D input (batch, seq_len, features). - /// - /// Normalizes over the last dimension (features). - pub fn forward(&self, input: &Tensor<3, D, B>) -> Tensor<3, D, Concrete> + /// Normalizes the last dimension of an input tensor. + pub fn forward( + &self, + input: &Tensor, + ) -> Tensor> where D: std::ops::Add + std::ops::Sub @@ -143,11 +117,14 @@ where DivOp: SimdBinaryOp, SumOp: SimdReduceOp, SqrtOp: SimdUnaryOp, - B: Fusion<3, D>, + B: Fusion, + Concrete: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, { - // Broadcast weight to input shape - let weight_broadcast: Tensor<3, D, _> = self.weight.broadcast_as(input.shape()); - let bias_broadcast: Option> = + let weight_broadcast: Tensor = self.weight.broadcast_as(input.shape()); + let bias_broadcast: Option> = self.bias.as_ref().map(|b| b.broadcast_as(input.shape())); input.layer_norm( &weight_broadcast, @@ -239,24 +216,6 @@ where self.eps } - /// Forward pass for 2D input. - pub fn forward_2d(&self, input: &Tensor<2, D, B>) -> Tensor<2, D, Concrete> - where - D: std::ops::Add - + std::ops::Sub - + std::ops::Mul - + std::ops::Div, - crate::AddOp: SimdBinaryOp, - SubOp: SimdBinaryOp, - MulOp: SimdBinaryOp, - DivOp: SimdBinaryOp, - SumOp: SimdReduceOp, - SqrtOp: SimdUnaryOp, - B: Fusion<2, D>, - { - self.forward(input) - } - /// Forward pass for any input rank. `OUT_RANK` equals `N - 1`. pub fn forward( &self, diff --git a/fusor-ml/fusor/src/layers/linear/autograd.rs b/fusor-ml/fusor/src/layers/linear/autograd.rs new file mode 100644 index 000000000..88571df26 --- /dev/null +++ b/fusor-ml/fusor/src/layers/linear/autograd.rs @@ -0,0 +1,112 @@ +//! Trainable linear layer implementation. + +use crate::autograd::{AutogradElement, Graph, Tensor}; + +/// A trainable linear (fully connected) layer. +/// +/// Computes `output = input @ weight.T + bias`. +pub struct Linear { + weight: Tensor<2, T>, + bias: Option>, +} + +impl Linear +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, +{ + /// Create a new Linear layer with the given weight and optional bias. + /// + /// Weight shape: (out_features, in_features) + /// Bias shape: (out_features,) + pub fn new(weight: Tensor<2, T>, bias: Option>) -> Self { + Self { weight, bias } + } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor<2, T> { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor<1, T>> { + self.bias.as_ref() + } + + /// Get the input features size. + pub fn in_features(&self) -> usize { + self.weight.shape()[1] + } + + /// Get the output features size. + pub fn out_features(&self) -> usize { + self.weight.shape()[0] + } +} + +impl Linear { + /// Import an inference layer as trainable graph leaves, dequantizing the weight to f32. + pub fn from_inference(graph: &Graph, layer: &crate::layers::Linear) -> Self { + let weight = graph.leaf(layer.weight().dequantize::<2>().into_concrete()); + let bias = layer.bias().map(|bias| graph.leaf(bias.clone())); + Self { weight, bias } + } +} + +impl Linear +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, +{ + /// Applies the linear projection to the last dimension of an input tensor. + pub fn forward(&self, input: &Tensor) -> Tensor { + assert!(R >= 2, "linear forward requires rank >= 2"); + + let input_shape = input.shape(); + let rows = input_shape[..R - 1].iter().product(); + let input_2d = input.reshape([rows, input_shape[R - 1]]); + let output_2d = input_2d.mat_mul_transposed_rhs(&self.weight); + let output_2d = if let Some(bias) = &self.bias { + output_2d.add_(bias) + } else { + output_2d + }; + let output_shape = std::array::from_fn(|axis| { + if axis == R - 1 { + self.out_features() + } else { + input_shape[axis] + } + }); + output_2d.reshape(output_shape) + } +} diff --git a/fusor-ml/fusor/src/layers/linear.rs b/fusor-ml/fusor/src/layers/linear/mod.rs similarity index 62% rename from fusor-ml/fusor/src/layers/linear.rs rename to fusor-ml/fusor/src/layers/linear/mod.rs index a0b6aa13e..5cedef0fa 100644 --- a/fusor-ml/fusor/src/layers/linear.rs +++ b/fusor-ml/fusor/src/layers/linear/mod.rs @@ -1,5 +1,7 @@ //! Linear layer implementation. +pub(crate) mod autograd; + use crate::{ CastTensor, CastTo, DataType, Device, Fusion, GgmlType, QMatrix, SimdElement, Tensor, VarBuilder, @@ -59,20 +61,6 @@ impl Linear { } } -/// Run `q_mat_mul(weight)` over a 3D input by flattening (batch, seq, in) to -/// (batch * seq, in), multiplying, then reshaping back. Shared by both the f32 -/// and the cast-via-f32 forward paths. -fn q_mat_mul_3d(input: &Tensor<3, f32>, weight: &QMatrix) -> Tensor<3, f32> { - let [batch, seq_len, in_features] = input.shape(); - let out_features = weight.shape()[0]; - - let input_2d: Tensor<2, f32> = input.reshape([batch * seq_len, in_features]).to_concrete(); - let output_2d = input_2d.q_mat_mul(weight); - output_2d - .reshape([batch, seq_len, out_features]) - .to_concrete() -} - // f32-specific implementations for loading and forward impl Linear { /// Load a Linear layer from a VarBuilder. @@ -86,34 +74,16 @@ impl Linear { Ok(Self { weight, bias }) } - /// Forward pass for 3D input (batch, seq_len, in_features) - /// - /// Input shape: (batch, seq_len, in_features) - /// Output shape: (batch, seq_len, out_features) - pub fn forward(&self, input: &Tensor<3, f32, B>) -> Tensor<3, f32> - where - B: Fusion<3, f32>, - { - let input_concrete = input.to_concrete(); - let output = q_mat_mul_3d(&input_concrete, &self.weight); - if let Some(bias) = &self.bias { - output.add_(bias) - } else { - output - } - } - /// Forward pass for 2D input (batch, in_features) - /// - /// Input shape: (batch, in_features) - /// Output shape: (batch, out_features) - pub fn forward_2d(&self, input: &Tensor<2, f32, B>) -> Tensor<2, f32> + /// Applies the linear projection to the last dimension of an input tensor. + pub fn forward(&self, input: &Tensor) -> Tensor where - B: Fusion<2, f32>, + B: Fusion, + (crate::gpu::Tensor, crate::gpu::Tensor<1, f32>): crate::gpu::MaxRank, + (crate::ConcreteTensor, crate::ConcreteTensor): crate::cpu::MaxRank, { let output = input.q_mat_mul(&self.weight); - if let Some(bias) = &self.bias { - output.add_(bias) + output.add_::<1, R, _>(bias) } else { output } @@ -127,17 +97,18 @@ where T: CastTo + CastTensor, f32: CastTo + CastTensor, { - /// Forward pass for 3D input with generic type. - /// Converts input to f32 for computation, then converts back. - pub fn forward_generic(&self, input: &Tensor<3, T, B>) -> Tensor<3, T> + /// Applies the linear projection after converting the input to f32. + pub fn forward_generic(&self, input: &Tensor) -> Tensor where - B: Fusion<3, T>, + B: Fusion, + (crate::gpu::Tensor, crate::gpu::Tensor<1, f32>): crate::gpu::MaxRank, + (crate::ConcreteTensor, crate::ConcreteTensor): crate::cpu::MaxRank, { let input_f32 = input.cast::(); - let output_f32 = q_mat_mul_3d(&input_f32, &self.weight); + let output_f32 = input_f32.q_mat_mul(&self.weight); let output_f32 = if let Some(bias) = &self.bias { let bias_f32: Tensor<1, f32> = bias.cast(); - output_f32.add_(&bias_f32) + output_f32.add_::<1, R, _>(&bias_f32) } else { output_f32 }; diff --git a/fusor-ml/fusor/src/layers/mod.rs b/fusor-ml/fusor/src/layers/mod.rs index 3a0a4f697..775868d22 100644 --- a/fusor-ml/fusor/src/layers/mod.rs +++ b/fusor-ml/fusor/src/layers/mod.rs @@ -4,15 +4,13 @@ //! //! All layers support loading from GGUF files via `VarBuilder` for f32 types. -mod conv; -mod conv1d; -mod embedding; -mod layer_norm; -mod linear; -mod rms_norm; +pub(crate) mod conv; +pub(crate) mod embedding; +pub(crate) mod layer_norm; +pub(crate) mod linear; +pub(crate) mod rms_norm; pub use conv::{ConvNd, ConvNdConfig}; -pub use conv1d::{Conv1d, Conv1dConfig}; pub use embedding::Embedding; pub use layer_norm::{LayerNorm, LayerNormNd}; pub use linear::Linear; diff --git a/fusor-ml/fusor/src/layers/rms_norm/autograd.rs b/fusor-ml/fusor/src/layers/rms_norm/autograd.rs new file mode 100644 index 000000000..9c8f42d64 --- /dev/null +++ b/fusor-ml/fusor/src/layers/rms_norm/autograd.rs @@ -0,0 +1,107 @@ +//! Trainable RMS normalization implementation. + +use crate::autograd::{AutogradElement, Graph, Tensor}; + +/// Root Mean Square Normalization. +/// +/// Normalizes the input over the last dimension without centering. +/// Formula: output = input / sqrt(mean(x^2) + eps) * weight +pub struct RmsNorm { + weight: Tensor, + bias: Option>, + eps: f32, +} + +impl RmsNorm { + /// Create a new RmsNorm layer. + /// + /// Weight should have shape matching the normalized dimension. + pub fn new(weight: Tensor, bias: Option>, eps: f32) -> Self { + Self { weight, bias, eps } + } + + /// Get the weight tensor. + pub fn weight(&self) -> &Tensor { + &self.weight + } + + /// Get the bias tensor if present. + pub fn bias(&self) -> Option<&Tensor> { + self.bias.as_ref() + } + + /// Get the epsilon value. + pub fn eps(&self) -> f32 { + self.eps + } +} + +impl RmsNorm { + /// Import a raw inference layer's weights as trainable graph leaves. + pub fn from_inference(graph: &Graph, layer: &crate::layers::RmsNorm) -> Self { + Self { + weight: graph.leaf(layer.weight().clone()), + bias: layer.bias().map(|bias| graph.leaf(bias.clone())), + eps: layer.eps(), + } + } +} + +impl RmsNorm<1, T> +where + crate::cpu::AddOp: crate::cpu::SimdBinaryOp, + crate::cpu::SubOp: crate::cpu::SimdBinaryOp, + crate::cpu::MulOp: crate::cpu::SimdBinaryOp, + crate::cpu::DivOp: crate::cpu::SimdBinaryOp, + crate::cpu::EqOp: crate::cpu::SimdBinaryOp, + crate::cpu::NeOp: crate::cpu::SimdBinaryOp, + crate::cpu::LtOp: crate::cpu::SimdBinaryOp, + crate::cpu::LteOp: crate::cpu::SimdBinaryOp, + crate::cpu::GtOp: crate::cpu::SimdBinaryOp, + crate::cpu::GteOp: crate::cpu::SimdBinaryOp, + crate::cpu::SumOp: crate::cpu::SimdReduceOp, + crate::cpu::MaxOp: crate::cpu::SimdReduceOp, + crate::cpu::MinOp: crate::cpu::SimdReduceOp, + crate::cpu::NegOp: crate::cpu::SimdUnaryOp, + crate::cpu::AbsOp: crate::cpu::SimdUnaryOp, + crate::cpu::SqrtOp: crate::cpu::SimdUnaryOp, + crate::cpu::ExpOp: crate::cpu::SimdUnaryOp, + crate::cpu::Exp2Op: crate::cpu::SimdUnaryOp, + crate::cpu::LogOp: crate::cpu::SimdUnaryOp, + crate::cpu::Log2Op: crate::cpu::SimdUnaryOp, + crate::cpu::SinOp: crate::cpu::SimdUnaryOp, + crate::cpu::CosOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanOp: crate::cpu::SimdUnaryOp, + crate::cpu::TanhOp: crate::cpu::SimdUnaryOp, + crate::cpu::SinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::CoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcosOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanOp: crate::cpu::SimdUnaryOp, + crate::cpu::AsinhOp: crate::cpu::SimdUnaryOp, + crate::cpu::AcoshOp: crate::cpu::SimdUnaryOp, + crate::cpu::AtanhOp: crate::cpu::SimdUnaryOp, + T: crate::CastTensor, + f32: crate::CastTensor, +{ + /// Normalizes the last dimension of an input tensor. + pub fn forward( + &self, + input: &Tensor, + ) -> Tensor + where + crate::ConcreteTensor: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + (crate::gpu::Tensor, crate::gpu::Tensor<1, T>): crate::gpu::MaxRank, + { + input.rms_norm_fused::<1, OUT_RANK>(&self.weight, self.bias.as_ref(), self.eps) + } + + /// Forward pass for `input + residual` followed by RMSNorm. + pub fn forward_residual(&self, input: &Tensor<3, T>, residual: &Tensor<3, T>) -> Tensor<3, T> { + input.rms_norm_residual_fused::<1, 2>(residual, &self.weight, self.bias.as_ref(), self.eps) + } +} diff --git a/fusor-ml/fusor/src/layers/rms_norm.rs b/fusor-ml/fusor/src/layers/rms_norm/mod.rs similarity index 69% rename from fusor-ml/fusor/src/layers/rms_norm.rs rename to fusor-ml/fusor/src/layers/rms_norm/mod.rs index e902f77ed..54410c801 100644 --- a/fusor-ml/fusor/src/layers/rms_norm.rs +++ b/fusor-ml/fusor/src/layers/rms_norm/mod.rs @@ -1,5 +1,7 @@ //! RMS normalization implementation. +pub(crate) mod autograd; + use crate::fusion::Concrete; use crate::{CastTensor, CastTo, DataType, Device, Fusion, SimdElement, Tensor, VarBuilder}; @@ -63,31 +65,21 @@ impl RmsNorm { } } -// Forward implementations for specific ranks (2D, 3D, 4D inputs) -// This avoids the complex trait bounds while still being useful impl RmsNorm<1, f32> { - /// Forward pass for 2D input (batch, features). - pub fn forward_2d(&self, input: &Tensor<2, f32, B>) -> Tensor<2, f32> - where - B: Fusion<2, f32>, - { - input.rms_norm_fused::<1, 1>(&self.weight, self.bias.as_ref(), self.eps) - } - - /// Forward pass for 3D input (batch, seq_len, features). - pub fn forward(&self, input: &Tensor<3, f32, B>) -> Tensor<3, f32> - where - B: Fusion<3, f32>, - { - input.rms_norm_fused::<1, 2>(&self.weight, self.bias.as_ref(), self.eps) - } - - /// Forward pass for 4D input (batch, heads, seq_len, features). - pub fn forward_4d(&self, input: &Tensor<4, f32, B>) -> Tensor<4, f32> + /// Normalizes the last dimension of an input tensor. + pub fn forward( + &self, + input: &Tensor, + ) -> Tensor where - B: Fusion<4, f32>, + B: Fusion, + Concrete: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + (crate::gpu::Tensor, crate::gpu::Tensor<1, f32>): crate::gpu::MaxRank, { - input.rms_norm_fused::<1, 3>(&self.weight, self.bias.as_ref(), self.eps) + input.rms_norm_fused::<1, OUT_RANK>(&self.weight, self.bias.as_ref(), self.eps) } } @@ -98,11 +90,18 @@ where T: CastTo + CastTensor, f32: CastTo + CastTensor, { - /// Forward pass for 3D input with generic type. - /// Converts input to f32 for computation, then converts back. - pub fn forward_generic(&self, input: &Tensor<3, T, B>) -> Tensor<3, T> + /// Normalizes after converting the input to f32, then converts it back. + pub fn forward_generic( + &self, + input: &Tensor, + ) -> Tensor where - B: Fusion<3, T>, + B: Fusion, + Concrete: crate::cpu::LastRank, + crate::gpu::Tensor: crate::gpu::LastRank, + as crate::gpu::LastRankInner>::LastRank: + crate::gpu::NextRankInner>, + (crate::gpu::Tensor, crate::gpu::Tensor<1, f32>): crate::gpu::MaxRank, { // Cast input and weights to f32 let input_f32 = input.cast::(); @@ -110,7 +109,8 @@ where let bias_f32: Option> = self.bias.as_ref().map(|b| b.cast()); // Compute RMS norm in f32 - let result_f32 = input_f32.rms_norm_fused::<1, 2>(&weight_f32, bias_f32.as_ref(), self.eps); + let result_f32 = + input_f32.rms_norm_fused::<1, OUT_RANK>(&weight_f32, bias_f32.as_ref(), self.eps); // Cast back to T result_f32.cast() @@ -158,18 +158,4 @@ where .rms_norm_residual_fused::<1, 2, _>(residual, &weight_f32, bias_f32.as_ref(), self.eps) .cast() } - - /// Forward pass for 4D input with generic type. - pub fn forward_generic_4d(&self, input: &Tensor<4, T, B>) -> Tensor<4, T> - where - B: Fusion<4, T>, - { - let input_f32 = input.cast::(); - let weight_f32: Tensor<1, f32> = self.weight.cast(); - let bias_f32: Option> = self.bias.as_ref().map(|b| b.cast()); - - let result_f32 = input_f32.rms_norm_fused::<1, 3>(&weight_f32, bias_f32.as_ref(), self.eps); - - result_f32.cast() - } } diff --git a/fusor-ml/fusor/src/lib.rs b/fusor-ml/fusor/src/lib.rs index 1d3360641..da820bfc1 100644 --- a/fusor-ml/fusor/src/lib.rs +++ b/fusor-ml/fusor/src/lib.rs @@ -11,6 +11,7 @@ #[cfg(not(any(feature = "cpu", feature = "gpu")))] compile_error!("fusor requires at least one backend feature: `cpu` or `gpu`."); +pub mod autograd; pub mod cache; mod composite; mod cpu; @@ -29,12 +30,13 @@ pub use quantized::{CpuF32Tensor, QMatrix}; use std::ops::{Deref, Range}; pub use composite::{ - MaskKind, RopeCache, ToVec, ToVec1, ToVec2, ToVec3, arange, arange_step, - base_inverse_frequency, cat, stack, + MaskKind, RopeCache, ToVec, arange, arange_step, base_inverse_frequency, cat, stack, }; pub use device::Device; pub use error::Error; pub use fusion::{Concrete, Fusion}; +#[cfg(feature = "gpu")] +pub use fusor_core::FusorConfig; pub use fusor_types::{D, Dim, FromArray, Layout, StrideSpec}; /// Result type for fusor operations. @@ -73,7 +75,8 @@ pub(crate) use crate::gpu::Tensor as GpuTensor; #[allow(unused_imports)] pub use crate::gpu::{ - CastTensor, DataType, FloatDataType, GgufReadError, NodeIndex, WasmNotSend, WasmNotSync, + CastTensor, DataType, FloatDataType, GgufReadError, KernelProfile, KernelProfileRow, NodeIndex, + WasmNotSend, WasmNotSync, }; pub use crate::gpu::{ @@ -382,6 +385,21 @@ where } } + /// Materialize the tensor while consuming its current handle. + /// + /// This is equivalent to [`Tensor::to_concrete`], but avoids creating a + /// redundant GPU graph alias when the caller no longer needs `self`. + pub fn into_concrete(self) -> Tensor + where + B: TensorBacking, + D: SimdElement, + { + match self { + Tensor::Cpu(t) => Tensor::Cpu(t.to_concrete()), + Tensor::Gpu(t) => Tensor::Gpu(t), + } + } + /// Resolve any pending work for this tensor without downloading it. pub async fn materialize(&self) where diff --git a/fusor-ml/tile-ir-kernels/examples/tile_ir.rs b/fusor-ml/tile-ir-kernels/examples/tile_ir.rs index a15146d4b..e3f8b21a5 100644 --- a/fusor-ml/tile-ir-kernels/examples/tile_ir.rs +++ b/fusor-ml/tile-ir-kernels/examples/tile_ir.rs @@ -128,6 +128,7 @@ fn qgemv_ir(format: GgmlQuantFormat, n: u32, k: u32) -> KernelIr { &y, 1, subgroup_config(32), + tile_ir_kernels::qgemv_selected_shape(format, k, n), Option::<&tile_ir_kernels::UnaryEpilogue>::None, ); }) @@ -151,6 +152,7 @@ fn qgemv_with_silu_epilogue_ir(format: GgmlQuantFormat, n: u32, k: u32) -> Kerne &y, 1, subgroup_config(32), + tile_ir_kernels::qgemv_selected_shape(format, k, n), Some(&silu), ); }) diff --git a/fusor-ml/tile-ir-kernels/src/dispatch.rs b/fusor-ml/tile-ir-kernels/src/dispatch.rs index 03e4d6a5e..aa2f63ce6 100644 --- a/fusor-ml/tile-ir-kernels/src/dispatch.rs +++ b/fusor-ml/tile-ir-kernels/src/dispatch.rs @@ -7,80 +7,13 @@ //! the literal generic //! arguments. //! -//! The mapping from environment variables and `(rows, cols)` heuristics to -//! ShapeKeys mirrors the table that previously lived inline in the builder -//! methods. Snapshot tests in this module pin the current behavior so that -//! the move from inline shape-tables to policy functions is observably a -//! no-op. -//! -//! Env vars (all preserved): -//! - `FUSOR_Q4K_MID_TILE` (rows<=4096, 4096<=cols<8192) -//! - `FUSOR_Q4K_LARGE_TILE` (rows<=4096, cols>=8192) -//! - `FUSOR_Q4K_TALL_TILE` (rows>4096, cols<=4096) -//! - `FUSOR_Q6K_LARGE_TILE` (rows<=4096, cols>=8192) -//! - `FUSOR_Q6K_TALL_TILE` (rows>4096, cols<=4096) +//! Shape selection is automatic and deterministic. Tests exercise the pure +//! policy functions directly; production has no environment-forced geometry. use fusor_tile_ir::{GgmlQuantFormat, SubgroupToken}; // ===== qgemv shapes (Q4K and Q6K ggml paths) ===== -const fn is_q4k_family(format: GgmlQuantFormat) -> bool { - format.is_q4k_family() -} - -const fn is_q6k_family(format: GgmlQuantFormat) -> bool { - format.is_q6k_family() -} - -/// Default qgemv output columns handled by one workgroup for `format`. -pub const fn qgemv_cols_per_workgroup(format: GgmlQuantFormat) -> u32 { - qgemv_subgroups_per_workgroup(format) * qgemv_cols_per_subgroup(format) -} - -/// Shape-aware qgemv output columns handled by one workgroup. -/// -/// This includes the Q4K/Q6K GGML specializations whose column grouping -/// depends on both K (`rows`) and N (`cols`). -pub fn qgemv_cols_per_workgroup_for_shape(format: GgmlQuantFormat, rows: u32, cols: u32) -> u32 { - if is_q4k_family(format) && rows <= 4096 && (4096..8192).contains(&cols) { - return q4k_mid_override(q4k_default_mid(rows, cols)).cols_per_workgroup(); - } - - if is_q4k_family(format) && rows <= 4096 && cols >= 8192 { - return q4k_large_override(q4k_default_large(rows, cols)).cols_per_workgroup(); - } - - if is_q4k_family(format) && rows > 4096 && cols <= 4096 { - return q4k_tall_override(q4k_default_tall(rows, cols)).cols_per_workgroup(); - } - - if is_q6k_family(format) && rows <= 4096 && cols >= 8192 { - return q6k_large_override(q6k_default_large(rows, cols)).cols_per_workgroup(); - } - - if is_q6k_family(format) && rows > 4096 && cols <= 4096 { - return q6k_tall_override(q6k_default_tall(rows, cols)).cols_per_workgroup(); - } - - qgemv_subgroups_per_workgroup_for_shape(format, rows, cols) * qgemv_cols_per_subgroup(format) -} - -pub(crate) const fn qgemv_cols_per_subgroup(format: GgmlQuantFormat) -> u32 { - match format { - GgmlQuantFormat::Q2K => 4, - GgmlQuantFormat::Q4_0 - | GgmlQuantFormat::Q4_0Native - | GgmlQuantFormat::Q4_1 - | GgmlQuantFormat::Q5_1 => 4, - GgmlQuantFormat::Q5_0 | GgmlQuantFormat::Q5_0Native => 4, - GgmlQuantFormat::Q3K | GgmlQuantFormat::Q8K => 2, - GgmlQuantFormat::Q4K | GgmlQuantFormat::Q4KNative => 8, - GgmlQuantFormat::Q6K | GgmlQuantFormat::Q6KNative => 4, - GgmlQuantFormat::Q8_0 | GgmlQuantFormat::Q8_0Native | GgmlQuantFormat::Q8_1 => 4, - GgmlQuantFormat::Q5K | GgmlQuantFormat::Q5KNative => 1, - } -} - pub(crate) const fn qgemv_subgroups_per_workgroup(format: GgmlQuantFormat) -> u32 { match format { GgmlQuantFormat::Q4K @@ -95,7 +28,7 @@ pub(crate) const fn qgemv_subgroups_per_workgroup(format: GgmlQuantFormat) -> u3 } /// Shape-aware subgroup count used by the qgemv dispatch policy. -pub const fn qgemv_subgroups_per_workgroup_for_shape( +const fn qgemv_subgroups_per_workgroup_for_shape( format: GgmlQuantFormat, rows: u32, _cols: u32, @@ -107,7 +40,7 @@ pub const fn qgemv_subgroups_per_workgroup_for_shape( } #[derive(Copy, Clone, Debug, PartialEq, Eq)] -pub(crate) struct QgemvShape { +pub struct QgemvShape { pub subgroups: u32, pub cols_per_subgroup: u32, } @@ -120,11 +53,65 @@ impl QgemvShape { } } - pub(crate) const fn cols_per_workgroup(self) -> u32 { + pub const fn cols_per_workgroup(self) -> u32 { self.subgroups * self.cols_per_subgroup } } +/// The workgroup geometry the qgemv builder will emit for `(format, rows, +/// output_cols)` — the single source for both the kernel body and the +/// dispatch grid. `rows` is the contraction depth (K) and `output_cols` the +/// epilogue-adjusted output width (N). Callers computing a dispatch must use +/// this (never a re-derived approximation): the launched grid and the +/// kernel's internal `qgemv_grid` agree by construction only when both come +/// from here. +pub fn qgemv_selected_shape(format: GgmlQuantFormat, rows: u32, output_cols: u32) -> QgemvShape { + match format { + GgmlQuantFormat::Q8_0 | GgmlQuantFormat::Q8_0Native => { + if output_cols >= 8192 { + QgemvShape::new(4, 8) + } else { + QgemvShape::new(4, 4) + } + } + GgmlQuantFormat::Q8_1 => QgemvShape::new(4, 4), + GgmlQuantFormat::Q4K | GgmlQuantFormat::Q4KNative => { + if rows <= 4096 && (4096..8192).contains(&output_cols) { + q4k_default_mid(rows, output_cols) + } else if rows <= 4096 && output_cols <= 4096 { + QgemvShape::new(8, 4) + } else if rows <= 4096 && output_cols >= 8192 { + q4k_default_large(rows, output_cols) + } else if rows > 4096 && output_cols <= 4096 { + q4k_default_tall(rows, output_cols) + } else if qgemv_subgroups_per_workgroup_for_shape(format, rows, output_cols) == 8 { + QgemvShape::new(8, 8) + } else { + QgemvShape::new(4, 8) + } + } + GgmlQuantFormat::Q5_0 | GgmlQuantFormat::Q5_0Native => QgemvShape::new(2, 4), + GgmlQuantFormat::Q4_0 + | GgmlQuantFormat::Q4_0Native + | GgmlQuantFormat::Q4_1 + | GgmlQuantFormat::Q5_1 + | GgmlQuantFormat::Q2K => QgemvShape::new(2, 4), + GgmlQuantFormat::Q3K | GgmlQuantFormat::Q8K => QgemvShape::new(2, 2), + GgmlQuantFormat::Q5K | GgmlQuantFormat::Q5KNative => QgemvShape::new(2, 1), + GgmlQuantFormat::Q6K | GgmlQuantFormat::Q6KNative => { + if rows <= 4096 && output_cols >= 8192 { + q6k_default_large(rows, output_cols) + } else if rows > 4096 && output_cols <= 4096 { + q6k_default_tall(rows, output_cols) + } else if qgemv_subgroups_per_workgroup_for_shape(format, rows, output_cols) == 4 { + QgemvShape::new(4, 4) + } else { + QgemvShape::new(8, 4) + } + } + } +} + /// Subgroup-width range advertised by the target adapter for one generated /// kernel. The generated shader reads the actual runtime subgroup size; the /// max is only used for the fixed workgroup allocation passed to WGSL. @@ -190,66 +177,6 @@ pub(crate) const fn q4k_default_mid(_rows: u32, cols: u32) -> QgemvShape { QgemvShape::new(2, 2) } -/// Apply `FUSOR_Q4K_MID_TILE` if set; otherwise return the default. The set -/// of accepted env values is exactly the inline `qgemv_ggml_env!` table that -/// used to live in `qgemv_tile`. -/// Per-context env override tables. Each entry maps an env-var token to the -/// `QgemvShape` it selects. Different contexts (mid / large / tall, Q4K vs -/// Q6K) accept different subsets of the 14 total shapes. -const Q4K_MID_TILES: &[(&str, QgemvShape)] = &[ - ("ggml_2x2", QgemvShape::new(2, 2)), - ("ggml_2x3", QgemvShape::new(2, 3)), - ("ggml_2x4", QgemvShape::new(2, 4)), - ("ggml_2x8", QgemvShape::new(2, 8)), - ("ggml_4x2", QgemvShape::new(4, 2)), - ("ggml_4x3", QgemvShape::new(4, 3)), - ("ggml_4x4", QgemvShape::new(4, 4)), - ("ggml_4x8", QgemvShape::new(4, 8)), - ("ggml_8x2", QgemvShape::new(8, 2)), - ("ggml_8x4", QgemvShape::new(8, 4)), -]; - -const Q4K_LARGE_TILES: &[(&str, QgemvShape)] = &[ - ("ggml_1x4", QgemvShape::new(1, 4)), - ("ggml_1x8", QgemvShape::new(1, 8)), - ("ggml_2x2", QgemvShape::new(2, 2)), - ("ggml_2x4", QgemvShape::new(2, 4)), - ("ggml_2x8", QgemvShape::new(2, 8)), - ("ggml_4x1", QgemvShape::new(4, 1)), - ("ggml_4x2", QgemvShape::new(4, 2)), - ("ggml_4x4", QgemvShape::new(4, 4)), - ("ggml_4x8", QgemvShape::new(4, 8)), - ("ggml_8x1", QgemvShape::new(8, 1)), - ("ggml_8x2", QgemvShape::new(8, 2)), - ("ggml_8x4", QgemvShape::new(8, 4)), -]; - -const STANDARD_8_TILES: &[(&str, QgemvShape)] = &[ - ("ggml_2x2", QgemvShape::new(2, 2)), - ("ggml_2x4", QgemvShape::new(2, 4)), - ("ggml_2x8", QgemvShape::new(2, 8)), - ("ggml_4x2", QgemvShape::new(4, 2)), - ("ggml_4x4", QgemvShape::new(4, 4)), - ("ggml_4x8", QgemvShape::new(4, 8)), - ("ggml_8x2", QgemvShape::new(8, 2)), - ("ggml_8x4", QgemvShape::new(8, 4)), -]; - -fn env_tile_override(var: &str, table: &[(&str, QgemvShape)], default: QgemvShape) -> QgemvShape { - let Ok(value) = std::env::var(var) else { - return default; - }; - table - .iter() - .find(|(name, _)| *name == value) - .map(|(_, shape)| *shape) - .unwrap_or(default) -} - -pub(crate) fn q4k_mid_override(default: QgemvShape) -> QgemvShape { - env_tile_override("FUSOR_Q4K_MID_TILE", Q4K_MID_TILES, default) -} - // ----- Q4K large (rows<=4096, cols>=8192) ----- /// Default Q4K large-shape: cols<=16_384 → 8x4, else 2x4. @@ -261,12 +188,6 @@ pub(crate) const fn q4k_default_large(_rows: u32, cols: u32) -> QgemvShape { } } -/// Apply `FUSOR_Q4K_LARGE_TILE` if set. Carries the same tile list as the -/// inline macro: adds 1x4/1x8/4x1/8x1 (no 2x3/4x3 entries). -pub(crate) fn q4k_large_override(default: QgemvShape) -> QgemvShape { - env_tile_override("FUSOR_Q4K_LARGE_TILE", Q4K_LARGE_TILES, default) -} - // ----- Q4K tall (rows>4096, cols<=4096) ----- /// Default Q4K tall-shape: 4x2. @@ -274,11 +195,6 @@ pub(crate) const fn q4k_default_tall(_rows: u32, _cols: u32) -> QgemvShape { QgemvShape::new(4, 2) } -/// Apply `FUSOR_Q4K_TALL_TILE` if set. Standard 8-tile set. -pub(crate) fn q4k_tall_override(default: QgemvShape) -> QgemvShape { - env_tile_override("FUSOR_Q4K_TALL_TILE", STANDARD_8_TILES, default) -} - // ----- Q6K large (rows<=4096, cols>=8192) ----- /// Default Q6K large-shape: cols<=16_384 → 2x2, else 2x4. @@ -290,11 +206,6 @@ pub(crate) const fn q6k_default_large(_rows: u32, cols: u32) -> QgemvShape { } } -/// Apply `FUSOR_Q6K_LARGE_TILE` if set. Standard 8-tile set. -pub(crate) fn q6k_large_override(default: QgemvShape) -> QgemvShape { - env_tile_override("FUSOR_Q6K_LARGE_TILE", STANDARD_8_TILES, default) -} - // ----- Q6K tall (rows>4096, cols<=4096) ----- /// Default Q6K tall-shape: 2x2. @@ -302,48 +213,36 @@ pub(crate) const fn q6k_default_tall(_rows: u32, _cols: u32) -> QgemvShape { QgemvShape::new(2, 2) } -/// Apply `FUSOR_Q6K_TALL_TILE` if set. Standard 8-tile set. -pub(crate) fn q6k_tall_override(default: QgemvShape) -> QgemvShape { - env_tile_override("FUSOR_Q6K_TALL_TILE", STANDARD_8_TILES, default) -} - #[cfg(test)] mod tests { - //! Snapshot tests pinning the current `(format, rows, cols, env) → - //! ShapeKey` mapping. These must continue to pass after the inline - //! `qgemv_ggml_env!` invocations and `if b.cols == ...` heuristics in - //! `kernels/qgemv.rs` are replaced with calls into this module. - //! - //! Env-var tests use a serial mutex because `std::env::set_var` is - //! process-global. They also unset the variable on entry to avoid - //! cross-test contamination from a developer's shell. + //! Snapshot tests pinning the automatic `(format, rows, cols) → ShapeKey` + //! mapping. use super::*; - use std::sync::Mutex; - use std::sync::OnceLock; - fn env_lock() -> &'static Mutex<()> { - static LOCK: OnceLock> = OnceLock::new(); - LOCK.get_or_init(|| Mutex::new(())) - } - fn with_env(var: &str, value: Option<&str>, f: impl FnOnce() -> R) -> R { - let _guard = env_lock().lock().unwrap_or_else(|e| e.into_inner()); - let prior = std::env::var(var).ok(); - // SAFETY: tests serialize on env_lock above; no other thread reads or - // mutates this var while the closure runs. - unsafe { - match value { - Some(v) => std::env::set_var(var, v), - None => std::env::remove_var(var), - } - } - let out = f(); - unsafe { - match prior { - Some(p) => std::env::set_var(var, p), - None => std::env::remove_var(var), - } - } - out + /// The selected shape is the single source for dispatch and kernel + /// geometry; these cells pin it where the deleted core-side re-derivation + /// used to disagree with the builder (over-dispatching masked workgroups + /// and permanently missing the prebuilt-pipeline fast path). + #[test] + fn selected_shape_pins_previously_desynced_cells() { + use GgmlQuantFormat as F; + let cols = |f, k, n| qgemv_selected_shape(f, k, n).cols_per_workgroup(); + // Q4K large: the old dispatch ladder said 8; the builder emits 32. + assert_eq!(cols(F::Q4K, 2048, 8192), 32); + assert_eq!(cols(F::Q4K, 4096, 11008), 32); + // Q4K mid: old ladder said 4 for every mid shape; the builder varies. + assert_eq!(cols(F::Q4K, 4096, 4097), 4); + assert_eq!(cols(F::Q4K, 4096, 5120), 12); + assert_eq!(cols(F::Q4K, 4096, 6144), 16); + // Q6K large <=16384: old ladder said 8; the builder emits 4. + assert_eq!(cols(F::Q6K, 2048, 8192), 4); + assert_eq!(cols(F::Q6K, 2048, 16385), 8); + // Q8_0 wide: old ladder said 32; the builder emits 32 (4x8). + assert_eq!(cols(F::Q8_0, 1024, 8192), 32); + // SmolLM2 decode cells. + assert_eq!(cols(F::Q4K, 576, 1536), 32); + assert_eq!(cols(F::Q4K, 1536, 576), 32); + assert_eq!(cols(F::Q6K, 576, 49152), 8); } #[test] @@ -385,138 +284,4 @@ mod tests { // Constant 2x2 from kernels/qgemv.rs. assert_eq!(q6k_default_tall(8192, 4096), QgemvShape::new(2, 2)); } - - #[test] - fn q4k_mid_override_table_unchanged() { - with_env("FUSOR_Q4K_MID_TILE", None, || { - assert_eq!( - q4k_mid_override(QgemvShape::new(2, 2)), - QgemvShape::new(2, 2) - ); - }); - let cases = [ - ("ggml_2x2", QgemvShape::new(2, 2)), - ("ggml_2x3", QgemvShape::new(2, 3)), - ("ggml_2x4", QgemvShape::new(2, 4)), - ("ggml_2x8", QgemvShape::new(2, 8)), - ("ggml_4x2", QgemvShape::new(4, 2)), - ("ggml_4x3", QgemvShape::new(4, 3)), - ("ggml_4x4", QgemvShape::new(4, 4)), - ("ggml_4x8", QgemvShape::new(4, 8)), - ("ggml_8x2", QgemvShape::new(8, 2)), - ("ggml_8x4", QgemvShape::new(8, 4)), - ]; - for (val, expect) in cases { - with_env("FUSOR_Q4K_MID_TILE", Some(val), || { - assert_eq!( - q4k_mid_override(QgemvShape::new(4, 4)), - expect, - "FUSOR_Q4K_MID_TILE={val}" - ); - }); - } - // Unrecognized value falls through to default. - with_env("FUSOR_Q4K_MID_TILE", Some("nonsense"), || { - assert_eq!( - q4k_mid_override(QgemvShape::new(4, 4)), - QgemvShape::new(4, 4) - ); - }); - } - - #[test] - fn q4k_large_override_table_unchanged() { - let cases = [ - ("ggml_1x4", QgemvShape::new(1, 4)), - ("ggml_1x8", QgemvShape::new(1, 8)), - ("ggml_2x2", QgemvShape::new(2, 2)), - ("ggml_2x4", QgemvShape::new(2, 4)), - ("ggml_2x8", QgemvShape::new(2, 8)), - ("ggml_4x1", QgemvShape::new(4, 1)), - ("ggml_4x2", QgemvShape::new(4, 2)), - ("ggml_4x4", QgemvShape::new(4, 4)), - ("ggml_4x8", QgemvShape::new(4, 8)), - ("ggml_8x1", QgemvShape::new(8, 1)), - ("ggml_8x2", QgemvShape::new(8, 2)), - ("ggml_8x4", QgemvShape::new(8, 4)), - ]; - for (val, expect) in cases { - with_env("FUSOR_Q4K_LARGE_TILE", Some(val), || { - assert_eq!( - q4k_large_override(QgemvShape::new(4, 4)), - expect, - "FUSOR_Q4K_LARGE_TILE={val}" - ); - }); - } - } - - #[test] - fn q4k_tall_override_table_unchanged() { - let cases = [ - ("ggml_2x2", QgemvShape::new(2, 2)), - ("ggml_2x4", QgemvShape::new(2, 4)), - ("ggml_2x8", QgemvShape::new(2, 8)), - ("ggml_4x2", QgemvShape::new(4, 2)), - ("ggml_4x4", QgemvShape::new(4, 4)), - ("ggml_4x8", QgemvShape::new(4, 8)), - ("ggml_8x2", QgemvShape::new(8, 2)), - ("ggml_8x4", QgemvShape::new(8, 4)), - ]; - for (val, expect) in cases { - with_env("FUSOR_Q4K_TALL_TILE", Some(val), || { - assert_eq!( - q4k_tall_override(QgemvShape::new(4, 2)), - expect, - "FUSOR_Q4K_TALL_TILE={val}" - ); - }); - } - } - - #[test] - fn q6k_large_override_table_unchanged() { - let cases = [ - ("ggml_2x2", QgemvShape::new(2, 2)), - ("ggml_2x4", QgemvShape::new(2, 4)), - ("ggml_2x8", QgemvShape::new(2, 8)), - ("ggml_4x2", QgemvShape::new(4, 2)), - ("ggml_4x4", QgemvShape::new(4, 4)), - ("ggml_4x8", QgemvShape::new(4, 8)), - ("ggml_8x2", QgemvShape::new(8, 2)), - ("ggml_8x4", QgemvShape::new(8, 4)), - ]; - for (val, expect) in cases { - with_env("FUSOR_Q6K_LARGE_TILE", Some(val), || { - assert_eq!( - q6k_large_override(QgemvShape::new(2, 2)), - expect, - "FUSOR_Q6K_LARGE_TILE={val}" - ); - }); - } - } - - #[test] - fn q6k_tall_override_table_unchanged() { - let cases = [ - ("ggml_2x2", QgemvShape::new(2, 2)), - ("ggml_2x4", QgemvShape::new(2, 4)), - ("ggml_2x8", QgemvShape::new(2, 8)), - ("ggml_4x2", QgemvShape::new(4, 2)), - ("ggml_4x4", QgemvShape::new(4, 4)), - ("ggml_4x8", QgemvShape::new(4, 8)), - ("ggml_8x2", QgemvShape::new(8, 2)), - ("ggml_8x4", QgemvShape::new(8, 4)), - ]; - for (val, expect) in cases { - with_env("FUSOR_Q6K_TALL_TILE", Some(val), || { - assert_eq!( - q6k_tall_override(QgemvShape::new(2, 2)), - expect, - "FUSOR_Q6K_TALL_TILE={val}" - ); - }); - } - } } diff --git a/fusor-ml/tile-ir-kernels/src/kernels.rs b/fusor-ml/tile-ir-kernels/src/kernels.rs index 9ec3f17bd..5c9729c32 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels.rs @@ -1,5 +1,6 @@ use fusor_tile_ir::{Layout, MemoryLevel, Shape}; +mod attention; mod helpers; mod matmul; mod qgemv; @@ -9,18 +10,22 @@ mod qmatmul; mod qmatmul_workgroup; mod quantized_matrix; +pub use attention::{ + flash_attention_bwd_supported, flash_attention_dispatch, flash_attention_f32, + flash_attention_supported, flash_attention_workgroup_bytes, flash_bwd_kv_dispatch, + flash_bwd_kv_f32, flash_bwd_q_dispatch, flash_bwd_q_f32, flash_lse_dispatch, flash_lse_f32, + FlashAttentionLayouts, FlashAttentionShape, FlashBwdLayouts, FlashKvOutputs, FlashMaskLayout, + FlashOperandLayout, FlashRowLayout, +}; pub use helpers::AccumCast; pub use matmul::{ - try_batched_coop_matmul, DenseCoopMatmulConfig, DenseCoopMatmulTile, DenseMatmulShape, - DenseMatmulTensors, -}; + DEFAULT_SWIZZLE_GROUP_M, + coop_tile_entries, merged_split_k_combine, split_k_combine, try_batched_coop_matmul, + try_batched_coop_matmul_split_k, try_merged_coop_matmul, CoopTileEntry, + DenseCoopMatmulConfig, DenseCoopMatmulTile, DenseMatmulShape, DenseMatmulTensors,}; pub use qgemv::{qgemv_with_epilogue, IntoQgemvEpilogues}; pub use qmatmul::qmatmul_with_epilogue; -pub use qmatmul_workgroup::{ - qgemv_workgroup_f16_with_epilogue, qgemv_workgroup_storage_f16_with_epilogue, - qgemv_workgroup_with_epilogue, qmatmul_workgroup_f16_with_epilogues, - qmatmul_workgroup_storage_f16_with_epilogues, qmatmul_workgroup_with_epilogues, -}; +pub use qmatmul_workgroup::qmatmul_workgroup_with_epilogues; pub use quantized_matrix::{quantized_matrix, quantized_matrix_for}; /// The default rank-1 unit-stride layout used by tile-ir's pre-built kernels diff --git a/fusor-ml/tile-ir-kernels/src/kernels/attention.rs b/fusor-ml/tile-ir-kernels/src/kernels/attention.rs new file mode 100644 index 000000000..87144684f --- /dev/null +++ b/fusor-ml/tile-ir-kernels/src/kernels/attention.rs @@ -0,0 +1,1519 @@ +//! Fused scaled-dot-product attention over cooperative matrices. +//! +//! One workgroup owns a `BR`-row query tile of one (batch, head) and streams +//! the KV axis in `BC`-wide tiles: QKᵀ and P·V run on simdgroup matrices +//! staged through workgroup memory, the online-softmax row statistics run on +//! per-row lanes over the staged score tile, and the output accumulates in +//! per-lane registers rescaled by the running max. Unlike the attention row +//! program (one query row per workgroup — decode's shape), K/V tiles are +//! shared across all `BR` rows in the workgroup, so prefill/training shapes +//! keep full data reuse. +//! +//! Operands are addressed through explicit per-axis element strides baked at +//! build time, so arbitrary strided rank-4 views (transposes, offsets) and +//! grouped-query K/V (`kv_groups > 1`, KV head = head / groups) need no +//! materialization. + +use fusor_tile_ir::tile::{Program, Storage, Tile, TileBlock}; +use fusor_tile_ir::{CoopMatrixToken, ElementType, ScalarElement, WorkgroupAxis}; + +use crate::dispatch::SubgroupConfig; +use crate::kernels::helpers::{ + scalar_of, + coop_load_a_fragments, coop_load_b_fragments, coop_load_b_fragments_transposed, + coop_mma_grid, dispatch_grid_1d, zero_coop_acc_grid, +}; + +/// Query rows per workgroup. +const BR: u32 = 32; +/// KV positions per streamed tile. +const BC: u32 = 16; +/// Cooperative fragment side. +const COOP_DIM: u32 = 8; +/// Element the Q/K/V/dO operand tiles stage in: f16 halves their threadgroup +/// footprint and the MMAs run f16 fragments against f32 accumulators at full +/// rate. Tiles that receive f32 accumulator stores (scores, dP) stay f32. + +/// Finite stand-in for -inf: masked scores exp to zero without the +/// `(-inf) - (-inf)` NaN when a row is entirely masked so far. +const MASKED_SCORE: f32 = -3.0e38; + +/// The kernel's math runs in f32; convert at the storage boundary when the +/// operand tensors are f16 (per-lane stores don't convert element types). +fn cast_to_storage(dst: &Storage, value: Tile) -> Tile { + match scalar_of(dst.element()) { + ScalarElement::F32 => value, + elem => value.cast(elem.element()), + } +} + +/// Widen a storage-loaded value into the kernel's f32 math. +fn cast_from_storage(src: &Storage, value: Tile) -> Tile { + match scalar_of(src.element()) { + ScalarElement::F32 => value, + _ => value.cast(ElementType::F32), + } +} + +/// Shape of one fused attention dispatch over rank-4 +/// `[batch, heads, seq, head_dim]` operands. +#[derive(Clone, Copy, Debug)] +pub struct FlashAttentionShape { + /// Batch size. + pub batch: u32, + /// Query heads. + pub heads: u32, + /// Query heads per KV head (1 = multi-head attention). K/V are addressed + /// at head `h / kv_groups`. + pub kv_groups: u32, + /// Query sequence length. + pub q_len: u32, + /// KV sequence length. + pub kv_len: u32, + /// Head dimension. + pub head_dim: u32, + /// Score scale (typically `head_dim^-0.5`). + pub scale: f32, + /// Causal masking via index comparison (mutually exclusive with `mask`). + pub causal: bool, +} + +/// Element strides of one rank-4 `[batch, heads, seq, head_dim]` operand into +/// its linear storage. K/V use their own head count (`heads / kv_groups`). +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct FlashOperandLayout { + /// Element offset of `[0, 0, 0, 0]`. + pub offset: u32, + /// Elements between consecutive batches. + pub batch_stride: u32, + /// Elements between consecutive heads. + pub head_stride: u32, + /// Elements between consecutive sequence positions. + pub seq_stride: u32, + /// Elements between consecutive head-dim positions. + pub dim_stride: u32, +} + +impl FlashOperandLayout { + /// The contiguous `[batch, heads, seq, head_dim]` layout. + pub fn contiguous(heads: u32, seq: u32, head_dim: u32) -> Self { + Self { + offset: 0, + batch_stride: heads * seq * head_dim, + head_stride: seq * head_dim, + seq_stride: head_dim, + dim_stride: 1, + } + } +} + +/// Element strides of the additive rank-2 `[q_len, kv_len]` mask, broadcast +/// over every (batch, head). +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct FlashMaskLayout { + /// Element offset of `[0, 0]`. + pub offset: u32, + /// Elements between consecutive query positions. + pub q_stride: u32, + /// Elements between consecutive KV positions. + pub kv_stride: u32, +} + +/// Per-operand layouts for one flash-attention dispatch. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct FlashAttentionLayouts { + pub q: FlashOperandLayout, + pub k: FlashOperandLayout, + pub v: FlashOperandLayout, + pub o: FlashOperandLayout, +} + +/// Whether [`flash_attention_f32`] can host this shape. +/// +/// The P·V pass splits the head dim across the four subgroups in 8-wide +/// fragments, so `head_dim` must be a multiple of 32. +pub fn flash_attention_supported(shape: &FlashAttentionShape, subgroups: SubgroupConfig) -> bool { + let block = subgroups.block_for_subgroups(4); + subgroups.is_fixed() + && shape.q_len % BR == 0 + && shape.kv_len % BC == 0 + && shape.head_dim % 32 == 0 + && shape.head_dim >= 32 + && shape.head_dim <= 80 + && block % BR == 0 + && shape.head_dim % (block / BR) == 0 + && shape.batch > 0 + && shape.heads > 0 + && shape.kv_groups > 0 + && shape.heads % shape.kv_groups == 0 +} + +/// Workgroup-memory footprint of [`flash_attention_f32`] in bytes for one +/// head dim and stage element: the staged Q/KV/P operand tiles in `stage` +/// (K and V share one tile) plus the f32 score tile and the three per-row +/// statistic arrays, each tile row carrying one pad element against bank +/// conflicts. Asserted equal to the lowered IR's `workgroup_bytes` in +/// `tests/footprint.rs`. +pub const fn flash_attention_workgroup_bytes(head_dim: u32, stage: ScalarElement) -> u64 { + let d = head_dim as u64; + let (br, bc) = (BR as u64, BC as u64); + let s_cols = if d > bc { d } else { bc }; + // A padded tile spans `rows * (cols + 1) - 1` elements: the pad after + // its last row is never addressed and is not allocated. + let stage_elements = (br * (d + 1) - 1) + (bc * (d + 1) - 1) + (br * (bc + 1) - 1); + let f32_elements = (br * (s_cols + 1) - 1) + 3 * br; + stage_elements * stage.byte_size() + f32_elements * ScalarElement::F32.byte_size() +} + +/// Decompose an operand's (batch, head) base into dynamic (row, col) origins +/// over a `[seq_stride, 1]`-strided rank-2 view of the same buffer: base +/// components divisible by the sequence stride advance whole rows, the rest +/// lands in the unit-stride column axis. Only unit-`dim_stride` operands +/// qualify (checked by the caller); this keeps `fill_tile`'s contiguous vec4 +/// staging path for every practical layout. +fn rank2_origins( + program: &mut TileBlock, + layout: FlashOperandLayout, + b: &Tile, + h: &Tile, +) -> (Tile, Tile) { + let ss = layout.seq_stride; + let mut row = Tile::u32(0); + let mut col = Tile::u32(0); + for (coeff, stride) in [(b, layout.batch_stride), (h, layout.head_stride)] { + if stride % ss == 0 { + row = row + coeff.clone() * (stride / ss); + } else { + col = col + coeff.clone() * stride; + } + } + if layout.offset % ss == 0 { + row = row + layout.offset / ss; + } else { + col = col + layout.offset; + } + (program.bind(row), program.bind(col)) +} + +/// The workgroup grid [`flash_attention_f32`] dispatches over: one workgroup +/// per `BR`-row query tile of one (batch, head). +pub fn flash_attention_dispatch( + shape: &FlashAttentionShape, + max_workgroups_per_dimension: u32, +) -> [u32; 3] { + let total_tiles = shape.batch * shape.heads * (shape.q_len / BR); + dispatch_grid_1d(total_tiles, max_workgroups_per_dimension) +} + +/// Emit the fused attention kernel. Returns `false` (program untouched) when +/// the shape fails [`flash_attention_supported`]. +#[allow(clippy::too_many_arguments)] +pub fn flash_attention_f32( + program: &mut Program, + q: &Storage, + k: &Storage, + v: &Storage, + mask: Option<(&Storage, FlashMaskLayout)>, + o: &Storage, + layouts: &FlashAttentionLayouts, + shape: FlashAttentionShape, + subgroups: SubgroupConfig, + coop: CoopMatrixToken, + max_workgroups_per_dimension: u32, +) -> bool { + if !flash_attention_supported(&shape, subgroups) || (shape.causal && mask.is_some()) { + return false; + } + let block = subgroups.block_for_subgroups(4); + let d = shape.head_dim; + let q_tiles = shape.q_len / BR; + let kv_tiles = shape.kv_len / BC; + let scalar = ScalarElement::F32; + let (lq, lk, lv, lo) = (layouts.q, layouts.k, layouts.v, layouts.o); + // Per-lane output slice: `block / BR` lanes share a row, each owning + // `d / (block / BR)` contiguous columns in registers. + let lanes_per_row = block / BR; + let cols_per_lane = d / lanes_per_row; + + // Staged tiles (+1 pad on the inner stride against bank conflicts). + // Operands stage in f16; `s_tile` receives the f32 QKᵀ accumulators (cols + // 0..BC) and, after a barrier, the f32 P·V partial (cols 0..d) — + // probabilities live in the small f16 `p_tile` so they can feed the P·V + // MMA as f16 A-fragments. + // Stage operands in their own element type: casting f32 operands to + // f16 tiles injects ~2e-4 noise per attention op, which measurably + // degrades training and compounds to NaN within a few hundred steps. + let stage = scalar_of(q.element()); + let q_tile = program.alloc_workgroup_tile_padded(stage, BR, d, 1); + // K and V share one tile: K is dead once the score MMA's post-store + // barrier passes, and V is not read until the barrier before the P*V + // MMA, so V stages into K's slot between the two existing barriers. + // The sharing costs nothing and holds the f16 kernel's footprint at + // 16.0 KB - the two-workgroups-per-core residency boundary. + let kv_tile = program.alloc_workgroup_tile_padded(stage, BC, d, 1); + let s_tile = program.alloc_workgroup_tile_padded(scalar, BR, d.max(BC), 1); + let p_tile = program.alloc_workgroup_tile_padded(stage, BR, BC, 1); + let m_arr = program.alloc_workgroup_array(scalar, BR); + let l_arr = program.alloc_workgroup_array(scalar, BR); + let alpha_arr = program.alloc_workgroup_array(scalar, BR); + + let s_stride = d.max(BC) + 1; + let p_stride = BC + 1; + + let kv_heads = shape.heads / shape.kv_groups; + let q_fast = fast_rows_view(q, lq, shape.batch * shape.heads * shape.q_len, d); + let k_fast = fast_rows_view(k, lk, shape.batch * kv_heads * shape.kv_len, d); + let v_fast = fast_rows_view(v, lv, shape.batch * kv_heads * shape.kv_len, d); + + let grid = flash_attention_dispatch(&shape, max_workgroups_per_dimension); + program.program_grid(block, grid, |program| { + let tile_id = program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; + let bh = tile_id.clone() / q_tiles; + let qt = program.bind(tile_id % q_tiles); + let b = program.bind(bh.clone() / shape.heads); + let h = program.bind(bh % shape.heads); + let kv_h = program.bind(h.clone() / shape.kv_groups); + let o_base = program.bind( + b.clone() * lo.batch_stride + + h.clone() * lo.head_stride + + qt.clone() * (BR * lo.seq_stride) + + lo.offset, + ); + // Position of this workgroup's first query row within the sequence + // (mask/causal indices are per-sequence, not per-storage-row). + let q_pos_base = program.bind(qt.clone() * BR); + + let lane = program.lane(); + // Zero the running statistics and the output registers. + program.if_then(lane.clone().lt(BR), |program| { + program.store_workgroup(&m_arr, lane.clone(), Tile::f32(MASKED_SCORE)); + program.store_workgroup(&l_arr, lane.clone(), Tile::f32(0.0)); + }); + let o_regs: Vec<_> = (0..cols_per_lane) + .map(|_| { + let reg = program.private(ElementType::F32); + program.store_local(®, Tile::f32(0.0)); + reg + }) + .collect(); + let o_row = program.bind(lane.clone() / lanes_per_row); + let o_col_base = program.bind((lane.clone() % lanes_per_row) * cols_per_lane); + + // Q tile is loop-invariant: staged once. + stage_rows( + program, q, &q_fast, lq, &q_tile, &b, &h, &q_pos_base, BR, d, &lane, block, + ); + + program.loop_range(kv_tiles, |program, kv_t| { + let kv_pos_base = program.bind(kv_t * BC); + if shape.causal { + // Tiles entirely past the query diagonal contribute nothing. + program.break_if(kv_pos_base.clone().gt(q_pos_base.clone() + (BR - 1))); + } + + // Stage K row-major; QKᵀ reads it through transposed fragment + // loads. V stages into the same tile after the score phase. + stage_rows( + program, k, &k_fast, lk, &kv_tile, &b, &kv_h, &kv_pos_base, BC, d, &lane, block, + ); + program.workgroup_barrier(); + + // S = Q · Kᵀ on fragments: 2×2 subgroup grid, each owning + // BR/2 rows × BC/2 columns of the score tile. + let subgroup_id = subgroups.token().subgroup_id(program); + let sg_row = program.bind(subgroup_id.clone() / 2); + let sg_col = program.bind(subgroup_id % 2); + let s_rows = BR / 2 / COOP_DIM; + let s_cols = BC / 2 / COOP_DIM; + let sg_row_base = program.bind(sg_row * (BR / 2)); + let sg_col_base = program.bind(sg_col.clone() * (BC / 2)); + let s_accs = zero_coop_acc_grid(program, coop, scalar, s_rows, s_cols); + for kk in 0..d / COOP_DIM { + let a_frags = coop_load_a_fragments( + program, + coop, + &q_tile, + &sg_row_base, + kk, + s_rows, + stage, + ); + let b_frags = coop_load_b_fragments_transposed( + program, + coop, + &kv_tile, + &sg_col_base, + kk, + s_cols, + stage, + ); + coop_mma_grid(program, coop, &s_accs, &a_frags, &b_frags); + } + for (r, row_accs) in s_accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + &s_tile, + sg_row_base.clone() + r as u32 * COOP_DIM, + sg_col_base.clone() + c as u32 * COOP_DIM, + ); + } + } + program.workgroup_barrier(); + + // Every subgroup's score MMA reads of K completed before the + // barrier above, so V now stages into the shared tile while the + // BR softmax lanes work the staged scores; the barrier below + // gates the P·V reads of V. + stage_rows( + program, v, &v_fast, lv, &kv_tile, &b, &kv_h, &kv_pos_base, BC, d, &lane, block, + ); + + // Online softmax over the staged scores: one lane per query row. + // Probabilities land in the f16 `p_tile` — the A operand of the + // P·V MMA below. + program.if_then(lane.clone().lt(BR), |program| { + let row = lane.clone(); + let q_pos = program.bind(q_pos_base.clone() + row.clone()); + let m_old = program.bind(program.load_workgroup(&m_arr, row.clone())); + let vals: Vec = (0..BC) + .map(|c| { + let raw = program.load_workgroup(&s_tile, row.clone() * s_stride + c) + * shape.scale; + let masked = if shape.causal { + let allowed = + (kv_pos_base.clone() + c).le(q_pos.clone()); + Tile::select(allowed, raw, Tile::f32(MASKED_SCORE)) + } else if let Some((mask, lm)) = &mask { + raw + cast_from_storage( + mask, + program.load( + mask.at(q_pos.clone() * lm.q_stride + + (kv_pos_base.clone() + c) * lm.kv_stride + + lm.offset), + Tile::all(), + 0.0, + ), + ) + } else { + raw + }; + program.bind(masked) + }) + .collect(); + let mut m_tile = vals[0].clone(); + for val in &vals[1..] { + m_tile = m_tile.max(val.clone()); + } + let m_new = program.bind(m_old.clone().max(m_tile)); + let alpha = program.bind((m_old - m_new.clone()).exp()); + let mut row_sum = Tile::f32(0.0); + for (c, val) in vals.iter().enumerate() { + let p = program.bind((val.clone() - m_new.clone()).exp()); + program.store_workgroup( + &p_tile, + row.clone() * p_stride + c as u32, + p.clone(), + ); + row_sum = row_sum + p; + } + let l_old = program.load_workgroup(&l_arr, row.clone()); + program.store_workgroup(&l_arr, row.clone(), l_old * alpha.clone() + row_sum); + program.store_workgroup(&m_arr, row.clone(), m_new); + program.store_workgroup(&alpha_arr, row, alpha); + }); + program.workgroup_barrier(); + + // P·V on fragments: each subgroup owns d/4 output columns across + // all BR rows. Accumulators stay in registers past the barrier + // below, so overwriting the score region is safe. + let pv_rows = BR / COOP_DIM; + let pv_cols = d / 4 / COOP_DIM; + let sg_d_base = program.bind(subgroups.token().subgroup_id(program) * (d / 4)); + let pv_accs = zero_coop_acc_grid(program, coop, scalar, pv_rows, pv_cols); + let zero_row = Tile::u32(0); + for kk in 0..BC / COOP_DIM { + let a_frags = coop_load_a_fragments( + program, + coop, + &p_tile, + &zero_row, + kk, + pv_rows, + stage, + ); + let b_frags = coop_load_b_fragments( + program, + coop, + &kv_tile, + &sg_d_base, + kk, + pv_cols, + stage, + ); + coop_mma_grid(program, coop, &pv_accs, &a_frags, &b_frags); + } + // Every subgroup finishes reading P before any overwrites it. + program.workgroup_barrier(); + for (r, row_accs) in pv_accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + &s_tile, + Tile::u32(r as u32 * COOP_DIM), + sg_d_base.clone() + c as u32 * COOP_DIM, + ); + } + } + program.workgroup_barrier(); + + // Fold the P·V partial into the output registers, rescaling the + // previous accumulation by this tile's alpha. + let alpha = program.bind(program.load_workgroup(&alpha_arr, o_row.clone())); + for (i, reg) in o_regs.iter().enumerate() { + let partial = program.load_workgroup( + &s_tile, + o_row.clone() * s_stride + o_col_base.clone() + i as u32, + ); + let folded = program.load_local(reg) * alpha.clone() + partial; + program.store_local(reg, folded); + } + // The trailing barrier gates this iteration's `s_tile`/`alpha_arr` + // reads against the next iteration's shared-memory writes. An + // attempted elision here (reasoning the post-staging barrier + // covers it) produced run-to-run nondeterminism and eventual NaNs + // at training step ~300 — the next iteration touches shared + // state before its first collective barrier. + program.workgroup_barrier(); + }); + + // Normalize by the softmax denominator and store. + let inv_l = program.bind(Tile::f32(1.0) / program.load_workgroup(&l_arr, o_row.clone())); + for (i, reg) in o_regs.iter().enumerate() { + program.store( + o.at(o_base.clone() + + o_row.clone() * lo.seq_stride + + (o_col_base.clone() + i as u32) * lo.dim_stride), + cast_to_storage(o, program.load_local(reg) * inv_l.clone()), + Tile::all(), + ); + } + }); + true +} + +// ---- backward ------------------------------------------------------------- + +/// Element strides of one rank-3 `[batch, heads, seq]` row statistic (the +/// forward's log-sum-exp, the backward's `rowsum(dO ∘ O)`). +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct FlashRowLayout { + /// Element offset of `[0, 0, 0]`. + pub offset: u32, + /// Elements between consecutive batches. + pub batch_stride: u32, + /// Elements between consecutive heads. + pub head_stride: u32, + /// Elements between consecutive sequence positions. + pub seq_stride: u32, +} + +/// Per-operand layouts for the backward kernels. `out` is dq for +/// [`flash_bwd_q_f32`]; for [`flash_bwd_kv_f32`] it is the combined dk/dv +/// tensor whose sequence axis spans `2 * kv_len` (dk rows first, dv rows at +/// `kv_len + position`). +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct FlashBwdLayouts { + pub q: FlashOperandLayout, + pub k: FlashOperandLayout, + pub v: FlashOperandLayout, + pub grad_o: FlashOperandLayout, + pub lse: FlashRowLayout, + pub dsum: FlashRowLayout, + pub out: FlashOperandLayout, +} + +/// Whether the backward family ([`flash_lse_f32`], [`flash_bwd_q_f32`], +/// [`flash_bwd_kv_f32`]) can host this shape. Grouped-query K/V is not yet +/// implemented on the kv pass, so the whole family requires `kv_groups == 1`. +pub fn flash_attention_bwd_supported( + shape: &FlashAttentionShape, + subgroups: SubgroupConfig, +) -> bool { + flash_attention_supported(shape, subgroups) + && shape.q_len % BR == 0 + && shape.kv_len % BR == 0 + && shape.kv_groups == 1 +} + +/// One workgroup per `BR`-row query tile of one (batch, head). +pub fn flash_lse_dispatch( + shape: &FlashAttentionShape, + max_workgroups_per_dimension: u32, +) -> [u32; 3] { + dispatch_grid_1d( + shape.batch * shape.heads * (shape.q_len / BR), + max_workgroups_per_dimension, + ) +} + +/// One workgroup per `BR`-row query tile of one (batch, head). +pub fn flash_bwd_q_dispatch( + shape: &FlashAttentionShape, + max_workgroups_per_dimension: u32, +) -> [u32; 3] { + dispatch_grid_1d( + shape.batch * shape.heads * (shape.q_len / BR), + max_workgroups_per_dimension, + ) +} + +/// One workgroup per `BR`-row KV tile of one (batch, head). +pub fn flash_bwd_kv_dispatch( + shape: &FlashAttentionShape, + max_workgroups_per_dimension: u32, +) -> [u32; 3] { + dispatch_grid_1d( + shape.batch * shape.heads * (shape.kv_len / BR), + max_workgroups_per_dimension, + ) +} + +/// The vectorized-staging rank-2 view of an operand, when its head dim is +/// unit-stride. +fn fast_rows_view( + storage: &Storage, + layout: FlashOperandLayout, + total_rows: u32, + d: u32, +) -> Option { + (layout.dim_stride == 1 && layout.seq_stride > 0) + .then(|| storage.restride([total_rows, d], [layout.seq_stride, 1])) +} + +/// Stage `rows` consecutive sequence rows into a `[rows][d + 1]` workgroup +/// tile: `fill_tile`'s vectorized path via the rank-2 view when available, +/// otherwise a scalar strided loop. +#[allow(clippy::too_many_arguments)] +fn stage_rows( + program: &mut TileBlock, + src: &Storage, + fast: &Option, + layout: FlashOperandLayout, + tile: &fusor_tile_ir::tile::WorkgroupTile, + b: &Tile, + h: &Tile, + seq_base: &Tile, + rows: u32, + d: u32, + lane: &Tile, + block: u32, +) { + if let Some(view) = fast { + let (row0, col0) = rank2_origins(program, layout, b, h); + program.fill_tile(tile, view, row0 + seq_base.clone(), col0); + return; + } + let base = program.bind( + b.clone() * layout.batch_stride + + h.clone() * layout.head_stride + + seq_base.clone() * layout.seq_stride + + layout.offset, + ); + let elems = rows * d; + debug_assert_eq!(elems % block, 0); + let stride = d + 1; + for i in 0..elems / block { + let flat = program.bind(lane.clone() + i * block); + let r = program.bind(flat.clone() / d); + let c = program.bind(flat % d); + let value = program.load( + src.at(base.clone() + r.clone() * layout.seq_stride + c.clone() * layout.dim_stride), + Tile::all(), + 0.0, + ); + program.store_workgroup(tile, r * stride + c, value); + } +} + +/// Load `rows` per-row statistics into a workgroup array (one lane each). +#[allow(clippy::too_many_arguments)] +fn load_row_stats( + program: &mut TileBlock, + src: &Storage, + layout: FlashRowLayout, + arr: &fusor_tile_ir::tile::WorkgroupTile, + b: &Tile, + h: &Tile, + seq_base: &Tile, + rows: u32, + lane: &Tile, +) { + let base = program.bind( + b.clone() * layout.batch_stride + + h.clone() * layout.head_stride + + seq_base.clone() * layout.seq_stride + + layout.offset, + ); + let seq_stride = layout.seq_stride; + program.if_then(lane.clone().lt(rows), |program| { + let value = program.load( + src.at(base.clone() + lane.clone() * seq_stride), + Tile::all(), + 0.0, + ); + program.store_workgroup(arr, lane.clone(), value); + }); +} + +/// Store a `[rows][d]` f32 accumulator grid — each of the four subgroups +/// owning the `d/4`-column slice at `subgroup_id * d/4` — to `out` at +/// `out_base + r * seq_stride + col * dim_stride`, staged through the given +/// dead f32 `[rows][BC + 1]` tiles. One subgroup stages per tile per round +/// (guarded coop stores are subgroup-uniform), then every lane copies out. +/// Requires `d / 4 <= BC` so a subgroup slice fits a tile; the last loop +/// iteration's trailing barrier covers the first round's tile reuse. +#[allow(clippy::too_many_arguments)] +fn store_acc_grid_chunked( + program: &mut TileBlock, + coop: CoopMatrixToken, + subgroups: SubgroupConfig, + accs: &[Vec], + tiles: &[&fusor_tile_ir::tile::WorkgroupTile], + out: &Storage, + layout: FlashOperandLayout, + out_base: &Tile, + rows: u32, + d: u32, + lane: &Tile, + block: u32, +) { + let chunk_cols = d / 4; + debug_assert!(chunk_cols <= BC && (rows * chunk_cols).is_multiple_of(block)); + let stride = BC + 1; + let sgid = subgroups.token().subgroup_id(program); + let mut sg = 0u32; + while sg < 4 { + if sg > 0 { + // The previous round's copies finish before the tiles are reused. + program.workgroup_barrier(); + } + let in_round = (tiles.len() as u32).min(4 - sg); + for i in 0..in_round { + let owner = sg + i; + program.if_then(sgid.clone().eq(owner), |program| { + for (r, row_accs) in accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + tiles[i as usize], + Tile::u32(r as u32 * COOP_DIM), + Tile::u32(c as u32 * COOP_DIM), + ); + } + } + }); + } + program.workgroup_barrier(); + for i in 0..in_round { + let col_base = (sg + i) * chunk_cols; + for pass in 0..rows * chunk_cols / block { + let flat = program.bind(lane.clone() + pass * block); + let r = program.bind(flat.clone() / chunk_cols); + let c = program.bind(flat % chunk_cols); + let value = + program.load_workgroup(tiles[i as usize], r.clone() * stride + c.clone()); + program.store( + out.at(out_base.clone() + + r * layout.seq_stride + + (c + col_base) * layout.dim_stride), + cast_to_storage(out, value), + Tile::all(), + ); + } + } + sg += in_round; + } +} + +/// Store a `[rows][d]` f32 accumulator grid to `out` staged through one dead +/// f32 `[rows][d + 1]` tile — the wide-head fallback when a subgroup's +/// column slice exceeds the chunk tiles. +#[allow(clippy::too_many_arguments)] +fn store_acc_grid_whole( + program: &mut TileBlock, + coop: CoopMatrixToken, + subgroups: SubgroupConfig, + accs: &[Vec], + tile: &fusor_tile_ir::tile::WorkgroupTile, + out: &Storage, + layout: FlashOperandLayout, + out_base: &Tile, + rows: u32, + d: u32, + lane: &Tile, + block: u32, +) { + let sg_d_base = program.bind(subgroups.token().subgroup_id(program) * (d / 4)); + for (r, row_accs) in accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + tile, + Tile::u32(r as u32 * COOP_DIM), + sg_d_base.clone() + c as u32 * COOP_DIM, + ); + } + } + program.workgroup_barrier(); + let lanes_per_row = block / rows; + let cols_per_lane = d / lanes_per_row; + let o_row = program.bind(lane.clone() / lanes_per_row); + let o_col_base = program.bind((lane.clone() % lanes_per_row) * cols_per_lane); + let stride = d + 1; + for i in 0..cols_per_lane { + let value = + program.load_workgroup(tile, o_row.clone() * stride + o_col_base.clone() + i); + program.store( + out.at(out_base.clone() + + o_row.clone() * layout.seq_stride + + (o_col_base.clone() + i) * layout.dim_stride), + cast_to_storage(out, value), + Tile::all(), + ); + } +} + +/// Emit the row log-sum-exp kernel: `lse[row] = m + ln Σ exp(scale·q·kᵀ +/// [+ mask] − m)` — the forward statistic that reconstructs probabilities +/// per tile. Returns `false` when the shape fails +/// [`flash_attention_bwd_supported`]. +#[allow(clippy::too_many_arguments)] +pub fn flash_lse_f32( + program: &mut Program, + q: &Storage, + k: &Storage, + mask: Option<(&Storage, FlashMaskLayout)>, + lse_out: &Storage, + q_layout: FlashOperandLayout, + k_layout: FlashOperandLayout, + lse_layout: FlashRowLayout, + shape: FlashAttentionShape, + subgroups: SubgroupConfig, + coop: CoopMatrixToken, + max_workgroups_per_dimension: u32, +) -> bool { + if !flash_attention_bwd_supported(&shape, subgroups) || (shape.causal && mask.is_some()) { + return false; + } + let block = subgroups.block_for_subgroups(4); + let d = shape.head_dim; + let q_tiles = shape.q_len / BR; + let kv_tiles = shape.kv_len / BC; + let scalar = ScalarElement::F32; + + // Stage operands in their own element type: casting f32 operands to + // f16 tiles injects ~2e-4 noise per attention op, which measurably + // degrades training and compounds to NaN within a few hundred steps. + let stage = scalar_of(q.element()); + let q_tile = program.alloc_workgroup_tile_padded(stage, BR, d, 1); + // Stage operands in their own element type: casting f32 operands to + // f16 tiles injects ~2e-4 noise per attention op, which measurably + // degrades training and compounds to NaN within a few hundred steps. + let stage = scalar_of(q.element()); + let k_tile = program.alloc_workgroup_tile_padded(stage, BC, d, 1); + let s_tile = program.alloc_workgroup_tile_padded(scalar, BR, BC, 1); + let m_arr = program.alloc_workgroup_array(scalar, BR); + let l_arr = program.alloc_workgroup_array(scalar, BR); + let s_stride = BC + 1; + + let q_fast = fast_rows_view(q, q_layout, shape.batch * shape.heads * shape.q_len, d); + let k_fast = fast_rows_view(k, k_layout, shape.batch * shape.heads * shape.kv_len, d); + + let grid = flash_lse_dispatch(&shape, max_workgroups_per_dimension); + program.program_grid(block, grid, |program| { + let tile_id = program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; + let bh = tile_id.clone() / q_tiles; + let qt = program.bind(tile_id % q_tiles); + let b = program.bind(bh.clone() / shape.heads); + let h = program.bind(bh % shape.heads); + let kv_h = program.bind(h.clone() / shape.kv_groups); + let q_pos_base = program.bind(qt * BR); + let lane = program.lane(); + + program.if_then(lane.clone().lt(BR), |program| { + program.store_workgroup(&m_arr, lane.clone(), Tile::f32(MASKED_SCORE)); + program.store_workgroup(&l_arr, lane.clone(), Tile::f32(0.0)); + }); + stage_rows( + program, q, &q_fast, q_layout, &q_tile, &b, &h, &q_pos_base, BR, d, &lane, block, + ); + + program.loop_range(kv_tiles, |program, kv_t| { + let kv_pos_base = program.bind(kv_t * BC); + if shape.causal { + program.break_if(kv_pos_base.clone().gt(q_pos_base.clone() + (BR - 1))); + } + stage_rows( + program, k, &k_fast, k_layout, &k_tile, &b, &kv_h, &kv_pos_base, BC, d, &lane, + block, + ); + program.workgroup_barrier(); + + let subgroup_id = subgroups.token().subgroup_id(program); + let sg_row = program.bind(subgroup_id.clone() / 2); + let sg_col = program.bind(subgroup_id % 2); + let s_rows = BR / 2 / COOP_DIM; + let s_cols = BC / 2 / COOP_DIM; + let sg_row_base = program.bind(sg_row * (BR / 2)); + let sg_col_base = program.bind(sg_col * (BC / 2)); + let s_accs = zero_coop_acc_grid(program, coop, scalar, s_rows, s_cols); + for kk in 0..d / COOP_DIM { + let a_frags = + coop_load_a_fragments(program, coop, &q_tile, &sg_row_base, kk, s_rows, stage); + let b_frags = coop_load_b_fragments_transposed( + program, coop, &k_tile, &sg_col_base, kk, s_cols, stage, + ); + coop_mma_grid(program, coop, &s_accs, &a_frags, &b_frags); + } + for (r, row_accs) in s_accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + &s_tile, + sg_row_base.clone() + r as u32 * COOP_DIM, + sg_col_base.clone() + c as u32 * COOP_DIM, + ); + } + } + program.workgroup_barrier(); + + program.if_then(lane.clone().lt(BR), |program| { + let row = lane.clone(); + let q_pos = program.bind(q_pos_base.clone() + row.clone()); + let m_old = program.bind(program.load_workgroup(&m_arr, row.clone())); + let vals: Vec = (0..BC) + .map(|c| { + let raw = program.load_workgroup(&s_tile, row.clone() * s_stride + c) + * shape.scale; + let masked = if shape.causal { + let allowed = (kv_pos_base.clone() + c).le(q_pos.clone()); + Tile::select(allowed, raw, Tile::f32(MASKED_SCORE)) + } else if let Some((mask, lm)) = &mask { + raw + cast_from_storage( + mask, + program.load( + mask.at(q_pos.clone() * lm.q_stride + + (kv_pos_base.clone() + c) * lm.kv_stride + + lm.offset), + Tile::all(), + 0.0, + ), + ) + } else { + raw + }; + program.bind(masked) + }) + .collect(); + let mut m_tile = vals[0].clone(); + for val in &vals[1..] { + m_tile = m_tile.max(val.clone()); + } + let m_new = program.bind(m_old.clone().max(m_tile)); + let alpha = program.bind((m_old - m_new.clone()).exp()); + let mut row_sum = Tile::f32(0.0); + for val in &vals { + row_sum = row_sum + (val.clone() - m_new.clone()).exp(); + } + let l_old = program.load_workgroup(&l_arr, row.clone()); + program.store_workgroup(&l_arr, row.clone(), l_old * alpha + row_sum); + program.store_workgroup(&m_arr, row, m_new); + }); + program.workgroup_barrier(); + }); + + let lse_base = program.bind( + b * lse_layout.batch_stride + + h * lse_layout.head_stride + + q_pos_base * lse_layout.seq_stride + + lse_layout.offset, + ); + program.if_then(lane.clone().lt(BR), |program| { + let m = program.load_workgroup(&m_arr, lane.clone()); + let l = program.load_workgroup(&l_arr, lane.clone()); + program.store( + lse_out.at(lse_base.clone() + lane.clone() * lse_layout.seq_stride), + cast_to_storage(lse_out, m + l.log()), + Tile::all(), + ); + }); + }); + true +} + +/// Emit the dQ kernel: per `BR`-row query tile, stream KV tiles +/// reconstructing `P = exp(scale·q·kᵀ [+ mask] − lse)`, form +/// `dS = P ∘ (dO·vᵀ − dsum) · scale`, and accumulate `dq = Σ dS·k` in +/// cooperative registers. Returns `false` when the shape fails +/// [`flash_attention_bwd_supported`]. +#[allow(clippy::too_many_arguments)] +pub fn flash_bwd_q_f32( + program: &mut Program, + q: &Storage, + k: &Storage, + v: &Storage, + grad_o: &Storage, + lse: &Storage, + dsum: &Storage, + mask: Option<(&Storage, FlashMaskLayout)>, + dq_out: &Storage, + layouts: &FlashBwdLayouts, + shape: FlashAttentionShape, + subgroups: SubgroupConfig, + coop: CoopMatrixToken, + max_workgroups_per_dimension: u32, +) -> bool { + if !flash_attention_bwd_supported(&shape, subgroups) || (shape.causal && mask.is_some()) { + return false; + } + let block = subgroups.block_for_subgroups(4); + let d = shape.head_dim; + let q_tiles = shape.q_len / BR; + let kv_tiles = shape.kv_len / BC; + let scalar = ScalarElement::F32; + let (lq, lk, lv, ldo) = (layouts.q, layouts.k, layouts.v, layouts.grad_o); + + // Stage operands in their own element type: casting f32 operands to + // f16 tiles injects ~2e-4 noise per attention op, which measurably + // degrades training and compounds to NaN within a few hundred steps. + let stage = scalar_of(q.element()); + let q_tile = program.alloc_workgroup_tile_padded(stage, BR, d, 1); + let do_tile = program.alloc_workgroup_tile_padded(stage, BR, d, 1); + // Stage operands in their own element type: casting f32 operands to + // f16 tiles injects ~2e-4 noise per attention op, which measurably + // degrades training and compounds to NaN within a few hundred steps. + let stage = scalar_of(q.element()); + let k_tile = program.alloc_workgroup_tile_padded(stage, BC, d, 1); + let v_tile = program.alloc_workgroup_tile_padded(stage, BC, d, 1); + let s_tile = program.alloc_workgroup_tile_padded(scalar, BR, BC, 1); + let dp_tile = program.alloc_workgroup_tile_padded(scalar, BR, BC, 1); + // dS feeds the dq MMA as f16 A-fragments; the f32 score/dP tiles stay the + // accumulator staging (and later the chunked output staging). + let ds_tile = program.alloc_workgroup_tile_padded(stage, BR, BC, 1); + let lse_arr = program.alloc_workgroup_array(scalar, BR); + let d_arr = program.alloc_workgroup_array(scalar, BR); + let s_stride = BC + 1; + // Output staging: dq columns leave through the dead f32 score/dP tiles + // when a subgroup's `d/4` slice fits; wider heads stage through one + // dedicated f32 tile. + let out_whole = + (d / 4 > BC).then(|| program.alloc_workgroup_tile_padded(scalar, BR, d, 1)); + + let q_fast = fast_rows_view(q, lq, shape.batch * shape.heads * shape.q_len, d); + let do_fast = fast_rows_view(grad_o, ldo, shape.batch * shape.heads * shape.q_len, d); + let k_fast = fast_rows_view(k, lk, shape.batch * shape.heads * shape.kv_len, d); + let v_fast = fast_rows_view(v, lv, shape.batch * shape.heads * shape.kv_len, d); + + let grid = flash_bwd_q_dispatch(&shape, max_workgroups_per_dimension); + program.program_grid(block, grid, |program| { + let tile_id = program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; + let bh = tile_id.clone() / q_tiles; + let qt = program.bind(tile_id % q_tiles); + let b = program.bind(bh.clone() / shape.heads); + let h = program.bind(bh % shape.heads); + let q_pos_base = program.bind(qt * BR); + let lane = program.lane(); + + stage_rows( + program, q, &q_fast, lq, &q_tile, &b, &h, &q_pos_base, BR, d, &lane, block, + ); + stage_rows( + program, grad_o, &do_fast, ldo, &do_tile, &b, &h, &q_pos_base, BR, d, &lane, block, + ); + load_row_stats( + program, lse, layouts.lse, &lse_arr, &b, &h, &q_pos_base, BR, &lane, + ); + load_row_stats( + program, dsum, layouts.dsum, &d_arr, &b, &h, &q_pos_base, BR, &lane, + ); + + // Each subgroup owns `d / 4` output columns across all rows. + let sg_d_base = program.bind(subgroups.token().subgroup_id(program) * (d / 4)); + let dq_rows = BR / COOP_DIM; + let dq_cols = d / 4 / COOP_DIM; + let dq_accs = zero_coop_acc_grid(program, coop, scalar, dq_rows, dq_cols); + let zero_row = Tile::u32(0); + + program.loop_range(kv_tiles, |program, kv_t| { + let kv_pos_base = program.bind(kv_t * BC); + if shape.causal { + program.break_if(kv_pos_base.clone().gt(q_pos_base.clone() + (BR - 1))); + } + stage_rows( + program, k, &k_fast, lk, &k_tile, &b, &h, &kv_pos_base, BC, d, &lane, block, + ); + stage_rows( + program, v, &v_fast, lv, &v_tile, &b, &h, &kv_pos_base, BC, d, &lane, block, + ); + program.workgroup_barrier(); + + // S = q·kᵀ and dP = dO·vᵀ on one 2×2 subgroup grid, K/V read + // through transposed fragment loads. + let subgroup_id = subgroups.token().subgroup_id(program); + let sg_row = program.bind(subgroup_id.clone() / 2); + let sg_col = program.bind(subgroup_id % 2); + let sg_row_base = program.bind(sg_row * (BR / 2)); + let sg_col_base = program.bind(sg_col * (BC / 2)); + let s_rows = BR / 2 / COOP_DIM; + let s_cols = BC / 2 / COOP_DIM; + let s_accs = zero_coop_acc_grid(program, coop, scalar, s_rows, s_cols); + let dp_accs = zero_coop_acc_grid(program, coop, scalar, s_rows, s_cols); + for kk in 0..d / COOP_DIM { + let a_frags = + coop_load_a_fragments(program, coop, &q_tile, &sg_row_base, kk, s_rows, stage); + let b_frags = coop_load_b_fragments_transposed( + program, coop, &k_tile, &sg_col_base, kk, s_cols, stage, + ); + coop_mma_grid(program, coop, &s_accs, &a_frags, &b_frags); + let da_frags = + coop_load_a_fragments(program, coop, &do_tile, &sg_row_base, kk, s_rows, stage); + let db_frags = coop_load_b_fragments_transposed( + program, coop, &v_tile, &sg_col_base, kk, s_cols, stage, + ); + coop_mma_grid(program, coop, &dp_accs, &da_frags, &db_frags); + } + for (r, row_accs) in s_accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + &s_tile, + sg_row_base.clone() + r as u32 * COOP_DIM, + sg_col_base.clone() + c as u32 * COOP_DIM, + ); + } + } + for (r, row_accs) in dp_accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + &dp_tile, + sg_row_base.clone() + r as u32 * COOP_DIM, + sg_col_base.clone() + c as u32 * COOP_DIM, + ); + } + } + program.workgroup_barrier(); + + // dS = P ∘ (dP − dsum) · scale, written to the f16 dS tile (the + // A operand of the dq MMA). + program.if_then(lane.clone().lt(BR), |program| { + let row = lane.clone(); + let q_pos = program.bind(q_pos_base.clone() + row.clone()); + let lse_row = program.bind(program.load_workgroup(&lse_arr, row.clone())); + let d_row = program.bind(program.load_workgroup(&d_arr, row.clone())); + for c in 0..BC { + let raw = program.load_workgroup(&s_tile, row.clone() * s_stride + c) + * shape.scale; + let masked = if shape.causal { + let allowed = (kv_pos_base.clone() + c).le(q_pos.clone()); + Tile::select(allowed, raw, Tile::f32(MASKED_SCORE)) + } else if let Some((mask, lm)) = &mask { + raw + cast_from_storage( + mask, + program.load( + mask.at(q_pos.clone() * lm.q_stride + + (kv_pos_base.clone() + c) * lm.kv_stride + + lm.offset), + Tile::all(), + 0.0, + ), + ) + } else { + raw + }; + let p = program.bind((masked - lse_row.clone()).exp()); + let dp = program.load_workgroup(&dp_tile, row.clone() * s_stride + c); + let ds = p * (dp - d_row.clone()) * shape.scale; + program.store_workgroup(&ds_tile, row.clone() * s_stride + c, ds); + } + }); + program.workgroup_barrier(); + + // dq += dS·k. + for kk in 0..BC / COOP_DIM { + let a_frags = + coop_load_a_fragments(program, coop, &ds_tile, &zero_row, kk, dq_rows, stage); + let b_frags = + coop_load_b_fragments(program, coop, &k_tile, &sg_d_base, kk, dq_cols, stage); + coop_mma_grid(program, coop, &dq_accs, &a_frags, &b_frags); + } + program.workgroup_barrier(); + }); + + // Stage the accumulators through dead f32 tiles and store. + let lo = layouts.out; + let out_base = program.bind( + b * lo.batch_stride + + h * lo.head_stride + + q_pos_base * lo.seq_stride + + lo.offset, + ); + if let Some(out_whole) = &out_whole { + store_acc_grid_whole( + program, coop, subgroups, &dq_accs, out_whole, dq_out, lo, &out_base, BR, d, + &lane, block, + ); + } else { + store_acc_grid_chunked( + program, + coop, + subgroups, + &dq_accs, + &[&s_tile, &dp_tile], + dq_out, + lo, + &out_base, + BR, + d, + &lane, + block, + ); + } + }); + true +} + +/// Which gradients [`flash_bwd_kv_f32`] emits. The independent modes exist so +/// each recognized contraction can lower alone; a later horizontal merge can +/// combine two single-output dispatches back into `Both`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum FlashKvOutputs { + /// dk and dv into one tensor whose sequence axis spans `2·kv_len` + /// (dk rows first, dv rows at `kv_len + position`). + Both, + /// dk only, into a `[batch, heads, kv_len, head_dim]` tensor. + Dk, + /// dv only, into a `[batch, heads, kv_len, head_dim]` tensor. + Dv, +} + +/// Emit the dK/dV kernel: per `BR`-row KV tile, stream query tiles +/// (descending when causal so the break skips the empty upper triangle), +/// reconstruct `Pᵀ` from the forward statistics, and accumulate +/// `dv = Σ Pᵀ·dO` and/or `dk = Σ dSᵀ·q` in cooperative registers per +/// `outputs`. Returns `false` when the shape fails +/// [`flash_attention_bwd_supported`]. +#[allow(clippy::too_many_arguments)] +pub fn flash_bwd_kv_f32( + program: &mut Program, + q: &Storage, + k: &Storage, + v: Option<&Storage>, + grad_o: &Storage, + lse: &Storage, + dsum: Option<&Storage>, + mask: Option<(&Storage, FlashMaskLayout)>, + dkv_out: &Storage, + layouts: &FlashBwdLayouts, + outputs: FlashKvOutputs, + shape: FlashAttentionShape, + subgroups: SubgroupConfig, + coop: CoopMatrixToken, + max_workgroups_per_dimension: u32, +) -> bool { + if !flash_attention_bwd_supported(&shape, subgroups) || (shape.causal && mask.is_some()) { + return false; + } + let emit_dk = outputs != FlashKvOutputs::Dv; + let emit_dv = outputs != FlashKvOutputs::Dk; + let block = subgroups.block_for_subgroups(4); + let d = shape.head_dim; + let kv_tiles = shape.kv_len / BR; + let q_tiles = shape.q_len / BC; + let scalar = ScalarElement::F32; + let (lq, lk, lv, ldo) = (layouts.q, layouts.k, layouts.v, layouts.grad_o); + + // Stage operands in their own element type: casting f32 operands to + // f16 tiles injects ~2e-4 noise per attention op, which measurably + // degrades training and compounds to NaN within a few hundred steps. + let stage = scalar_of(q.element()); + let k_tile = program.alloc_workgroup_tile_padded(stage, BR, d, 1); + // dPᵀ = v·dOᵀ is only needed for dk. + let v_tile = emit_dk.then(|| program.alloc_workgroup_tile_padded(stage, BR, d, 1)); + // Stage operands in their own element type: casting f32 operands to + // f16 tiles injects ~2e-4 noise per attention op, which measurably + // degrades training and compounds to NaN within a few hundred steps. + let stage = scalar_of(q.element()); + let q_tile = program.alloc_workgroup_tile_padded(stage, BC, d, 1); + let do_tile = program.alloc_workgroup_tile_padded(stage, BC, d, 1); + let st_tile = program.alloc_workgroup_tile_padded(scalar, BR, BC, 1); + let dpt_tile = emit_dk.then(|| program.alloc_workgroup_tile_padded(scalar, BR, BC, 1)); + // Pᵀ and dSᵀ feed the dv/dk MMAs as f16 A-fragments; the f32 tiles above + // keep the accumulator staging (and later the chunked output staging). + let pt_tile = emit_dv.then(|| program.alloc_workgroup_tile_padded(stage, BR, BC, 1)); + let dst_tile = emit_dk.then(|| program.alloc_workgroup_tile_padded(stage, BR, BC, 1)); + let lse_arr = program.alloc_workgroup_array(scalar, BC); + let d_arr = emit_dk.then(|| program.alloc_workgroup_array(scalar, BC)); + let st_stride = BC + 1; + // Wide-head output staging fallback, as in the dq kernel. + let out_whole = + (d / 4 > BC).then(|| program.alloc_workgroup_tile_padded(scalar, BR, d, 1)); + + let q_fast = fast_rows_view(q, lq, shape.batch * shape.heads * shape.q_len, d); + let do_fast = fast_rows_view(grad_o, ldo, shape.batch * shape.heads * shape.q_len, d); + let k_fast = fast_rows_view(k, lk, shape.batch * shape.heads * shape.kv_len, d); + let v_fast = v.and_then(|v| fast_rows_view(v, lv, shape.batch * shape.heads * shape.kv_len, d)); + + let grid = flash_bwd_kv_dispatch(&shape, max_workgroups_per_dimension); + program.program_grid(block, grid, |program| { + let tile_id = program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; + let bh = tile_id.clone() / kv_tiles; + let kvt = program.bind(tile_id % kv_tiles); + let b = program.bind(bh.clone() / shape.heads); + let h = program.bind(bh % shape.heads); + let kv_pos_base = program.bind(kvt * BR); + let lane = program.lane(); + + stage_rows( + program, k, &k_fast, lk, &k_tile, &b, &h, &kv_pos_base, BR, d, &lane, block, + ); + if let Some(v_tile) = &v_tile { + let v = v.expect("dk emission requires the values operand"); + stage_rows( + program, v, &v_fast, lv, v_tile, &b, &h, &kv_pos_base, BR, d, &lane, block, + ); + } + + let sg_d_base = program.bind(subgroups.token().subgroup_id(program) * (d / 4)); + let acc_rows = BR / COOP_DIM; + let acc_cols = d / 4 / COOP_DIM; + let dk_accs = + emit_dk.then(|| zero_coop_acc_grid(program, coop, scalar, acc_rows, acc_cols)); + let dv_accs = + emit_dv.then(|| zero_coop_acc_grid(program, coop, scalar, acc_rows, acc_cols)); + let zero_row = Tile::u32(0); + + program.loop_range(q_tiles, |program, i| { + // Causal iterates query tiles descending: once a tile is fully + // above the diagonal every remaining one is too. + let q_t = if shape.causal { + program.bind(Tile::u32(q_tiles - 1) - i) + } else { + program.bind(i) + }; + let q_pos_base = program.bind(q_t * BC); + if shape.causal { + program.break_if((q_pos_base.clone() + (BC - 1)).lt(kv_pos_base.clone())); + } + stage_rows( + program, q, &q_fast, lq, &q_tile, &b, &h, &q_pos_base, BC, d, &lane, block, + ); + stage_rows( + program, grad_o, &do_fast, ldo, &do_tile, &b, &h, &q_pos_base, BC, d, &lane, + block, + ); + load_row_stats( + program, lse, layouts.lse, &lse_arr, &b, &h, &q_pos_base, BC, &lane, + ); + if let Some(d_arr) = &d_arr { + let dsum = dsum.expect("dk emission requires the dsum operand"); + load_row_stats( + program, dsum, layouts.dsum, d_arr, &b, &h, &q_pos_base, BC, &lane, + ); + } + program.workgroup_barrier(); + + // S̃ = k·qᵀ (and dPᵀ = v·dOᵀ when dk is emitted) on one 2×2 + // subgroup grid; Q/dO read through transposed fragment loads. + let subgroup_id = subgroups.token().subgroup_id(program); + let sg_row = program.bind(subgroup_id.clone() / 2); + let sg_col = program.bind(subgroup_id % 2); + let sg_row_base = program.bind(sg_row * (BR / 2)); + let sg_col_base = program.bind(sg_col * (BC / 2)); + let s_rows = BR / 2 / COOP_DIM; + let s_cols = BC / 2 / COOP_DIM; + let st_accs = zero_coop_acc_grid(program, coop, scalar, s_rows, s_cols); + let dpt_accs = + emit_dk.then(|| zero_coop_acc_grid(program, coop, scalar, s_rows, s_cols)); + for kk in 0..d / COOP_DIM { + let a_frags = + coop_load_a_fragments(program, coop, &k_tile, &sg_row_base, kk, s_rows, stage); + let b_frags = coop_load_b_fragments_transposed( + program, coop, &q_tile, &sg_col_base, kk, s_cols, stage, + ); + coop_mma_grid(program, coop, &st_accs, &a_frags, &b_frags); + if let (Some(dpt_accs), Some(v_tile)) = (&dpt_accs, &v_tile) { + let da_frags = coop_load_a_fragments( + program, coop, v_tile, &sg_row_base, kk, s_rows, stage, + ); + let db_frags = coop_load_b_fragments_transposed( + program, coop, &do_tile, &sg_col_base, kk, s_cols, stage, + ); + coop_mma_grid(program, coop, dpt_accs, &da_frags, &db_frags); + } + } + for (r, row_accs) in st_accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + &st_tile, + sg_row_base.clone() + r as u32 * COOP_DIM, + sg_col_base.clone() + c as u32 * COOP_DIM, + ); + } + } + if let (Some(dpt_accs), Some(dpt_tile)) = (&dpt_accs, &dpt_tile) { + for (r, row_accs) in dpt_accs.iter().enumerate() { + for (c, acc) in row_accs.iter().enumerate() { + coop.coop_store_tile( + program, + acc, + dpt_tile, + sg_row_base.clone() + r as u32 * COOP_DIM, + sg_col_base.clone() + c as u32 * COOP_DIM, + ); + } + } + } + program.workgroup_barrier(); + + // Pᵀ over the score tile (and dSᵀ over the dPᵀ tile when dk is + // emitted). Row = KV position, column = query position, + // statistics indexed by query. + program.if_then(lane.clone().lt(BR), |program| { + let kv_row = lane.clone(); + let kv_pos = program.bind(kv_pos_base.clone() + kv_row.clone()); + for c in 0..BC { + let q_pos = program.bind(q_pos_base.clone() + c); + let raw = program + .load_workgroup(&st_tile, kv_row.clone() * st_stride + c) + * shape.scale; + let masked = if shape.causal { + let allowed = kv_pos.clone().le(q_pos.clone()); + Tile::select(allowed, raw, Tile::f32(MASKED_SCORE)) + } else if let Some((mask, lm)) = &mask { + raw + cast_from_storage( + mask, + program.load( + mask.at(q_pos.clone() * lm.q_stride + + kv_pos.clone() * lm.kv_stride + + lm.offset), + Tile::all(), + 0.0, + ), + ) + } else { + raw + }; + let lse_col = program.load_workgroup(&lse_arr, Tile::u32(c)); + let p = program.bind((masked - lse_col).exp()); + if let (Some(d_arr), Some(dpt_tile), Some(dst_tile)) = + (&d_arr, &dpt_tile, &dst_tile) + { + let d_col = program.load_workgroup(d_arr, Tile::u32(c)); + let dp = + program.load_workgroup(dpt_tile, kv_row.clone() * st_stride + c); + let ds = p.clone() * (dp - d_col) * shape.scale; + program.store_workgroup(dst_tile, kv_row.clone() * st_stride + c, ds); + } + if let Some(pt_tile) = &pt_tile { + program.store_workgroup(pt_tile, kv_row.clone() * st_stride + c, p); + } + } + }); + program.workgroup_barrier(); + + // dv += Pᵀ·dO and dk += dSᵀ·q. + for kk in 0..BC / COOP_DIM { + if let (Some(dv_accs), Some(pt_tile)) = (&dv_accs, &pt_tile) { + let a_frags = coop_load_a_fragments( + program, coop, pt_tile, &zero_row, kk, acc_rows, stage, + ); + let b_frags = coop_load_b_fragments( + program, coop, &do_tile, &sg_d_base, kk, acc_cols, stage, + ); + coop_mma_grid(program, coop, dv_accs, &a_frags, &b_frags); + } + if let (Some(dk_accs), Some(dst_tile)) = (&dk_accs, &dst_tile) { + let ka_frags = coop_load_a_fragments( + program, coop, dst_tile, &zero_row, kk, acc_rows, stage, + ); + let kb_frags = coop_load_b_fragments( + program, coop, &q_tile, &sg_d_base, kk, acc_cols, stage, + ); + coop_mma_grid(program, coop, dk_accs, &ka_frags, &kb_frags); + } + } + program.workgroup_barrier(); + }); + + // Stage each accumulator grid through dead f32 tiles and store. In + // `Both` mode dk lands at sequence `kv_pos` and dv at + // `kv_len + kv_pos`; single-output modes write at `kv_pos`. + let dv_seq_offset = if emit_dk { shape.kv_len } else { 0 }; + let lo = layouts.out; + let mut stores: Vec<(u32, &Vec>)> = Vec::new(); + if let Some(dk_accs) = &dk_accs { + stores.push((0, dk_accs)); + } + if let Some(dv_accs) = &dv_accs { + stores.push((dv_seq_offset, dv_accs)); + } + let mut chunk_tiles: Vec<&fusor_tile_ir::tile::WorkgroupTile> = vec![&st_tile]; + if let Some(dpt_tile) = &dpt_tile { + chunk_tiles.push(dpt_tile); + } + for (i, (seq_offset, accs)) in stores.iter().enumerate() { + if i > 0 { + // The first output's copies finish before its staging tiles + // are reused. + program.workgroup_barrier(); + } + let out_base = program.bind( + b.clone() * lo.batch_stride + + h.clone() * lo.head_stride + + (kv_pos_base.clone() + *seq_offset) * lo.seq_stride + + lo.offset, + ); + if let Some(out_whole) = &out_whole { + store_acc_grid_whole( + program, coop, subgroups, accs, out_whole, dkv_out, lo, &out_base, BR, d, + &lane, block, + ); + } else { + store_acc_grid_chunked( + program, + coop, + subgroups, + accs, + &chunk_tiles, + dkv_out, + lo, + &out_base, + BR, + d, + &lane, + block, + ); + } + } + }); + true +} diff --git a/fusor-ml/tile-ir-kernels/src/kernels/helpers.rs b/fusor-ml/tile-ir-kernels/src/kernels/helpers.rs index 1892d0225..805c06c12 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/helpers.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/helpers.rs @@ -7,8 +7,8 @@ use crate::types::QmatmulExtra; /// accumulator elements are [`ScalarElement`] data. The `F32 -> F32` / /// `F16 -> F16` cases are identity; the `F16 -> F32` case inserts the cast /// pair that lets F16 storage be loaded into F32 accumulators and stored back. -/// Used by the unified `batched_matmul_with_epilogues` / `batched_gemv_*` -/// kernels so we don't have to duplicate every body per (storage, accum) pair. +/// Used by the workgroup-tiled quantized kernels so we don't have to duplicate +/// every body per (storage, accum) pair. #[derive(Copy, Clone, Debug, PartialEq, Eq)] pub struct AccumCast { storage: ScalarElement, @@ -69,7 +69,7 @@ impl AccumCast { } } -fn zero_literal(element: ScalarElement) -> TileLiteral { +pub(super) fn zero_literal(element: ScalarElement) -> TileLiteral { match element { ScalarElement::F32 => TileLiteral::f32(0.0), ScalarElement::F16 => TileLiteral::F16(0), @@ -200,6 +200,33 @@ pub(super) fn coop_load_b_fragments( .collect() } +/// Cooperatively load `cols` B-role 8x8 fragments of a workgroup tile's +/// TRANSPOSE: the tile holds the operand row-major and each fragment reads +/// its transpose, so no staged transpose copy is needed. +pub(super) fn coop_load_b_fragments_transposed( + program: &TileBlock<'_>, + coop: CoopMatrixToken, + tile: &WorkgroupTile, + sg_col_base: &Tile, + kk: u32, + cols: u32, + scalar: ScalarElement, +) -> Vec { + (0..cols) + .map(|c| { + coop.coop_load_b_transposed( + program, + tile, + kk * COOP_DIM, + sg_col_base.clone() + c * COOP_DIM, + scalar, + COOP_DIM, + COOP_DIM, + ) + }) + .collect() +} + /// Cooperatively load `cols` C-role fragments from a rank-1 column vector, /// broadcasting each 8-column slice across the fragment rows. pub(super) fn coop_load_c_broadcast_fragments( @@ -247,9 +274,56 @@ pub(super) fn coop_mma_grid( /// 1D-logical workgroup count dispatched as a 3D grid clamped to /// `max_per_dim` in each axis. Shared by dense and quantized matmul /// dispatch paths. +/// Clamp a spread grid's overhang workgroup ids to the last valid unit of +/// work. Cooperative ops require uniform control flow, so an early return +/// fails naga validation; the clamp is branch-free — overhang workgroups +/// redundantly recompute the final tile and store identical values, which +/// is deterministic. Emitted only when the grid over-covers, keeping +/// exact-grid kernels' IR unchanged. +pub(super) fn clamp_grid_overhang( + program: &mut TileBlock<'_>, + id: Tile, + total: u32, + grid: [u32; 3], +) -> Tile { + let covered = grid[0] as u64 * grid[1] as u64 * grid[2] as u64; + if covered == u64::from(total) { + return id; + } + program.bind(id.min(Tile::u32(total - 1))) +} + pub(super) fn dispatch_grid_1d(total_workgroups: u32, max_per_dim: u32) -> [u32; 3] { assert!(total_workgroups > 0, "matmul dispatch must have workgroups"); assert!(max_per_dim > 0, "max_per_dim must be non-zero"); + if total_workgroups <= max_per_dim { + return [total_workgroups, 1, 1]; + } + // Prefer an exact factorization: tile totals are products of tile-grid + // dims and batch, so a divisor at or under the cap almost always exists + // (65536 tiles = 256 x 256), and an exact grid needs no overhang clamp. + let mut x = max_per_dim; + while x > 1 { + if total_workgroups.is_multiple_of(x) { + let rest = total_workgroups / x; + if rest <= max_per_dim { + return [x, rest, 1]; + } + let mut y = max_per_dim; + while y > 1 { + if rest.is_multiple_of(y) && rest / y <= max_per_dim { + return [x, y, rest / y]; + } + y -= 1; + } + } + // Bound the scan: totals this composite always factor quickly; a + // pathological near-prime falls through to the covering grid. + if max_per_dim - x > 4096 { + break; + } + x -= 1; + } let x = total_workgroups.min(max_per_dim); let y_needed = total_workgroups.div_ceil(x); let y = y_needed.min(max_per_dim); diff --git a/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs b/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs index 022b211a5..d51d888b8 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/matmul.rs @@ -5,13 +5,21 @@ use fusor_tile_ir::{CoopMatrixToken, ScalarElement, SubgroupToken, TileLiteral, use crate::{ dispatch::SubgroupConfig, + kernels::helpers::zero_coop_acc_grid, kernels::helpers::{ coop_load_a_fragments, coop_load_b_fragments, coop_mma_grid, coop_store_acc_grid, - dispatch_grid_1d, scalar_of, zero_coop_acc_grid, + clamp_grid_overhang, dispatch_grid_1d, scalar_of, zero_literal, + }, + types::{ + DenseMatmulEpilogues, UnaryEpilogue, apply_optional_epilogue, + cooperative_store_layout_supported, }, - types::{cooperative_store_layout_supported, DenseMatmulEpilogues}, }; +/// Cooperative-matrix fragment side: every per-subgroup fragment grid +/// counts whole 8x8 fragments. +const COOP_DIM: u32 = 8; + /// Logical shape for flattened batched dense matmul views. #[derive(Clone, Copy, Debug)] pub struct DenseMatmulShape { @@ -44,31 +52,129 @@ pub struct DenseCoopMatmulTile { pub bk: u32, } +impl DenseCoopMatmulTile { + /// Elements in one staged A/B workgroup-tile pair: a `bm x bk` A tile + /// plus a `bk x (bn / n_passes)` B tile, each row carrying one pad + /// element against shared-memory bank conflicts. A padded tile spans + /// `rows * (cols + 1) - 1` elements — the pad after its last row is + /// never addressed and is not allocated. + pub const fn stage_pair_elements(self, n_passes: u32) -> u64 { + let bn_pass = (self.bn / n_passes) as u64; + let a_tile = self.bm as u64 * (self.bk as u64 + 1) - 1; + let b_tile = self.bk as u64 * (bn_pass + 1) - 1; + a_tile + b_tile + } +} + /// Capability and tile selection for a cooperative dense matmul attempt. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct DenseCoopMatmulConfig { pub coop: CoopMatrixToken, pub subgroups: SubgroupConfig, pub tile: DenseCoopMatmulTile, + /// How the workgroup's subgroups tile the output block. Derived by + /// [`CoopTileEntry::subgroup_split`]; it travels on the config so a + /// forced-geometry experiment can override it without editing the table. + pub row_groups: u32, + pub col_groups: u32, + /// Staged operand tile pairs, 1 or 2. Two pairs let the next K-tile fill + /// while the current one's MMAs issue; one pair halves the workgroup's + /// threadgroup footprint, so more workgroups sit resident per core and + /// hide each other's epilogue drain. Which wins is a per-contraction + /// trade the caller scores; `single_buffered` entries force 1. + pub stage_buffers: u32, + /// Traversal-order parameter: the tile swizzle walks the grid in + /// super-blocks of this many M-lines so a resident wavefront covers a + /// near-square output patch (see [`DEFAULT_SWIZZLE_GROUP_M`]). + pub swizzle_group_m: u32, + /// Stage operands through workgroup tiles of this element instead of the + /// storage element. `Some(F16)` over f32 storage halves the staged bytes + /// and shared-memory footprint while accumulating in f32 — operands + /// round to f16, products do not. Ignored unless the storage is f32. + pub staging: Option, } #[derive(Clone, Copy)] -struct CoopTileEntry { - tile: DenseCoopMatmulTile, - row_groups: u32, - col_groups: u32, - n_passes: u32, - single_buffered: bool, +pub struct CoopTileEntry { + pub tile: DenseCoopMatmulTile, + /// Subgroups per workgroup. Sets `block`; a compiled kernel property. + /// Its factorization into rows x cols is not data — see + /// [`Self::subgroup_split`]. + pub subgroups: u32, + pub n_passes: u32, + pub single_buffered: bool, } impl CoopTileEntry { + /// `(row_groups, col_groups)`: how this entry's subgroups tile the output + /// block. At a fixed subgroup count `S` the per-subgroup fragment grid has + /// invariant area (`bm * bn_pass / (64 * S)`) and the staged operand bytes + /// do not depend on the factorization at all, so the only quantity it + /// moves is threadgroup fragment loads: a subgroup issues `tr + tc` loads + /// to feed `tr * tc` MMAs per kk-step, which summed over the workgroup is + /// `cg * bm + rg * bn_pass`. Minimize that subject to both fragment sides + /// staying whole multiples of `COOP_DIM`; ties keep the smaller `rg`, + /// which reproduces the table's hand-set 64x128 row. + pub const fn subgroup_split(self) -> (u32, u32) { + let bn_pass = self.tile.bn / self.n_passes; + let mut best_rg = 0; + let mut best_loads = 0; + let mut rg = 1; + while rg <= self.subgroups { + let cg = self.subgroups / rg; + if self.subgroups % rg == 0 + && self.tile.bm % (COOP_DIM * rg) == 0 + && bn_pass % (COOP_DIM * cg) == 0 + { + let loads = cg * self.tile.bm + rg * bn_pass; + if best_rg == 0 || loads < best_loads { + best_rg = rg; + best_loads = loads; + } + } + rg += 1; + } + // Total for a hypothetical row whose sides admit no legal split; the + // kernels' own divisibility asserts catch it at build time. + if best_rg == 0 { + (1, self.subgroups) + } else { + (best_rg, self.subgroups / best_rg) + } + } + const fn block(self, subgroups: SubgroupConfig) -> u32 { - subgroups.block_for_subgroups(self.row_groups * self.col_groups) + subgroups.block_for_subgroups(self.subgroups) + } + + /// Workgroup-memory footprint of this entry's single-pass kernel in + /// bytes: one staged A/B pair of the given stage element, doubled unless + /// the entry is single-buffered. Asserted equal to the lowered IR's + /// `workgroup_bytes` per entry in `tests/footprint.rs`. + pub const fn workgroup_bytes(self, stage: ScalarElement) -> u64 { + self.workgroup_bytes_at(stage, if self.single_buffered { 1 } else { 2 }) + } + + /// The same footprint at an explicit staging depth. The residency term + /// reads this: `max_workgroup_storage / workgroup_bytes_at(..)` is how + /// many workgroups a core can hold. + pub const fn workgroup_bytes_at(self, stage: ScalarElement, buffers: u32) -> u64 { + self.tile.stage_pair_elements(self.n_passes) * stage.byte_size() * buffers as u64 } } -/// Try to emit a fast cooperative-matrix batched matmul. Returns false -/// when shape/layout/epilogues require the generic path. The storage element +/// The full cooperative-matrix tile candidate set, geometry plus the static +/// execution properties a selection cost model scores over. This table is +/// the single source of truth for coop tile geometry; selection layers must +/// derive from it rather than duplicating rows. +pub fn coop_tile_entries() -> &'static [CoopTileEntry] { + COOP_TILE_TABLE +} + +/// Try to emit a fast cooperative-matrix batched matmul. Optional unary +/// pre-epilogues run while staging A/B; a post-epilogue runs over the +/// workgroup's output tile after the cooperative store. Returns false when +/// shape/layout requirements need the generic path. The storage element /// travels in the bound [`Storage`] views, so both F32 and F16 use the same /// runtime dispatch table. pub fn try_batched_coop_matmul( @@ -84,6 +190,11 @@ pub fn try_batched_coop_matmul( coop, subgroups, tile, + row_groups, + col_groups, + staging, + stage_buffers, + swizzle_group_m, } = config; let subgroup = subgroups.token(); let DenseCoopMatmulTile { bm, bn, bk } = tile; @@ -92,114 +203,17 @@ pub fn try_batched_coop_matmul( // to `ceil(m / bm) * bm` per batch and its columns to `ceil(n / bn) * bn` // (the stores cover whole tiles; the pad region holds garbage the // logical view never reads). - if !subgroups.is_fixed() - || epilogues.pre_a.is_some() - || epilogues.pre_b.is_some() - || epilogues.post.is_some() - || !cooperative_store_layout_supported(y.layout()) - { + if !subgroups.is_fixed() || !cooperative_store_layout_supported(y.layout()) { return false; } - let total_tiles = shape.batch * shape.m.div_ceil(bm) * shape.n.div_ceil(bn); - if total_tiles > max_workgroups_per_dimension { - return false; - } - - // Tile geometry per supported (bm, bn, bk). bk=16 across the board keeps - // the double-buffered workgroup tile footprint inside Apple's 32 KB - // threadgroup-memory limit; with bk=32 the per-WG shared memory for the - // bigger BM/BN variants overflows (e.g. Tile128x64 bk=32 double-buffer - // = ~50 KB). The (256, 256, 16) entry runs single-buffered because the - // 256×K A tile would exceed the limit when doubled; its single-buffer - // overhead is amortized by halving global A reads vs (128, 512, 16). - // - // Schema: (bm, bn, bk, row_groups, col_groups, n_passes, single_buffered). - const COOP_TILE_TABLE: &[CoopTileEntry] = &[ - CoopTileEntry { - tile: DenseCoopMatmulTile { - bm: 256, - bn: 256, - bk: 16, - }, - row_groups: 8, - col_groups: 1, - n_passes: 8, - single_buffered: true, - }, - CoopTileEntry { - tile: DenseCoopMatmulTile { - bm: 128, - bn: 512, - bk: 16, - }, - row_groups: 4, - col_groups: 2, - n_passes: 8, - single_buffered: false, - }, - CoopTileEntry { - tile: DenseCoopMatmulTile { - bm: 128, - bn: 256, - bk: 16, - }, - row_groups: 4, - col_groups: 2, - n_passes: 4, - single_buffered: false, - }, - CoopTileEntry { - tile: DenseCoopMatmulTile { - bm: 128, - bn: 128, - bk: 16, - }, - row_groups: 4, - col_groups: 4, - n_passes: 2, - single_buffered: false, - }, - CoopTileEntry { - tile: DenseCoopMatmulTile { - bm: 128, - bn: 64, - bk: 16, - }, - row_groups: 4, - col_groups: 2, - n_passes: 1, - single_buffered: false, - }, - CoopTileEntry { - tile: DenseCoopMatmulTile { - bm: 64, - bn: 128, - bk: 16, - }, - row_groups: 2, - col_groups: 4, - n_passes: 2, - single_buffered: false, - }, - CoopTileEntry { - tile: DenseCoopMatmulTile { - bm: 64, - bn: 64, - bk: 16, - }, - row_groups: 2, - col_groups: 2, - n_passes: 1, - single_buffered: false, - }, - ]; - let Some(entry) = COOP_TILE_TABLE.iter().find(|entry| entry.tile == tile) else { + let Some(entry) = coop_tile_entry(tile) else { return false; }; let block = entry.block(subgroups); - if entry.single_buffered { + if entry.single_buffered || stage_buffers <= 1 { batched_coop_matmul_perf_single( program, + staging, a, b, y, @@ -211,14 +225,16 @@ pub fn try_batched_coop_matmul( bm, bn, bk, - entry.row_groups, - entry.col_groups, + row_groups, + col_groups, entry.n_passes, subgroups, + epilogues, ); } else { batched_coop_matmul_perf( program, + staging, a, b, y, @@ -230,18 +246,857 @@ pub fn try_batched_coop_matmul( bm, bn, bk, - entry.row_groups, - entry.col_groups, + row_groups, + col_groups, entry.n_passes, subgroups, + epilogues, + swizzle_group_m, ); } true } +/// Tile geometry per supported (bm, bn, bk). Each entry's workgroup-memory +/// footprint is [`CoopTileEntry::workgroup_bytes`] over +/// [`DenseCoopMatmulTile::stage_pair_elements`], asserted equal to the +/// lowered IR per entry in `tests/footprint.rs`. bk=16 across the board +/// keeps every double-buffered f32 entry inside Apple's 32 KB +/// threadgroup-memory limit (bk=32 overflows the bigger BM/BN variants), +/// and `single_buffered` is exactly "two f32 pairs would exceed that limit" +/// (also asserted there): the (256, 256, 16) entry trades load/MMA overlap +/// for fitting, amortized by halving global A reads vs (128, 512, 16). +const COOP_TILE_TABLE: &[CoopTileEntry] = &[ + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 256, + bn: 256, + bk: 16, + }, + subgroups: 8, + n_passes: 8, + single_buffered: true, + }, + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 128, + bn: 512, + bk: 16, + }, + subgroups: 8, + n_passes: 8, + single_buffered: false, + }, + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 128, + bn: 256, + bk: 16, + }, + subgroups: 8, + n_passes: 4, + single_buffered: false, + }, + // The original (4, 4) profile — the table's only 16-subgroup, 512-lane + // configuration — miscomputed (all-zero output even on aligned shapes; + // caught by `coop_tile_conformance`). Re-profiled into the proven + // 8-subgroup family: per pass this is exactly the 128x256 entry's + // per-subgroup geometry with half the passes. + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 128, + bn: 128, + bk: 16, + }, + subgroups: 8, + n_passes: 2, + single_buffered: false, + }, + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 128, + bn: 64, + bk: 16, + }, + subgroups: 8, + n_passes: 1, + single_buffered: false, + }, + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 64, + bn: 128, + bk: 16, + }, + subgroups: 8, + n_passes: 2, + single_buffered: false, + }, + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 64, + bn: 64, + bk: 16, + }, + subgroups: 4, + n_passes: 1, + single_buffered: false, + }, + // Small-side tiles for contractions with a 16-wide M or N side + // (attention head_dim contractions, narrow-vocab heads). Fragment + // sides stay multiples of COOP_DIM=8: 64×16 splits into 32×8 + // per-subgroup fragments, 16×64 into 8×32, and so on. + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 64, + bn: 16, + bk: 16, + }, + subgroups: 4, + n_passes: 1, + single_buffered: false, + }, + CoopTileEntry { + tile: DenseCoopMatmulTile { + bm: 16, + bn: 64, + bk: 16, + }, + subgroups: 4, + n_passes: 1, + single_buffered: false, + }, +]; + +fn coop_tile_entry(tile: DenseCoopMatmulTile) -> Option<&'static CoopTileEntry> { + COOP_TILE_TABLE.iter().find(|entry| entry.tile == tile) +} + +/// Split-K partials for a starved cooperative-matrix tile grid: dispatch +/// `splits × total_tiles` workgroups, each running the coop K loop over one +/// contiguous span of `ceil(k_iterations / splits)` K-tiles and storing its +/// partial accumulator to `y` (the scratch buffer) at split-major rows — +/// row `(split · batch + b) · m_padded + m`. A combine kernel +/// ([`split_k_combine`]) then folds the `splits` partials into the real +/// output. Only the sum order changes versus the single-pass kernel. +/// +/// Returns false when the tile geometry is unsupported (unknown or +/// single-buffered table entries) or the grid exceeds the dispatch limit. +/// +/// K bounds are skipped automatically when the spans partition K exactly. +pub fn try_batched_coop_matmul_split_k( + program: &mut Program, + tensors: DenseMatmulTensors<'_>, + shape: DenseMatmulShape, + splits: u32, + max_workgroups_per_dimension: u32, + config: DenseCoopMatmulConfig, +) -> bool { + let DenseMatmulTensors { a, b, y } = tensors; + let DenseCoopMatmulConfig { + coop, + subgroups, + tile, + row_groups, + col_groups, + staging, + // The partials body always stages one pair; see `buffers` below. + stage_buffers: _, + // The split grid is starved by construction; traversal order has no + // resident wavefront to shape. + swizzle_group_m: _, + } = config; + let subgroup = subgroups.token(); + let DenseCoopMatmulTile { bm, bn, bk } = tile; + if !subgroups.is_fixed() || splits < 2 || !cooperative_store_layout_supported(y.layout()) { + return false; + } + let Some(entry) = coop_tile_entry(tile) else { + return false; + }; + // The split path targets tiny tile grids, which never select the + // single-buffered (256, 256) geometry. + if entry.single_buffered { + return false; + } + let tiles_m = shape.m.div_ceil(bm); + let tiles_n = shape.n.div_ceil(bn); + let total_tiles = shape.batch * tiles_m * tiles_n; + let Some(total_workgroups) = splits.checked_mul(total_tiles) else { + return false; + }; + + let block = entry.block(subgroups); + let bn_pass: u32 = bn / entry.n_passes; + let subgroup_rows: u32 = bm / row_groups; + let subgroup_cols_per_pass: u32 = bn_pass / col_groups; + let tile_rows_per_sg: u32 = subgroup_rows / COOP_DIM; + let tile_cols_per_sg: u32 = subgroup_cols_per_pass / COOP_DIM; + let scalar = scalar_of(a.element()); + let stage_scalar = staging + .filter(|_| scalar == ScalarElement::F32) + .unwrap_or(scalar); + + let k_iterations = shape.k.div_ceil(bk); + let span_iters = k_iterations.div_ceil(splits); + let m_padded = tiles_m * bm; + + // The partials body always stages one pair. A split grid exists to raise + // occupancy, and a second pair halves how many of its workgroups a core + // can hold: measured directly, double-buffering this loop costs 55% of + // wall time on 64x2048x256 at identical tile, splits and grid. + let buffers = 1; + let a_tiles: Vec<_> = (0..buffers) + .map(|_| program.alloc_workgroup_tile_padded(stage_scalar, bm, bk, 1)) + .collect(); + let b_tiles: Vec<_> = (0..buffers) + .map(|_| program.alloc_workgroup_tile_padded(stage_scalar, bk, bn_pass, 1)) + .collect(); + + let grid = dispatch_grid_1d(total_workgroups, max_workgroups_per_dimension); + program.program_grid(block, grid, |program| { + let wg_id = program.bind( + program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1], + ); + let wg_id = clamp_grid_overhang(program, wg_id, total_workgroups, grid); + let split = program.bind(wg_id.clone() / total_tiles); + let tile_id = wg_id % total_tiles; + let batch = tile_id.clone() / (tiles_m * tiles_n); + let local_tile = tile_id % (tiles_m * tiles_n); + let m_tile = local_tile.clone() / tiles_n; + let n_tile = local_tile % tiles_n; + let row_base = m_tile * bm; + let col_base = n_tile * bn; + let a_batch_base = batch.clone() * shape.m; + let b_batch_base = batch.clone() * shape.k; + // Split-major scratch rows: slice `split` holds one full padded + // [batch · m_padded, n_padded] partial. + let y_batch_base = (split.clone() * shape.batch + batch) * m_padded; + // The K bound is live only when a span can overrun the logical K + // extent (K not dividing the tile, or the spans not covering K + // exactly). A live bound forces the tile fills onto the scalar + // per-element path and off the vec4 staging fast path. + let k_spans_aligned = shape.k.is_multiple_of(bk) && k_iterations.is_multiple_of(splits); + let a_bounds: [Option; 2] = [ + (!shape.m.is_multiple_of(bm)).then(|| a_batch_base.clone() + shape.m), + (!k_spans_aligned).then(|| Tile::literal(TileLiteral::U32(shape.k))), + ]; + let b_bounds: [Option; 2] = [ + (!k_spans_aligned).then(|| b_batch_base.clone() + shape.k), + (!shape.n.is_multiple_of(bn)).then(|| Tile::literal(TileLiteral::U32(shape.n))), + ]; + + let subgroup_id = subgroup.subgroup_id(program); + let sg_row = subgroup_id.clone() / col_groups; + let sg_col = subgroup_id % col_groups; + let sg_row_base = sg_row * subgroup_rows; + let sg_col_base_in_pass = sg_col * subgroup_cols_per_pass; + + let span_base = program.bind(split * span_iters); + coop_perf_pass_loop( + program, + coop, + scalar, + entry.n_passes, + bn_pass, + tile_rows_per_sg, + tile_cols_per_sg, + y, + &y_batch_base, + &row_base, + &col_base, + &sg_row_base, + &sg_col_base_in_pass, + bm, + block, + None, + |program, pass_col_base, accs| { + coop_k_span_pairs( + program, + &span_base, + span_iters, + bk, + buffers, + |program, k_base, buffer| { + coop_stage_and_mma( + program, + coop, + a, + b, + &a_tiles[buffer], + &b_tiles[buffer], + &a_batch_base, + &b_batch_base, + &row_base, + pass_col_base, + k_base, + &sg_row_base, + &sg_col_base_in_pass, + &a_bounds, + &b_bounds, + accs, + tile_rows_per_sg, + tile_cols_per_sg, + bk, + COOP_DIM, + stage_scalar, + block, + bm, + bn_pass, + None, + None, + ); + }, + ); + }, + ); + }); + true +} + +/// One kernel running several independent same-shape cooperative-matrix +/// matmuls: each segment owns a contiguous range of workgroups guarded by a +/// uniform linear-workgroup-id range compare (the same discipline as the +/// merged n-ary and row-program kernels), and runs the standard coop tile +/// body over its own `a`/`b`/`y` bindings. All segments share one logical +/// `shape`, tile geometry, and split factor, so the guarded bodies differ +/// only in their storage bindings and the workgroup tiles are allocated +/// once and reused by every branch (the guards are workgroup-uniform). +/// +/// `splits == 1` runs each segment as the single-pass double-buffered body +/// (numerics identical to [`try_batched_coop_matmul`]); `splits >= 2` runs +/// each segment as the split-K partials body (numerics identical to +/// [`try_batched_coop_matmul_split_k`]) — the caller must follow with +/// [`merged_split_k_combine`] over the same segment order. +/// +/// Returns false when the tile geometry is unsupported or the grid exceeds +/// the dispatch limit; callers then fall back to per-segment kernels. +#[allow(clippy::too_many_arguments)] +pub fn try_merged_coop_matmul( + program: &mut Program, + segments: &[DenseMatmulTensors<'_>], + shape: DenseMatmulShape, + splits: u32, + max_workgroups_per_dimension: u32, + config: DenseCoopMatmulConfig, +) -> bool { + let DenseCoopMatmulConfig { + coop, + subgroups, + tile, + row_groups, + col_groups, + staging, + stage_buffers, + // Merged segments walk per-segment grids; the swizzle applies to the + // standalone dense path. + swizzle_group_m: _, + } = config; + let subgroup = subgroups.token(); + let DenseCoopMatmulTile { bm, bn, bk } = tile; + if segments.is_empty() || !subgroups.is_fixed() { + return false; + } + if segments + .iter() + .any(|segment| !cooperative_store_layout_supported(segment.y.layout())) + { + return false; + } + let Some(entry) = coop_tile_entry(tile) else { + return false; + }; + // Merged bodies stay double-buffer-table only, like the split path. + if entry.single_buffered { + return false; + } + let tiles_m = shape.m.div_ceil(bm); + let tiles_n = shape.n.div_ceil(bn); + let total_tiles = shape.batch * tiles_m * tiles_n; + let Some(per_segment) = splits.max(1).checked_mul(total_tiles) else { + return false; + }; + let Some(total_workgroups) = per_segment.checked_mul(segments.len() as u32) else { + return false; + }; + if total_workgroups > max_workgroups_per_dimension { + return false; + } + + let block = entry.block(subgroups); + let bn_pass: u32 = bn / entry.n_passes; + let subgroup_rows: u32 = bm / row_groups; + let subgroup_cols_per_pass: u32 = bn_pass / col_groups; + let tile_rows_per_sg: u32 = subgroup_rows / COOP_DIM; + let tile_cols_per_sg: u32 = subgroup_cols_per_pass / COOP_DIM; + let scalar = scalar_of(segments[0].a.element()); + let stage_scalar = scalar; + let _ = staging; + + let k_iterations = shape.k.div_ceil(bk); + let m_padded = tiles_m * bm; + let split_k = splits >= 2; + let span_iters = k_iterations.div_ceil(splits.max(1)); + let k_pairs = k_iterations / 2; + let k_remainder = k_iterations % 2; + + // Shared workgroup tiles: every guarded branch has the same geometry. + // Both bodies double-buffer with two pairs, matching the standalone + // kernels. + // Split partials always stage one pair (see `try_batched_coop_matmul_split_k`); + // the single-pass body takes the scored depth. + let buffers = if split_k { 1 } else { stage_buffers.clamp(1, 2) as usize }; + let a_tiles: Vec<_> = (0..buffers) + .map(|_| program.alloc_workgroup_tile_padded(stage_scalar, bm, bk, 1)) + .collect(); + let b_tiles: Vec<_> = (0..buffers) + .map(|_| program.alloc_workgroup_tile_padded(stage_scalar, bk, bn_pass, 1)) + .collect(); + + let grid = dispatch_grid_1d(total_workgroups, max_workgroups_per_dimension); + program.program_grid(block, grid, |program| { + // Keep the flat workgroup id a raw builtin expression: `bind` routes + // through a function-space local, whose loads naga's uniformity + // analysis marks non-uniform — and the segment guards below must be + // uniform control flow for the coop ops (and barriers) inside. + let wg_id = program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; + for (index, segment) in segments.iter().enumerate() { + let DenseMatmulTensors { a, b, y } = *segment; + let base = index as u32 * per_segment; + let in_segment = wg_id.clone().ge(base) & wg_id.clone().lt(base + per_segment); + program.if_then(in_segment, |program| { + let local = program.bind(wg_id.clone() - base); + let (split, tile_id) = if split_k { + ( + Some(program.bind(local.clone() / total_tiles)), + local % total_tiles, + ) + } else { + (None, local) + }; + let batch = tile_id.clone() / (tiles_m * tiles_n); + let local_tile = tile_id % (tiles_m * tiles_n); + let m_tile = local_tile.clone() / tiles_n; + let n_tile = local_tile % tiles_n; + let row_base = m_tile * bm; + let col_base = n_tile * bn; + let a_batch_base = batch.clone() * shape.m; + let b_batch_base = batch.clone() * shape.k; + // Split partials land at split-major scratch rows; the + // single-pass output lands at the padded batch rows. + let y_batch_base = match &split { + Some(split) => (split.clone() * shape.batch + batch) * m_padded, + None => batch * m_padded, + }; + // Bounds mirror the standalone kernels exactly: the split + // path may elide aligned K bounds (vec4 staging fast path), + // the single-pass path keeps K live only for ragged K. + let k_spans_aligned = + split_k && shape.k.is_multiple_of(bk) && k_iterations.is_multiple_of(splits); + let k_bound_live = if split_k { + !k_spans_aligned + } else { + !shape.k.is_multiple_of(bk) + }; + let a_bounds: [Option; 2] = [ + (!shape.m.is_multiple_of(bm)).then(|| a_batch_base.clone() + shape.m), + k_bound_live.then(|| Tile::literal(TileLiteral::U32(shape.k))), + ]; + let b_bounds: [Option; 2] = [ + k_bound_live.then(|| b_batch_base.clone() + shape.k), + (!shape.n.is_multiple_of(bn)).then(|| Tile::literal(TileLiteral::U32(shape.n))), + ]; + + let subgroup_id = subgroup.subgroup_id(program); + let sg_row = subgroup_id.clone() / col_groups; + let sg_col = subgroup_id % col_groups; + let sg_row_base = sg_row * subgroup_rows; + let sg_col_base_in_pass = sg_col * subgroup_cols_per_pass; + + let span_base = split.map(|split| program.bind(split * span_iters)); + coop_perf_pass_loop( + program, + coop, + scalar, + entry.n_passes, + bn_pass, + tile_rows_per_sg, + tile_cols_per_sg, + y, + &y_batch_base, + &row_base, + &col_base, + &sg_row_base, + &sg_col_base_in_pass, + bm, + block, + None, + |program, pass_col_base, accs| { + if let Some(span_base) = &span_base { + // Split-K span: the same double-buffered K-pair + // loop over this split's contiguous K-tile range. + coop_k_span_pairs( + program, + span_base, + span_iters, + bk, + buffers, + |program, k_base, buffer| { + coop_stage_and_mma( + program, + coop, + a, + b, + &a_tiles[buffer], + &b_tiles[buffer], + &a_batch_base, + &b_batch_base, + &row_base, + pass_col_base, + k_base, + &sg_row_base, + &sg_col_base_in_pass, + &a_bounds, + &b_bounds, + accs, + tile_rows_per_sg, + tile_cols_per_sg, + bk, + COOP_DIM, + stage_scalar, + block, + bm, + bn_pass, + None, + None, + ); + }, + ); + return; + } + if buffers == 1 { + program.loop_range(k_iterations, |program, iter_idx| { + let k_base = iter_idx * bk; + coop_stage_and_mma( + program, coop, a, b, &a_tiles[0], &b_tiles[0], + &a_batch_base, &b_batch_base, &row_base, + pass_col_base, &k_base, &sg_row_base, + &sg_col_base_in_pass, &a_bounds, &b_bounds, accs, + tile_rows_per_sg, tile_cols_per_sg, bk, COOP_DIM, + stage_scalar, block, bm, bn_pass, None, None, + ); + program.workgroup_barrier(); + }); + return; + } + // Single-pass: the same K-pair loop written out + // against literal K bases (no span offset to add). + let (a_tile_1, b_tile_1) = (&a_tiles[1], &b_tiles[1]); + if k_pairs > 0 { + program.loop_range(k_pairs, |program, pair_idx| { + let k_base_0 = pair_idx.clone() * (2 * bk); + let k_base_1 = pair_idx * (2 * bk) + bk; + coop_stage_and_mma( + program, + coop, + a, + b, + &a_tiles[0], + &b_tiles[0], + &a_batch_base, + &b_batch_base, + &row_base, + pass_col_base, + &k_base_0, + &sg_row_base, + &sg_col_base_in_pass, + &a_bounds, + &b_bounds, + accs, + tile_rows_per_sg, + tile_cols_per_sg, + bk, + COOP_DIM, + stage_scalar, + block, + bm, + bn_pass, + None, + None, + ); + coop_stage_and_mma( + program, + coop, + a, + b, + a_tile_1, + b_tile_1, + &a_batch_base, + &b_batch_base, + &row_base, + pass_col_base, + &k_base_1, + &sg_row_base, + &sg_col_base_in_pass, + &a_bounds, + &b_bounds, + accs, + tile_rows_per_sg, + tile_cols_per_sg, + bk, + COOP_DIM, + stage_scalar, + block, + bm, + bn_pass, + None, + None, + ); + }); + } + if k_remainder == 1 { + let k_base_epi = + Tile::literal(TileLiteral::U32((k_iterations - 1) * bk)); + coop_stage_and_mma( + program, + coop, + a, + b, + &a_tiles[0], + &b_tiles[0], + &a_batch_base, + &b_batch_base, + &row_base, + pass_col_base, + &k_base_epi, + &sg_row_base, + &sg_col_base_in_pass, + &a_bounds, + &b_bounds, + accs, + tile_rows_per_sg, + tile_cols_per_sg, + bk, + COOP_DIM, + stage_scalar, + block, + bm, + bn_pass, + None, + None, + ); + program.workgroup_barrier(); + } + }, + ); + }); + } + }); + true +} + +/// The merged counterpart of [`split_k_combine`]: one kernel folding the +/// split-K partials of several same-shape segments, each `y` a read-write +/// view of that segment's whole `(1 + splits)`-slice buffer, each segment +/// guarded by its linear-workgroup-id range in the same segment order as +/// [`try_merged_coop_matmul`]. +pub fn merged_split_k_combine( + program: &mut Program, + ys: &[&Storage], + rows: u32, + cols: u32, + splits: u32, + max_workgroups_per_dimension: u32, +) { + const BLOCK: u32 = 256; + let total = rows * cols; + let per_segment = total.div_ceil(BLOCK); + let total_workgroups = per_segment * ys.len() as u32; + let scalar = scalar_of(ys[0].element()); + let zero = zero_literal(scalar); + let grid = dispatch_grid_1d(total_workgroups, max_workgroups_per_dimension); + program.program_grid(BLOCK, grid, |program| { + let wg_id = program.bind( + program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1], + ); + for (index, y) in ys.iter().enumerate() { + let base = index as u32 * per_segment; + let in_segment = wg_id.clone().ge(base) & wg_id.clone().lt(base + per_segment); + program.if_then(in_segment, |program| { + let local = wg_id.clone() - base; + let index = program.bind(local * BLOCK + program.lane()); + let active = index.clone().lt(total); + let row = program.bind(index.clone() / cols); + let col = program.bind(index % cols); + let mut acc = program.load( + y.at((row.clone() + rows, col.clone())), + active.clone(), + zero, + ); + for split in 2..=splits { + acc = acc + + program.load( + y.at((row.clone() + split * rows, col.clone())), + active.clone(), + zero, + ); + } + program.store(y.at((row, col)), acc, active); + }); + } + }); +} + +/// Fold the split-K partials into the output. `y` is one read-write view of +/// the whole `(1 + splits)` -slice buffer — `[(1 + splits) · rows, cols]` +/// where rows `0..rows` are the real (padded) output and slice `s ∈ +/// 1..=splits` holds one partial at rows `s · rows..`. One lane per output +/// element sums the partials and stores slice 0. A single binding keeps the +/// buffer from being bound with conflicting access modes. +pub fn split_k_combine( + program: &mut Program, + y: &Storage, + rows: u32, + cols: u32, + splits: u32, + max_workgroups_per_dimension: u32, +) { + const BLOCK: u32 = 256; + let total = rows * cols; + let scalar = scalar_of(y.element()); + let zero = zero_literal(scalar); + let grid = dispatch_grid_1d(total.div_ceil(BLOCK), max_workgroups_per_dimension); + program.program_grid(BLOCK, grid, |program| { + let wg_id = program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; + let index = program.bind(wg_id * BLOCK + program.lane()); + let active = index.clone().lt(total); + let row = program.bind(index.clone() / cols); + let col = program.bind(index % cols); + let mut acc = program.load( + y.at((row.clone() + rows, col.clone())), + active.clone(), + zero, + ); + for split in 2..=splits { + acc = acc + + program.load( + y.at((row.clone() + split * rows, col.clone())), + active.clone(), + zero, + ); + } + program.store(y.at((row, col)), acc, active); + }); +} + /// Stage one `BK`-tile of A and B into `a_tile`/`b_tile`, barrier, then run the /// `kk` MMA sweep into the accumulator grid. The caller decides the trailing /// barrier; the K-pair shape elides it between halves. +#[allow(clippy::too_many_arguments)] +fn fill_tile_bounded_with_epilogue( + program: &mut TileBlock<'_>, + dst: &fusor_tile_ir::tile::WorkgroupTile, + src: &Storage, + row_base: Tile, + col_base: Tile, + bounds: [Option; 2], + rows: u32, + cols: u32, + padded_stride: u32, + lanes: u32, + epilogue: Option<&UnaryEpilogue>, + scalar: ScalarElement, +) { + let Some(epilogue) = epilogue else { + program.fill_tile_bounded(dst, src, row_base, col_base, bounds); + return; + }; + + let total = rows * cols; + let passes = total.div_ceil(lanes); + for pass in 0..passes { + let flat = program.lane() + pass * lanes; + let local_row = flat.clone() / cols; + let local_col = flat.clone() % cols; + let global_row = row_base.clone() + local_row.clone(); + let global_col = col_base.clone() + local_col.clone(); + let within_tile = flat.lt(total); + let mut active = within_tile.clone(); + if let Some(bound) = &bounds[0] { + active = active & global_row.clone().lt(bound.clone()); + } + if let Some(bound) = &bounds[1] { + active = active & global_col.clone().lt(bound.clone()); + } + let zero = zero_literal(scalar); + let loaded = program.load(src.at((global_row, global_col)), active.clone(), zero); + let transformed = apply_optional_epilogue(Some(epilogue), loaded); + let value = Tile::select(active, transformed, Tile::literal(zero)); + let tile_index = local_row * padded_stride + local_col; + program.if_then(within_tile, |program| { + program.store_workgroup(dst, tile_index, value); + }); + } +} + +#[allow(clippy::too_many_arguments)] +fn apply_post_epilogue_in_place( + program: &mut TileBlock<'_>, + y: &Storage, + y_batch_base: &Tile, + row_base: &Tile, + col_base: &Tile, + rows: u32, + cols: u32, + lanes: u32, + epilogue: Option<&UnaryEpilogue>, + scalar: ScalarElement, +) { + let Some(epilogue) = epilogue else { + return; + }; + + // Cooperative accumulator fragments are opaque to scalar tile-IR. Store + // them first, synchronize storage visibility within the workgroup, then + // map the epilogue over this workgroup's disjoint output tile in place. + program.storage_barrier(); + let y_scalar = scalar_of(y.element()); + let total = rows * cols; + let passes = total.div_ceil(lanes); + for pass in 0..passes { + let flat = program.lane() + pass * lanes; + let local_row = flat.clone() / cols; + let local_col = flat.clone() % cols; + let row = y_batch_base.clone() + row_base.clone() + local_row; + let col = col_base.clone() + local_col; + let active = flat.lt(total); + let loaded = program.load( + y.at((row.clone(), col.clone())), + active.clone(), + zero_literal(y_scalar), + ); + // A dtype-changing chain reads the matmul in its operand dtype while + // the store landed in the chain's (wider) output dtype: rounding the + // exact stored accumulator back down here reproduces the unfused + // matmul's output bit-for-bit before the chain transforms it. + let loaded = if y_scalar != scalar { + loaded.cast(scalar.element()) + } else { + loaded + }; + let value = apply_optional_epilogue(Some(epilogue), loaded); + program.store(y.at((row, col)), value, active); + } +} + #[allow(clippy::too_many_arguments)] fn coop_stage_and_mma( program: &mut TileBlock<'_>, @@ -265,20 +1120,39 @@ fn coop_stage_and_mma( bk: u32, coop_dim: u32, scalar: ScalarElement, + block: u32, + bm: u32, + bn_pass: u32, + pre_a: Option<&UnaryEpilogue>, + pre_b: Option<&UnaryEpilogue>, ) { - program.fill_tile_bounded( + fill_tile_bounded_with_epilogue( + program, a_tile, a, a_batch_base.clone() + row_base.clone(), - k_base, + k_base.clone(), a_bounds.clone(), + bm, + bk, + bk + 1, + block, + pre_a, + scalar, ); - program.fill_tile_bounded( + fill_tile_bounded_with_epilogue( + program, b_tile, b, b_batch_base.clone() + k_base.clone(), - pass_col_base, + pass_col_base.clone(), b_bounds.clone(), + bk, + bn_pass, + bn_pass + 1, + block, + pre_b, + scalar, ); program.workgroup_barrier(); @@ -306,6 +1180,53 @@ fn coop_stage_and_mma( } } +/// A double-buffered K loop over `iters` consecutive `bk`-tiles starting at +/// K-tile index `span_base`. `stage(program, k_base, buffer)` emits one +/// stage-and-MMA against tile pair 0 or 1. +/// +/// Alternating pairs is what exposes the global fill to the MMAs: the only +/// synchronization [`coop_stage_and_mma`] emits is the barrier between its +/// fill and its cooperative loads, so the second half of a K-pair fills +/// while the first half's MMAs are still issuing. Correctness rests on that +/// same barrier — a thread entering the next pair's fill of buffer 0 has +/// passed the barrier the previous stage put after its fill of buffer 1, +/// and every thread finished its buffer-0 MMAs before reaching it. +fn coop_k_span_pairs( + program: &mut TileBlock<'_>, + span_base: &Tile, + iters: u32, + bk: u32, + buffers: usize, + mut stage: impl FnMut(&mut TileBlock<'_>, &Tile, usize), +) { + if buffers == 1 { + program.loop_range(iters, |program, iter_idx| { + let k_base = (span_base.clone() + iter_idx) * bk; + stage(program, &k_base, 0); + // Trailing barrier: the next iteration overwrites the tiles this + // one just read through the coop loads. + program.workgroup_barrier(); + }); + return; + } + let pairs = iters / 2; + if pairs > 0 { + program.loop_range(pairs, |program, pair_idx| { + let first = (span_base.clone() + pair_idx.clone() * 2) * bk; + stage(program, &first, 0); + let second = (span_base.clone() + (pair_idx * 2 + 1)) * bk; + stage(program, &second, 1); + }); + } + if iters % 2 == 1 { + let last = (span_base.clone() + (iters - 1)) * bk; + stage(program, &last, 0); + // The odd tile closes the span on buffer 0, so nothing behind it + // separates it from whatever reuses the tiles next. + program.workgroup_barrier(); + } +} + /// Shared pass-loop scaffolding for the coop-perf matmul variants (single- /// and double-buffered). For each of `N_PASSES` column sub-passes, allocates /// a fresh accumulator grid, runs the caller-supplied K-loop body, then @@ -327,13 +1248,21 @@ fn coop_perf_pass_loop( col_base: &Tile, sg_row_base: &Tile, sg_col_base_in_pass: &Tile, + bm: u32, + block: u32, + post: Option<&UnaryEpilogue>, mut k_body: F, ) where F: FnMut(&mut TileBlock<'_>, &Tile, &[Vec]), { + // Always accumulate in f32, matching the composed contraction's + // accumulator (`as_fused_reduce` upgrades f16 accumulation to f32). + // f16 operands run the mixed f16xf16->f32 MMA at full rate, and an f16 + // output converts the fragment per thread at the store. + let acc_scalar = ScalarElement::F32; for n_pass in 0..n_passes { let pass_col_base = col_base.clone() + n_pass * bn_pass; - let accs = zero_coop_acc_grid(program, coop, scalar, tile_rows_per_sg, tile_cols_per_sg); + let accs = zero_coop_acc_grid(program, coop, acc_scalar, tile_rows_per_sg, tile_cols_per_sg); k_body(program, &pass_col_base, &accs); @@ -348,6 +1277,18 @@ fn coop_perf_pass_loop( sg_row_base, sg_col_base_in_pass, ); + apply_post_epilogue_in_place( + program, + y, + y_batch_base, + row_base, + &pass_col_base, + bm, + bn_pass, + block, + post, + scalar, + ); } } @@ -359,6 +1300,7 @@ fn coop_perf_pass_loop( #[allow(clippy::too_many_arguments)] fn batched_coop_matmul_perf_single( program: &mut Program, + staging: Option, a: &Storage, b: &Storage, y: &Storage, @@ -374,8 +1316,8 @@ fn batched_coop_matmul_perf_single( col_groups: u32, n_passes: u32, subgroups: SubgroupConfig, + epilogues: &DenseMatmulEpilogues<'_>, ) { - const COOP_DIM: u32 = 8; debug_assert!(n_passes >= 1); debug_assert_eq!(bn % n_passes, 0); let bn_pass: u32 = bn / n_passes; @@ -393,6 +1335,10 @@ fn batched_coop_matmul_perf_single( let tile_cols_per_sg: u32 = subgroup_cols_per_pass / COOP_DIM; let scalar = scalar_of(a.element()); + // The single-buffered body loads fragments straight off its persistent + // tiles; it has not been taught mixed staging. + let stage_scalar = scalar; + let _ = staging; let tiles_m = shape.m.div_ceil(bm); let tiles_n = shape.n.div_ceil(bn); @@ -403,8 +1349,8 @@ fn batched_coop_matmul_perf_single( // extents. let m_padded = tiles_m * bm; - let a_tile = program.alloc_workgroup_tile_padded(scalar, bm, bk, 1); - let b_tile = program.alloc_workgroup_tile_padded(scalar, bk, bn_pass, 1); + let a_tile = program.alloc_workgroup_tile_padded(stage_scalar, bm, bk, 1); + let b_tile = program.alloc_workgroup_tile_padded(stage_scalar, bk, bn_pass, 1); let grid = dispatch_grid_1d(total_tiles, max_workgroups_per_dimension); program.program_grid(block, grid, |program| { @@ -449,6 +1395,9 @@ fn batched_coop_matmul_perf_single( &col_base, &sg_row_base, &sg_col_base_in_pass, + bm, + block, + epilogues.post, |program, pass_col_base, accs| { program.loop_range(k_iterations, |program, iter_idx| { let k_base = iter_idx * bk; @@ -473,7 +1422,12 @@ fn batched_coop_matmul_perf_single( tile_cols_per_sg, bk, COOP_DIM, - scalar, + stage_scalar, + block, + bm, + bn_pass, + epilogues.pre_a, + epilogues.pre_b, ); // Trailing barrier required: next iter overwrites the same // tile that this iter just finished reading via coop loads. @@ -484,6 +1438,88 @@ fn batched_coop_matmul_perf_single( }); } +/// Consecutive tile ids that share one B column-slab under the L2 tile-order +/// swizzle in [`batched_coop_matmul_perf`]. +/// Default swizzle M-group: the measured optimum on Apple M-series, and the +/// power of two nearest the square root of the concurrently-resident +/// workgroup count there (a near-square resident output patch minimizes the +/// wavefront's combined operand footprint). Selection derives the value per +/// device; fixed-geometry callers (labs, tests) use this default. +pub const DEFAULT_SWIZZLE_GROUP_M: u32 = 8; + +/// Remap one batch's linear tile index into super-blocked `(m_tile, n_tile)` +/// coordinates so concurrently-resident workgroups share operand slabs +/// (threadblock swizzling for L2 reuse). +/// +/// Row-major order walks a whole row of `tiles_n` output tiles before +/// advancing `m`, so the resident wavefront shares one A row-slab but +/// streams a distinct B column-slab (`k * bn` bytes) per workgroup and +/// re-streams the full B operand once per tile row. The swizzle instead +/// walks the grid in super-blocks of `SWIZZLE_GROUP_M` M-lines, +/// M-fastest: `SWIZZLE_GROUP_M` consecutive workgroups share one B +/// column-slab while touching only `SWIZZLE_GROUP_M` A row-slabs, so a +/// resident wavefront of `R` workgroups covers a near-square +/// `SWIZZLE_GROUP_M x (R / SWIZZLE_GROUP_M)` patch of the output whose +/// operand k-window footprint is minimal — both operands get cache reuse +/// instead of one. +/// +/// The map stays a bijection on `[0, tiles_m * tiles_n)`: when `tiles_m` is +/// not a multiple of the group size, the ragged tail walks its remaining +/// `tiles_m % SWIZZLE_GROUP_M` M-lines in the same order. All divisors are +/// build-time u32 constants (the constant-divisor lowering is the proven +/// path on Apple GPUs; runtime divisors are not). +fn swizzled_tile_coords( + program: &mut TileBlock<'_>, + local_tile: Tile, + tiles_m: u32, + tiles_n: u32, + group: u32, +) -> (Tile, Tile) { + if tiles_m <= 1 || tiles_n <= 1 { + // Degenerate grids: the swizzle is a no-op; keep the plain row-major + // decomposition. + let m_tile = local_tile.clone() / tiles_n; + let n_tile = local_tile % tiles_n; + return (m_tile, n_tile); + } + // Ids [0, threshold) cover the full super-blocks: each spans `group` + // consecutive M-lines by all `tiles_n` N-lines, walked M-fastest. Ids + // [threshold, ..) cover the ragged tail of `tail` M-lines the same way. + let span = group * tiles_n; + let num_full = tiles_m / group; + let tail = tiles_m % group; + let threshold = num_full * span; + let local = program.bind(local_tile); + + let full = (num_full > 0).then(|| { + let group_idx = local.clone() / span; + let in_group = program.bind(local.clone() % span); + let m_tile = group_idx * group + in_group.clone() % group; + let n_tile = in_group / group; + (m_tile, n_tile) + }); + let tail_coords = (tail > 0).then(|| { + // `max` keeps the discarded branch's operand in range instead of + // wrapping below zero when `local < threshold`. + let rem = program.bind(local.clone().max(threshold) - threshold); + let m_tile = rem.clone() % tail + num_full * group; + let n_tile = rem / tail; + (m_tile, n_tile) + }); + match (full, tail_coords) { + (Some(full), None) => full, + (None, Some(tail)) => tail, + (Some((m_full, n_full)), Some((m_tail, n_tail))) => { + let in_full = local.lt(threshold); + ( + Tile::select(in_full.clone(), m_full, m_tail), + Tile::select(in_full, n_full, n_tail), + ) + } + (None, None) => unreachable!("tiles_m > 1 yields a full block or a tail"), + } +} + /// Cooperative-matrix batched matmul. /// /// Per-workgroup output tile is `BM × BN`. The N axis is split into @@ -497,6 +1533,7 @@ fn batched_coop_matmul_perf_single( #[allow(clippy::too_many_arguments)] fn batched_coop_matmul_perf( program: &mut Program, + staging: Option, a: &Storage, b: &Storage, y: &Storage, @@ -512,8 +1549,9 @@ fn batched_coop_matmul_perf( col_groups: u32, n_passes: u32, subgroups: SubgroupConfig, + epilogues: &DenseMatmulEpilogues<'_>, + swizzle_group_m: u32, ) { - const COOP_DIM: u32 = 8; debug_assert!(n_passes >= 1, "n_passes must be at least 1"); debug_assert_eq!(bn % n_passes, 0, "bn must be divisible by n_passes"); let bn_pass: u32 = bn / n_passes; @@ -531,6 +1569,16 @@ fn batched_coop_matmul_perf( let tile_cols_per_sg: u32 = subgroup_cols_per_pass / COOP_DIM; let scalar = scalar_of(a.element()); + let stage_scalar = staging + .filter(|_| { + scalar == ScalarElement::F32 + && epilogues.pre_a.is_none() + && epilogues.pre_b.is_none() + }) + .unwrap_or(scalar); + // bk stays at the table's 16 even though half-width tiles would fit a + // 32-deep slab: the doubled footprint (25.2KB) drops threadgroup + // residency from two workgroups to one and measured 26-36% slower. let tiles_m = shape.m.div_ceil(bm); let tiles_n = shape.n.div_ceil(bn); @@ -547,20 +1595,23 @@ fn batched_coop_matmul_perf( // conflicts on the inner stride (matches `stride_a = block_k + 1` in // `coop_gemm.rs`). Two A and two B tiles let the K loop issue both halves // of a K-pair before barriering. - let a_tile_0 = program.alloc_workgroup_tile_padded(scalar, bm, bk, 1); - let a_tile_1 = program.alloc_workgroup_tile_padded(scalar, bm, bk, 1); - let b_tile_0 = program.alloc_workgroup_tile_padded(scalar, bk, bn_pass, 1); - let b_tile_1 = program.alloc_workgroup_tile_padded(scalar, bk, bn_pass, 1); + let a_tile_0 = program.alloc_workgroup_tile_padded(stage_scalar, bm, bk, 1); + let a_tile_1 = program.alloc_workgroup_tile_padded(stage_scalar, bm, bk, 1); + let b_tile_0 = program.alloc_workgroup_tile_padded(stage_scalar, bk, bn_pass, 1); + let b_tile_1 = program.alloc_workgroup_tile_padded(stage_scalar, bk, bn_pass, 1); let grid = dispatch_grid_1d(total_tiles, max_workgroups_per_dimension); program.program_grid(block, grid, |program| { - let tile_id = program.program_id(WorkgroupAxis::X) - + program.program_id(WorkgroupAxis::Y) * grid[0] - + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; + let tile_id = program.bind( + program.program_id(WorkgroupAxis::X) + + program.program_id(WorkgroupAxis::Y) * grid[0] + + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1], + ); + let tile_id = clamp_grid_overhang(program, tile_id, total_tiles, grid); let batch = tile_id.clone() / (tiles_m * tiles_n); let local_tile = tile_id % (tiles_m * tiles_n); - let m_tile = local_tile.clone() / tiles_n; - let n_tile = local_tile % tiles_n; + let (m_tile, n_tile) = + swizzled_tile_coords(program, local_tile, tiles_m, tiles_n, swizzle_group_m); let row_base = m_tile * bm; let col_base = n_tile * bn; let a_batch_base = batch.clone() * shape.m; @@ -595,6 +1646,9 @@ fn batched_coop_matmul_perf( &col_base, &sg_row_base, &sg_col_base_in_pass, + bm, + block, + epilogues.post, |program, pass_col_base, accs| { if k_pairs > 0 { program.loop_range(k_pairs, |program, pair_idx| { @@ -628,7 +1682,12 @@ fn batched_coop_matmul_perf( tile_cols_per_sg, bk, COOP_DIM, - scalar, + stage_scalar, + block, + bm, + bn_pass, + epilogues.pre_a, + epilogues.pre_b, ); coop_stage_and_mma( @@ -652,7 +1711,12 @@ fn batched_coop_matmul_perf( tile_cols_per_sg, bk, COOP_DIM, - scalar, + stage_scalar, + block, + bm, + bn_pass, + epilogues.pre_a, + epilogues.pre_b, ); // No trailing barrier: next iter writes to tile_0 first // (different from MMA-tile_1 reads above) — barrier-2 of @@ -685,7 +1749,12 @@ fn batched_coop_matmul_perf( tile_cols_per_sg, bk, COOP_DIM, - scalar, + stage_scalar, + block, + bm, + bn_pass, + epilogues.pre_a, + epilogues.pre_b, ); program.workgroup_barrier(); } diff --git a/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs b/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs index b2133e3cb..c7f2d0a02 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/qgemv.rs @@ -18,11 +18,7 @@ use fusor_tile_ir::tile::{range, Mask, Program, Storage, Tile, TileBlock}; use fusor_tile_ir::{GgmlQuantFormat, QuantizedMatrix, TileLiteral}; -use crate::dispatch::{ - q4k_default_large, q4k_default_mid, q4k_default_tall, q4k_large_override, q4k_mid_override, - q4k_tall_override, q6k_default_large, q6k_default_tall, q6k_large_override, q6k_tall_override, - qgemv_subgroups_per_workgroup_for_shape, QgemvShape, SubgroupConfig, -}; +use crate::dispatch::{qgemv_selected_shape, QgemvShape, SubgroupConfig}; use crate::grid::{ dot4_sum, qgemv_grid, qgemv_has_no_packed_load_tails, qgemv_program_scope, store_qgemv_sums_with_epilogue, QgemvStoreTarget, @@ -98,6 +94,7 @@ fn qgemv_shape(subgroups: u32, cols_per_subgroup: u32) -> QgemvShape { /// &y, /// 1, /// fusor_tile_ir_kernels::SubgroupConfig::fixed(subgroup, 32), +/// fusor_tile_ir_kernels::qgemv_selected_shape(GgmlQuantFormat::Q4K, 256, 128), /// Option::<&UnaryEpilogue>::None, /// ); /// }); @@ -110,10 +107,11 @@ pub fn qgemv_with_epilogue<'a>( y: &Storage, workgroups_x: u32, subgroups: SubgroupConfig, + shape: QgemvShape, epilogues: impl IntoQgemvEpilogues<'a>, ) { let epilogues = epilogues.into_qgemv_epilogues(); - qgemv_tile_with_epilogue(program, a, b, y, workgroups_x, subgroups, &epilogues); + qgemv_tile_with_epilogue(program, a, b, y, workgroups_x, subgroups, shape, &epilogues); } /// Format-dispatched qgemv body with optional pre/post unary epilogues. @@ -132,35 +130,21 @@ pub(crate) fn qgemv_tile_with_epilogue( y: &Storage, workgroups_x: u32, subgroups: SubgroupConfig, + shape: QgemvShape, ep: &QmatmulEpilogues<'_>, ) { let [m, _] = matrix_shape(a.layout()); assert_eq!(m, 1, "qgemv requires a single input row"); + debug_assert_eq!( + shape, + qgemv_selected_shape(b.format, b.rows, ep.post_output_cols(b.cols)), + "caller-provided qgemv shape must be the selected shape for this matrix", + ); let tensors = QgemvTensors { a, b, y }; - let output_cols = ep.post_output_cols(b.cols); match b.format { GgmlQuantFormat::Q8_0 | GgmlQuantFormat::Q8_0Native => { - if output_cols >= 8192 { - return qgemv_perf_with_epilogue( - program, - tensors, - workgroups_x, - subgroups, - ep, - qgemv_shape(4, 8), - 8, - ); - } - qgemv_perf_with_epilogue( - program, - tensors, - workgroups_x, - subgroups, - ep, - qgemv_shape(4, 4), - 8, - ) + qgemv_perf_with_epilogue(program, tensors, workgroups_x, subgroups, ep, shape, 8) } GgmlQuantFormat::Q8_1 => qgemv_perf_with_epilogue( program, @@ -172,19 +156,6 @@ pub(crate) fn qgemv_tile_with_epilogue( 8, ), GgmlQuantFormat::Q4K | GgmlQuantFormat::Q4KNative => { - let shape = if b.rows <= 4096 && (4096..8192).contains(&output_cols) { - q4k_mid_override(q4k_default_mid(b.rows, output_cols)) - } else if b.rows <= 4096 && output_cols <= 4096 { - qgemv_shape(8, 4) - } else if b.rows <= 4096 && output_cols >= 8192 { - q4k_large_override(q4k_default_large(b.rows, output_cols)) - } else if b.rows > 4096 && output_cols <= 4096 { - q4k_tall_override(q4k_default_tall(b.rows, output_cols)) - } else if qgemv_subgroups_per_workgroup_for_shape(b.format, b.rows, output_cols) == 8 { - qgemv_shape(8, 8) - } else { - qgemv_shape(4, 8) - }; // The decode matmuls (no pre-epilogue) take the ggml super-block- // amortized dot, which decodes each 256-element super-block's // scale/min once per lane instead of re-decoding per 16-element @@ -248,36 +219,9 @@ pub(crate) fn qgemv_tile_with_epilogue( 8, ), GgmlQuantFormat::Q6K | GgmlQuantFormat::Q6KNative => { - if b.rows <= 4096 && output_cols >= 8192 { - let shape = q6k_large_override(q6k_default_large(b.rows, output_cols)); - return qgemv_perf_with_epilogue( - program, - tensors, - workgroups_x, - subgroups, - ep, - shape, - 8, - ); - } - if b.rows > 4096 && output_cols <= 4096 { - let shape = q6k_tall_override(q6k_default_tall(b.rows, output_cols)); - return qgemv_perf_with_epilogue( - program, - tensors, - workgroups_x, - subgroups, - ep, - shape, - 8, - ); - } - let (shape, values_per_lane) = - if qgemv_subgroups_per_workgroup_for_shape(b.format, b.rows, output_cols) == 4 { - (qgemv_shape(4, 4), 8) - } else { - (qgemv_shape(8, 4), 16) - }; + // Only the default (8, 4) regime uses the 16-value lanes; the + // large/tall/4-subgroup regimes all stage 8 values per lane. + let values_per_lane = if shape == qgemv_shape(8, 4) { 16 } else { 8 }; // The decode matmuls (no pre-epilogue) take the ggml super-block- // amortized dot, which decodes each 256-element super-block's `d` // and sub-block scales once per 16-element lane region instead of diff --git a/fusor-ml/tile-ir-kernels/src/kernels/qmatmul.rs b/fusor-ml/tile-ir-kernels/src/kernels/qmatmul.rs index 9b8432212..f0c40f956 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/qmatmul.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/qmatmul.rs @@ -68,7 +68,12 @@ pub fn qmatmul_with_epilogue( assert_eq!(b.cols, y_n, "qmatmul output column count must match B"); if m == 1 { - super::qgemv::qgemv_with_epilogue(program, a, b, y, 1, subgroups, epilogues); + let shape = crate::dispatch::qgemv_selected_shape( + b.format, + b.rows, + epilogues.post_output_cols(b.cols), + ); + super::qgemv::qgemv_tile_with_epilogue(program, a, b, y, 1, subgroups, shape, epilogues); } else { qmatmul_tile_with_epilogue(program, a, b, y, epilogues, coop, subgroups, bm, bn, bk); } @@ -198,9 +203,6 @@ pub(crate) fn qmatmul_try_coop( bn: u32, bk: u32, ) -> bool { - if std::env::var_os("FUSOR_DIAG_DISABLE_COOP").is_some() { - return false; - } if b.format.is_q4k_family() || b.format.is_q6k_family() { return false; } diff --git a/fusor-ml/tile-ir-kernels/src/kernels/qmatmul_workgroup.rs b/fusor-ml/tile-ir-kernels/src/kernels/qmatmul_workgroup.rs index 501ad617f..150ce6b49 100644 --- a/fusor-ml/tile-ir-kernels/src/kernels/qmatmul_workgroup.rs +++ b/fusor-ml/tile-ir-kernels/src/kernels/qmatmul_workgroup.rs @@ -1,69 +1,80 @@ -//! Workgroup-tiled quantized matmul / gemv kernels for GPUs that don't -//! expose `Features::SUBGROUP`. +//! One register-tile template for the quantized matmul / gemv family that +//! runs where the subgroup paths can't: adapters without +//! `Features::SUBGROUP`, plus every f16-activation quantized matmul (the +//! selector routes those here on any adapter). //! //! The subgroup-based qmatmul/qgemv paths in this crate partition lanes by //! `subgroup_id` and reduce via `subgroup_reduce_*`, which `Mesa lavapipe` //! (Linux CI's software Vulkan) and other adapters without the SUBGROUP -//! feature can't validate. The kernels below mirror the dense -//! `batched_matmul_with_epilogues` strategy — stage A and a dequantized B -//! into workgroup memory, then have each lane accumulate a `TM x TN` -//! register sub-tile — so they only use `program.lane()` and -//! `workgroup_barrier()`. They're cooperative across the workgroup, never -//! the subgroup. +//! feature can't validate. The template below stages A and a dequantized B +//! into workgroup memory and has each lane accumulate a `tm x tn` register +//! sub-tile, so it only uses `program.lane()` and `workgroup_barrier()`. +//! It's cooperative across the workgroup, never the subgroup. +//! +//! Two geometries instantiate it: [`MATMUL_SHAPE`] for the general case and +//! the single-row [`GEMV_SHAPE`] for `m == 1`. A single-row grid has no M +//! axis — every row index folds to the constant 0, all lanes fan across N, +//! and the accumulators fan across the post epilogue's matrix-column offsets. //! //! Storage and staging element types travel as [`ScalarElement`] data through -//! [`AccumCast`]. The bodies share three building blocks: +//! [`AccumCast`]. The body shares three building blocks: //! - [`stage_storage_tile_with_pre`] — cooperative per-lane staging of a dense //! source into a workgroup tile, applying the optional pre-activation -//! epilogue per element. Used for A in both kernels. +//! epilogue per element. Used for A. //! - [`TileBlock::fill_tile_quantized`] — per-lane dequantize-into-workgroup- //! tile for B. This path is not coop-forcing, preserving the lavapipe //! invariant that this kernel requests neither `SUBGROUP` nor //! `COOPERATIVE_MATRIX`. //! - [`accumulate_register_tile_from_workgroup`] — per-lane register -//! accumulation reading both staged tiles. Parameterized over the register -//! tile shape (`tm`, `tn`), so the matmul body uses 4x4 and the gemv body -//! uses 1x1. +//! accumulation reading both staged tiles. use fusor_tile_ir::tile::{Mask, Program, Storage, Tile, TileBlock, WorkgroupTile}; -use fusor_tile_ir::{QuantizedMatrix, ScalarElement, TileLiteral, WorkgroupAxis}; +use fusor_tile_ir::{QuantizedMatrix, ScalarElement, WorkgroupAxis}; use crate::kernels::helpers::{dispatch_grid_1d, load_qmatmul_extra, scalar_of, AccumCast}; use crate::types::{ - apply_qmatmul_post_epilogue, apply_qmatmul_post_epilogue_values, apply_qmatmul_pre_epilogue, - matrix_shape, QmatmulEpilogues, + apply_qmatmul_post_epilogue_values, apply_qmatmul_pre_epilogue, matrix_shape, QmatmulEpilogues, }; -const QMATMUL_LANES: u32 = 64; -const QGEMV_LANES: u32 = 64; -const QMATMUL_TM: u32 = 4; -const QMATMUL_TN: u32 = 4; -const QGEMV_TN: u32 = 1; - -struct RegisterTileWorkgroups<'a> { - a: &'a WorkgroupTile, - b: &'a WorkgroupTile, -} - -struct RegisterTileLane<'a> { - row: &'a Tile, - col: &'a Tile, -} +const LANES: u32 = 64; -struct RegisterTileShape { +/// One instantiation's geometry: a workgroup covers a `bm x bn` output tile, +/// staging `bk` K-elements of A and B per pass, and every lane accumulates a +/// `tm x tn` register sub-tile. +#[derive(Clone, Copy)] +struct WorkgroupTileShape { + bm: u32, bn: u32, bk: u32, tm: u32, tn: u32, } +/// `bm`/`bn` are pinned to the 4x4-register-tile geometry across [`LANES`]. +const MATMUL_SHAPE: WorkgroupTileShape = WorkgroupTileShape { + bm: 32, + bn: 32, + bk: 8, + tm: 4, + tn: 4, +}; + +/// Single output row, one output column per lane. +const GEMV_SHAPE: WorkgroupTileShape = WorkgroupTileShape { + bm: 1, + bn: LANES, + bk: 8, + tm: 1, + tn: 1, +}; + /// Stage `src` rows in `[row_base, row_base + rows)` and cols in /// `[col_base, col_base + cols)` into the workgroup tile `dst`, applying -/// `pre` per element. Cooperative across all `lanes` workgroup lanes. Pads +/// `pre` per element. Cooperative across all [`LANES`] workgroup lanes. Pads /// out-of-bound source positions with zero, and guards the workgroup-tile /// store so lanes with `flat >= rows * cols` don't write past the tile -/// (qgemv passes a 1xBK tile to a 64-lane workgroup; the unused lanes -/// would otherwise corrupt adjacent workgroup memory). +/// (the gemv shape passes a 1xBK tile to a 64-lane workgroup; the unused +/// lanes would otherwise corrupt adjacent workgroup memory). /// /// `stor_cast` promotes storage loads into the f32 accumulator; `staging_cast` /// demotes the post-pre-epilogue f32 value back to the staged tile element. @@ -82,12 +93,11 @@ fn stage_storage_tile_with_pre( epilogues: &QmatmulEpilogues<'_>, rows: u32, cols: u32, - lanes: u32, ) { let tile_elements = rows * cols; - let passes = (rows * cols).div_ceil(lanes); + let passes = (rows * cols).div_ceil(LANES); for pass in 0..passes { - let flat = program.lane() + pass * lanes; + let flat = program.lane() + pass * LANES; let local_row = flat.clone() / cols; let local_col = flat.clone() % cols; let global_row = row_base.clone() + local_row.clone(); @@ -111,7 +121,7 @@ fn stage_storage_tile_with_pre( let value = Tile::select( in_bounds, apply_qmatmul_pre_epilogue(epilogues, loaded, pre_extras), - Tile::literal(TileLiteral::f32(0.0)), + Tile::f32(0.0), ); let value = staging_cast.from_accum(value); // Re-use the same flat index but only emit the store on lanes that @@ -127,29 +137,31 @@ fn stage_storage_tile_with_pre( /// sub-tile rooted at `(lane_row * tm, lane_col * tn)` in the workgroup /// tiles. Caller is responsible for the surrounding `workgroup_barrier()`s. /// -/// Layout: `A_tile` is row-major `BM x BK` (index = row*BK + k), `B_tile` is -/// row-major `BK x BN` (index = k*BN + col). `staging_cast` promotes the +/// Layout: `A_tile` is row-major `bm x bk` (index = row*bk + k), `B_tile` is +/// row-major `bk x bn` (index = k*bn + col). `staging_cast` promotes the /// staged-tile loads to the f32 accumulator. fn accumulate_register_tile_from_workgroup( program: &mut TileBlock<'_>, - tiles: RegisterTileWorkgroups<'_>, + a_tile: &WorkgroupTile, + b_tile: &WorkgroupTile, staging_cast: &AccumCast, - lane: RegisterTileLane<'_>, - shape: RegisterTileShape, + lane_row: &Tile, + lane_col: &Tile, + shape: WorkgroupTileShape, ) -> Vec { - let RegisterTileShape { bn, bk, tm, tn } = shape; + let WorkgroupTileShape { bn, bk, tm, tn, .. } = shape; (0..tm * tn) .map(|idx| { let r = idx / tn; let c = idx % tn; - let local_row = lane.row.clone() * tm + r; - let local_col = lane.col.clone() * tn + c; - let mut sum = Tile::literal(TileLiteral::f32(0.0)); + let local_row = lane_row.clone() * tm + r; + let local_col = lane_col.clone() * tn + c; + let mut sum = Tile::f32(0.0); for kk in 0..bk { let a_value = staging_cast - .into_accum(program.load_workgroup(tiles.a, local_row.clone() * bk + kk)); + .into_accum(program.load_workgroup(a_tile, local_row.clone() * bk + kk)); let b_value = staging_cast - .into_accum(program.load_workgroup(tiles.b, local_col.clone() + kk * bn)); + .into_accum(program.load_workgroup(b_tile, local_col.clone() + kk * bn)); sum = sum + a_value * b_value; } sum @@ -157,76 +169,13 @@ fn accumulate_register_tile_from_workgroup( .collect() } -/// Workgroup-tiled quantized matmul. Each workgroup produces a `BM x BN` -/// output tile by staging A and a dequantized B into workgroup memory and -/// having every lane accumulate a `TM x TN` register sub-tile. No subgroup -/// ops — uses only `program.lane()` and `workgroup_barrier()`. -/// -/// `BM` and `BN` must equal 32 (matches the `4x4` register tile across 64 -/// lanes). `BK` is the K-axis staging chunk. +/// Workgroup-tiled quantized matmul staged through `staging_element` (f16 +/// staging requires shader-f16 support; accumulation stays f32 and the output +/// keeps the storage element). Each workgroup produces one output tile by +/// staging A and a dequantized B into workgroup memory and having every lane +/// accumulate its register sub-tile. No subgroup ops — uses only +/// `program.lane()` and `workgroup_barrier()`. pub fn qmatmul_workgroup_with_epilogues( - program: &mut Program, - a: &Storage, - b: &QuantizedMatrix, - y: &Storage, - epilogues: &QmatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, -) { - qmatmul_workgroup_with_epilogues_impl( - program, - a, - b, - y, - ScalarElement::F32, - epilogues, - max_workgroups_per_dimension, - ); -} - -/// F16-staged variant of [`qmatmul_workgroup_with_epilogues`]. This requires -/// shader-f16 support but otherwise shares the f32 implementation and keeps -/// f32 accumulation/output. -pub fn qmatmul_workgroup_f16_with_epilogues( - program: &mut Program, - a: &Storage, - b: &QuantizedMatrix, - y: &Storage, - epilogues: &QmatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, -) { - qmatmul_workgroup_with_epilogues_impl( - program, - a, - b, - y, - ScalarElement::F16, - epilogues, - max_workgroups_per_dimension, - ); -} - -/// F16-storage and F16-staged variant of [`qmatmul_workgroup_with_epilogues`]. -/// Accumulates in f32 and writes f16 outputs directly. -pub fn qmatmul_workgroup_storage_f16_with_epilogues( - program: &mut Program, - a: &Storage, - b: &QuantizedMatrix, - y: &Storage, - epilogues: &QmatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, -) { - qmatmul_workgroup_with_epilogues_impl( - program, - a, - b, - y, - ScalarElement::F16, - epilogues, - max_workgroups_per_dimension, - ); -} - -fn qmatmul_workgroup_with_epilogues_impl( program: &mut Program, a: &Storage, b: &QuantizedMatrix, @@ -235,55 +184,76 @@ fn qmatmul_workgroup_with_epilogues_impl( epilogues: &QmatmulEpilogues<'_>, max_workgroups_per_dimension: u32, ) { - // BM/BN are pinned to the 4x4-register-tile geometry across QMATMUL_LANES. - // BK is the K-axis staging chunk per pass. - const BM: u32 = 32; - const BN: u32 = 32; - const BK: u32 = 8; - let bk = BK; + let [m, k] = matrix_shape(&a.view().layout); + let shape = if m == 1 { GEMV_SHAPE } else { MATMUL_SHAPE }; + let WorkgroupTileShape { bm, bn, bk, tm, tn } = shape; + let single_row = bm == 1; // F32 accumulation throughout; storage / staging elements are runtime data. - let stor_scalar = scalar_of(a.element()); - let stor_cast = AccumCast::new(stor_scalar, ScalarElement::F32); + let stor_cast = AccumCast::new(scalar_of(a.element()), ScalarElement::F32); let staging_cast = AccumCast::new(staging_element, ScalarElement::F32); - let [m, k] = matrix_shape(&a.view().layout); - let n = b.cols; + let n = if single_row { + epilogues.post_output_cols(b.cols) + } else { + b.cols + }; assert_eq!(k, b.rows, "qmatmul K dimensions must match"); let [y_m, y_n] = matrix_shape(&y.view().layout); assert_eq!(m, y_m, "qmatmul output row count must match A"); assert_eq!(n, y_n, "qmatmul output column count must match B"); - let tiles_m = m.div_ceil(BM); - let tiles_n = n.div_ceil(BN); + // The single-row shape accumulates one value per matrix-column offset of + // the post epilogue (the default offset list is a single `0`); the tiled + // shape holds one accumulator per output element at the tile's own base. + let column_offsets: Vec> = if single_row { + epilogues + .post_accumulator_offsets() + .iter() + .copied() + .map(Some) + .collect() + } else { + vec![None] + }; + let value_count = column_offsets.len(); + + let tiles_m = m.div_ceil(bm); + let tiles_n = n.div_ceil(bn); let total_tiles = tiles_m * tiles_n; let k_tiles = k.div_ceil(bk); let grid = dispatch_grid_1d(total_tiles, max_workgroups_per_dimension); - let a_tile = program.alloc_workgroup_tile(staging_element, BM, bk); - let b_tile = program.alloc_workgroup_tile(staging_element, bk, BN); - let b_clone = b.clone(); + let a_tile = program.alloc_workgroup_tile(staging_element, bm, bk); + let b_tile = program.alloc_workgroup_tile(staging_element, bk, bn); + let b = b.clone(); - program.program_grid(QMATMUL_LANES, grid, |program| { + program.program_grid(LANES, grid, |program| { let tile_id = program.program_id(WorkgroupAxis::X) + program.program_id(WorkgroupAxis::Y) * grid[0] + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; let tile_active = tile_id.clone().lt(total_tiles); - let m_tile = tile_id.clone() / tiles_n; - let n_tile = tile_id % tiles_n; - let lane = program.lane(); - let lane_row = lane.clone() / (BN / QMATMUL_TN); - let lane_col = lane % (BN / QMATMUL_TN); - let m_tile_base = m_tile * BM; - let n_tile_base = n_tile * BN; - let row_base = m_tile_base.clone() + lane_row.clone() * QMATMUL_TM; - let col_base = n_tile_base.clone() + lane_col.clone() * QMATMUL_TN; + let (m_tile_base, n_tile) = if single_row { + (Tile::u32(0), tile_id) + } else { + (tile_id.clone() / tiles_n * bm, tile_id % tiles_n) + }; + let (lane_row, lane_col) = if single_row { + (Tile::u32(0), lane) + } else { + (lane.clone() / (bn / tn), lane % (bn / tn)) + }; + let n_tile_base = n_tile * bn; + let row_base = if single_row { + Tile::u32(0) + } else { + m_tile_base.clone() + lane_row.clone() * tm + }; + let col_base = n_tile_base.clone() + lane_col.clone() * tn; - let init: [Tile; (QMATMUL_TM * QMATMUL_TN) as usize] = - std::array::from_fn(|_| Tile::literal(TileLiteral::f32(0.0))); - let sums = program.fold( + let sums = program.fold_vec( fusor_tile_ir::tile::range(k_tiles), - init, + vec![Tile::f32(0.0); (tm * tn) as usize * value_count], |program, k_tile, accs| { let k_base = k_tile * bk; stage_storage_tile_with_pre( @@ -298,228 +268,30 @@ fn qmatmul_workgroup_with_epilogues_impl( m, k, epilogues, - BM, - bk, - QMATMUL_LANES, - ); - program.fill_tile_quantized(&b_tile, &b_clone, k_base, n_tile_base.clone()); - program.workgroup_barrier(); - - let chunk_vec = accumulate_register_tile_from_workgroup( - program, - RegisterTileWorkgroups { - a: &a_tile, - b: &b_tile, - }, - &staging_cast, - RegisterTileLane { - row: &lane_row, - col: &lane_col, - }, - RegisterTileShape { - bn: BN, - bk, - tm: QMATMUL_TM, - tn: QMATMUL_TN, - }, - ); - let mut chunk_iter = chunk_vec.into_iter(); - let next: [Tile; (QMATMUL_TM * QMATMUL_TN) as usize] = std::array::from_fn(|idx| { - let chunk = - program.bind(chunk_iter.next().expect("register tile size matches")); - accs[idx].clone() + chunk - }); - program.workgroup_barrier(); - next - }, - ); - - for (idx, sum) in sums.into_iter().enumerate() { - let r = idx as u32 / QMATMUL_TN; - let c = idx as u32 % QMATMUL_TN; - let row = row_base.clone() + r; - let col = col_base.clone() + c; - let extras = epilogues - .post_extra_inputs - .iter() - .map(|extra| load_qmatmul_extra(program, extra, &row, &col, n)) - .collect::>(); - let value = apply_qmatmul_post_epilogue(epilogues, sum, extras); - let value = stor_cast.from_accum(value); - let mask = tile_active.clone() & row.clone().lt(m) & col.clone().lt(n); - program.store(y.at((row, col)), value, mask); - } - }); -} - -/// Workgroup-tiled quantized GEMV (`m == 1`) for adapters without subgroups. -/// All `QGEMV_LANES` lanes fan out across the BN columns of one output tile. -/// Stages A's single row into workgroup memory and reuses -/// [`accumulate_register_tile_from_workgroup`] with `tm = 1`, `tn = 1`. -pub fn qgemv_workgroup_with_epilogue( - program: &mut Program, - a: &Storage, - b: &QuantizedMatrix, - y: &Storage, - epilogues: &QmatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, -) { - qgemv_workgroup_with_epilogue_impl( - program, - a, - b, - y, - ScalarElement::F32, - epilogues, - max_workgroups_per_dimension, - ); -} - -/// F16-staged variant of [`qgemv_workgroup_with_epilogue`]. Requires -/// shader-f16 support; accumulation and output remain f32. -pub fn qgemv_workgroup_f16_with_epilogue( - program: &mut Program, - a: &Storage, - b: &QuantizedMatrix, - y: &Storage, - epilogues: &QmatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, -) { - qgemv_workgroup_with_epilogue_impl( - program, - a, - b, - y, - ScalarElement::F16, - epilogues, - max_workgroups_per_dimension, - ); -} - -/// F16-storage and F16-staged variant of [`qgemv_workgroup_with_epilogue`]. -/// Accumulates in f32 and writes f16 outputs directly. -pub fn qgemv_workgroup_storage_f16_with_epilogue( - program: &mut Program, - a: &Storage, - b: &QuantizedMatrix, - y: &Storage, - epilogues: &QmatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, -) { - qgemv_workgroup_with_epilogue_impl( - program, - a, - b, - y, - ScalarElement::F16, - epilogues, - max_workgroups_per_dimension, - ); -} - -fn qgemv_workgroup_with_epilogue_impl( - program: &mut Program, - a: &Storage, - b: &QuantizedMatrix, - y: &Storage, - staging_element: ScalarElement, - epilogues: &QmatmulEpilogues<'_>, - max_workgroups_per_dimension: u32, -) { - // BN is pinned to QGEMV_LANES (one column per lane). BK is the K-axis - // staging chunk per pass. - const BN: u32 = QGEMV_LANES; - const BK: u32 = 8; - let bk = BK; - - let stor_scalar = scalar_of(a.element()); - let stor_cast = AccumCast::new(stor_scalar, ScalarElement::F32); - let staging_cast = AccumCast::new(staging_element, ScalarElement::F32); - - let [m, k] = matrix_shape(&a.view().layout); - let n = epilogues.post_output_cols(b.cols); - assert_eq!(m, 1, "qgemv_workgroup expects a single input row"); - assert_eq!(k, b.rows, "qgemv K dimensions must match"); - let [y_m, y_n] = matrix_shape(&y.view().layout); - assert_eq!(y_m, 1, "qgemv output must be single-row"); - assert_eq!(n, y_n, "qgemv output column count must match B"); - - let tiles_n = n.div_ceil(BN); - let k_tiles = k.div_ceil(bk); - let grid = dispatch_grid_1d(tiles_n, max_workgroups_per_dimension); - // BM=1 logical row tile. Reuse the stager with rows=1. - let a_tile = program.alloc_workgroup_tile(staging_element, 1, bk); - let b_tile = program.alloc_workgroup_tile(staging_element, bk, BN); - let b_clone = b.clone(); - - program.program_grid(QGEMV_LANES, grid, |program| { - let tile_id = program.program_id(WorkgroupAxis::X) - + program.program_id(WorkgroupAxis::Y) * grid[0] - + program.program_id(WorkgroupAxis::Z) * grid[0] * grid[1]; - let tile_active = tile_id.clone().lt(tiles_n); - let n_tile = tile_id; - let n_tile_base = n_tile * BN; - - let lane = program.lane(); - let lane_row = Tile::literal(TileLiteral::U32(0)); - let lane_col = lane; - let row_base = Tile::literal(TileLiteral::U32(0)); - let col_base = n_tile_base.clone() + lane_col.clone() * QGEMV_TN; - - let post_accumulator_offsets = epilogues.post_accumulator_offsets().to_vec(); - let post_value_arity = post_accumulator_offsets.len(); - let sums = program.fold_vec( - fusor_tile_ir::tile::range(k_tiles), - vec![Tile::literal(TileLiteral::f32(0.0)); QGEMV_TN as usize * post_value_arity], - |program, k_tile, accs| { - let k_base = k_tile * bk; - stage_storage_tile_with_pre( - program, - &a_tile, - a, - &stor_cast, - &staging_cast, - &row_base, - &k_base, - &tile_active, - 1, - k, - epilogues, - 1, + bm, bk, - QGEMV_LANES, ); let mut next = accs; - for (value_idx, offset) in post_accumulator_offsets.iter().copied().enumerate() { - program.fill_tile_quantized( - &b_tile, - &b_clone, - k_base.clone(), - n_tile_base.clone() + offset, - ); + for (value_idx, offset) in column_offsets.iter().enumerate() { + let b_col_base = match offset { + Some(offset) => n_tile_base.clone() + *offset, + None => n_tile_base.clone(), + }; + program.fill_tile_quantized(&b_tile, &b, k_base.clone(), b_col_base); program.workgroup_barrier(); - let chunk_vec = accumulate_register_tile_from_workgroup( + let chunks = accumulate_register_tile_from_workgroup( program, - RegisterTileWorkgroups { - a: &a_tile, - b: &b_tile, - }, + &a_tile, + &b_tile, &staging_cast, - RegisterTileLane { - row: &lane_row, - col: &lane_col, - }, - RegisterTileShape { - bn: BN, - bk, - tm: 1, - tn: QGEMV_TN, - }, + &lane_row, + &lane_col, + shape, ); - for (idx, chunk) in chunk_vec.into_iter().enumerate() { - let accum_idx = idx * post_value_arity + value_idx; - next[accum_idx] = next[accum_idx].clone() + program.bind(chunk); + for (idx, chunk) in chunks.into_iter().enumerate() { + let slot = idx * value_count + value_idx; + next[slot] = next[slot].clone() + program.bind(chunk); } program.workgroup_barrier(); } @@ -527,22 +299,28 @@ fn qgemv_workgroup_with_epilogue_impl( }, ); - for (idx, values) in sums.chunks(post_value_arity).enumerate() { - let row = Tile::literal(TileLiteral::U32(0)); - let col = col_base.clone() + idx as u32; + for (idx, values) in sums.chunks(value_count).enumerate() { + let r = idx as u32 / tn; + let c = idx as u32 % tn; + let row = if single_row { + Tile::u32(0) + } else { + row_base.clone() + r + }; + let col = col_base.clone() + c; let extras = epilogues .post_extra_inputs .iter() .map(|extra| load_qmatmul_extra(program, extra, &row, &col, n)) .collect::>(); - let value = if values.len() == 1 { - apply_qmatmul_post_epilogue(epilogues, values[0].clone(), extras) + let value = apply_qmatmul_post_epilogue_values(epilogues, values.to_vec(), extras); + let value = stor_cast.from_accum(value); + let mask = if single_row { + tile_active.clone() & col.clone().lt(n) } else { - apply_qmatmul_post_epilogue_values(epilogues, values.to_vec(), extras) + tile_active.clone() & row.clone().lt(m) & col.clone().lt(n) }; - let value = stor_cast.from_accum(value); - let mask = tile_active.clone() & col.clone().lt(n); - program.store(y.at((0u32, col)), value, mask); + program.store(y.at((row, col)), value, mask); } }); } diff --git a/fusor-ml/tile-ir-kernels/src/lib.rs b/fusor-ml/tile-ir-kernels/src/lib.rs index 4af8c0d15..8464463d2 100644 --- a/fusor-ml/tile-ir-kernels/src/lib.rs +++ b/fusor-ml/tile-ir-kernels/src/lib.rs @@ -10,17 +10,21 @@ mod kernels; mod types; pub use dispatch::{ - qgemv_cols_per_workgroup, qgemv_cols_per_workgroup_for_shape, - qgemv_subgroups_per_workgroup_for_shape, SubgroupConfig, + qgemv_selected_shape, + SubgroupConfig, }; pub use kernels::{ - linear_storage_layout, qgemv_with_epilogue, qgemv_workgroup_f16_with_epilogue, - qgemv_workgroup_storage_f16_with_epilogue, qgemv_workgroup_with_epilogue, - qmatmul_with_epilogue, qmatmul_workgroup_f16_with_epilogues, - qmatmul_workgroup_storage_f16_with_epilogues, qmatmul_workgroup_with_epilogues, - quantized_matrix, quantized_matrix_for, try_batched_coop_matmul, AccumCast, - DenseCoopMatmulConfig, DenseCoopMatmulTile, DenseMatmulShape, DenseMatmulTensors, - IntoQgemvEpilogues, + coop_tile_entries, DEFAULT_SWIZZLE_GROUP_M, flash_attention_bwd_supported, flash_attention_dispatch, + flash_attention_f32, flash_attention_supported, flash_attention_workgroup_bytes, + flash_bwd_kv_dispatch, flash_bwd_kv_f32, + flash_bwd_q_dispatch, flash_bwd_q_f32, flash_lse_dispatch, flash_lse_f32, + linear_storage_layout, qgemv_with_epilogue, qmatmul_with_epilogue, + qmatmul_workgroup_with_epilogues, + merged_split_k_combine, quantized_matrix, quantized_matrix_for, split_k_combine, + try_batched_coop_matmul, try_batched_coop_matmul_split_k, try_merged_coop_matmul, + AccumCast, CoopTileEntry, DenseCoopMatmulConfig, DenseCoopMatmulTile, DenseMatmulShape, + DenseMatmulTensors, FlashAttentionLayouts, FlashAttentionShape, FlashBwdLayouts, + FlashKvOutputs, FlashMaskLayout, FlashOperandLayout, FlashRowLayout, IntoQgemvEpilogues, }; pub use types::{ cooperative_store_layout_supported, DenseMatmulEpilogues, QmatmulEpilogues, QmatmulExtra, diff --git a/fusor-ml/tile-ir-kernels/src/types.rs b/fusor-ml/tile-ir-kernels/src/types.rs index 3c9becd46..1ea442c76 100644 --- a/fusor-ml/tile-ir-kernels/src/types.rs +++ b/fusor-ml/tile-ir-kernels/src/types.rs @@ -186,7 +186,7 @@ pub(crate) fn apply_epilogue_values_with_extras( } } -/// Bundle of pre- and post-reduce epilogues for dense F32 matmul kernels. +/// Bundle of pre- and post-reduce epilogues for dense F32/F16 matmul kernels. #[derive(Clone, Default)] pub struct DenseMatmulEpilogues<'a> { /// Optional transform applied to each loaded lhs value before the product. diff --git a/fusor-ml/tile-ir-kernels/tests/footprint.rs b/fusor-ml/tile-ir-kernels/tests/footprint.rs new file mode 100644 index 000000000..22cf7cdb9 --- /dev/null +++ b/fusor-ml/tile-ir-kernels/tests/footprint.rs @@ -0,0 +1,320 @@ +//! Deviceless IR-vs-analytic workgroup-footprint checks: every +//! `COOP_TILE_TABLE` entry and the flash-attention forward kernel must lower +//! to exactly the byte count their `workgroup_bytes` formulas report, so the +//! selection layers can plan occupancy from pure arithmetic. + +use fusor_tile_ir::{tile, ScalarElement, Shape}; +use fusor_tile_ir_kernels::{ + coop_tile_entries, flash_attention_f32, flash_attention_supported, + flash_attention_workgroup_bytes, try_batched_coop_matmul, try_batched_coop_matmul_split_k, + CoopTileEntry, DenseCoopMatmulConfig, DenseMatmulEpilogues, DenseMatmulShape, + DenseMatmulTensors, FlashAttentionLayouts, FlashAttentionShape, FlashOperandLayout, + SubgroupConfig, DEFAULT_SWIZZLE_GROUP_M, +}; + +fn subgroup_config() -> SubgroupConfig { + SubgroupConfig::fixed(fusor_tile_ir::SubgroupToken::new_unchecked(), 32) +} + +/// Build one table entry's kernel IR without buffers: the standard path for +/// `splits: None` (at the requested staging depth, or one pair when the entry +/// forces it), the split-K partials kernel otherwise. +fn coop_matmul_ir( + entry: &CoopTileEntry, + storage: ScalarElement, + staging: Option, + stage_buffers: u32, + splits: Option, +) -> fusor_tile_ir::KernelIr { + let geometry = entry.tile; + tile::build(|program| { + let shape = DenseMatmulShape { + batch: 1, + m: geometry.bm, + k: geometry.bk * 4, + n: geometry.bn, + }; + let a = program.storage_read(storage.element(), Shape::new([shape.m, shape.k])); + let b = program.storage_read(storage.element(), Shape::new([shape.k, shape.n])); + // Split-K over-allocates the output with one scratch slice per split. + let y_rows = shape.m * (splits.unwrap_or(0) + 1); + let y = program.storage_write(storage.element(), Shape::new([y_rows, shape.n])); + let tensors = DenseMatmulTensors { + a: &a, + b: &b, + y: &y, + }; + let (row_groups, col_groups) = entry.subgroup_split(); + let config = DenseCoopMatmulConfig { + coop: fusor_tile_ir::CoopMatrixToken::new_unchecked(), + subgroups: subgroup_config(), + tile: geometry, + row_groups, + col_groups, + staging, + stage_buffers, + swizzle_group_m: DEFAULT_SWIZZLE_GROUP_M, + }; + let emitted = match splits { + Some(splits) => { + try_batched_coop_matmul_split_k(program, tensors, shape, splits, 65_535, config) + } + None => try_batched_coop_matmul( + program, + tensors, + shape, + &DenseMatmulEpilogues::empty(), + 65_535, + config, + ), + }; + assert!(emitted, "kernel declined {geometry:?}"); + }) +} + +/// Every table entry's lowered footprint equals the analytic formula, for +/// both storage elements, through both the double-buffered perf body and +/// the single-buffered body (the entry's flag routes it). +#[test] +fn coop_table_footprints_match_ir() { + for entry in coop_tile_entries() { + for storage in [ScalarElement::F32, ScalarElement::F16] { + let ir = coop_matmul_ir(entry, storage, None, 2, None); + assert_eq!( + ir.workgroup_bytes(), + entry.workgroup_bytes(storage), + "{:?} {storage:?}", + entry.tile, + ); + } + } +} + +/// The split-K partials kernel stages exactly one tile pair, whatever depth +/// the config asks for: a split grid exists to raise occupancy and a second +/// pair would halve how many of its workgroups a core holds. +#[test] +fn split_k_footprints_match_ir() { + for entry in coop_tile_entries() { + // The split path declines single-buffered geometry. + if entry.single_buffered { + continue; + } + let ir = coop_matmul_ir(entry, ScalarElement::F32, None, 2, Some(2)); + assert_eq!( + ir.workgroup_bytes(), + entry.workgroup_bytes_at(ScalarElement::F32, 1), + "{:?} split-k", + entry.tile, + ); + } +} + +/// `staging: Some(F16)` over f32 storage stages the whole pair set in f16 — +/// the formula's stage axis, validated without enabling staging in any +/// production path. +#[test] +fn f16_staging_over_f32_storage_matches_ir() { + for entry in coop_tile_entries() { + // The single-buffered body ignores staging. + if entry.single_buffered { + continue; + } + let ir = coop_matmul_ir(entry, ScalarElement::F32, Some(ScalarElement::F16), 2, None); + assert_eq!( + ir.workgroup_bytes(), + entry.workgroup_bytes(ScalarElement::F16), + "{:?} staged f16", + entry.tile, + ); + } +} + +/// `stage_buffers: 1` on a double-bufferable entry lowers to exactly one +/// staged pair. `DispatchPolicy::core_workgroup_slots` divides the workgroup +/// storage limit by this number, so a wrong footprint here would mis-price +/// the whole staging-depth choice. +#[test] +fn single_pair_staging_halves_the_footprint() { + for entry in coop_tile_entries() { + if entry.single_buffered { + continue; + } + for storage in [ScalarElement::F32, ScalarElement::F16] { + let ir = coop_matmul_ir(entry, storage, None, 1, None); + assert_eq!( + ir.workgroup_bytes(), + entry.workgroup_bytes_at(storage, 1), + "{:?} {storage:?} single pair", + entry.tile, + ); + } + } +} + +/// `single_buffered` is a derived property, not a tuning choice: set exactly +/// when two f32 pairs would overflow Apple's 32 KB threadgroup-memory limit. +#[test] +fn single_buffered_is_exactly_the_32kb_overflow() { + for entry in coop_tile_entries() { + let two_pair_f32 = + 2 * entry.tile.stage_pair_elements(entry.n_passes) * ScalarElement::F32.byte_size(); + assert_eq!( + entry.single_buffered, + two_pair_f32 > 32 * 1024, + "{:?}: two-pair f32 footprint {two_pair_f32}", + entry.tile, + ); + } +} + +fn flash_shape(head_dim: u32) -> FlashAttentionShape { + FlashAttentionShape { + batch: 1, + heads: 1, + kv_groups: 1, + q_len: 32, + kv_len: 32, + head_dim, + scale: 1.0, + causal: false, + } +} + +fn flash_ir(head_dim: u32, storage: ScalarElement) -> fusor_tile_ir::KernelIr { + tile::build(|program| { + let shape = flash_shape(head_dim); + let q_elems = shape.q_len * head_dim; + let kv_elems = shape.kv_len * head_dim; + let q = program.storage_read(storage.element(), Shape::new([q_elems])); + let k = program.storage_read(storage.element(), Shape::new([kv_elems])); + let v = program.storage_read(storage.element(), Shape::new([kv_elems])); + let o = program.storage_write(storage.element(), Shape::new([q_elems])); + let layouts = FlashAttentionLayouts { + q: FlashOperandLayout::contiguous(1, shape.q_len, head_dim), + k: FlashOperandLayout::contiguous(1, shape.kv_len, head_dim), + v: FlashOperandLayout::contiguous(1, shape.kv_len, head_dim), + o: FlashOperandLayout::contiguous(1, shape.q_len, head_dim), + }; + assert!(flash_attention_f32( + program, + &q, + &k, + &v, + None, + &o, + &layouts, + shape, + subgroup_config(), + fusor_tile_ir::CoopMatrixToken::new_unchecked(), + 65_535, + )); + }) +} + +/// The forward kernel's lowered footprint equals the analytic formula for +/// every supported head dim and stage element. +#[test] +fn flash_forward_footprints_match_ir() { + let mut covered = 0; + for head_dim in [32, 64, 80] { + if !flash_attention_supported(&flash_shape(head_dim), subgroup_config()) { + continue; + } + for storage in [ScalarElement::F32, ScalarElement::F16] { + let ir = flash_ir(head_dim, storage); + assert_eq!( + ir.workgroup_bytes(), + flash_attention_workgroup_bytes(head_dim, storage), + "d={head_dim} {storage:?}", + ); + covered += 1; + } + } + assert!(covered >= 4, "supported-shape sweep collapsed: {covered}"); +} + +/// The f16 d=64 forward kernel sits at 16.0 KB — the two-workgroups-per-core +/// residency boundary on Apple's 32 KB budget, inside WebGPU's 16 KB default +/// workgroup-storage limit. +#[test] +fn flash_f16_d64_holds_the_residency_boundary() { + let bytes = flash_attention_workgroup_bytes(64, ScalarElement::F16); + assert_eq!(bytes, 16_022); + assert!(bytes <= 16 << 10); +} + +fn flash_ir_byte_arena(head_dim: u32, storage: ScalarElement) -> fusor_tile_ir::KernelIr { + tile::build(|program| { + program.enable_byte_arena(fusor_tile_ir::ByteArenaToken::new_unchecked()); + let shape = flash_shape(head_dim); + let q_elems = shape.q_len * head_dim; + let kv_elems = shape.kv_len * head_dim; + let q = program.storage_read(storage.element(), Shape::new([q_elems])); + let k = program.storage_read(storage.element(), Shape::new([kv_elems])); + let v = program.storage_read(storage.element(), Shape::new([kv_elems])); + let o = program.storage_write(storage.element(), Shape::new([q_elems])); + let layouts = FlashAttentionLayouts { + q: FlashOperandLayout::contiguous(1, shape.q_len, head_dim), + k: FlashOperandLayout::contiguous(1, shape.kv_len, head_dim), + v: FlashOperandLayout::contiguous(1, shape.kv_len, head_dim), + o: FlashOperandLayout::contiguous(1, shape.q_len, head_dim), + }; + assert!(flash_attention_f32( + program, + &q, + &k, + &v, + None, + &o, + &layouts, + shape, + subgroup_config(), + fusor_tile_ir::CoopMatrixToken::new_unchecked(), + 65_535, + )); + }) +} + +/// The byte arena self-selects: it only replaces typed regions when +/// cross-stride reuse actually shrinks the footprint. The f16 forward +/// kernel's tiles are all live across the KV loop, so today the arena finds +/// nothing and the footprint must stay exactly the regions number. +#[test] +fn flash_f16_byte_arena_footprint() { + let regions = flash_ir(64, ScalarElement::F16).workgroup_bytes(); + let packed = flash_ir_byte_arena(64, ScalarElement::F16).workgroup_bytes(); + assert_eq!(regions, 16_022); + assert_eq!(packed, regions); +} + +/// `subgroup_split` is a derivation, not a table column: it reproduces the +/// hand-set factorization on seven of the nine rows and on the eighth +/// through its documented tie-break, and it deliberately disagrees on the +/// two 16-wide rows — (64,16,16) 2x2 -> 4x1 and (16,64,16) 2x2 -> 1x4, each +/// going from 1.25 to 1.00 threadgroup fragment loads per MMA at identical +/// MMA count, staged bytes and workgroup footprint. +#[test] +fn subgroup_split_derives_the_table() { + for entry in coop_tile_entries() { + let (bm, bn) = (entry.tile.bm, entry.tile.bn); + // The nine hand-set factorizations the table carried before the + // derivation replaced them. + let expected = match (bm, bn) { + (256, 256) => (8, 1), + (64, 128) => (2, 4), + (64, 64) => (2, 2), + // The two rows the derivation deliberately moves: 2x2 wasted a + // fragment load per MMA on both. + (64, 16) => (4, 1), + (16, 64) => (1, 4), + _ => (4, 2), + }; + assert_eq!(entry.subgroup_split(), expected, "{bm}x{bn}"); + // Both fragment sides stay whole 8x8 fragment counts. + let (rg, cg) = entry.subgroup_split(); + assert_eq!(rg * cg, entry.subgroups, "{bm}x{bn} factorization"); + assert_eq!(bm % (8 * rg), 0, "{bm}x{bn} A side"); + assert_eq!((bn / entry.n_passes) % (8 * cg), 0, "{bm}x{bn} B side"); + } +} diff --git a/fusor-ml/tile-ir-kernels/tests/goldens/qmatmul_workgroup_ir.txt b/fusor-ml/tile-ir-kernels/tests/goldens/qmatmul_workgroup_ir.txt new file mode 100644 index 000000000..0b9c04f3f --- /dev/null +++ b/fusor-ml/tile-ir-kernels/tests/goldens/qmatmul_workgroup_ir.txt @@ -0,0 +1,38 @@ +m=32 k=256 n=32 Q4K storage=F32 staging=F32 None ir 0x531b563df342dbdd naga 0xed4b78ebea21f504 +m=33 k=260 n=33 Q4K storage=F32 staging=F32 None ir 0x92337b6a3cf86308 naga 0x7ada05866079f027 +m=96 k=512 n=160 Q4K storage=F32 staging=F32 None ir 0x3f2fb3a660da99e8 naga 0x3faec31cb90b871d +m=1 k=256 n=128 Q4K storage=F32 staging=F32 None ir 0x45ad8f4ca71a81c3 naga 0x5f2cb934edde5c28 +m=1 k=260 n=130 Q4K storage=F32 staging=F32 None ir 0x3cbf77230e107cd7 naga 0x5fa90d2f573e75c7 +m=1 k=512 n=64 Q4K storage=F32 staging=F32 None ir 0x137e7f9e513159bd naga 0xb18d24552f4c8bd9 +m=32 k=256 n=32 Q8_0 storage=F32 staging=F32 None ir 0x23b9d797956bb267 naga 0xe6062acea692c703 +m=32 k=256 n=32 Q4KNative storage=F32 staging=F32 None ir 0x83f9c212baed8850 naga 0x482f14c8ce078dde +m=32 k=256 n=32 Q6K storage=F32 staging=F32 None ir 0xf5e9797d529c415b naga 0x414e8f5e791d4311 +m=32 k=256 n=32 Q4K storage=F32 staging=F32 Post ir 0xf601a337dca984cd naga 0x1699279a9ac717ea +m=32 k=256 n=32 Q4K storage=F32 staging=F32 PostColumn ir 0x03d6ca07268f905d naga 0x711bc930b0de5331 +m=32 k=256 n=32 Q4K storage=F32 staging=F32 PostPointwise ir 0xa0cca3ddf7e9f55d naga 0xb168fd59fc9c4a55 +m=32 k=256 n=32 Q4K storage=F32 staging=F32 Pre ir 0x5819dee2562927ba naga 0x3595b485ffa0bcd8 +m=32 k=256 n=32 Q4K storage=F32 staging=F32 PrePointwise ir 0x3f289c24c3461be4 naga 0xdda7c6d83422df20 +m=32 k=256 n=32 Q4K storage=F32 staging=F16 None ir 0xab6e9db5e0fe522f naga 0xfc5fb2008fe787f3 +m=32 k=256 n=32 Q4K storage=F32 staging=F16 Post ir 0xacfdd143ebcd365f naga 0x59002fe09b8eb237 +m=32 k=256 n=32 Q4K storage=F32 staging=F16 PostColumn ir 0xc72db0614b1d75af naga 0x8586b4c3caffe590 +m=32 k=256 n=32 Q4K storage=F32 staging=F16 PostPointwise ir 0x293e291a207c0ecf naga 0x7a46007224d17c4d +m=32 k=256 n=32 Q4K storage=F32 staging=F16 Pre ir 0x753c0578c0517414 naga 0xa92f033b9fe07bac +m=32 k=256 n=32 Q4K storage=F32 staging=F16 PrePointwise ir 0x3afd05d0b1b4cb46 naga 0x2e2538f818b0f487 +m=32 k=256 n=32 Q4K storage=F16 staging=F16 None ir 0xdf842c829059730b naga 0xfdda6d8158347231 +m=1 k=256 n=128 Q8_0 storage=F32 staging=F32 None ir 0x36ce47c25500d83f naga 0x5fed1f36a8cc1692 +m=1 k=256 n=128 Q4KNative storage=F32 staging=F32 None ir 0x882e3f6cde455338 naga 0x21a4e8db762bef4e +m=1 k=256 n=128 Q6K storage=F32 staging=F32 None ir 0x6342357159a9bec5 naga 0x050e6b5dc9f7e25c +m=1 k=256 n=128 Q4K storage=F32 staging=F32 Post ir 0x1333f460867e3199 naga 0xe04e5c019e057f51 +m=1 k=256 n=128 Q4K storage=F32 staging=F32 PostColumn ir 0xbbd7d63a6a95caac naga 0xebb023f0eb8c7a2b +m=1 k=256 n=128 Q4K storage=F32 staging=F32 PostPointwise ir 0xc95ae4ffc6fef71e naga 0x2fd5ba8789980780 +m=1 k=256 n=128 Q4K storage=F32 staging=F32 Pre ir 0xdb436d049750c09e naga 0x30e67ffd136ec461 +m=1 k=256 n=128 Q4K storage=F32 staging=F32 PrePointwise ir 0x44ad7b6a793b297e naga 0xde19a55a3bd1c6e6 +m=1 k=256 n=128 Q4K storage=F32 staging=F32 PairedOffsets ir 0x8b67b3ea2902d276 naga 0xf48c5c97e94a4c21 +m=1 k=256 n=128 Q4K storage=F32 staging=F16 None ir 0xd35fe8c2c917861b naga 0xeb3146dcc842db8a +m=1 k=256 n=128 Q4K storage=F32 staging=F16 Post ir 0x172545bc88f1c551 naga 0xfed86680c1eb5f41 +m=1 k=256 n=128 Q4K storage=F32 staging=F16 PostColumn ir 0x70e18a7439eeee24 naga 0x4f5292252f1b57bc +m=1 k=256 n=128 Q4K storage=F32 staging=F16 PostPointwise ir 0xe0d46b523f60eca6 naga 0x909d53f509989ea7 +m=1 k=256 n=128 Q4K storage=F32 staging=F16 Pre ir 0x186c9a2f0747acd8 naga 0x7ae618ca5e0b0a2b +m=1 k=256 n=128 Q4K storage=F32 staging=F16 PrePointwise ir 0xb1abc97138fb69b4 naga 0x3ab986d7d50f20b1 +m=1 k=256 n=128 Q4K storage=F32 staging=F16 PairedOffsets ir 0xeb505c79f8fc7b4e naga 0xf9b8af4216ce7a93 +m=1 k=256 n=128 Q4K storage=F16 staging=F16 None ir 0xaa5ab31c627c8da5 naga 0xf2d877744782675f diff --git a/fusor-ml/tile-ir-kernels/tests/lowering.rs b/fusor-ml/tile-ir-kernels/tests/lowering.rs index b8eb3cbd3..dce0f7974 100644 --- a/fusor-ml/tile-ir-kernels/tests/lowering.rs +++ b/fusor-ml/tile-ir-kernels/tests/lowering.rs @@ -1,10 +1,10 @@ use fusor_tile_ir::{tile, GgmlQuantFormat, NagaKernel, ScalarElement, Shape}; use fusor_tile_ir_kernels::{ - qgemv_with_epilogue, qgemv_workgroup_f16_with_epilogue, qgemv_workgroup_with_epilogue, - qmatmul_with_epilogue, qmatmul_workgroup_f16_with_epilogues, qmatmul_workgroup_with_epilogues, - quantized_matrix, try_batched_coop_matmul, DenseCoopMatmulConfig, DenseCoopMatmulTile, - DenseMatmulEpilogues, DenseMatmulShape, DenseMatmulTensors, QmatmulEpilogues, SubgroupConfig, - UnaryEpilogue, UnaryEpilogueWithExtras, + DEFAULT_SWIZZLE_GROUP_M, + qgemv_with_epilogue, qmatmul_with_epilogue, qmatmul_workgroup_with_epilogues, quantized_matrix, + try_batched_coop_matmul, DenseCoopMatmulConfig, DenseCoopMatmulTile, DenseMatmulEpilogues, + DenseMatmulShape, DenseMatmulTensors, QmatmulEpilogues, SubgroupConfig, UnaryEpilogue, + UnaryEpilogueWithExtras, }; fn lower_or_fail(ir: &fusor_tile_ir::KernelIr, label: &str) -> NagaKernel { @@ -41,6 +41,7 @@ fn qgemv_ir_with_subgroup_size( &y, 1, subgroup_config(subgroup_size), + fusor_tile_ir_kernels::qgemv_selected_shape(format, rows, cols), Option::<&UnaryEpilogue>::None, ); }) @@ -175,12 +176,74 @@ fn cooperative_dense_f32_matmul_lowers() { bn: 64, bk: 16, }, + row_groups: 2, + col_groups: 2, + staging: None, + stage_buffers: 2, + swizzle_group_m: DEFAULT_SWIZZLE_GROUP_M, }, )); }); lower_or_fail(&ir, "cooperative dense f32 matmul"); } +#[test] +fn cooperative_dense_f32_matmul_with_pre_and_post_epilogues_lowers() { + let pre = UnaryEpilogue::new("test_scale", |tile| tile * tile::Tile::f32(0.5)); + let post = UnaryEpilogue::new("test_tanh", |tile| tile.tanh()); + let ir = tile::build(|program| { + let shape = DenseMatmulShape { + batch: 1, + m: 61, + k: 63, + n: 59, + }; + let a = program.storage_read( + ScalarElement::F32.element(), + Shape::new([shape.batch * shape.m, shape.k]), + ); + let b = program.storage_read( + ScalarElement::F32.element(), + Shape::new([shape.batch * shape.k, shape.n]), + ); + let y = program.storage_write( + ScalarElement::F32.element(), + // Cooperative stores cover the whole selected output tile. + Shape::new([64, 64]), + ); + assert!(try_batched_coop_matmul( + program, + DenseMatmulTensors { + a: &a, + b: &b, + y: &y, + }, + shape, + &DenseMatmulEpilogues { + pre_a: Some(&pre), + pre_b: None, + post: Some(&post), + }, + 65_535, + DenseCoopMatmulConfig { + coop: coop_token(), + subgroups: subgroup_config(32), + tile: DenseCoopMatmulTile { + bm: 64, + bn: 64, + bk: 16, + }, + row_groups: 2, + col_groups: 2, + staging: None, + stage_buffers: 2, + swizzle_group_m: DEFAULT_SWIZZLE_GROUP_M, + }, + )); + }); + lower_or_fail(&ir, "cooperative dense f32 matmul with epilogues"); +} + #[test] fn cooperative_dense_f16_matmul_lowers() { let ir = tile::build(|program| { @@ -220,6 +283,11 @@ fn cooperative_dense_f16_matmul_lowers() { bn: 64, bk: 16, }, + row_groups: 2, + col_groups: 2, + staging: None, + stage_buffers: 2, + swizzle_group_m: DEFAULT_SWIZZLE_GROUP_M, }, )); }); @@ -265,6 +333,11 @@ fn cooperative_dense_f32_matmul_128x128_lowers() { bn: 128, bk: 16, }, + row_groups: 4, + col_groups: 2, + staging: None, + stage_buffers: 2, + swizzle_group_m: DEFAULT_SWIZZLE_GROUP_M, }, )); }); @@ -310,6 +383,11 @@ fn cooperative_dense_f32_matmul_128x64_lowers() { bn: 64, bk: 16, }, + row_groups: 4, + col_groups: 2, + staging: None, + stage_buffers: 2, + swizzle_group_m: DEFAULT_SWIZZLE_GROUP_M, }, )); }); @@ -357,6 +435,11 @@ fn cooperative_dense_f32_matmul_128x256_npass_lowers() { bn: 256, bk: 16, }, + row_groups: 4, + col_groups: 2, + staging: None, + stage_buffers: 2, + swizzle_group_m: DEFAULT_SWIZZLE_GROUP_M, }, )); }); @@ -425,117 +508,83 @@ fn module_uses_tanh(module: &naga::Module) -> bool { } #[test] -fn workgroup_qmatmul_lowers_without_subgroups() { - let ir = tile::build(|program| { - let a = program.storage_read(ScalarElement::F32.element(), Shape::new([32, 256])); - let b = quantized_matrix(program, GgmlQuantFormat::Q8_0, 256, 32); - let y = program.storage_write(ScalarElement::F32.element(), Shape::new([32, 32])); - qmatmul_workgroup_with_epilogues(program, &a, &b, &y, &QmatmulEpilogues::empty(), 65_535); - }); - let lowered = lower_or_fail(&ir, "workgroup qmatmul"); - assert!( - !module_uses_subgroup(lowered.module()), - "workgroup qmatmul emitted subgroup ops" - ); -} - -#[test] -fn f16_staged_workgroup_qmatmul_lowers_without_subgroups() { - let ir = tile::build(|program| { - let a = program.storage_read(ScalarElement::F32.element(), Shape::new([32, 256])); - let b = quantized_matrix(program, GgmlQuantFormat::Q4KNative, 256, 32); - let y = program.storage_write(ScalarElement::F32.element(), Shape::new([32, 32])); - qmatmul_workgroup_f16_with_epilogues( - program, - &a, - &b, - &y, - &QmatmulEpilogues::empty(), - 65_535, - ); - }); - let lowered = lower_or_fail(&ir, "f16 staged workgroup qmatmul"); - assert!( - module_uses_f16(lowered.module()), - "f16 staged workgroup qmatmul did not allocate f16 scratch" - ); - assert!( - !module_uses_subgroup(lowered.module()), - "f16 staged workgroup qmatmul emitted subgroup ops" - ); -} - -#[test] -fn workgroup_qgemv_lowers_without_subgroups() { - let ir = tile::build(|program| { - let a = program.storage_read(ScalarElement::F32.element(), Shape::new([1, 256])); - let b = quantized_matrix(program, GgmlQuantFormat::Q4K, 256, 128); - let y = program.storage_write(ScalarElement::F32.element(), Shape::new([1, 128])); - qgemv_workgroup_with_epilogue(program, &a, &b, &y, &QmatmulEpilogues::empty(), 65_535); - }); - let lowered = lower_or_fail(&ir, "workgroup qgemv"); - assert!( - !module_uses_subgroup(lowered.module()), - "workgroup qgemv emitted subgroup ops" - ); -} - -#[test] -fn f16_staged_workgroup_qgemv_lowers_without_subgroups() { - let ir = tile::build(|program| { - let a = program.storage_read(ScalarElement::F32.element(), Shape::new([1, 256])); - let b = quantized_matrix(program, GgmlQuantFormat::Q4KNative, 256, 128); - let y = program.storage_write(ScalarElement::F32.element(), Shape::new([1, 128])); - qgemv_workgroup_f16_with_epilogue(program, &a, &b, &y, &QmatmulEpilogues::empty(), 65_535); - }); - let lowered = lower_or_fail(&ir, "f16 staged workgroup qgemv"); - assert!( - module_uses_f16(lowered.module()), - "f16 staged workgroup qgemv did not allocate f16 scratch" - ); - assert!( - !module_uses_subgroup(lowered.module()), - "f16 staged workgroup qgemv emitted subgroup ops" - ); -} - -#[test] -fn q4k_native_workgroup_qgemv_lowers_without_subgroups() { - let ir = tile::build(|program| { - let a = program.storage_read(ScalarElement::F32.element(), Shape::new([1, 256])); - let b = quantized_matrix(program, GgmlQuantFormat::Q4KNative, 256, 128); - let y = program.storage_write(ScalarElement::F32.element(), Shape::new([1, 128])); - qgemv_workgroup_with_epilogue(program, &a, &b, &y, &QmatmulEpilogues::empty(), 65_535); - }); - let lowered = lower_or_fail(&ir, "q4k native workgroup qgemv"); - assert!( - !module_uses_subgroup(lowered.module()), - "native workgroup qgemv emitted subgroup ops" - ); -} - -#[test] -fn workgroup_qgemv_accumulator_offsets_lower_without_subgroups() { - let post = UnaryEpilogueWithExtras::new_with_value_arity("paired_product", 2, 0, |values| { +fn workgroup_qmatmul_family_lowers_without_subgroups() { + let paired = UnaryEpilogueWithExtras::new_with_value_arity("paired_product", 2, 0, |values| { values[0].clone() * values[1].clone() }); let offsets = [0, 64]; - let epilogues = QmatmulEpilogues { - post_with_extras: Some(&post), + let paired_epilogues = QmatmulEpilogues { + post_with_extras: Some(&paired), post_accumulator_offsets: &offsets, ..QmatmulEpilogues::empty() }; - let ir = tile::build(|program| { - let a = program.storage_read(ScalarElement::F32.element(), Shape::new([1, 256])); - let b = quantized_matrix(program, GgmlQuantFormat::Q4K, 256, 128); - let y = program.storage_write(ScalarElement::F32.element(), Shape::new([1, 64])); - qgemv_workgroup_with_epilogue(program, &a, &b, &y, &epilogues, 65_535); - }); - let lowered = lower_or_fail(&ir, "workgroup qgemv accumulator offsets"); - assert!( - !module_uses_subgroup(lowered.module()), - "workgroup qgemv accumulator offsets emitted subgroup ops" - ); + let empty = QmatmulEpilogues::empty(); + // (label, [m, k, n], format, staging element, epilogues) + let cases = [ + ( + "workgroup qmatmul", + [32, 256, 32], + GgmlQuantFormat::Q8_0, + ScalarElement::F32, + &empty, + ), + ( + "f16 staged workgroup qmatmul", + [32, 256, 32], + GgmlQuantFormat::Q4KNative, + ScalarElement::F16, + &empty, + ), + ( + "workgroup qgemv", + [1, 256, 128], + GgmlQuantFormat::Q4K, + ScalarElement::F32, + &empty, + ), + ( + "f16 staged workgroup qgemv", + [1, 256, 128], + GgmlQuantFormat::Q4KNative, + ScalarElement::F16, + &empty, + ), + ( + "q4k native workgroup qgemv", + [1, 256, 128], + GgmlQuantFormat::Q4KNative, + ScalarElement::F32, + &empty, + ), + ( + "workgroup qgemv accumulator offsets", + [1, 256, 64], + GgmlQuantFormat::Q4K, + ScalarElement::F32, + &paired_epilogues, + ), + ]; + for (label, [m, k, n], format, staging, epilogues) in cases { + // The paired epilogue reads two accumulators per output column. + let matrix_cols = n * epilogues.post_value_arity() as u32; + let ir = tile::build(|program| { + let a = program.storage_read(ScalarElement::F32.element(), Shape::new([m, k])); + let b = quantized_matrix(program, format, k, matrix_cols); + let y = program.storage_write(ScalarElement::F32.element(), Shape::new([m, n])); + qmatmul_workgroup_with_epilogues(program, &a, &b, &y, staging, epilogues, 65_535); + }); + let lowered = lower_or_fail(&ir, label); + assert!( + !module_uses_subgroup(lowered.module()), + "{label} emitted subgroup ops" + ); + if staging == ScalarElement::F16 { + assert!( + module_uses_f16(lowered.module()), + "{label} did not allocate f16 scratch" + ); + } + } } fn module_uses_subgroup(module: &naga::Module) -> bool { diff --git a/fusor-ml/tile-ir-kernels/tests/qmatmul_workgroup_ir.rs b/fusor-ml/tile-ir-kernels/tests/qmatmul_workgroup_ir.rs new file mode 100644 index 000000000..becff960b --- /dev/null +++ b/fusor-ml/tile-ir-kernels/tests/qmatmul_workgroup_ir.rs @@ -0,0 +1,244 @@ +//! Machine-pinned IR goldens for the workgroup-tiled quantized matmul family. +//! +//! The family is emitted by one register-tile template (`qmatmul_workgroup`), +//! and runs where the subgroup paths can't: adapters without +//! `Features::SUBGROUP`, plus every f16-activation quantized matmul. That +//! makes it hard to cover with the local GPU suites, so each config's tile IR +//! and lowered Naga module are hashed and pinned here. The digests were +//! captured from the hand-written kernels the template replaced, which the +//! template reproduced bit-for-bit across the full +//! shape x format x storage/staging x epilogue cross product (336 configs); +//! this file keeps the union of one-axis sweeps around the two base shapes. +//! +//! An intentional codegen change re-captures the goldens from the failure +//! output. + +use fusor_tile_ir::{tile, ElementType, GgmlQuantFormat, ScalarElement, Shape}; +use fusor_tile_ir_kernels::{ + qmatmul_workgroup_with_epilogues, quantized_matrix, QmatmulEpilogues, QmatmulExtra, + UnaryEpilogue, UnaryEpilogueWithExtras, +}; + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum Epilogue { + None, + Post, + PostColumn, + PostPointwise, + Pre, + PrePointwise, + PairedOffsets, +} + +#[derive(Clone, Copy, Debug)] +struct Case { + m: u32, + k: u32, + n: u32, + format: GgmlQuantFormat, + storage: ScalarElement, + staging: ScalarElement, + epilogue: Epilogue, +} + +impl Case { + fn label(&self) -> String { + format!( + "m={} k={} n={} {:?} storage={:?} staging={:?} {:?}", + self.m, self.k, self.n, self.format, self.storage, self.staging, self.epilogue + ) + } + + /// Matrix columns backing `n` output columns: the paired-offset epilogue + /// reads two accumulators per output column. + fn matrix_cols(&self) -> u32 { + match self.epilogue { + Epilogue::PairedOffsets => self.n * 2, + _ => self.n, + } + } +} + +fn build_ir(case: Case) -> fusor_tile_ir::KernelIr { + tile::build(move |program| { + let a = program.storage_read(case.storage.element(), Shape::new([case.m, case.k])); + let b = quantized_matrix(program, case.format, case.k, case.matrix_cols()); + let column = program.storage_read(ElementType::F32, Shape::new([case.n])); + let pointwise = program.storage_read(ElementType::F32, Shape::new([case.m, case.n])); + let pre_pointwise = program.storage_read(ElementType::F32, Shape::new([case.m, case.k])); + let y = program.storage_write(case.storage.element(), Shape::new([case.m, case.n])); + + let post = UnaryEpilogue::new("test_post", |tile| tile.tanh()); + let pre = UnaryEpilogue::new("test_pre", |tile| tile.silu()); + let post_extras = UnaryEpilogueWithExtras::new("test_post_extras", 1, |tiles| { + tiles[0].clone() * tiles[1].clone() + }); + let pre_extras = UnaryEpilogueWithExtras::new("test_pre_extras", 1, |tiles| { + tiles[0].clone() + tiles[1].clone() + }); + let paired = + UnaryEpilogueWithExtras::new_with_value_arity("test_paired_product", 2, 0, |values| { + values[0].clone().silu() * values[1].clone() + }); + let column_extra = [QmatmulExtra::Column(&column)]; + let pointwise_extra = [QmatmulExtra::Pointwise(&pointwise)]; + let pre_pointwise_extra = [QmatmulExtra::Pointwise(&pre_pointwise)]; + let offsets = [0, case.matrix_cols() / 2]; + + let epilogues = match case.epilogue { + Epilogue::None => QmatmulEpilogues::empty(), + Epilogue::Post => QmatmulEpilogues::post(&post), + Epilogue::PostColumn => QmatmulEpilogues { + post_with_extras: Some(&post_extras), + post_extra_inputs: &column_extra, + ..QmatmulEpilogues::empty() + }, + Epilogue::PostPointwise => QmatmulEpilogues { + post_with_extras: Some(&post_extras), + post_extra_inputs: &pointwise_extra, + ..QmatmulEpilogues::empty() + }, + Epilogue::Pre => QmatmulEpilogues::pre(&pre), + Epilogue::PrePointwise => QmatmulEpilogues { + pre_with_extras: Some(&pre_extras), + pre_extra_inputs: &pre_pointwise_extra, + ..QmatmulEpilogues::empty() + }, + Epilogue::PairedOffsets => QmatmulEpilogues { + post_with_extras: Some(&paired), + post_accumulator_offsets: &offsets, + ..QmatmulEpilogues::empty() + }, + }; + + qmatmul_workgroup_with_epilogues(program, &a, &b, &y, case.staging, &epilogues, 65_535); + }) +} + +/// Union of one-axis sweeps around the two base shapes: every shape, quant +/// format, storage/staging pair and epilogue form appears, without paying the +/// full cross product's dump-formatting cost. +fn cases() -> Vec { + const BASE: [(u32, u32, u32); 2] = [(32, 256, 32), (1, 256, 128)]; + let mut cases: Vec = Vec::new(); + let mut push = |case: Case| { + if !cases.iter().any(|seen| seen.label() == case.label()) { + cases.push(case); + } + }; + let base_case = |(m, k, n): (u32, u32, u32)| Case { + m, + k, + n, + format: GgmlQuantFormat::Q4K, + storage: ScalarElement::F32, + staging: ScalarElement::F32, + epilogue: Epilogue::None, + }; + // Aligned and ragged shapes for both geometries: the tiled family covers + // 32x32 output tiles over an 8-deep K chunk, the single-row family 1x64. + for shape in [ + (32, 256, 32), + (33, 260, 33), + (96, 512, 160), + (1, 256, 128), + (1, 260, 130), + (1, 512, 64), + ] { + push(base_case(shape)); + } + for shape in BASE { + for format in [ + GgmlQuantFormat::Q8_0, + GgmlQuantFormat::Q4K, + GgmlQuantFormat::Q4KNative, + GgmlQuantFormat::Q6K, + ] { + push(Case { + format, + ..base_case(shape) + }); + } + for (storage, staging) in [ + (ScalarElement::F32, ScalarElement::F32), + (ScalarElement::F32, ScalarElement::F16), + (ScalarElement::F16, ScalarElement::F16), + ] { + push(Case { + storage, + staging, + ..base_case(shape) + }); + // The paired-accumulator epilogue is a single-row form, and f16 + // storage only ever runs without epilogues. + if storage == ScalarElement::F16 { + continue; + } + for epilogue in [ + Epilogue::Post, + Epilogue::PostColumn, + Epilogue::PostPointwise, + Epilogue::Pre, + Epilogue::PrePointwise, + Epilogue::PairedOffsets, + ] { + if epilogue == Epilogue::PairedOffsets && shape.0 != 1 { + continue; + } + push(Case { + storage, + staging, + epilogue, + ..base_case(shape) + }); + } + } + } + cases +} + +fn fnv1a(bytes: &[u8]) -> u64 { + let mut hash = 0xcbf29ce484222325u64; + for &byte in bytes { + hash ^= byte as u64; + hash = hash.wrapping_mul(0x100000001b3); + } + hash +} + +/// `Expr`'s cached `hash` mixes the `Rc` address of every referenced local, so +/// it differs between two builds of the same program. Everything else in the +/// dump is structural; local identity is pinned by the lowered Naga module, +/// whose handles are arena indices. +fn structural_dump(ir: &fusor_tile_ir::KernelIr) -> String { + format!("{ir:#?}") + .lines() + .filter(|line| !line.trim_start().starts_with("hash: ")) + .collect::>() + .join("\n") +} + +#[test] +fn workgroup_qmatmul_ir_matches_golden() { + let measured = cases() + .into_iter() + .map(|case| { + let ir = build_ir(case); + let ir_digest = fnv1a(structural_dump(&ir).as_bytes()); + let lowered = ir + .lower_to_naga() + .unwrap_or_else(|error| panic!("lowering failed for {}: {error}", case.label())); + let naga_digest = fnv1a(format!("{:#?}", lowered.module()).as_bytes()); + format!( + "{} ir {ir_digest:#018x} naga {naga_digest:#018x}", + case.label() + ) + }) + .collect::>() + .join("\n"); + let golden = include_str!("goldens/qmatmul_workgroup_ir.txt"); + assert!( + golden.trim() == measured.trim(), + "workgroup qmatmul IR golden mismatch; measured values:\n{measured}" + ); +} diff --git a/fusor-ml/tile-ir-runtime/Cargo.toml b/fusor-ml/tile-ir-runtime/Cargo.toml index 20e1b65f8..0b1c3704b 100644 --- a/fusor-ml/tile-ir-runtime/Cargo.toml +++ b/fusor-ml/tile-ir-runtime/Cargo.toml @@ -13,9 +13,13 @@ lru = { version = "0.14.0", default-features = false } parking_lot = "0.12.3" rustc-hash = "2.1.1" tracing = "0.1.41" +serde = { version = "1", features = ["derive"] } +bincode = "1" +naga = { git = "https://github.com/ealmloff/wgpu", branch = "yield-now", features = ["serialize", "deserialize"] } [target.'cfg(target_vendor = "apple")'.dependencies] wgpu = { git = "https://github.com/ealmloff/wgpu", branch = "yield-now", default-features = false, features = ["metal"] } +libc = "0.2" [target.'cfg(windows)'.dependencies] wgpu = { git = "https://github.com/ealmloff/wgpu", branch = "yield-now", default-features = false, features = ["dx12"] } diff --git a/fusor-ml/tile-ir-runtime/src/buffer_pool.rs b/fusor-ml/tile-ir-runtime/src/buffer_pool.rs index 289962df6..b16e7e504 100644 --- a/fusor-ml/tile-ir-runtime/src/buffer_pool.rs +++ b/fusor-ml/tile-ir-runtime/src/buffer_pool.rs @@ -2,7 +2,7 @@ use std::{ num::NonZeroU64, sync::{ Arc, - atomic::{AtomicBool, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, }, }; @@ -52,7 +52,7 @@ impl CachedBuffer { } } -fn prune_cached_buffers(buffers: &mut Vec) { +fn prune_cached_buffers(buffers: &mut Vec, live_bytes: &AtomicU64) { let mut kept_free_buffers = 0; buffers.retain(|cached| { let is_free = Arc::strong_count(&cached.buffer) == 1; @@ -64,11 +64,63 @@ fn prune_cached_buffers(buffers: &mut Vec) { kept_free_buffers += 1; true } else { + // The pool holds the last reference: dropping it deallocates. + live_bytes.fetch_sub(cached.buffer.size(), Ordering::Relaxed); false } }); } +/// The in-flight GPU memory cap for pool allocations. +/// +/// Exceeding physical unified memory with GPU-referenced allocations does not +/// fail gracefully on macOS: the OS wires the memory and the machine panics +/// (`watchdogd` starvation / `IOGPUMemory` asserts) rather than returning an +/// error. Apple's contract is `MTLDevice.recommendedMaxWorkingSetSize` (about +/// two thirds of RAM below 64GB); PyTorch MPS and MLX both enforce a +/// framework-level watermark against it for the same reason. wgpu does not +/// expose the Metal value, so this reproduces the same formula from physical +/// RAM. [`crate::FusorConfig::max_gpu_memory_bytes`] overrides; non-Apple +/// targets default to uncapped (VRAM exhaustion fails with driver errors, not +/// a dead OS). +fn gpu_memory_cap(config_override: Option) -> u64 { + if let Some(cap) = config_override { + return cap.max(1); + } + #[cfg(target_vendor = "apple")] + { + let mut memsize: u64 = 0; + let mut len = std::mem::size_of::(); + let name = c"hw.memsize"; + // SAFETY: standard sysctlbyname read of a u64 with matching size. + let ok = unsafe { + libc::sysctlbyname( + name.as_ptr(), + &mut memsize as *mut u64 as *mut libc::c_void, + &mut len, + std::ptr::null_mut(), + 0, + ) + }; + if ok == 0 && memsize > 0 { + return memsize / 3 * 2; + } + } + u64::MAX +} + +/// Cumulative allocation statistics for a [`BufferPool`]. `requested` counts +/// every buffer handed out; `created` counts only the ones that missed the +/// pool cache and hit the wgpu allocator. Snapshot before/after a step and +/// diff to measure allocations per step. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct BufferPoolCounters { + pub requested: u64, + pub created: u64, + /// Bytes of pool-created buffers still tracked (free-list + handed out). + pub live_bytes: u64, +} + /// Per-device buffer pool keyed by `(size, usage)`. Reuses freed buffer /// storage so common tensor allocations skip the wgpu allocator. pub struct BufferPool { @@ -78,6 +130,14 @@ pub struct BufferPool { RwLock, FxBuildHasher>>, initialized_buffers_dirty: AtomicBool, initialized_buffer_keys: Mutex>, + buffers_requested: AtomicU64, + buffers_created: AtomicU64, + /// Bytes of pool-created buffers still tracked. Slightly overcounts when + /// the LRU evicts a whole bucket (those buffers stop being tracked + /// without a decrement) — the safe direction for a memory cap. + live_bytes: AtomicU64, + /// In-flight allocation cap; see [`gpu_memory_cap`]. + memory_cap: u64, } impl std::fmt::Debug for BufferPool { @@ -87,7 +147,11 @@ impl std::fmt::Debug for BufferPool { } impl BufferPool { - pub fn new(device: Arc, queue: Arc) -> Self { + pub fn new( + device: Arc, + queue: Arc, + config: &crate::FusorConfig, + ) -> Self { let buffer_allocation_cache = RwLock::new(LruCache::with_hasher( const { std::num::NonZeroUsize::new(BUFFER_ALLOCATION_CACHE_SIZE).unwrap() }, Default::default(), @@ -98,9 +162,34 @@ impl BufferPool { buffer_allocation_cache, initialized_buffers_dirty: AtomicBool::new(false), initialized_buffer_keys: Mutex::new(Vec::new()), + buffers_requested: AtomicU64::new(0), + buffers_created: AtomicU64::new(0), + live_bytes: AtomicU64::new(0), + memory_cap: gpu_memory_cap(config.max_gpu_memory_bytes), + } + } + + /// Snapshot the cumulative allocation counters. + pub fn counters(&self) -> BufferPoolCounters { + BufferPoolCounters { + requested: self.buffers_requested.load(Ordering::Relaxed), + created: self.buffers_created.load(Ordering::Relaxed), + live_bytes: self.live_bytes.load(Ordering::Relaxed), } } + /// Whether `buffer` is one of the pool's tracked buffers in the + /// `(size, usage)` bucket — i.e. the pool holds its own strong clone of + /// it. Liveness accounting (allocation-reuse ledger) uses this to + /// enumerate the pool as an expected `Arc` holder. Read-only: does not + /// touch LRU order. + pub fn is_tracked(&self, size: u64, usage: BufferUsages, buffer: &Arc) -> bool { + let cache = self.buffer_allocation_cache.read(); + cache + .peek(&(size, usage)) + .is_some_and(|buffers| buffers.iter().any(|c| Arc::ptr_eq(&c.buffer, buffer))) + } + /// Reset the initialized flag on all cached buffers. pub fn reset_initialized_buffers(&self) { if !self.initialized_buffers_dirty.swap(false, Ordering::AcqRel) { @@ -116,7 +205,7 @@ impl BufferPool { for buffer in buffers.iter_mut() { buffer.writen = false; } - prune_cached_buffers(buffers); + prune_cached_buffers(buffers, &self.live_bytes); } } } @@ -161,27 +250,10 @@ impl BufferPool { .store(true, Ordering::Release); self.initialized_buffer_keys.lock().push((size, usage)); } + self.buffers_requested.fetch_add(1, Ordering::Relaxed); let buffer = self .get_cached_buffer(size, usage, to_initilize) - .unwrap_or_else(|| { - let new_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { - label: Some("Tensor Buffer"), - size, - usage, - mapped_at_creation: false, - }); - - let buffer = Arc::new(new_buffer); - self.buffer_allocation_cache - .write() - .get_or_insert_mut((size, usage), Vec::new) - .push(CachedBuffer::new(buffer.clone(), to_initilize)); - if let Some(buffers) = self.buffer_allocation_cache.write().get_mut(&(size, usage)) - { - prune_cached_buffers(buffers); - } - buffer - }); + .unwrap_or_else(|| self.create_uncached_buffer(size, usage, to_initilize)); // Buffers created with init data are fully overwritten by the caller, so // only the to-be-written-by-a-kernel buffers need poisoning to surface // zero-initialization assumptions. @@ -191,6 +263,52 @@ impl BufferPool { buffer } + /// Cache-miss path: create a fresh wgpu buffer, enforcing the in-flight + /// memory cap. Growing past physical unified memory does not fail + /// gracefully on macOS — the OS panics — so when a new allocation would + /// cross the cap this first blocks on the GPU (completed submissions + /// release their buffers back to the free lists) and retries the cache; + /// only a working set that genuinely exceeds the cap fails, loudly, on + /// the host. + fn create_uncached_buffer( + &self, + size: u64, + usage: wgpu::BufferUsages, + to_initilize: bool, + ) -> Arc { + let cap = self.memory_cap; + if self.live_bytes.load(Ordering::Relaxed).saturating_add(size) > cap { + let _ = self.device.poll(wgpu::PollType::wait_indefinitely()); + if let Some(buffer) = self.get_cached_buffer(size, usage, to_initilize) { + return buffer; + } + let live = self.live_bytes.load(Ordering::Relaxed); + if live.saturating_add(size) > cap { + panic!( + "fusor: allocating {size} more bytes of GPU memory would exceed the \ + in-flight cap ({live} bytes live, cap {cap}). This working set does not \ + fit safely in unified memory; reduce the batch/model size, read results \ + back more often, or raise FUSOR_MAX_GPU_MEMORY_BYTES." + ); + } + } + self.buffers_created.fetch_add(1, Ordering::Relaxed); + self.live_bytes.fetch_add(size, Ordering::Relaxed); + let new_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { + label: Some("Tensor Buffer"), + size, + usage, + mapped_at_creation: false, + }); + + let buffer = Arc::new(new_buffer); + let mut cache = self.buffer_allocation_cache.write(); + let buffers = cache.get_or_insert_mut((size, usage), Vec::new); + buffers.push(CachedBuffer::new(buffer.clone(), to_initilize)); + prune_cached_buffers(buffers, &self.live_bytes); + buffer + } + /// Overwrite a buffer with [`DIRTY_FILL_BYTE`] so a later kernel that reads /// an unwritten region sees poison instead of zeros. Only storage buffers /// that can be a copy destination are poisoned; readback/staging buffers diff --git a/fusor-ml/tile-ir-runtime/src/cache.rs b/fusor-ml/tile-ir-runtime/src/cache.rs index 1e0920ee2..f7fa92cff 100644 --- a/fusor-ml/tile-ir-runtime/src/cache.rs +++ b/fusor-ml/tile-ir-runtime/src/cache.rs @@ -13,7 +13,7 @@ use parking_lot::RwLock; use rustc_hash::{FxBuildHasher, FxHasher}; use wgpu::{BindGroupLayout, PipelineLayout}; -use crate::DirectPlanCache; +use crate::KernelPlanCache; #[cfg(not(target_arch = "wasm32"))] const KERNEL_CACHE_SIZE: usize = 4096; @@ -26,17 +26,25 @@ const DIRECT_DYNAMIC_BIND_GROUP_CACHE_SIZE: usize = 512; /// Content-addressed key used to dedupe compiled kernel modules, shader /// modules, and pipelines across dispatches of the same kernel. +/// +/// Built on the canonical two-lane hash (see [`crate::two_lane_salted`]); +/// trusted without exact verification, so the hashed inputs must cover every +/// fact that changes generated source or binding layout — a collision or an +/// omitted field both mean dispatching the wrong pipeline. #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub struct KernelCacheKey([u64; 2]); impl KernelCacheKey { + pub(crate) fn parts(&self) -> [u64; 2] { + self.0 + } + + pub(crate) fn from_parts(parts: [u64; 2]) -> Self { + Self(parts) + } + pub fn from_hash_inputs(hash_inputs: impl Fn(&mut FxHasher)) -> Self { - Self(std::array::from_fn(|salt| { - let mut hasher = FxHasher::default(); - (salt as u64).hash(&mut hasher); - hash_inputs(&mut hasher); - hasher.finish() - })) + Self(crate::two_lane_salted(hash_inputs)) } } @@ -144,12 +152,13 @@ impl DirectDynamicBindGroupKey { /// pipeline cache. pub struct KernelCache { pub(crate) device: Arc, + config: Arc, pub(crate) wgpu_cache: Option, cache_file: Option, pub(crate) kernels: RwLock, FxBuildHasher>>, pub(crate) direct_dynamic_bind_group_cache: RwLock>, - direct_plan_cache: DirectPlanCache, + kernel_plan_cache: KernelPlanCache, direct_three_buffer_bind_group_layout: OnceLock, direct_three_buffer_pipeline_layout: OnceLock, } @@ -168,8 +177,15 @@ fn make_lru(size: usize) -> RwLock, adapter: &wgpu::Adapter) -> Self { + pub fn new( + device: Arc, + adapter: &wgpu::Adapter, + config: Arc, + ) -> Self { use wgpu::PipelineCacheDescriptor; + // tile-ir cannot see the config (the dependency points this way); + // push the liveness/arena trace flag down at device creation. + fusor_tile_ir::set_liveness_trace(config.trace_arena); let filename = wgpu::util::pipeline_cache_key(&adapter.get_info()); let (wgpu_cache, cache_file) = if let Some(filename) = filename.filter(|_| device.features().contains(wgpu::Features::PIPELINE_CACHE)) @@ -189,13 +205,20 @@ impl KernelCache { (None, None) }; + let kernel_plan_cache = KernelPlanCache::new(config.trace_resolve_host); + kernel_plan_cache.attach_disk( + device_capability_fingerprint(&device), + config.kernel_cache_dir.clone(), + ); + Self { device, + config, wgpu_cache, cache_file, kernels: make_lru(KERNEL_CACHE_SIZE), direct_dynamic_bind_group_cache: make_lru(DIRECT_DYNAMIC_BIND_GROUP_CACHE_SIZE), - direct_plan_cache: DirectPlanCache::new(), + kernel_plan_cache, direct_three_buffer_bind_group_layout: OnceLock::new(), direct_three_buffer_pipeline_layout: OnceLock::new(), } @@ -205,8 +228,13 @@ impl KernelCache { &self.device } - pub fn direct_plan_cache(&self) -> &DirectPlanCache { - &self.direct_plan_cache + /// The process configuration this cache was constructed with. + pub fn config(&self) -> &crate::FusorConfig { + &self.config + } + + pub fn kernel_plan_cache(&self) -> &KernelPlanCache { + &self.kernel_plan_cache } pub fn direct_three_buffer_bind_group_layout(&self) -> BindGroupLayout { @@ -267,7 +295,11 @@ impl KernelCache { } pub fn create_naga_shader_module(&self, kernel: &NagaKernel) -> wgpu::ShaderModule { - crate::note_compile("shader"); + crate::note_compile(&self.config, "shader"); + #[cfg(not(target_arch = "wasm32"))] + if let Some(dir) = &self.config.dump_shaders { + dump_shader(kernel, dir); + } // SAFETY: all kernels avoid out-of-bounds memory access and unbounded loops. unsafe { self.device.create_shader_module_trusted( @@ -304,6 +336,20 @@ fn shader_source(kernel: &NagaKernel) -> wgpu::ShaderSource<'static> { wgpu::ShaderSource::Naga(Cow::Owned(kernel.module().clone())) } +/// Debug aid: with `FUSOR_DUMP_SHADERS=`, every compiled kernel is also +/// serialized to WGSL in that directory, named by a running counter. +#[cfg(not(target_arch = "wasm32"))] +fn dump_shader(kernel: &NagaKernel, dir: &std::path::Path) { + use std::sync::atomic::{AtomicUsize, Ordering}; + static COUNTER: AtomicUsize = AtomicUsize::new(0); + let n = COUNTER.fetch_add(1, Ordering::Relaxed); + let _ = std::fs::create_dir_all(dir); + let _ = std::fs::write( + dir.join(format!("shader_{n:03}.ir.txt")), + format!("{:#?}", kernel.module()), + ); +} + #[cfg(target_arch = "wasm32")] fn shader_source(kernel: &NagaKernel) -> wgpu::ShaderSource<'static> { let mut wgsl = String::from(kernel.wgsl_extension_prelude()); @@ -329,3 +375,27 @@ impl Drop for KernelCache { } } } + +/// Everything device-side that steers kernel codegen: feature bits, the +/// limits the lowerer consults, and the codegen-altering environment +/// switches. Persistent plans are salted by this so a capability change can +/// never replay a mismatched kernel. +fn device_capability_fingerprint(device: &wgpu::Device) -> u64 { + use std::hash::{Hash, Hasher}; + let mut hasher = FxHasher::default(); + format!("{:?}", device.features()).hash(&mut hasher); + let limits = device.limits(); + limits.max_compute_workgroup_size_x.hash(&mut hasher); + limits.max_compute_workgroup_size_y.hash(&mut hasher); + limits.max_compute_workgroup_size_z.hash(&mut hasher); + limits + .max_compute_invocations_per_workgroup + .hash(&mut hasher); + limits + .max_compute_workgroups_per_dimension + .hash(&mut hasher); + limits + .max_storage_buffers_per_shader_stage + .hash(&mut hasher); + hasher.finish() +} diff --git a/fusor-ml/tile-ir-runtime/src/config.rs b/fusor-ml/tile-ir-runtime/src/config.rs new file mode 100644 index 000000000..115d7c50f --- /dev/null +++ b/fusor-ml/tile-ir-runtime/src/config.rs @@ -0,0 +1,178 @@ +//! Process configuration, parsed from the environment exactly once. +//! +//! Every runtime knob and trace flag lives here instead of being read from +//! the process environment at its point of use. [`FusorConfig::from_env`] +//! runs once at device creation; the value is threaded through the +//! constructors that need it ([`crate::KernelCache`], [`crate::BufferPool`], +//! and the device layer above), so configuration is explicit, inspectable, +//! and settable programmatically without touching the environment. + +use std::path::PathBuf; + +/// All Fusor runtime knobs and trace flags. +/// +/// `Default` disables every flag and leaves every knob at its built-in +/// policy; `from_env` reads the documented `FUSOR_*` variables. +#[derive(Debug, Clone, Default)] +pub struct FusorConfig { + /// Log resolver dispatch-category counts (`FUSOR_TRACE_RESOLVE`). + pub trace_resolve: bool, + /// Log host-side resolver pass timings (`FUSOR_TRACE_RESOLVE_HOST`). + pub trace_resolve_host: bool, + /// Log decode dispatch counts (`FUSOR_TRACE_DECODE`). + pub trace_decode: bool, + /// Log per-kernel dispatch names (`FUSOR_TRACE_DECODE_NAMES`). + pub trace_decode_names: bool, + /// Log decode timing at the model layer (`FUSOR_TRACE_DECODE_TIMING`; + /// `KALOSM_TRACE_DECODE_TIMING` is honored for compatibility). + pub trace_decode_timing: bool, + /// Request GPU timestamp queries and print per-kernel GPU timings + /// (`FUSOR_TRACE_GPU_KERNELS`). + pub trace_gpu_kernels: bool, + /// Log sampler pipeline decisions (`FUSOR_TRACE_SAMPLER`). + pub trace_sampler: bool, + /// Log split-K matmul selection (`FUSOR_TRACE_SPLITK`). + pub trace_splitk: bool, + /// Log horizontal matmul-merge decisions (`FUSOR_TRACE_MATMUL_MERGE`). + pub trace_matmul_merge: bool, + /// Log row-program fusion decisions (`FUSOR_TRACE_ROW_FUSION`). + pub trace_row_fusion: bool, + /// Log tiled-reduce lowering decisions (`FUSOR_TRACE_REDUCE_TILED`). + pub trace_reduce_tiled: bool, + /// Log per-kernel build times (`FUSOR_TRACE_BUILD_TIMES`). + pub trace_build_times: bool, + /// Log workgroup-tile liveness, arena packing, and barrier elision + /// (`FUSOR_TRACE_ARENA`). Pushed into tile-ir via + /// [`fusor_tile_ir::set_liveness_trace`] when the config is applied. + pub trace_arena: bool, + /// Log every shader-module / pipeline compilation + /// (`FUSOR_TRACE_PIPELINE_COMPILES`). + pub trace_pipeline_compiles: bool, + /// Validate sampler outputs against a CPU reference + /// (`FUSOR_DEBUG_SAMPLER`). + pub debug_sampler: bool, + /// Cross-check structurally shared fusion plans against fresh planning + /// (`FUSOR_VERIFY_PLAN_SHARING`). + pub verify_plan_sharing: bool, + /// Log the per-resolve ingest and window-capture ledgers of the + /// recognition-hoisting spike (`FUSOR_SPIKE_HOISTING`; see + /// `compute_graph/resolve/egraph/HOISTING_SPIKE.md`). Measurement only: + /// the ledgers change no decision. + pub spike_hoisting: bool, + /// Skip the pre-ingest recognition sweep for resolves with at most this + /// many execution nodes (`FUSOR_SPIKE_NO_RECOGNITION`), so the e-graph + /// ingests the un-preshrunk graph. This is the cost side of hoisting + /// every recognizer into a fusion generator; it is scoped by graph size + /// because an un-preshrunk training step does not fit in unified memory. + pub spike_no_recognition: Option, + /// Override the structural fusion-plan window horizon + /// (`FUSOR_SPIKE_WINDOW_DEPTH`); unset keeps the built-in stub depth. + pub spike_window_depth: Option, + /// Log what the extraction duplication gate refused + /// (`FUSOR_SPIKE_DUP_LEDGER`): the producers a fusion candidate would + /// have inlined while they still materialize for someone else. This is + /// the rematerialization surface. Measurement only; changes no decision. + pub spike_dup_ledger: bool, + /// Score the extraction byte term as total traffic — input reads plus the + /// output write — instead of the output write alone + /// (`FUSOR_SPIKE_READ_TRAFFIC`). Without it, inlining a producer looks + /// free because only the deleted write is counted and the added reads of + /// that producer's own inputs are not. + pub spike_read_traffic: bool, + /// Rewrite every built-in row-phase combine into an equivalent general + /// combine before emission (`FUSOR_SPIKE_GENERAL_COMBINE`). The results + /// must be bit-identical: this is how the general fold path is validated + /// against the closed-operator path on real reductions. + pub spike_general_combine: bool, + /// Compare extraction costs on one clock — dispatches, bytes and work + /// converted to nanoseconds by measured roofline constants — instead of + /// the lexicographic tuple (`FUSOR_SPIKE_SCALAR_COST`). Implies + /// `spike_read_traffic`. Under the tuple, dispatch count is effectively + /// infinite and no amount of bandwidth or arithmetic can outweigh it. + pub spike_scalar_cost: bool, + /// Let candidates that duplicate a live producer through to the cost + /// model instead of rejecting them outright (`FUSOR_SPIKE_NO_DUP_GATE`). + /// Measures whether the lexicographic cost tuple can price duplication + /// on its own. Experimental: the cost model does not count input reads, + /// so this is expected to over-fuse. + pub spike_no_dup_gate: bool, + /// Write every generated shader to this directory (`FUSOR_DUMP_SHADERS`). + pub dump_shaders: Option, + /// Write a Graphviz digraph of the execution graph after each resolver + /// stage to this directory (`FUSOR_DUMP_STAGES`). Requires the `graphvis` + /// feature; one `resolve-.dot` per stage per resolve. + pub dump_stages: Option, + /// Override the lazy graph's auto-flush node threshold + /// (`FUSOR_GRAPH_FLUSH_THRESHOLD`; 0 disables auto-flush). + pub graph_flush_threshold: Option, + /// Override dispatches recorded per compute pass on giant graphs + /// (`FUSOR_RESOLVE_DISPATCHES_PER_PASS`). + pub resolve_dispatches_per_pass: Option, + /// Override dispatches per queue submit on giant graphs + /// (`FUSOR_RESOLVE_DISPATCHES_PER_SUBMIT`). + pub resolve_dispatches_per_submit: Option, + /// Override the top-k chunking policy's minimum candidates per chunk + /// (`FUSOR_TOP_K_MIN_CANDIDATES_PER_CHUNK`). + pub top_k_min_candidates_per_chunk: Option, + /// Cap on pooled GPU memory before allocation panics; used to catch + /// runaway graphs in tests (`FUSOR_MAX_GPU_MEMORY_BYTES`). + pub max_gpu_memory_bytes: Option, + /// Override the on-disk kernel-plan cache directory + /// (`FUSOR_KERNEL_CACHE_DIR`); platform cache conventions apply when + /// unset. + pub kernel_cache_dir: Option, + /// Preferred wgpu adapter substring match (`WGPU_ADAPTER_NAME`). + pub adapter_name: Option, +} + +fn flag(name: &str) -> bool { + std::env::var_os(name).is_some() +} + +fn parse(name: &str) -> Option { + std::env::var(name).ok()?.parse().ok() +} + +impl FusorConfig { + /// Read every documented variable from the process environment. + pub fn from_env() -> Self { + Self { + trace_resolve: flag("FUSOR_TRACE_RESOLVE"), + trace_resolve_host: flag("FUSOR_TRACE_RESOLVE_HOST"), + trace_decode: flag("FUSOR_TRACE_DECODE"), + trace_decode_names: flag("FUSOR_TRACE_DECODE_NAMES"), + trace_decode_timing: flag("FUSOR_TRACE_DECODE_TIMING") + || flag("KALOSM_TRACE_DECODE_TIMING"), + trace_gpu_kernels: flag("FUSOR_TRACE_GPU_KERNELS"), + trace_sampler: flag("FUSOR_TRACE_SAMPLER"), + trace_splitk: flag("FUSOR_TRACE_SPLITK"), + trace_matmul_merge: flag("FUSOR_TRACE_MATMUL_MERGE"), + trace_row_fusion: flag("FUSOR_TRACE_ROW_FUSION"), + trace_reduce_tiled: flag("FUSOR_TRACE_REDUCE_TILED"), + trace_build_times: flag("FUSOR_TRACE_BUILD_TIMES"), + trace_arena: flag("FUSOR_TRACE_ARENA"), + trace_pipeline_compiles: flag("FUSOR_TRACE_PIPELINE_COMPILES"), + debug_sampler: flag("FUSOR_DEBUG_SAMPLER"), + verify_plan_sharing: flag("FUSOR_VERIFY_PLAN_SHARING"), + spike_hoisting: flag("FUSOR_SPIKE_HOISTING"), + spike_no_recognition: parse("FUSOR_SPIKE_NO_RECOGNITION"), + spike_window_depth: parse("FUSOR_SPIKE_WINDOW_DEPTH"), + spike_dup_ledger: flag("FUSOR_SPIKE_DUP_LEDGER"), + spike_read_traffic: flag("FUSOR_SPIKE_READ_TRAFFIC"), + spike_scalar_cost: flag("FUSOR_SPIKE_SCALAR_COST"), + spike_general_combine: flag("FUSOR_SPIKE_GENERAL_COMBINE"), + spike_no_dup_gate: flag("FUSOR_SPIKE_NO_DUP_GATE"), + dump_shaders: std::env::var_os("FUSOR_DUMP_SHADERS").map(PathBuf::from), + dump_stages: std::env::var_os("FUSOR_DUMP_STAGES").map(PathBuf::from), + graph_flush_threshold: parse("FUSOR_GRAPH_FLUSH_THRESHOLD"), + resolve_dispatches_per_pass: parse("FUSOR_RESOLVE_DISPATCHES_PER_PASS") + .filter(|&v: &usize| v > 0), + resolve_dispatches_per_submit: parse("FUSOR_RESOLVE_DISPATCHES_PER_SUBMIT") + .filter(|&v: &usize| v > 0), + top_k_min_candidates_per_chunk: parse("FUSOR_TOP_K_MIN_CANDIDATES_PER_CHUNK"), + max_gpu_memory_bytes: parse("FUSOR_MAX_GPU_MEMORY_BYTES"), + kernel_cache_dir: std::env::var_os("FUSOR_KERNEL_CACHE_DIR").map(PathBuf::from), + adapter_name: std::env::var("WGPU_ADAPTER_NAME").ok(), + } + } +} diff --git a/fusor-ml/tile-ir-runtime/src/direct_kernel.rs b/fusor-ml/tile-ir-runtime/src/direct_kernel.rs index fa6349ade..716083418 100644 --- a/fusor-ml/tile-ir-runtime/src/direct_kernel.rs +++ b/fusor-ml/tile-ir-runtime/src/direct_kernel.rs @@ -25,6 +25,10 @@ enum DirectKernelKind { /// quantized matrix) so dispatch skips the kernel-cache LRU entirely. Storage3 { pipeline: wgpu::ComputePipeline, + /// The lowered kernel behind the pipeline, when the construction + /// site still had it (the per-matrix decode pipeline cache keeps + /// only the pipeline). Plans need it to persist across processes. + cached: Option>, input: Arc, weight: Arc, output: Arc, @@ -55,6 +59,7 @@ enum DirectKernelTemplateKind { }, Storage3 { pipeline: wgpu::ComputePipeline, + cached: Option>, }, Sequence(Vec), } @@ -71,6 +76,95 @@ pub struct DirectKernelTemplate { kind: DirectKernelTemplateKind, } +impl DirectKernelTemplate { + /// The serializable form of this template, or `None` for kernels whose + /// pipeline cannot be rebuilt from a module alone. + pub(crate) fn to_disk(&self) -> Option { + let kind = match &self.kind { + DirectKernelTemplateKind::Dynamic { cached, bindings } => { + crate::disk_cache::DiskTemplateKind::Dynamic { + module: cached.kernel.module().clone(), + subgroups: cached.kernel.subgroups(), + bindings: bindings + .iter() + .map(|binding| (binding.binding, binding.read_only)) + .collect(), + } + } + DirectKernelTemplateKind::Storage3 { cached, .. } => { + let cached = cached.as_ref()?; + crate::disk_cache::DiskTemplateKind::Storage3 { + module: cached.kernel.module().clone(), + subgroups: cached.kernel.subgroups(), + } + } + DirectKernelTemplateKind::Sequence(templates) => { + crate::disk_cache::DiskTemplateKind::Sequence( + templates + .iter() + .map(|template| template.to_disk()) + .collect::>>()?, + ) + } + }; + Some(crate::disk_cache::DiskTemplate { + name: self.name.clone(), + dispatch_size: self.dispatch_size, + kind, + }) + } + + /// Rebuild a template from its serialized form; `None` (a cache miss) + /// when the stored module no longer validates. + pub(crate) fn from_disk( + disk: crate::disk_cache::DiskTemplate, + cache: &KernelCache, + ) -> Option { + let kind = match disk.kind { + crate::disk_cache::DiskTemplateKind::Dynamic { + module, + subgroups, + bindings, + } => { + let kernel = fusor_tile_ir::NagaKernel::from_module(module, subgroups).ok()?; + DirectKernelTemplateKind::Dynamic { + cached: Arc::new(CachedKernel::new(Arc::new(kernel))), + bindings: bindings + .into_iter() + .map(|(binding, read_only)| DirectKernelTemplateBinding { + binding, + read_only, + }) + .collect(), + } + } + crate::disk_cache::DiskTemplateKind::Storage3 { module, subgroups } => { + let kernel = fusor_tile_ir::NagaKernel::from_module(module, subgroups).ok()?; + let cached = Arc::new(CachedKernel::new(Arc::new(kernel))); + let pipeline = + crate::dispatch::prepare_three_buffer_pipeline(cache, &disk.name, &cached); + DirectKernelTemplateKind::Storage3 { + pipeline, + cached: Some(cached), + } + } + crate::disk_cache::DiskTemplateKind::Sequence(templates) => { + DirectKernelTemplateKind::Sequence( + templates + .into_iter() + .map(|template| Self::from_disk(template, cache)) + .collect::>>()?, + ) + } + }; + Some(Self { + name: disk.name, + dispatch_size: disk.dispatch_size, + kind, + }) + } +} + pub struct PreparedDirectDispatch { steps: Vec, _buffers: Vec>, @@ -99,6 +193,7 @@ impl DirectKernel { pub fn from_prepared_three_buffer_pipeline( name: impl Into, pipeline: wgpu::ComputePipeline, + cached: Option>, input: Arc, weight: Arc, output: Arc, @@ -108,6 +203,7 @@ impl DirectKernel { name: name.into(), dispatch_size, kind: DirectKernelKind::Storage3 { + cached, pipeline, input, weight, @@ -142,8 +238,11 @@ impl DirectKernel { }) .collect(), }, - DirectKernelKind::Storage3 { pipeline, .. } => DirectKernelTemplateKind::Storage3 { + DirectKernelKind::Storage3 { + pipeline, cached, .. + } => DirectKernelTemplateKind::Storage3 { pipeline: pipeline.clone(), + cached: cached.clone(), }, DirectKernelKind::Sequence(kernels) => DirectKernelTemplateKind::Sequence( kernels.iter().map(DirectKernel::to_template).collect(), @@ -187,6 +286,7 @@ impl DirectKernel { input, weight, output, + cached: _, } => { let bind_group_layout = cache.direct_three_buffer_bind_group_layout(); let bind_entries = [ @@ -218,32 +318,7 @@ impl DirectKernel { }) } DirectKernelKind::Dynamic { cached, bindings } => { - let bind_group_layout = cached - .dynamic_bind_group_layout - .get_or_init(|| { - let layout_entries = bindings - .iter() - .map(|binding| wgpu::BindGroupLayoutEntry { - binding: binding.binding, - visibility: wgpu::ShaderStages::COMPUTE, - ty: wgpu::BindingType::Buffer { - ty: wgpu::BufferBindingType::Storage { - read_only: binding.read_only, - }, - has_dynamic_offset: false, - min_binding_size: None, - }, - count: None, - }) - .collect::>(); - cache - .device - .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { - label: Some(&self.name), - entries: &layout_entries, - }) - }) - .clone(); + let (bind_group_layout, pipeline) = self.dynamic_pipeline(cache, cached, bindings); let bind_entries = bindings .iter() @@ -287,45 +362,6 @@ impl DirectKernel { .clone() }; - let pipeline_layout = cached - .dynamic_pipeline_layout - .get_or_init(|| { - cache - .device - .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some(&self.name), - bind_group_layouts: &[Some(&bind_group_layout)], - immediate_size: 0, - }) - }) - .clone(); - - let shader = cache.shader_for(cached); - let pipeline = cached - .pipeline - .get_or_init(|| { - crate::note_compile(&format!( - "pipeline name={} dispatch={:?} bindings={}", - self.name, - self.dispatch_size, - bindings.len() - )); - cache - .device - .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { - label: Some(&self.name), - layout: Some(&pipeline_layout), - module: shader, - entry_point: Some("main"), - cache: cache.wgpu_cache.as_ref(), - compilation_options: PipelineCompilationOptions { - zero_initialize_workgroup_memory: false, - ..Default::default() - }, - }) - }) - .clone(); - Some(PreparedDirectDispatch { steps: vec![PreparedDirectDispatchStep { pipeline, @@ -341,6 +377,86 @@ impl DirectKernel { } } + /// The buffer-independent compiled artifacts of a dynamic kernel: + /// bind-group layout and compute pipeline (plus the shader module and + /// pipeline layout behind them). Everything sits behind per-kernel + /// once-cells, so this is thread-safe and idempotent. + fn dynamic_pipeline( + &self, + cache: &KernelCache, + cached: &Arc, + bindings: &[DirectKernelBinding], + ) -> (wgpu::BindGroupLayout, wgpu::ComputePipeline) { + let bind_group_layout = cached + .dynamic_bind_group_layout + .get_or_init(|| { + let layout_entries = bindings + .iter() + .map(|binding| wgpu::BindGroupLayoutEntry { + binding: binding.binding, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { + read_only: binding.read_only, + }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }) + .collect::>(); + cache + .device + .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some(&self.name), + entries: &layout_entries, + }) + }) + .clone(); + let pipeline_layout = cached + .dynamic_pipeline_layout + .get_or_init(|| { + cache + .device + .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some(&self.name), + bind_group_layouts: &[Some(&bind_group_layout)], + immediate_size: 0, + }) + }) + .clone(); + + let shader = cache.shader_for(cached); + let pipeline = cached + .pipeline + .get_or_init(|| { + crate::note_compile( + cache.config(), + &format!( + "pipeline name={} dispatch={:?} bindings={}", + self.name, + self.dispatch_size, + bindings.len() + ), + ); + cache + .device + .create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some(&self.name), + layout: Some(&pipeline_layout), + module: shader, + entry_point: Some("main"), + cache: cache.wgpu_cache.as_ref(), + compilation_options: PipelineCompilationOptions { + zero_initialize_workgroup_memory: false, + ..Default::default() + }, + }) + }) + .clone(); + (bind_group_layout, pipeline) + } + pub fn bindings_for_test(&self) -> Vec { match &self.kind { DirectKernelKind::Dynamic { bindings, .. } => bindings.clone(), @@ -444,13 +560,16 @@ impl DirectKernel { bindings, } } - DirectKernelKind::Storage3 { pipeline, .. } => { + DirectKernelKind::Storage3 { + pipeline, cached, .. + } => { let input = new[*cursor].clone(); let weight = new[*cursor + 1].clone(); let output = new[*cursor + 2].clone(); *cursor += 3; DirectKernelKind::Storage3 { pipeline: pipeline.clone(), + cached: cached.clone(), input, weight, output, @@ -506,13 +625,14 @@ impl DirectKernelTemplate { bindings, } } - DirectKernelTemplateKind::Storage3 { pipeline } => { + DirectKernelTemplateKind::Storage3 { pipeline, cached } => { let input = new[*cursor].clone(); let weight = new[*cursor + 1].clone(); let output = new[*cursor + 2].clone(); *cursor += 3; DirectKernelKind::Storage3 { pipeline: pipeline.clone(), + cached: cached.clone(), input, weight, output, diff --git a/fusor-ml/tile-ir-runtime/src/disk_cache.rs b/fusor-ml/tile-ir-runtime/src/disk_cache.rs new file mode 100644 index 000000000..2646ec700 --- /dev/null +++ b/fusor-ml/tile-ir-runtime/src/disk_cache.rs @@ -0,0 +1,164 @@ +//! Persistent cache of lowered direct-kernel plans. +//! +//! Kernel construction (tile-ir program building and Naga lowering) dominates +//! first-sight resolves; the resulting plans are pure functions of the +//! 128-bit structural [`KernelCacheKey`], so they can be reused across +//! processes. Plans are stored bufferless (the in-memory plan cache rebinds +//! the caller's buffers positionally) as one file per key under a salt +//! directory that encodes the executable identity and the device capability +//! fingerprint — any compiler change or capability change starts a cold +//! cache rather than risking a stale kernel. +//! +//! Every failure path (missing file, decode error, version or key mismatch, +//! revalidation failure) falls back to rebuilding the kernel, so the cache +//! can only miss, never corrupt. + +use std::hash::{Hash, Hasher}; +use std::path::PathBuf; + +use rustc_hash::FxHasher; +use serde::{Deserialize, Serialize}; + +use crate::cache::KernelCacheKey; + +pub(crate) const DISK_PLAN_FORMAT_VERSION: u32 = 3; +/// Salt directories untouched for this long are removed on open: they belong +/// to executables that have since been rebuilt. +const STALE_SALT_AGE: std::time::Duration = std::time::Duration::from_secs(14 * 24 * 60 * 60); + +#[derive(Serialize, Deserialize)] +pub(crate) struct DiskPlanFile { + pub(crate) format: u32, + pub(crate) key: [u64; 2], + pub(crate) plans: Vec, +} + +#[derive(Serialize, Deserialize)] +pub(crate) struct DiskPlan { + /// Caller-buffer index per kernel binding slot. + pub(crate) permutation: Vec, + /// For each caller-buffer position, the first position holding the same + /// buffer at record time. + pub(crate) alias_class: Vec, + pub(crate) template: DiskTemplate, +} + +#[derive(Serialize, Deserialize)] +pub(crate) struct DiskTemplate { + pub(crate) name: String, + pub(crate) dispatch_size: [u32; 3], + pub(crate) kind: DiskTemplateKind, +} + +#[derive(Serialize, Deserialize)] +pub(crate) enum DiskTemplateKind { + Dynamic { + module: naga::Module, + subgroups: bool, + /// (binding index, read only) per buffer, in binding order. + bindings: Vec<(u32, bool)>, + }, + /// The singleton three-buffer (input, weight, output) fast-path layout; + /// the pipeline is rebuilt from the module on revival. + Storage3 { + module: naga::Module, + subgroups: bool, + }, + Sequence(Vec), +} + +pub(crate) struct DiskPlanCache { + dir: PathBuf, +} + +impl DiskPlanCache { + /// Open (creating if needed) the plan directory for this executable and + /// device fingerprint, or `None` when its location is unresolvable. + /// `dir_override` replaces the platform cache directory when set. + pub(crate) fn open(device_fingerprint: u64, dir_override: Option) -> Option { + let base = match dir_override { + Some(dir) => dir, + None => default_cache_dir()?, + }; + remove_stale_salts(&base); + let dir = base.join(format!("{:016x}", salt(device_fingerprint)?)); + std::fs::create_dir_all(&dir).ok()?; + Some(Self { dir }) + } + + fn path(&self, key: KernelCacheKey) -> PathBuf { + let [a, b] = key.parts(); + self.dir.join(format!("{a:016x}{b:016x}.plan")) + } + + pub(crate) fn load(&self, key: KernelCacheKey) -> Option { + let bytes = std::fs::read(self.path(key)).ok()?; + let file: DiskPlanFile = bincode::deserialize(&bytes).ok()?; + (file.format == DISK_PLAN_FORMAT_VERSION && file.key == key.parts()).then_some(file) + } + + /// Persist a plan. The write is synchronous — callers run on the + /// parallel kernel-build workers, and a detached write racing process + /// exit would silently drop exactly the largest (most valuable) plans. + /// It is atomic (temp file + rename) so concurrent processes see either + /// the whole file or none. + pub(crate) fn store(&self, file: DiskPlanFile) { + let path = self.path(KernelCacheKey::from_parts(file.key)); + let Ok(bytes) = bincode::serialize(&file) else { + return; + }; + let tmp = path.with_extension(format!("tmp{}", std::process::id())); + if std::fs::write(&tmp, bytes).is_ok() { + let _ = std::fs::rename(&tmp, &path); + } + } +} + +/// A salt covering everything that changes generated kernels: the exact +/// executable (any rebuild may change kernel emission) and the device +/// capability fingerprint (features and limits steer codegen). +fn salt(device_fingerprint: u64) -> Option { + let exe = std::env::current_exe().ok()?; + let meta = std::fs::metadata(&exe).ok()?; + let mut hasher = FxHasher::default(); + exe.hash(&mut hasher); + meta.len().hash(&mut hasher); + meta.modified() + .ok()? + .duration_since(std::time::UNIX_EPOCH) + .ok()? + .as_nanos() + .hash(&mut hasher); + device_fingerprint.hash(&mut hasher); + Some(hasher.finish()) +} + +fn default_cache_dir() -> Option { + #[cfg(target_vendor = "apple")] + let base = std::env::var_os("HOME").map(|home| PathBuf::from(home).join("Library/Caches")); + #[cfg(not(target_vendor = "apple"))] + let base = std::env::var_os("XDG_CACHE_HOME") + .map(PathBuf::from) + .or_else(|| std::env::var_os("HOME").map(|home| PathBuf::from(home).join(".cache"))); + Some(base?.join("fusor-ml/kernel-plans")) +} + +fn remove_stale_salts(base: &std::path::Path) { + let Ok(entries) = std::fs::read_dir(base) else { + return; + }; + let now = std::time::SystemTime::now(); + for entry in entries.flatten() { + let stale = entry + .metadata() + .ok() + .and_then(|meta| meta.modified().ok()) + .is_some_and(|modified| { + now.duration_since(modified) + .is_ok_and(|age| age > STALE_SALT_AGE) + }); + if stale { + let _ = std::fs::remove_dir_all(entry.path()); + } + } +} diff --git a/fusor-ml/tile-ir-runtime/src/dispatch.rs b/fusor-ml/tile-ir-runtime/src/dispatch.rs index c13307ffe..548b42cc7 100644 --- a/fusor-ml/tile-ir-runtime/src/dispatch.rs +++ b/fusor-ml/tile-ir-runtime/src/dispatch.rs @@ -15,7 +15,18 @@ fn cached_kernel( if let Some(cached) = cache.kernels.write().get(&key) { return Some(cached.clone()); } - let kernel = Arc::new(build_ir()?.lower_to_naga().ok()?); + let ir = build_ir()?; + // The threadgroup footprint decides residency on Apple Silicon (two + // workgroups per core at or below half the 32 KB budget); record it per + // fresh build so selection work always has real numbers. + tracing::debug!("kernel_built workgroup_bytes={}", ir.workgroup_bytes()); + let lowered = ir.lower_to_naga(); + if cache.config().trace_matmul_merge + && let Err(error) = &lowered + { + eprintln!("tile_ir_lower_error: {error:?}"); + } + let kernel = Arc::new(lowered.ok()?); Some(cache.get_or_insert_kernel(key, || kernel)) } @@ -35,7 +46,11 @@ pub fn dynamic_kernel_from_ir( dispatch_size: [u32; 3], ) -> Option { let cached = cached_kernel(cache, cache_key, build_ir)?; - let bindings = bindings_from_naga(cached.kernel.module(), buffers)?; + let bindings = bindings_from_naga( + cached.kernel.module(), + buffers, + cache.config().trace_matmul_merge, + )?; Some(DirectKernel::from_cached( name, cached, @@ -90,7 +105,7 @@ pub fn run_direct_kernel( /// Build a compute pipeline using the singleton 3-buffer pipeline layout /// for an already-cached kernel. The shader is shared with the dynamic path. -fn prepare_three_buffer_pipeline( +pub(crate) fn prepare_three_buffer_pipeline( cache: &KernelCache, name: &str, cached: &Arc, @@ -122,9 +137,10 @@ pub fn three_buffer_pipeline_from_ir( name: &str, cache_key: KernelCacheKey, build_ir: impl FnOnce() -> Option, -) -> Option { +) -> Option<(wgpu::ComputePipeline, Arc)> { let cached = cached_kernel(cache, cache_key, build_ir)?; - Some(prepare_three_buffer_pipeline(cache, name, &cached)) + let pipeline = prepare_three_buffer_pipeline(cache, name, &cached); + Some((pipeline, cached)) } /// Read each storage `GlobalVariable` from the Naga module in `(group, binding)` @@ -134,6 +150,7 @@ pub fn three_buffer_pipeline_from_ir( fn bindings_from_naga( module: &wgpu::naga::Module, buffers: impl IntoIterator>, + trace_mismatch: bool, ) -> Option> { let mut storages: Vec<(u32, bool)> = module .global_variables @@ -151,6 +168,13 @@ fn bindings_from_naga( let buffers: Vec> = buffers.into_iter().collect(); if buffers.len() != storages.len() { + if trace_mismatch { + eprintln!( + "bindings_from_naga mismatch: buffers={} storages={}", + buffers.len(), + storages.len() + ); + } return None; } Some( diff --git a/fusor-ml/tile-ir-runtime/src/lib.rs b/fusor-ml/tile-ir-runtime/src/lib.rs index 32d01cd53..e88edfab4 100644 --- a/fusor-ml/tile-ir-runtime/src/lib.rs +++ b/fusor-ml/tile-ir-runtime/src/lib.rs @@ -7,30 +7,35 @@ mod buffer_pool; mod cache; +mod config; mod direct_kernel; +mod disk_cache; mod dispatch; mod plan_cache; +mod two_lane; -pub use buffer_pool::BufferPool; +pub use buffer_pool::{BufferPool, BufferPoolCounters}; pub use cache::{ CachedKernel, DirectDynamicBindGroupKey, KernelCache, KernelCacheKey, KernelVariantKey, }; +pub use config::FusorConfig; pub use direct_kernel::{ DirectKernel, DirectKernelBinding, DirectKernelTemplate, PreparedDirectDispatch, }; pub use dispatch::{ dynamic_kernel_from_ir, run_direct_kernel, run_kernel, three_buffer_pipeline_from_ir, }; -pub use plan_cache::DirectPlanCache; +pub use plan_cache::KernelPlanCache; +pub use two_lane::{TwoLaneHasher, single_lane, two_lane_salted}; /// Diagnostic: total shader-module / compute-pipeline compilations performed at /// runtime. Each WGSL shader-module and pipeline creation bumps this counter. -/// It is logged when `FUSOR_TRACE_PIPELINE_COMPILES` is set. +/// It is logged when [`FusorConfig::trace_pipeline_compiles`] is set. static COMPILES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); -pub(crate) fn note_compile(what: &str) { +pub(crate) fn note_compile(config: &FusorConfig, what: &str) { let n = COMPILES.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1; - if std::env::var_os("FUSOR_TRACE_PIPELINE_COMPILES").is_some() { + if config.trace_pipeline_compiles { tracing::info!("fusor_compile #{n} {what}"); } } diff --git a/fusor-ml/tile-ir-runtime/src/plan_cache.rs b/fusor-ml/tile-ir-runtime/src/plan_cache.rs index e98a5ad56..b6b307492 100644 --- a/fusor-ml/tile-ir-runtime/src/plan_cache.rs +++ b/fusor-ml/tile-ir-runtime/src/plan_cache.rs @@ -10,150 +10,352 @@ use lru::LruCache; use parking_lot::Mutex; use rustc_hash::FxBuildHasher; +use crate::cache::KernelCache; use crate::{DirectKernel, DirectKernelTemplate, KernelCacheKey}; #[cfg(not(target_arch = "wasm32"))] -const DIRECT_PLAN_CACHE_SIZE: usize = 4096; +const KERNEL_PLAN_CACHE_SIZE: usize = 4096; #[cfg(target_arch = "wasm32")] -const DIRECT_PLAN_CACHE_SIZE: usize = 512; +const KERNEL_PLAN_CACHE_SIZE: usize = 512; -/// Per-device cache for direct-kernel plans. +/// Per-device cache for lowered kernel plans. /// /// This stores bufferless direct-kernel templates and replays them with the /// caller-provided binding buffers for the current dispatch. The cache never /// infers binding provenance from pointer equality; callers must provide the /// buffers in the exact order returned by [`DirectKernel::binding_buffers`]. -pub struct DirectPlanCache { - enabled: bool, - plans: Mutex, FxBuildHasher>>, +pub struct KernelPlanCache { + plans: Mutex, FxBuildHasher>>, + /// Persistent plan store, attached once the device capability + /// fingerprint is known. + disk: std::sync::OnceLock>, hits: AtomicU64, misses: AtomicU64, + disk_hits: AtomicU64, + /// Log hit/miss totals on every event (wasm always logs). + trace_events: bool, } -struct CachedDirectKernelPlan { +struct CachedKernelPlan { template: DirectKernelTemplate, - binding_count: usize, + /// Caller-buffer index per kernel binding slot: kernels may bind the + /// caller's buffers in any order (or bind one buffer several times), so + /// rebinding routes `caller_buffers[permutation[slot]]` into each slot. + permutation: Vec, + /// For each caller-buffer position, the first position holding the same + /// buffer at record time. The kernel body is only correct for callers + /// with the *identical* aliasing pattern: a body built for distinct + /// buffers binds an aliased pair twice (wrong and rejected by wgpu), + /// and a body built over an alias (an in-place output) would clobber a + /// caller whose buffers are distinct. + alias_class: Vec, } -impl std::fmt::Debug for DirectPlanCache { +impl std::fmt::Debug for KernelPlanCache { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DirectPlanCache") - .field("enabled", &self.enabled) - .finish() + f.debug_struct("KernelPlanCache").finish() } } -impl Default for DirectPlanCache { +impl Default for KernelPlanCache { fn default() -> Self { - Self::new() + Self::new(false) } } -impl DirectPlanCache { - pub fn new() -> Self { +impl KernelPlanCache { + pub fn new(trace_events: bool) -> Self { Self { - enabled: std::env::var_os("FUSOR_DISABLE_DECODE_PLAN_CACHE").is_none(), plans: Mutex::new(LruCache::with_hasher( - NonZeroUsize::new(DIRECT_PLAN_CACHE_SIZE) - .expect("direct plan cache size must be non-zero"), + NonZeroUsize::new(KERNEL_PLAN_CACHE_SIZE) + .expect("kernel plan cache size must be non-zero"), Default::default(), )), + disk: std::sync::OnceLock::new(), hits: AtomicU64::new(0), misses: AtomicU64::new(0), + disk_hits: AtomicU64::new(0), + trace_events, } } - pub fn enabled(&self) -> bool { - self.enabled + /// Attach the persistent plan store. Kernel codegen depends on device + /// capabilities, so the store is salted by their fingerprint. + /// `dir_override` replaces the platform cache directory when set. + pub fn attach_disk(&self, device_fingerprint: u64, dir_override: Option) { + let _ = self.disk.set(crate::disk_cache::DiskPlanCache::open( + device_fingerprint, + dir_override, + )); } pub fn try_get_or_insert( &self, + cache: &KernelCache, key: KernelCacheKey, binding_buffers: &[Arc], build: impl FnOnce() -> Result, ) -> Result { - let mut kernels = self.try_get_or_insert_many(key, &[binding_buffers], || { + let mut kernels = self.try_get_or_insert_many(cache, key, &[binding_buffers], || { build().map(|kernel| vec![kernel]) })?; Ok(kernels .pop() - .expect("single direct plan cache result must contain one kernel")) + .expect("single kernel plan cache result must contain one kernel")) } pub fn try_get_or_insert_many( &self, + cache: &KernelCache, key: KernelCacheKey, binding_buffers: &[&[Arc]], build: impl FnOnce() -> Result, E>, ) -> Result, E> { - if !self.enabled { - return build(); + if let Some(kernels) = self.get_many(cache, key, binding_buffers) { + return Ok(kernels); } + let built = build()?; + self.insert_many(key, &built, binding_buffers); + Ok(built) + } + /// A cached plan (memory first, then the persistent store) rebound to + /// `binding_buffers`, or `None` on a miss. + pub fn get_many( + &self, + cache: &KernelCache, + key: KernelCacheKey, + binding_buffers: &[&[Arc]], + ) -> Option> { + if let Some(plan) = self.memory_plan(key) + && binding_shape_matches(&plan, binding_buffers) { - let mut plans = self.plans.lock(); - if let Some(plan) = plans.get(&key) - && binding_shape_matches(plan, binding_buffers) - { - let hit_total = self.hits.fetch_add(1, Ordering::Relaxed) + 1; - trace_cache_event(hit_total, self.misses.load(Ordering::Relaxed)); - return Ok(plan - .iter() - .zip(binding_buffers) - .map(|(plan, buffers)| plan.template.bind_buffers(buffers)) - .collect()); - } + let hit_total = self.hits.fetch_add(1, Ordering::Relaxed) + 1; + self.trace_cache_event(hit_total, self.misses.load(Ordering::Relaxed)); + return Some(bind_plan(&plan, binding_buffers)); + } + + if let Some(disk) = self.disk.get().and_then(Option::as_ref) + && let Some(file) = disk.load(key) + && let Some(plan) = plans_from_disk(file, cache) + && binding_shape_matches(&plan, binding_buffers) + { + let bound = bind_plan(&plan, binding_buffers); + self.plans.lock().put(key, plan.into()); + let disk_total = self.disk_hits.fetch_add(1, Ordering::Relaxed) + 1; + tracing::debug!("kernel_plan_disk_hit total={disk_total}"); + return Some(bound); } let miss_total = self.misses.fetch_add(1, Ordering::Relaxed) + 1; - trace_cache_event(self.hits.load(Ordering::Relaxed), miss_total); - let built = build()?; - if binding_buffers_match(&built, binding_buffers) { - self.plans.lock().put(key, record_plan(&built)); + self.trace_cache_event(self.hits.load(Ordering::Relaxed), miss_total); + None + } + + fn trace_cache_event(&self, hits: u64, misses: u64) { + if cfg!(target_arch = "wasm32") || self.trace_events { + tracing::info!("kernel_plan_cache hit={hits} miss={misses}"); } - Ok(built) } + + /// Clone the immutable plan snapshot while touching its LRU entry, then + /// release the cache lock before validating or binding caller buffers. + fn memory_plan(&self, key: KernelCacheKey) -> Option> { + self.plans.lock().get(&key).cloned() + } + + /// Record a built plan when its true binding order matches the caller's + /// buffer list; silently skips plans the positional rebind model cannot + /// express (internal scratch allocations, deduplicated bindings). + pub fn insert_many( + &self, + key: KernelCacheKey, + kernels: &[DirectKernel], + binding_buffers: &[&[Arc]], + ) { + let Some(plan) = record_plan(kernels, binding_buffers) else { + return; + }; + if let Some(disk) = self.disk.get().and_then(Option::as_ref) + && let Some(file) = plans_to_disk(key, &plan) + { + disk.store(file); + } + self.plans.lock().put(key, plan.into()); + } +} + +fn bind_plan( + plan: &[CachedKernelPlan], + binding_buffers: &[&[Arc]], +) -> Vec { + plan.iter() + .zip(binding_buffers) + .map(|(plan, buffers)| { + let routed: Vec> = plan + .permutation + .iter() + .map(|&index| buffers[index].clone()) + .collect(); + plan.template.bind_buffers(&routed) + }) + .collect() } -fn record_plan(kernels: &[DirectKernel]) -> Vec { +fn plans_from_disk( + file: crate::disk_cache::DiskPlanFile, + cache: &KernelCache, +) -> Option> { + file.plans + .into_iter() + .map(|plan| { + let template = DirectKernelTemplate::from_disk(plan.template, cache)?; + let len = plan.alias_class.len(); + (plan.permutation.iter().all(|&index| index < len) + && plan + .alias_class + .iter() + .enumerate() + .all(|(index, &class)| class <= index)) + .then_some(CachedKernelPlan { + template, + permutation: plan.permutation, + alias_class: plan.alias_class, + }) + }) + .collect() +} + +fn plans_to_disk( + key: KernelCacheKey, + plans: &[CachedKernelPlan], +) -> Option { + let plans = plans + .iter() + .map(|plan| { + Some(crate::disk_cache::DiskPlan { + permutation: plan.permutation.clone(), + alias_class: plan.alias_class.clone(), + template: plan.template.to_disk()?, + }) + }) + .collect::>>()?; + Some(crate::disk_cache::DiskPlanFile { + format: crate::disk_cache::DISK_PLAN_FORMAT_VERSION, + key: key.parts(), + plans, + }) +} + +/// Record templates plus the binding permutation against the caller's +/// buffers; `None` when a kernel binds a buffer the caller does not present +/// (an internal allocation the positional rebind model cannot express) or +/// when the caller's buffers alias each other. Aliasing makes the +/// permutation ambiguous: an in-place output recorded over its input would +/// permanently route the output binding into the input slot, corrupting any +/// later dispatch of the same structural key whose buffers do not alias. +fn record_plan( + kernels: &[DirectKernel], + binding_buffers: &[&[Arc]], +) -> Option> { + if kernels.len() != binding_buffers.len() { + return None; + } kernels .iter() - .map(|kernel| CachedDirectKernelPlan { - template: kernel.to_template(), - binding_count: kernel.binding_buffers().len(), + .zip(binding_buffers) + .map(|(kernel, expected)| { + let permutation = kernel + .binding_buffers() + .iter() + .map(|bound| { + expected + .iter() + .position(|buffer| Arc::ptr_eq(buffer, bound)) + }) + .collect::>>()?; + Some(CachedKernelPlan { + template: kernel.to_template(), + permutation, + alias_class: alias_classes(expected), + }) + }) + .collect() +} + +/// For each position, the first position holding the same buffer. +fn alias_classes(buffers: &[Arc]) -> Vec { + buffers + .iter() + .enumerate() + .map(|(index, buffer)| { + buffers[..index] + .iter() + .position(|earlier| Arc::ptr_eq(earlier, buffer)) + .unwrap_or(index) }) .collect() } +/// Whether the caller's buffers reproduce the recorded aliasing pattern +/// exactly (same positions aliased, same positions distinct). +fn alias_pattern_matches(recorded: &[usize], buffers: &[Arc]) -> bool { + recorded.len() == buffers.len() && alias_classes(buffers) == recorded +} + fn binding_shape_matches( - plan: &[CachedDirectKernelPlan], + plan: &[CachedKernelPlan], binding_buffers: &[&[Arc]], ) -> bool { plan.len() == binding_buffers.len() && plan .iter() .zip(binding_buffers) - .all(|(plan, buffers)| plan.binding_count == buffers.len()) + .all(|(plan, buffers)| alias_pattern_matches(&plan.alias_class, buffers)) } -fn binding_buffers_match(kernels: &[DirectKernel], expected: &[&[Arc]]) -> bool { - if kernels.len() != expected.len() { - return false; +#[cfg(test)] +mod tests { + use super::*; + + fn test_cache(capacity: usize) -> KernelPlanCache { + KernelPlanCache { + plans: Mutex::new(LruCache::with_hasher( + NonZeroUsize::new(capacity).unwrap(), + Default::default(), + )), + disk: std::sync::OnceLock::new(), + hits: AtomicU64::new(0), + misses: AtomicU64::new(0), + disk_hits: AtomicU64::new(0), + trace_events: false, + } + } + + fn empty_plan() -> Arc<[CachedKernelPlan]> { + Arc::from(Vec::new()) } - kernels.iter().zip(expected).all(|(kernel, expected)| { - let actual = kernel.binding_buffers(); - actual.len() == expected.len() - && actual - .iter() - .zip(*expected) - .all(|(actual, expected)| Arc::ptr_eq(actual, expected)) - }) -} + #[test] + fn memory_plan_releases_lock_and_touches_lru() { + let cache = test_cache(2); + let first = KernelCacheKey::from_parts([1, 1]); + let second = KernelCacheKey::from_parts([2, 2]); + let third = KernelCacheKey::from_parts([3, 3]); + { + let mut plans = cache.plans.lock(); + plans.put(first, empty_plan()); + plans.put(second, empty_plan()); + } + + let first_snapshot = cache.memory_plan(first).unwrap(); + assert!(cache.plans.try_lock().is_some()); -fn trace_cache_event(hits: u64, misses: u64) { - if cfg!(target_arch = "wasm32") || std::env::var_os("FUSOR_TRACE_RESOLVE_HOST").is_some() { - tracing::info!("direct_plan_cache hit={hits} miss={misses}"); + cache.plans.lock().put(third, empty_plan()); + let plans = cache.plans.lock(); + assert!(plans.peek(&first).is_some()); + assert!(plans.peek(&second).is_none()); + assert!(plans.peek(&third).is_some()); + assert!(first_snapshot.is_empty()); } } diff --git a/fusor-ml/tile-ir-runtime/src/two_lane.rs b/fusor-ml/tile-ir-runtime/src/two_lane.rs new file mode 100644 index 000000000..94c06a478 --- /dev/null +++ b/fusor-ml/tile-ir-runtime/src/two_lane.rs @@ -0,0 +1,75 @@ +//! The canonical structural-identity hash. +//! +//! Every "same structure?" question in the stack — kernel cache keys, kernel +//! plan keys, flush-replay fingerprints, fusion-plan window keys, semantic +//! payload identity — answers with a 128-bit key produced here, in one of +//! two flavors: +//! +//! - [`TwoLaneHasher`]: an accumulating hasher for streaming writes. Lane +//! `b` is fed a deterministic mix of the same 64-bit words as lane `a`, +//! so per-write entropy stays 64 bits (one FxHash); the second lane widens +//! the *accumulator* state to make cross-item cancellation collisions +//! harder, not the per-item hash. +//! - [`two_lane_salted`]: a one-shot flavor that re-runs the input closure +//! once per differently-seeded lane, giving the full 128 bits of per-key +//! entropy when the inputs are cheap to re-hash. +//! +//! Collision contract: consumers trust these keys without byte-exact +//! verification. Each key domain documents at its newtype what a collision +//! would cost and what secondary validation (replay validation, verify +//! flags, recorder poisoning) bounds the damage. + +use std::hash::{Hash, Hasher}; + +use rustc_hash::FxHasher; + +/// Two differently-seeded accumulator lanes over the same 64-bit words. +pub struct TwoLaneHasher { + a: FxHasher, + b: FxHasher, +} + +impl Default for TwoLaneHasher { + fn default() -> Self { + Self::new() + } +} + +impl TwoLaneHasher { + pub fn new() -> Self { + let mut a = FxHasher::default(); + 0u64.hash(&mut a); + let mut b = FxHasher::default(); + 1u64.hash(&mut b); + Self { a, b } + } + + pub fn write_u64(&mut self, value: u64) { + value.hash(&mut self.a); + (value.rotate_left(32) ^ 0x9E37_79B9_7F4A_7C15).hash(&mut self.b); + } + + pub fn finish(self) -> [u64; 2] { + [self.a.finish(), self.b.finish()] + } +} + +/// One-shot two-lane key: the closure's writes are hashed once per lane +/// under distinct seeds, so both lanes carry independent 64-bit digests of +/// the full input stream. +pub fn two_lane_salted(hash_inputs: impl Fn(&mut FxHasher)) -> [u64; 2] { + std::array::from_fn(|salt| { + let mut hasher = FxHasher::default(); + (salt as u64).hash(&mut hasher); + hash_inputs(&mut hasher); + hasher.finish() + }) +} + +/// Single-lane convenience for hashing a sub-structure into one word before +/// feeding it to a [`TwoLaneHasher`]. +pub fn single_lane(f: impl FnOnce(&mut FxHasher)) -> u64 { + let mut hasher = FxHasher::default(); + f(&mut hasher); + hasher.finish() +} diff --git a/fusor-ml/tile-ir/src/analysis/elide.rs b/fusor-ml/tile-ir/src/analysis/elide.rs new file mode 100644 index 000000000..f2fa615f8 --- /dev/null +++ b/fusor-ml/tile-ir/src/analysis/elide.rs @@ -0,0 +1,229 @@ +//! Barrier elision: remove workgroup barriers that separate nothing. +//! +//! Elision never judges absolute correctness — lane-ownership discipline is +//! the kernel author's (a tile whose rows are only ever touched by their +//! owning thread needs no barrier, and the analysis cannot see that). It +//! preserves the *separation structure* instead: a barrier is removable only +//! when every conservatively-hazardous access pair it currently separates is +//! also separated by another barrier. Under that rule removal can never +//! introduce a race that the original ordering excluded, and the packer +//! (which re-derives sharing legality from the post-elision body) can only +//! lose opportunities the removed barrier alone provided — which the rule +//! also forbids, because expanded-range endpoints are access positions. +//! +//! Hazard pairs are enumerated two ways: +//! - *Forward*: accesses `x < y` (any two tiles or one tile, at least one +//! write). A barrier separates the pair when it sits in `(x, y]` and every +//! loop between the barrier and the pair's innermost common loop is +//! guaranteed to complete (a skippable inner loop can skip the barrier). +//! - *Back edge*: accesses `x`, `y` inside a loop `L` (at least one write) +//! race from iteration `i`'s later access to iteration `i + 1`'s earlier +//! one. A barrier inside `L` separates the wrap when it covers the +//! circular interval `(y, L.end) ∪ (L.start, x]`; `Break` does not +//! invalidate it (taking the back edge means the full body executed), so +//! this holds even in break loops. + +use crate::ir::{KernelIr, Stmt}; + +use super::LivenessInfo; + +/// Remove every removable barrier from `ir`. Returns how many were removed. +pub(crate) fn elide_barriers(ir: &mut KernelIr) -> usize { + let has_barrier = { + fn any_barrier(stmts: &[Stmt]) -> bool { + stmts.iter().any(|stmt| match stmt { + Stmt::Barrier => true, + Stmt::Loop { body, .. } => any_barrier(body), + Stmt::If { accept, reject, .. } => any_barrier(accept) || any_barrier(reject), + _ => false, + }) + } + any_barrier(&ir.body) + }; + if !has_barrier { + return 0; + } + + let info = LivenessInfo::compute(ir); + if info.tiles.is_empty() { + // No workgroup tiles: control barriers still order storage traffic + // paired with storage barriers; leave them alone. + return 0; + } + let pairs = hazard_pairs(&info); + let mut alive = vec![true; info.barriers.len()]; + // Greedy forward order, re-checking against the remaining set so two + // barriers that mutually back each other up cannot both go. + for candidate in 0..info.barriers.len() { + alive[candidate] = false; + let preserved = pairs.iter().all(|pair| { + !separates(&info, candidate, pair) + || (0..info.barriers.len()) + .any(|other| alive[other] && separates(&info, other, pair)) + }); + if !preserved { + alive[candidate] = true; + } + } + + let removed: Vec<&[u32]> = info + .barriers + .iter() + .zip(&alive) + .filter(|(_, alive)| !**alive) + .map(|(barrier, _)| barrier.path.as_slice()) + .collect(); + if removed.is_empty() { + return 0; + } + if super::trace_enabled() { + eprintln!("arena-elide removing {} barrier(s)", removed.len()); + } + // Paths sorted descending so earlier removals never shift later ones. + let mut paths: Vec> = removed.iter().map(|path| path.to_vec()).collect(); + paths.sort_unstable_by(|a, b| b.cmp(a)); + for path in &paths { + remove_stmt(&mut ir.body, path); + } + paths.len() +} + +/// One conservatively-hazardous access pair at raw walk positions. +struct HazardPair { + /// Earlier access position (forward), or the wrap target (back edge). + x: u32, + /// Later access position (forward), or the wrap source (back edge). + y: u32, + /// Innermost common loop for forward pairs; the wrapped loop for back + /// edges. + scope: Option, + back_edge: bool, +} + +fn hazard_pairs(info: &LivenessInfo) -> Vec { + // Pairs collapse to position signatures: tile identity is irrelevant to + // separation, so dedupe keeps the greedy pass cheap on barrier-heavy + // kernels (elision runs on every eager kernel build). + let mut seen = rustc_hash::FxHashSet::default(); + let mut pairs = Vec::new(); + let accesses: Vec<_> = info + .order + .iter() + .flat_map(|key| info.tiles[key].accesses.iter().copied()) + .collect(); + for (index, a) in accesses.iter().enumerate() { + for b in &accesses[index + 1..] { + if !a.kind.writes() && !b.kind.writes() { + continue; + } + let (x, y) = if a.position <= b.position { + (a.position, b.position) + } else { + (b.position, a.position) + }; + let scope = info.innermost_common_loop(x, y); + if x != y && seen.insert((x, y, false)) { + pairs.push(HazardPair { + x, + y, + scope, + back_edge: false, + }); + } + // The wrap: both accesses inside a common loop race across the + // back edge (y@i races x@i+1), including x == y across + // iterations. + if scope.is_some() && seen.insert((x, y, true)) { + pairs.push(HazardPair { + x, + y, + scope, + back_edge: true, + }); + } + } + } + pairs +} + +/// Whether barrier `candidate` orders the pair. The barrier must execute on +/// every path between the two accesses: every loop enclosing the barrier +/// below the pair's scope must be guaranteed to complete. +fn separates(info: &LivenessInfo, candidate: usize, pair: &HazardPair) -> bool { + let barrier = &info.barriers[candidate]; + if !guaranteed_within(info, candidate, pair.scope) { + return false; + } + if pair.back_edge { + let scope = &info.loops[pair.scope.expect("back edges carry a scope") as usize].span; + // Inside the wrapped loop, covering (y, end) or (start, x]. + let position = barrier.position; + let inside = scope.first < position && position < scope.last; + inside && (position > pair.y || position <= pair.x) + } else { + barrier.position > pair.x && barrier.position <= pair.y + } +} + +/// Every loop enclosing the barrier strictly below `scope` completes every +/// pass, so the barrier executes whenever control flows from one end of the +/// scope to the other. +fn guaranteed_within(info: &LivenessInfo, candidate: usize, scope: Option) -> bool { + let barrier = &info.barriers[candidate]; + match scope { + None => barrier.guaranteed, + // A barrier positionally between two same-scope accesses is inside + // the scope by construction; `guaranteed_below` rejects the rest. + Some(scope) => info.guaranteed_below(barrier, scope) || barrier.guaranteed, + } +} + +fn remove_stmt(body: &mut Vec, path: &[u32]) { + let index = path[0] as usize; + match path.len() { + 1 => { + debug_assert!(matches!(body[index], Stmt::Barrier)); + body.remove(index); + } + _ => match &mut body[index] { + Stmt::Loop { body: inner, .. } => remove_stmt(inner, &path[1..]), + _ => unreachable!("barrier paths descend through loop bodies only"), + }, + } +} + +/// A profitable barrier insertion: one uniform barrier at a top-level +/// boundary lets otherwise-blocked tiles share an allocation. Policy — +/// whether the saving crosses an occupancy class worth a barrier — lives +/// with the caller; tile-ir only reports the delta. +pub struct BarrierSuggestion { + /// Root-level statement index to insert before. + index: usize, + /// Workgroup bytes saved by the insertion. + pub bytes_saved: u64, +} + +pub(crate) fn barrier_suggestions(ir: &KernelIr) -> Vec { + let current = crate::lower::workgroup_bytes(ir); + let mut suggestions = Vec::new(); + // Root boundaries are uniform by construction (every thread reaches + // them); simulate each insertion exactly rather than re-deriving the + // packer's position arithmetic. + for index in 1..ir.body.len() { + let mut candidate = ir.clone(); + candidate.body.insert(index, Stmt::Barrier); + let saved = current.saturating_sub(crate::lower::workgroup_bytes(&candidate)); + if saved > 0 { + suggestions.push(BarrierSuggestion { + index, + bytes_saved: saved, + }); + } + } + suggestions.sort_by_key(|suggestion| std::cmp::Reverse(suggestion.bytes_saved)); + suggestions +} + +pub(crate) fn apply_barrier_suggestion(ir: &mut KernelIr, suggestion: &BarrierSuggestion) { + ir.body.insert(suggestion.index, Stmt::Barrier); +} diff --git a/fusor-ml/tile-ir/src/analysis/mod.rs b/fusor-ml/tile-ir/src/analysis/mod.rs new file mode 100644 index 000000000..25ee116c2 --- /dev/null +++ b/fusor-ml/tile-ir/src/analysis/mod.rs @@ -0,0 +1,588 @@ +//! Workgroup-tile liveness analysis. +//! +//! One walk over a [`KernelIr`] body produces [`LivenessInfo`]: per-tile live +//! ranges, uniform-barrier positions, and loop spans. The lowering arena +//! consumes it to share workgroup allocations. +//! +//! Two workgroup tiles may share one allocation when their live ranges are +//! disjoint *and* a uniform workgroup barrier orders every thread's last +//! touch of the earlier tile before any thread's first touch of the later +//! one. Threads of a workgroup are not in lockstep, so plain program-order +//! disjointness is not enough — without the barrier a fast thread could +//! write the later tile while a slow thread still reads the earlier one. +//! +//! Loops add a wrap-around hazard: when both tiles live inside a common +//! loop, the later tile's last touch of iteration `i` races the earlier +//! tile's first touch of iteration `i + 1`. [`expand_ranges_over_loops`] +//! folds that in by widening every range to cover each loop body it +//! intersects, so two tiles sharing a loop always overlap and plain +//! interval disjointness plus one forward barrier is sound. +//! +//! Barriers inside `If` blocks are not uniform and never count. Barriers +//! inside loops that may break, return, or run zero dynamic iterations may +//! be skipped at runtime, so they never separate tiles either (they could +//! only ever separate tiles living wholly outside the loop — see +//! [`BarrierInfo::guaranteed`]). + +use std::sync::atomic::{AtomicU8, Ordering}; + +use rustc_hash::FxHashMap; + +use crate::ir::{Accumulator, Expr, KernelIr, Stmt, Tile, TileUse}; +use crate::{ElementType, MemoryLevel}; + +mod elide; +pub use elide::BarrierSuggestion; +pub(crate) use elide::{apply_barrier_suggestion, barrier_suggestions, elide_barriers}; + +mod verify; +pub(crate) use verify::verify_arena; + +mod trace { + use super::*; + + const UNSET: u8 = 2; + static TRACE: AtomicU8 = AtomicU8::new(UNSET); + + /// Enable or disable liveness/arena tracing at runtime. Called by the + /// runtime crate when `FusorConfig` is materialized; until then the + /// `FUSOR_TRACE_ARENA` env var is the fallback so standalone tile-ir + /// tests keep working. + pub fn set_liveness_trace(enabled: bool) { + TRACE.store(enabled as u8, Ordering::Relaxed); + } + + pub(crate) fn enabled() -> bool { + match TRACE.load(Ordering::Relaxed) { + UNSET => { + let on = std::env::var_os("FUSOR_TRACE_ARENA").is_some(); + TRACE.store(on as u8, Ordering::Relaxed); + on + } + value => value == 1, + } + } +} + +pub use trace::set_liveness_trace; +pub(crate) use trace::enabled as trace_enabled; + +#[derive(Clone, Copy)] +pub(crate) struct LiveRange { + pub first: u32, + pub last: u32, +} + +/// How a statement touches a tile. +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) enum AccessKind { + Read, + Write, + /// Collective read-modify-write (reduction scratch). + ReadWrite, +} + +impl AccessKind { + pub(crate) fn writes(self) -> bool { + !matches!(self, Self::Read) + } +} + +#[derive(Clone, Copy)] +pub(crate) struct TileAccess { + /// Position of the enclosing statement. + pub position: u32, + pub kind: AccessKind, +} + +pub(crate) struct TileLiveness { + /// Live range after loop expansion. + pub range: LiveRange, + pub element: ElementType, + /// Allocation extent in elements of `element`. + pub elements: u32, + /// Every touch in walk order, at raw (pre-expansion) positions. + pub accesses: Vec, + /// When every access lies inside one innermost loop: that loop and the + /// tile's per-iteration phase (raw positions expanded over loops nested + /// inside it). Enables sharing between in-loop tiles whose phases are + /// barrier-separated both forward and across the back edge. + pub scoped: Option<(u32, LiveRange)>, + /// Consumed as a raw cooperative-matrix pointer (`CoopLoad` / + /// `CoopStoreTile`): the emitted array type must equal the tile's + /// element, so its region never widens to a canonical type. + pub coop: bool, +} + +pub(crate) struct LoopInfo { + /// Positions spanned by the loop: `start` is the `Loop` statement itself, + /// `end` the synthetic position after the body. + pub span: LiveRange, + /// A `Break` statement is attributed to this loop (innermost frame). + pub has_break: bool, + /// A `Return` statement occurs anywhere in the body (exits every loop). + pub has_return: bool, + /// The loop count when it is a static literal. + pub static_count: Option, +} + +impl LoopInfo { + /// Whether every dynamic execution of this loop runs the full body at + /// least once: a positive static-literal count with no early exit. A + /// dynamic count may be zero at runtime, and a `Break`/`Return` can skip + /// the tail of the body — either way a barrier inside the loop is not + /// guaranteed to execute. + pub(crate) fn guaranteed_once(&self) -> bool { + self.static_count.is_some_and(|count| count > 0) && !self.has_break && !self.has_return + } +} + +pub(crate) struct BarrierInfo { + pub position: u32, + /// Statement indices from the body root to the barrier, descending only + /// through `Loop` bodies (uniform barriers are never inside `If`). + pub path: Vec, + /// Enclosing loop indices, outermost first. + pub enclosing_loops: Vec, + /// Every enclosing loop is [`LoopInfo::guaranteed_once`], so every + /// thread passes this barrier on every full pass of the enclosing body. + /// Only guaranteed barriers separate tile live ranges: an in-loop + /// barrier can only ever separate tiles living wholly outside the loop + /// (range expansion pins intersecting tiles to the span boundary), which + /// is exactly the case a zero-trip, `Break`, or `Return` can skip. + pub guaranteed: bool, +} + +pub(crate) struct LivenessInfo { + pub tiles: FxHashMap<*const (), TileLiveness>, + /// First-touch order of workgroup tiles (deterministic assignment order). + pub order: Vec<*const ()>, + /// Uniform workgroup barriers, in position order. + pub barriers: Vec, + /// Completed loop spans, indexed stably from frame push. + pub loops: Vec, +} + +impl LivenessInfo { + pub(crate) fn compute(ir: &KernelIr) -> Self { + let mut walk = Walk::default(); + walk.visit_stmts(&ir.body); + walk.expand_ranges_over_loops(); + for barrier in &mut walk.barriers { + barrier.guaranteed = barrier + .enclosing_loops + .iter() + .all(|&index| walk.loops[index as usize].guaranteed_once()); + } + let mut info = Self { + tiles: walk.tiles, + order: walk.order, + barriers: walk.barriers, + loops: walk.loops, + }; + info.compute_scoped_phases(); + if trace_enabled() { + for (index, key) in info.order.iter().enumerate() { + let tile = &info.tiles[key]; + eprintln!( + "arena-tile {index}: {:?} x{} range=({},{})", + tile.element, tile.elements, tile.range.first, tile.range.last + ); + } + eprintln!( + "arena-barriers {:?}", + info.barriers + .iter() + .map(|barrier| barrier.position) + .collect::>() + ); + for info in &info.loops { + eprintln!("arena-loop ({},{})", info.span.first, info.span.last); + } + } + info + } + + /// The innermost loop whose span strictly contains `[x, y]`. + pub(crate) fn innermost_common_loop(&self, x: u32, y: u32) -> Option { + let mut best: Option = None; + for (index, info) in self.loops.iter().enumerate() { + if info.span.first < x && y < info.span.last { + let tighter = match best { + None => true, + Some(previous) => { + let previous = &self.loops[previous as usize].span; + info.span.first >= previous.first && info.span.last <= previous.last + } + }; + if tighter { + best = Some(index as u32); + } + } + } + best + } + + /// Every loop enclosing `barrier` strictly below `scope` completes every + /// pass, so the barrier executes on every full pass of `scope`'s body. + /// `Break` in `scope` itself does not disqualify: taking the back edge + /// means the full body executed, and after an exit the loop's tiles are + /// touched no more. + pub(crate) fn guaranteed_below(&self, barrier: &BarrierInfo, scope: u32) -> bool { + match barrier + .enclosing_loops + .iter() + .position(|&index| index == scope) + { + None => false, + Some(position) => barrier.enclosing_loops[position + 1..] + .iter() + .all(|&index| self.loops[index as usize].guaranteed_once()), + } + } + + fn compute_scoped_phases(&mut self) { + let mut scoped: Vec<(*const (), Option<(u32, LiveRange)>)> = Vec::new(); + for &key in &self.order { + let tile = &self.tiles[&key]; + let first = tile.accesses.iter().map(|access| access.position).min(); + let last = tile.accesses.iter().map(|access| access.position).max(); + let (Some(first), Some(last)) = (first, last) else { + scoped.push((key, None)); + continue; + }; + let Some(home) = self.innermost_common_loop(first, last) else { + scoped.push((key, None)); + continue; + }; + // Expand the phase over loops nested inside the home loop, to + // fixpoint: a touch inside a nested loop recurs every nested + // iteration. + let home_span = self.loops[home as usize].span; + let mut phase = LiveRange { first, last }; + loop { + let mut changed = false; + for info in &self.loops { + let span = info.span; + let nested = span.first > home_span.first && span.last < home_span.last; + let intersects = phase.first < span.last && phase.last > span.first; + if nested + && intersects + && (phase.first > span.first || phase.last < span.last) + { + phase.first = phase.first.min(span.first); + phase.last = phase.last.max(span.last); + changed = true; + } + } + if !changed { + break; + } + } + scoped.push((key, Some((home, phase)))); + } + for (key, value) in scoped { + self.tiles.get_mut(&key).expect("walk-recorded tile").scoped = value; + } + } + + /// A barrier inside loop `scope` at a position satisfying `in_interval`, + /// executing on every full pass of the body. + fn scoped_barrier(&self, scope: u32, in_interval: impl Fn(u32) -> bool) -> bool { + let span = self.loops[scope as usize].span; + self.barriers.iter().any(|barrier| { + span.first < barrier.position + && barrier.position < span.last + && in_interval(barrier.position) + && self.guaranteed_below(barrier, scope) + }) + } + + /// Whether `later` may reuse memory whose previous occupant is + /// `earlier`, considering both the plain interval arm and the loop + /// phase arm. + pub(crate) fn can_follow_tiles(&self, earlier: &TileLiveness, later: &TileLiveness) -> bool { + if self.can_follow(earlier.range, later.range) { + return true; + } + // Phase arm: both tiles live only inside one common loop, with + // disjoint per-iteration phases, a barrier between the phases, and + // a barrier covering the wrap back to the earlier phase. + let (Some((home_a, phase_a)), Some((home_b, phase_b))) = (earlier.scoped, later.scoped) + else { + return false; + }; + if home_a != home_b { + return false; + } + let (first, second) = if phase_a.first <= phase_b.first { + (phase_a, phase_b) + } else { + (phase_b, phase_a) + }; + first.last < second.first + && self.scoped_barrier(home_a, |p| p > first.last && p <= second.first) + && self.scoped_barrier(home_a, |p| p > second.last || p <= first.first) + } + + /// A guaranteed uniform barrier strictly after `after` and at or before + /// `at`. Barriers inside loops that may break, return, or run zero + /// iterations are skippable at runtime and never separate. + pub(crate) fn separating_barrier(&self, after: u32, at: u32) -> bool { + self.barriers + .iter() + .any(|barrier| barrier.guaranteed && barrier.position > after && barrier.position <= at) + } + + /// Whether `later` may reuse memory whose previous occupant was + /// `earlier`: disjoint expanded ranges with a uniform barrier ordering + /// every thread's last touch of `earlier` before any first touch of + /// `later`. + pub(crate) fn can_follow(&self, earlier: LiveRange, later: LiveRange) -> bool { + earlier.last < later.first && self.separating_barrier(earlier.last, later.first) + } +} + +#[derive(Default)] +struct Walk { + position: u32, + tiles: FxHashMap<*const (), TileLiveness>, + order: Vec<*const ()>, + barriers: Vec, + loops: Vec, + /// Open loop frames as indices into `loops`. + loop_stack: Vec, + /// Statement indices from the body root, descending through `Loop` + /// bodies only. + path: Vec, + /// Kind attributed to the next `touch` (writes are statement-level, so + /// the statement arm sets this before visiting). + access_kind: AccessKind, + /// `If` nesting depth: barriers below a conditional are not uniform. + conditional_depth: u32, +} + +impl Default for AccessKind { + fn default() -> Self { + Self::Read + } +} + +impl Walk { + fn touch(&mut self, tile: &Tile, coop: bool) { + if tile.layout.memory_level() != MemoryLevel::Workgroup { + return; + } + let key = std::rc::Rc::as_ptr(tile) as *const (); + let position = self.position; + if !self.tiles.contains_key(&key) { + self.order.push(key); + self.tiles.insert( + key, + TileLiveness { + range: LiveRange { + first: position, + last: position, + }, + element: tile.element, + elements: tile.layout.allocation_element_count().get(), + accesses: Vec::new(), + scoped: None, + coop: false, + }, + ); + } + let liveness = self.tiles.get_mut(&key).expect("inserted above"); + liveness.range.last = position; + liveness.coop |= coop; + liveness.accesses.push(TileAccess { + position, + kind: self.access_kind, + }); + } + + fn visit_expr(&mut self, expr: &Expr) { + let mut touched: Vec<(Tile, TileUse)> = Vec::new(); + expr.kind().for_each_tile(&mut |tile, tile_use| { + touched.push((tile.clone(), tile_use)); + }); + for (tile, tile_use) in &touched { + self.access_kind = match tile_use { + TileUse::Read | TileUse::CoopRead => AccessKind::Read, + TileUse::ReadWrite => AccessKind::ReadWrite, + }; + self.touch(tile, matches!(tile_use, TileUse::CoopRead)); + } + self.access_kind = AccessKind::Read; + expr.kind().for_each_child(&mut |child| self.visit_expr(child)); + } + + fn visit_stmts(&mut self, stmts: &[Stmt]) { + for (index, stmt) in stmts.iter().enumerate() { + self.position += 1; + match stmt { + Stmt::Store { + addr, value, mask, .. + } => { + match addr { + crate::ir::Addr::Rc2 { row, col } => { + self.visit_expr(row); + self.visit_expr(col); + } + crate::ir::Addr::Linear(index) => self.visit_expr(index), + } + self.visit_expr(value); + self.visit_expr(mask); + } + Stmt::StoreLocal { value, .. } => self.visit_expr(value), + Stmt::StoreTile { dst, index, value } => { + self.access_kind = AccessKind::Write; + self.touch(dst, false); + self.access_kind = AccessKind::Read; + self.visit_expr(index); + self.visit_expr(value); + } + Stmt::FillTile { dst, value, bounds } => { + self.access_kind = AccessKind::Write; + self.touch(dst, false); + self.access_kind = AccessKind::Read; + self.visit_expr(value); + for bound in bounds.iter().flatten() { + self.visit_expr(bound); + } + } + Stmt::CoopStore { addr, .. } => match addr { + crate::ir::Addr::Rc2 { row, col } => { + self.visit_expr(row); + self.visit_expr(col); + } + crate::ir::Addr::Linear(index) => self.visit_expr(index), + }, + Stmt::CoopStoreTile { tile, row, col, .. } => { + self.access_kind = AccessKind::Write; + self.touch(tile, true); + self.access_kind = AccessKind::Read; + self.visit_expr(row); + self.visit_expr(col); + } + Stmt::If { + condition, + accept, + reject, + } => { + self.visit_expr(condition); + self.conditional_depth += 1; + self.visit_stmts(accept); + self.visit_stmts(reject); + self.conditional_depth -= 1; + } + Stmt::Loop { + count, + accumulators, + body, + .. + } => { + // Count and accumulator inits run once, before the loop: + // header position, outside the span. + if let Some(count) = count { + self.visit_expr(count); + } + for Accumulator { init, .. } in accumulators { + self.visit_expr(init); + } + let loop_index = self.loops.len() as u32; + self.loops.push(LoopInfo { + span: LiveRange { + first: self.position, + last: self.position, + }, + has_break: false, + has_return: false, + static_count: count.as_ref().and_then(literal_u32), + }); + self.loop_stack.push(loop_index); + if self.conditional_depth == 0 { + self.path.push(index as u32); + } + self.visit_stmts(body); + if self.conditional_depth == 0 { + self.path.pop(); + } + // Accumulator updates run at the end of EVERY iteration — + // after any in-loop barrier — so their tile touches are + // attributed inside the span and expand over the loop. + if !accumulators.is_empty() { + self.position += 1; + for Accumulator { update, .. } in accumulators { + self.visit_expr(update); + } + } + self.position += 1; + self.loop_stack.pop().expect("loop frame pushed above"); + self.loops[loop_index as usize].span.last = self.position; + } + Stmt::Break => { + if let Some(&frame) = self.loop_stack.last() { + self.loops[frame as usize].has_break = true; + } + } + Stmt::Return => { + for &frame in &self.loop_stack { + self.loops[frame as usize].has_return = true; + } + } + Stmt::Barrier => { + if self.conditional_depth == 0 { + let mut path = self.path.clone(); + path.push(index as u32); + self.barriers.push(BarrierInfo { + position: self.position, + path, + enclosing_loops: self.loop_stack.clone(), + // Finalized after the walk, once every enclosing + // loop's break/return/count facts are complete. + guaranteed: false, + }); + } + } + Stmt::StorageBarrier => {} + } + } + } + + /// Expand every tile's range to cover each loop body it intersects, + /// to fixpoint across nesting. A touch inside a loop recurs every + /// iteration, so for hazard purposes the tile is live across the whole + /// body — including the back edge. After expansion, two tiles sharing a + /// loop always overlap (never merge), and plain interval disjointness + /// plus one forward barrier is sound; `Break` only shortens executions + /// of ranges the expansion already covers. + fn expand_ranges_over_loops(&mut self) { + loop { + let mut changed = false; + for liveness in self.tiles.values_mut() { + let range = &mut liveness.range; + for info in &self.loops { + let span = info.span; + let intersects = range.first < span.last && range.last > span.first; + if intersects && (range.first > span.first || range.last < span.last) { + range.first = range.first.min(span.first); + range.last = range.last.max(span.last); + changed = true; + } + } + } + if !changed { + break; + } + } + } +} + +fn literal_u32(expr: &Expr) -> Option { + use crate::ir::{ExprKind, TileLiteral}; + match expr.kind() { + ExprKind::Literal(TileLiteral::U32(value)) => Some(*value), + _ => None, + } +} diff --git a/fusor-ml/tile-ir/src/analysis/verify.rs b/fusor-ml/tile-ir/src/analysis/verify.rs new file mode 100644 index 000000000..704587fa2 --- /dev/null +++ b/fusor-ml/tile-ir/src/analysis/verify.rs @@ -0,0 +1,195 @@ +//! Barrier-hazard verification of a tile-arena placement. +//! +//! Independent recheck of every pair of tiles whose bytes overlap: the +//! packer's first-fit consults only each allocation's most recent occupant +//! (sound by barrier transitivity), so the verifier deliberately re-derives +//! legality for ALL pairs from the liveness facts alone. Any two tiles +//! sharing bytes must have disjoint expanded live ranges with a guaranteed +//! uniform barrier between them — otherwise a fast thread can touch one +//! tile while a slow thread still touches the other, through the same +//! memory. A failure is a lowering-time error, never a runtime NaN hunt. + +use crate::lower::arena::{ArenaMode, Placement, TileArena}; + +use super::LivenessInfo; + +pub(crate) fn verify_arena(info: &LivenessInfo, arena: &TileArena) -> Result<(), String> { + // (byte interval, liveness) per placed tile; regions get disjoint + // synthetic base offsets so interval overlap means "same bytes" in both + // modes. + let mut placed: Vec<(u64, u64, *const ())> = Vec::with_capacity(info.order.len()); + for &key in &info.order { + let tile = &info.tiles[&key]; + let stride = tile + .element + .workgroup_array_stride() + .map(u64::from) + .unwrap_or_else(|| tile.element.byte_size()); + let extent = u64::from(tile.elements) * stride; + let base = match arena.assignment.get(&key) { + Some(Placement::Region { index }) => { + debug_assert!(matches!(arena.mode, ArenaMode::Regions)); + // Regions cannot overlap each other; give each a base far + // past any real allocation. + (*index as u64) << 40 + } + Some(Placement::Arena { byte_offset }) => u64::from(*byte_offset), + None => continue, + }; + placed.push((base, base + extent, key)); + } + + verify_overlaps(info, &placed) +} + +fn verify_overlaps( + info: &LivenessInfo, + placed: &[(u64, u64, *const ())], +) -> Result<(), String> { + for (index, &(a_start, a_end, a_key)) in placed.iter().enumerate() { + for &(b_start, b_end, b_key) in &placed[index + 1..] { + if a_start >= b_end || b_start >= a_end { + continue; + } + let a = &info.tiles[&a_key]; + let b = &info.tiles[&b_key]; + let ordered = + info.can_follow_tiles(a, b) || info.can_follow_tiles(b, a); + if !ordered { + return Err(format!( + "tiles share bytes without a guaranteed separating barrier: \ + {:?} x{} live ({},{}) at [{},{}) overlaps {:?} x{} live ({},{}) at [{},{})", + a.element, + a.elements, + a.range.first, + a.range.last, + a_start, + a_end, + b.element, + b.elements, + b.range.first, + b.range.last, + b_start, + b_end, + )); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use rustc_hash::FxHashMap; + + use super::super::{BarrierInfo, LiveRange, LivenessInfo, TileLiveness}; + use super::*; + use crate::ElementType; + + fn info_with( + ranges: &[(u32, u32)], + barriers: &[(u32, bool)], + ) -> (LivenessInfo, Vec<*const ()>) { + let keys: Vec<*const ()> = (1..=ranges.len()).map(|i| i as *const ()).collect(); + let mut tiles = FxHashMap::default(); + for (key, &(first, last)) in keys.iter().zip(ranges) { + tiles.insert( + *key, + TileLiveness { + range: LiveRange { first, last }, + element: ElementType::F32, + elements: 16, + accesses: Vec::new(), + scoped: None, + coop: false, + }, + ); + } + let info = LivenessInfo { + tiles, + order: keys.clone(), + barriers: barriers + .iter() + .map(|&(position, guaranteed)| BarrierInfo { + position, + path: Vec::new(), + enclosing_loops: Vec::new(), + guaranteed, + }) + .collect(), + loops: Vec::new(), + }; + (info, keys) + } + + fn same_region(keys: &[*const ()]) -> TileArena { + let mut assignment = FxHashMap::default(); + for &key in keys { + assignment.insert(key, Placement::Region { index: 0 }); + } + TileArena { + mode: ArenaMode::Regions, + regions: vec![crate::lower::arena::Region { + canonical: ElementType::F32, + elements: 16, + }], + arena_bytes: 0, + assignment, + } + } + + #[test] + fn rejects_overlapping_ranges_in_one_region() { + let (info, keys) = info_with(&[(1, 5), (3, 8)], &[(2, true)]); + assert!(verify_arena(&info, &same_region(&keys)).is_err()); + } + + #[test] + fn rejects_disjoint_ranges_without_barrier() { + let (info, keys) = info_with(&[(1, 3), (5, 8)], &[]); + assert!(verify_arena(&info, &same_region(&keys)).is_err()); + } + + #[test] + fn rejects_disjoint_ranges_with_only_poisoned_barrier() { + let (info, keys) = info_with(&[(1, 3), (5, 8)], &[(4, false)]); + assert!(verify_arena(&info, &same_region(&keys)).is_err()); + } + + #[test] + fn accepts_barrier_separated_chain() { + let (info, keys) = info_with(&[(1, 3), (5, 8), (10, 12)], &[(4, true), (9, true)]); + assert!(verify_arena(&info, &same_region(&keys)).is_ok()); + } + + #[test] + fn byte_arena_overlap_requires_barrier() { + let (info, keys) = info_with(&[(1, 3), (4, 8)], &[]); + let mut assignment = FxHashMap::default(); + // 16 f32 elements = 64 bytes each; offsets 0 and 32 overlap. + assignment.insert(keys[0], Placement::Arena { byte_offset: 0 }); + assignment.insert(keys[1], Placement::Arena { byte_offset: 32 }); + let arena = TileArena { + mode: ArenaMode::ByteArena, + regions: Vec::new(), + arena_bytes: 96, + assignment, + }; + assert!(verify_arena(&info, &arena).is_err()); + } + + #[test] + fn byte_arena_disjoint_intervals_need_no_barrier() { + let (info, keys) = info_with(&[(1, 3), (2, 8)], &[]); + let mut assignment = FxHashMap::default(); + assignment.insert(keys[0], Placement::Arena { byte_offset: 0 }); + assignment.insert(keys[1], Placement::Arena { byte_offset: 64 }); + let arena = TileArena { + mode: ArenaMode::ByteArena, + regions: Vec::new(), + arena_bytes: 128, + assignment, + }; + assert!(verify_arena(&info, &arena).is_ok()); + } +} diff --git a/fusor-ml/tile-ir/src/ir/element.rs b/fusor-ml/tile-ir/src/ir/element.rs index f436a5778..9e6239997 100644 --- a/fusor-ml/tile-ir/src/ir/element.rs +++ b/fusor-ml/tile-ir/src/ir/element.rs @@ -20,6 +20,14 @@ pub enum ScalarElement { } impl ScalarElement { + /// Byte size of one scalar as stored in memory. + pub const fn byte_size(self) -> u64 { + match self { + Self::F32 | Self::U32 | Self::Bool => 4, + Self::F16 => 2, + } + } + /// Element type for this scalar by itself. pub const fn element(self) -> ElementType { match self { @@ -98,6 +106,43 @@ impl ElementType { } } + /// Byte size of one element as allocated in an array of this type. + /// Cooperative fragments live in registers, not addressable arrays, and + /// report the size of their scalar so footprint sums stay conservative. + pub const fn byte_size(self) -> u64 { + match self { + Self::F32 | Self::U32 | Self::Bool => 4, + Self::F16 => 2, + Self::Vector { scalar, lanes } => scalar.byte_size() * lanes as u64, + Self::CoopMatrix { scalar, .. } => scalar.byte_size(), + } + } + + /// Array stride of one element in a workgroup array, or `None` for + /// elements that cannot back one (bool, cooperative fragments). The + /// single source of stride truth: allocation packing and Naga array + /// emission both read this, so they can never disagree. Differs from + /// [`Self::byte_size`] for vec3, which pads to the vec4 stride. + pub const fn workgroup_array_stride(self) -> Option { + match self { + Self::F32 | Self::U32 => Some(4), + Self::F16 => Some(2), + Self::Vector { scalar, lanes } => { + let size = match scalar { + ScalarElement::F32 | ScalarElement::U32 => 4, + ScalarElement::F16 => 2, + ScalarElement::Bool => return None, + }; + match lanes { + 2 => Some(2 * size), + 3 | 4 => Some(4 * size), + _ => None, + } + } + Self::Bool | Self::CoopMatrix { .. } => None, + } + } + /// Returns true when this element stores or computes with f16 data. pub const fn uses_f16(self) -> bool { matches!( diff --git a/fusor-ml/tile-ir/src/ir/expr.rs b/fusor-ml/tile-ir/src/ir/expr.rs index 9e73f9d3b..cb0c6788a 100644 --- a/fusor-ml/tile-ir/src/ir/expr.rs +++ b/fusor-ml/tile-ir/src/ir/expr.rs @@ -74,6 +74,9 @@ pub enum CoopSrc { row: Box, /// Column coordinate of the fragment origin. col: Box, + /// Load the transpose of the addressed region: fragment `(i, j)` + /// reads tile element `(col + j, row + i)`. + transposed: bool, }, /// A rank-1 storage vector broadcast across all fragment rows. BroadcastCol { @@ -278,6 +281,13 @@ pub enum ExprKind { /// Number of f32 lanes produced. lanes: u32, }, + /// Extract one component of a vector value (naga `AccessIndex`). + VecComponent { + /// The vector expression. + vector: Box, + /// Component index (< the vector's lane count). + component: u32, + }, /// Project lane `lane` out of a `Dequantize` (usually wrapped in `Shared`). LaneOf { /// The block being projected (a `Dequantize`, typically `Shared`). @@ -439,11 +449,17 @@ fn hash_addr(addr: &Addr, h: &mut FxHasher) { fn hash_coop_src(src: &CoopSrc, h: &mut FxHasher) { match src { - CoopSrc::TileRegion { tile, row, col } => { + CoopSrc::TileRegion { + tile, + row, + col, + transposed, + } => { h.write_u8(0); hash_ptr(tile, h); hash_expr(row, h); hash_expr(col, h); + h.write_u8(*transposed as u8); } CoopSrc::BroadcastCol { src, col } => { h.write_u8(1); @@ -483,6 +499,10 @@ fn hash_kind_into(kind: &ExprKind, h: &mut FxHasher) { std::mem::discriminant(kind).hash(h); match kind { ExprKind::Literal(lit) => lit.hash(h), + ExprKind::VecComponent { vector, component } => { + hash_kind_into(vector.kind(), h); + component.hash(h); + } ExprKind::Builtin(builtin) => builtin.hash(h), ExprKind::LoadLocal(local) => hash_ptr(local, h), ExprKind::Load { @@ -642,13 +662,15 @@ fn coop_src_eq(a: &CoopSrc, b: &CoopSrc) -> bool { tile: tx, row: rx, col: cx, + transposed: px, }, CoopSrc::TileRegion { tile: ty, row: ry, col: cy, + transposed: py, }, - ) => Rc::ptr_eq(tx, ty) && expr_eq(rx, ry) && expr_eq(cx, cy), + ) => Rc::ptr_eq(tx, ty) && expr_eq(rx, ry) && expr_eq(cx, cy) && px == py, ( CoopSrc::BroadcastCol { src: sx, col: cx }, CoopSrc::BroadcastCol { src: sy, col: cy }, @@ -891,3 +913,131 @@ fn kind_eq(a: &ExprKind, b: &ExprKind) -> bool { _ => false, } } + +impl ExprKind { + /// Visit every direct child expression of this node, exhaustively: new + /// variants must extend this or fail to compile, so structural passes + /// (liveness, footprint) cannot silently skip children. + pub(crate) fn for_each_child(&self, f: &mut dyn FnMut(&Expr)) { + let addr = |addr: &Addr, f: &mut dyn FnMut(&Expr)| match addr { + Addr::Rc2 { row, col } => { + f(row); + f(col); + } + Addr::Linear(index) => f(index), + }; + match self { + Self::Literal(_) | Self::Builtin(_) | Self::LoadLocal(_) => {} + Self::Load { + addr: a, + mask, + fill, + .. + } => { + addr(a, f); + f(mask); + f(fill); + } + Self::LoadTile { index, .. } => f(index), + Self::Unary { value, .. } => f(value), + Self::Binary { left, right, .. } | Self::Compare { left, right, .. } => { + f(left); + f(right); + } + Self::Cast { value, .. } | Self::Bitcast { value, .. } => f(value), + Self::Select { + condition, + accept, + reject, + } => { + f(condition); + f(accept); + f(reject); + } + Self::Vec { parts, .. } => { + for part in parts { + f(part); + } + } + Self::Dot { left, right } => { + f(left); + f(right); + } + Self::Reduce { value, .. } => f(value), + Self::CoopLoad { src, .. } => match src { + CoopSrc::TileRegion { row, col, .. } => { + f(row); + f(col); + } + CoopSrc::BroadcastCol { col, .. } => f(col), + }, + Self::CoopMma { a, b, c } => { + f(a); + f(b); + f(c); + } + Self::Dequantize { + k_base, + col, + mask, + fill, + .. + } => { + f(k_base); + f(col); + f(mask); + f(fill); + } + Self::VecComponent { vector, .. } => f(vector), + Self::LaneOf { block, .. } => f(block), + Self::QuantizedDot { + activations, + k_base, + col, + mask, + fill, + .. + } => { + for activation in activations { + f(activation); + } + f(k_base); + f(col); + f(mask); + f(fill); + } + Self::Shared(inner) => f(inner), + } + } + + /// Visit every workgroup/private tile this node references directly + /// (not through children): tile loads, cooperative tile-region loads, + /// and reduction scratch tiles. Loads are plain reads; reduction scratch + /// is a collective read-modify-write. + pub(crate) fn for_each_tile(&self, f: &mut dyn FnMut(&Tile, TileUse)) { + match self { + Self::LoadTile { tile, .. } => f(tile, TileUse::Read), + Self::CoopLoad { + src: CoopSrc::TileRegion { tile, .. }, + .. + } => f(tile, TileUse::CoopRead), + Self::Reduce { kind, .. } => match kind { + ReduceKind::Subgroup => {} + ReduceKind::Workgroup { scratch, .. } => f(scratch, TileUse::ReadWrite), + ReduceKind::Loop { scratch, .. } => f(scratch, TileUse::ReadWrite), + }, + _ => {} + } + } +} + +/// How an expression node uses a tile it references directly. +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) enum TileUse { + Read, + /// Read through a raw cooperative-matrix pointer: the tile's emitted + /// array type must equal its element type. + CoopRead, + /// Collective read-modify-write (reduction scratch). + ReadWrite, +} diff --git a/fusor-ml/tile-ir/src/ir/mod.rs b/fusor-ml/tile-ir/src/ir/mod.rs index 2d0157cee..81cae57ff 100644 --- a/fusor-ml/tile-ir/src/ir/mod.rs +++ b/fusor-ml/tile-ir/src/ir/mod.rs @@ -14,6 +14,7 @@ pub use storage::{ mod expr; pub use expr::{Addr, Builtin, CoopSrc, Expr, ExprKind, Node, QuantActivation, ReduceKind, Source}; +pub(crate) use expr::TileUse; mod program; pub use program::{Accumulator, KernelIr, Stmt}; diff --git a/fusor-ml/tile-ir/src/ir/program.rs b/fusor-ml/tile-ir/src/ir/program.rs index 3cea67412..1836e0c87 100644 --- a/fusor-ml/tile-ir/src/ir/program.rs +++ b/fusor-ml/tile-ir/src/ir/program.rs @@ -17,6 +17,10 @@ pub struct KernelIr { pub block: u32, /// Program statements. pub(crate) body: Vec, + /// The backend supports the workgroup-alias extension (proved by a + /// device token at build time): mixed-stride tiles may pack into one + /// byte arena instead of per-type regions. + pub(crate) byte_arena: bool, } impl Default for KernelIr { @@ -26,6 +30,7 @@ impl Default for KernelIr { grid: [1, 1, 1], block: 0, body: Vec::new(), + byte_arena: false, } } } @@ -35,6 +40,27 @@ impl KernelIr { pub fn lower_to_naga(&self) -> Result { crate::lower::lower_to_naga(self) } + + /// Total workgroup (threadgroup) memory this kernel allocates, in bytes. + /// + /// The residency input for occupancy-aware selection: Apple Silicon keeps + /// two workgroups per core resident only while a kernel's footprint stays + /// at or below half the 32 KB threadgroup budget. + pub fn workgroup_bytes(&self) -> u64 { + crate::lower::workgroup_bytes(self) + } + + /// Top-level barrier insertions that would shrink the workgroup + /// footprint, best first. The caller owns the policy: apply one only + /// when the saving buys an occupancy class worth the barrier's cost. + pub fn barrier_suggestions(&self) -> Vec { + crate::analysis::barrier_suggestions(self) + } + + /// Insert the suggested barrier. + pub fn apply_barrier_suggestion(&mut self, suggestion: &crate::BarrierSuggestion) { + crate::analysis::apply_barrier_suggestion(self, suggestion); + } } /// One accumulator carried by a counted `Stmt::Loop`. @@ -104,6 +130,19 @@ pub enum Stmt { /// Destination address. addr: Addr, }, + /// Cooperatively store an accumulator into a workgroup tile at + /// `(row, col)` — the staging step between fragment math and per-lane + /// passes over the same values (attention's softmax over staged scores). + CoopStoreTile { + /// Accumulator local. + acc: Local, + /// Destination workgroup tile. + tile: Tile, + /// Row coordinate of the fragment origin. + row: Box, + /// Column coordinate of the fragment origin. + col: Box, + }, /// Per-invocation control flow. If { /// Bool condition. @@ -130,6 +169,8 @@ pub enum Stmt { Break, /// Return from the kernel entry point. Return, - /// Workgroup-scope memory barrier. + /// Workgroup-memory barrier with workgroup execution scope. Barrier, + /// Storage-memory barrier with workgroup execution scope. + StorageBarrier, } diff --git a/fusor-ml/tile-ir/src/lib.rs b/fusor-ml/tile-ir/src/lib.rs index 980971361..008065376 100644 --- a/fusor-ml/tile-ir/src/lib.rs +++ b/fusor-ml/tile-ir/src/lib.rs @@ -25,6 +25,7 @@ //! # Ok::<(), Box>(()) //! ``` +mod analysis; mod ir; mod kernel_builder; mod lower; @@ -40,10 +41,11 @@ pub use ir::{ // LocalDecl, Node, ReduceKind, Source, Stmt, Tile, TileDecl) is intentionally not // re-exported: consumers only build via `tile`/`KernelBuilder` and lower to an // opaque `NagaKernel`. Internal code names these through `crate::ir::*`. +pub use analysis::{set_liveness_trace, BarrierSuggestion}; pub use kernel_builder::{KernelBuilder, KernelTensorRef}; pub use lower::{LowerError, NagaKernel}; pub use quantized::{GgmlQuantFormat, QuantizedMatrix}; -pub use tile::{CoopMatrixToken, SubgroupToken}; +pub use tile::{ByteArenaToken, CoopMatrixToken, SubgroupToken}; #[cfg(test)] mod tests; diff --git a/fusor-ml/tile-ir/src/lower/analysis.rs b/fusor-ml/tile-ir/src/lower/analysis.rs index 96359c75f..7b980c044 100644 --- a/fusor-ml/tile-ir/src/lower/analysis.rs +++ b/fusor-ml/tile-ir/src/lower/analysis.rs @@ -130,6 +130,18 @@ impl Analysis { self.note_view(dst); self.visit_addr(addr); } + Stmt::CoopStoreTile { + acc, + tile, + row, + col, + } => { + self.caps.uses_coop = true; + self.note_local(acc); + self.note_tile(tile); + self.visit_expr(row); + self.visit_expr(col); + } Stmt::If { condition, accept, @@ -161,7 +173,7 @@ impl Analysis { self.visit_stmt(s); } } - Stmt::Break | Stmt::Return | Stmt::Barrier => {} + Stmt::Break | Stmt::Return | Stmt::Barrier | Stmt::StorageBarrier => {} } } @@ -184,7 +196,12 @@ impl Analysis { fn visit_coop_src(&mut self, src: &CoopSrc) { match src { - CoopSrc::TileRegion { tile, row, col } => { + CoopSrc::TileRegion { + tile, + row, + col, + transposed: _, + } => { self.note_tile(tile); self.visit_expr(row); self.visit_expr(col); @@ -211,6 +228,7 @@ impl Analysis { self.note_element(expr.element()); match expr.kind() { ExprKind::Literal(_) => {} + ExprKind::VecComponent { vector, .. } => self.visit_expr(vector), ExprKind::Builtin(builtin) => match builtin { Builtin::SubgroupId => self.caps.subgroup_id = true, Builtin::SubgroupLane => self.caps.subgroup_lane = true, diff --git a/fusor-ml/tile-ir/src/lower/arena.rs b/fusor-ml/tile-ir/src/lower/arena.rs new file mode 100644 index 000000000..84c123499 --- /dev/null +++ b/fusor-ml/tile-ir/src/lower/arena.rs @@ -0,0 +1,300 @@ +//! Workgroup-tile allocation packing. +//! +//! Consumes [`crate::analysis::LivenessInfo`] and places each workgroup tile +//! into shared memory. Two modes: +//! +//! - [`ArenaMode::Regions`] (portable): tiles pack into per-region typed +//! arrays, every tile at offset 0. Tiles of the same stride class share a +//! region when their live ranges are barrier-separated; a region holding +//! more than one element type is emitted with a class-neutral canonical +//! type and every access bitcasts the value (never the address — within a +//! stride class, element index `i` names the same bytes for every type). +//! - [`ArenaMode::ByteArena`] (Metal fork): one byte arena, tiles at byte +//! offsets via interval strip-packing, so tiles of *different* strides +//! (f16 staging next to f32 accumulators) can reuse the same bytes. +//! +//! Sharing legality is [`LivenessInfo::can_follow`]: disjoint expanded live +//! ranges plus a guaranteed uniform barrier between them. In the byte arena +//! the "previous occupant" is per byte interval, tracked as a segment list; +//! transitivity of the barrier chain applies pointwise per byte. + +use rustc_hash::FxHashMap; + +use crate::ElementType; +use crate::analysis::{LivenessInfo, trace_enabled}; +use crate::ir::ScalarElement; + +/// How this kernel's workgroup tiles are laid out. +#[derive(Clone, Copy, PartialEq, Eq)] +pub(crate) enum ArenaMode { + /// One naga global per region; every tile at offset 0 in its region. + Regions, + /// One threadgroup byte arena; tiles at byte offsets via the workgroup + /// alias extension. Chosen only when the kernel proves backend support + /// (`KernelIr::byte_arena`) and mixes stride widths. + ByteArena, +} + +/// A stride-compatibility class. `lanes` is part of the key so vec3 (data +/// 12 B, stride 16 B) never mixes with vec4, and value bitcasts stay +/// per-component. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +struct StrideClass { + stride: u32, + lanes: u32, +} + +fn stride_class(element: ElementType) -> Option { + let stride = element.workgroup_array_stride()?; + let lanes = match element { + ElementType::Vector { lanes, .. } => lanes, + _ => 1, + }; + Some(StrideClass { stride, lanes }) +} + +/// Whether tiles of elements `a` and `b` may occupy one typed region: equal +/// types always; otherwise the same stride class with a value-level bitcast +/// between them. Only 32-bit-scalar casts qualify (`f32 <-> u32`, per +/// component for vectors): f16 has no same-width partner, so 2-byte tiles +/// never join a 4-byte region and no sub-word read-modify-write hazard can +/// arise. +fn bitcast_compatible(a: ElementType, b: ElementType) -> bool { + if a == b { + return true; + } + let (Some(class_a), Some(class_b)) = (stride_class(a), stride_class(b)) else { + return false; + }; + class_a == class_b && scalar_of(a).is_some_and(|scalar| scalar.byte_size() == 4) + && scalar_of(b).is_some_and(|scalar| scalar.byte_size() == 4) +} + +fn scalar_of(element: ElementType) -> Option { + match element { + ElementType::F32 => Some(ScalarElement::F32), + ElementType::U32 => Some(ScalarElement::U32), + ElementType::F16 => Some(ScalarElement::F16), + ElementType::Vector { scalar, .. } => Some(scalar), + ElementType::Bool | ElementType::CoopMatrix { .. } => None, + } +} + +/// The canonical emission type for a heterogeneous region of this class: +/// the u32-based type of the same shape. +fn neutral(class: StrideClass) -> ElementType { + if class.lanes == 1 { + ElementType::U32 + } else { + ElementType::vector(ScalarElement::U32, class.lanes) + } +} + +/// One time-shared typed allocation (Regions mode). +pub(crate) struct Region { + /// Emission element type of the backing array: the occupant element + /// while homogeneous (bit-identical emission to unshared lowering), + /// widened to the class-neutral u32 form on the first cross-type join. + pub canonical: ElementType, + /// Array length in canonical elements (stride-equal across the class, + /// so occupant extents compare directly). + pub elements: u32, +} + +#[derive(Clone, Copy)] +pub(crate) enum Placement { + Region { index: usize }, + Arena { byte_offset: u32 }, +} + +/// The computed tile placement for one kernel. +pub(crate) struct TileArena { + pub mode: ArenaMode, + /// Regions in creation order (empty in ByteArena mode). + pub regions: Vec, + /// Packed arena extent, 16-byte aligned (0 in Regions mode). + pub arena_bytes: u32, + /// Tile identity (`Rc` pointer) -> placement. + pub assignment: FxHashMap<*const (), Placement>, +} + +impl TileArena { + pub(crate) fn assign(info: &LivenessInfo, byte_arena: bool) -> Self { + let mut strides = Vec::new(); + for key in &info.order { + let tile = &info.tiles[key]; + match tile.element.workgroup_array_stride() { + Some(stride) if !strides.contains(&stride) => strides.push(stride), + _ => {} + } + } + let mixed_strides = strides.len() > 1; + let all_packable = info + .order + .iter() + .all(|key| stride_class(info.tiles[key].element).is_some()); + let regions = Self::assign_regions(info); + if byte_arena && mixed_strides && all_packable { + // The arena only wins when cross-stride reuse actually fires: + // without it, 16-byte rounding makes it a strict loss, so pick + // by measured footprint. + let packed = Self::assign_byte_arena(info); + if packed.total_bytes() < regions.total_bytes() { + return packed; + } + } + regions + } + + fn assign_regions(info: &LivenessInfo) -> Self { + let mut regions: Vec = Vec::new(); + // Every occupant per region. The plain interval arm would be sound + // checking only the most recent occupant (barrier transitivity), + // but the loop-phase arm does not compose: A->B and B->C phase + // separation does not imply the C->A wrap is covered. Regions hold + // a handful of tiles, so checking all occupants costs nothing. + let mut region_occupants: Vec> = Vec::new(); + // A coop-consumed occupant pins the region's type: widening the + // canonical would retype the raw pointer its CooperativeLoad/Store + // sees. + let mut region_coop: Vec = Vec::new(); + let mut assignment = FxHashMap::default(); + for &key in &info.order { + let tile = &info.tiles[&key]; + let (range, element, elements) = (tile.range, tile.element, tile.elements); + let reused = regions.iter().enumerate().position(|(index, region)| { + let type_ok = region.canonical == element + || (bitcast_compatible(region.canonical, element) + && !tile.coop + && !region_coop[index]); + type_ok + && region_occupants[index] + .iter() + .all(|occupant| info.can_follow_tiles(&info.tiles[occupant], tile)) + }); + let index = match reused { + Some(index) => { + if trace_enabled() { + eprintln!( + "arena-share region={index} element={:?} elems={elements} range=({},{}) occupants={}", + element, + range.first, + range.last, + region_occupants[index].len() + ); + } + let region = &mut regions[index]; + region.elements = region.elements.max(elements); + if region.canonical != element { + region.canonical = neutral( + stride_class(element).expect("bitcast-compatible implies a class"), + ); + } + region_coop[index] |= tile.coop; + region_occupants[index].push(key); + index + } + None => { + regions.push(Region { + canonical: element, + elements, + }); + region_occupants.push(vec![key]); + region_coop.push(tile.coop); + regions.len() - 1 + } + }; + assignment.insert(key, Placement::Region { index }); + } + Self { + mode: ArenaMode::Regions, + regions, + arena_bytes: 0, + assignment, + } + } + + fn assign_byte_arena(info: &LivenessInfo) -> Self { + // Full placement history: with the loop-phase arm, legality against + // only each byte's most recent occupant does not compose, so a + // candidate is checked against EVERY placement its bytes overlap. + struct Placed { + start: u32, + end: u32, + key: *const (), + } + let mut placed: Vec = Vec::new(); + let mut assignment = FxHashMap::default(); + let mut arena_end = 0u32; + for &key in &info.order { + let tile = &info.tiles[&key]; + let stride = tile + .element + .workgroup_array_stride() + .expect("mode selection requires packable elements"); + let extent = tile.elements * stride; + // Stride doubles as alignment: every supported element's Naga + // array stride is a power of two at least as large as its + // alignment (vec3 already padded to the vec4 stride). + let align = stride; + let align_up = |value: u32| value.div_ceil(align) * align; + let mut candidates: Vec = std::iter::once(0) + .chain(placed.iter().map(|entry| align_up(entry.end))) + .collect(); + candidates.sort_unstable(); + candidates.dedup(); + let offset = candidates + .into_iter() + .find(|&offset| { + placed + .iter() + .filter(|entry| entry.start < offset + extent && entry.end > offset) + .all(|entry| info.can_follow_tiles(&info.tiles[&entry.key], tile)) + }) + .expect("the offset past every placement always fits"); + if trace_enabled() { + eprintln!( + "arena-pack offset={offset} bytes={extent} element={:?} range=({},{})", + tile.element, tile.range.first, tile.range.last + ); + } + let end = offset + extent; + placed.push(Placed { + start: offset, + end, + key, + }); + arena_end = arena_end.max(end); + assignment.insert(key, Placement::Arena { + byte_offset: offset, + }); + } + Self { + mode: ArenaMode::ByteArena, + regions: Vec::new(), + arena_bytes: arena_end.div_ceil(16) * 16, + assignment, + } + } + + /// Total post-arena workgroup footprint in bytes. + pub(crate) fn total_bytes(&self) -> u64 { + match self.mode { + ArenaMode::Regions => self + .regions + .iter() + .map(|region| { + let stride = region + .canonical + .workgroup_array_stride() + .map(u64::from) + // Elements that cannot back an array (rejected at + // emission) still count their data size. + .unwrap_or_else(|| region.canonical.byte_size()); + u64::from(region.elements) * stride + }) + .sum(), + ArenaMode::ByteArena => u64::from(self.arena_bytes), + } + } +} diff --git a/fusor-ml/tile-ir/src/lower/coop.rs b/fusor-ml/tile-ir/src/lower/coop.rs index b1414a13c..47acc2922 100644 --- a/fusor-ml/tile-ir/src/lower/coop.rs +++ b/fusor-ml/tile-ir/src/lower/coop.rs @@ -32,10 +32,20 @@ impl<'a> Lowerer<'a> { self.lower_store_local(expressions, body, dst, value) } Stmt::StoreTile { dst, index, value } => { + let source = value.element(); let value = self.lower_expr(expressions, body, value)?; + // Mixed-precision staging: a per-lane f32 value stored into an + // f16 tile (or vice versa) converts on the way in, matching the + // FillTile fast path's cast. + let value = match (source, dst.element) { + (ElementType::F32, ElementType::F16) | (ElementType::F16, ElementType::F32) => { + self.cast_tile_value(expressions, body, value, source, dst.element) + } + _ => value, + }; let index = self.lower_expr(expressions, body, index)?; let pointer = self.tile_dynamic_pointer(expressions, dst, index, body)?; - body.push(Statement::Store { pointer, value }, Span::default()); + self.store_tile_value(expressions, body, dst, pointer, value); Ok(()) } Stmt::FillTile { dst, value, bounds } => { @@ -44,6 +54,12 @@ impl<'a> Lowerer<'a> { Stmt::CoopStore { acc, dst, addr } => { self.lower_store_coop_acc(expressions, body, acc, dst, addr) } + Stmt::CoopStoreTile { + acc, + tile, + row, + col, + } => self.lower_store_coop_tile(expressions, body, acc, tile, row, col), Stmt::If { condition, accept, @@ -111,6 +127,10 @@ impl<'a> Lowerer<'a> { ); Ok(()) } + Stmt::StorageBarrier => { + body.push(Statement::ControlBarrier(Barrier::STORAGE), Span::default()); + Ok(()) + } } } @@ -163,14 +183,28 @@ impl<'a> Lowerer<'a> { CoopMatrixRole::C => CooperativeRole::C, }; match src { - CoopSrc::TileRegion { tile, row, col } => { + CoopSrc::TileRegion { + tile, + row, + col, + transposed, + } => { let layout = self.tile_layout(tile); let stride_u = Self::row_major_tile_stride(layout)?; let _ = scalar; let row_h = self.lower_expr(expressions, body, row)?; let col_h = self.lower_expr(expressions, body, col)?; + // A transposed load addresses the fragment origin in the + // tile's memory coordinates (swapped) and flips the load + // orientation, so fragment (i, j) reads tile (col + j, + // row + i) without a staged transpose. + let (first, second) = if *transposed { + (col_h, row_h) + } else { + (row_h, col_h) + }; let index = - self.tile_matrix_index_inline(expressions, body, row_h, col_h, stride_u); + self.tile_matrix_index_inline(expressions, body, first, second, stride_u); let ptr = self.tile_dynamic_pointer(expressions, tile, index, body)?; let stride = self.u32(expressions, stride_u); Ok(self.emit( @@ -186,8 +220,9 @@ impl<'a> Lowerer<'a> { // Metal's simdgroup matrix orientation makes // row-major A/B fragments multiply as B * A. // Keep Fusor's logical A * B by holding coop - // fragments transposed internally. - row_major: false, + // fragments transposed internally; a transposed + // load flips that orientation. + row_major: *transposed, }, }, )) @@ -450,8 +485,14 @@ impl<'a> Lowerer<'a> { storage_index, &mut accept, )?; - values[i as usize] = - Some(Self::emit_load(expressions, &mut accept, storage_ptr)); + let loaded = Self::emit_load(expressions, &mut accept, storage_ptr); + values[i as usize] = Some(self.cast_tile_value( + expressions, + &mut accept, + loaded, + src.buffer.element, + dst.element, + )); } for i in 0..VEC { let tile_index = self.add_literal_u32_emitted( @@ -462,12 +503,12 @@ impl<'a> Lowerer<'a> { ); let tile_ptr = self.tile_dynamic_pointer(expressions, dst, tile_index, &mut accept)?; - accept.push( - Statement::Store { - pointer: tile_ptr, - value: values[i as usize].expect("loaded above"), - }, - Span::default(), + self.store_tile_value( + expressions, + &mut accept, + dst, + tile_ptr, + values[i as usize].expect("loaded above"), ); } Ok(accept) @@ -543,8 +584,14 @@ impl<'a> Lowerer<'a> { storage_index, &mut accept, )?; - values[i as usize] = - Some(Self::emit_load(expressions, &mut accept, storage_ptr)); + let loaded = Self::emit_load(expressions, &mut accept, storage_ptr); + values[i as usize] = Some(self.cast_tile_value( + expressions, + &mut accept, + loaded, + src.buffer.element, + dst.element, + )); } for i in 0..VEC { let tile_index = self.add_literal_u32_emitted( @@ -555,12 +602,12 @@ impl<'a> Lowerer<'a> { ); let tile_ptr = self.tile_dynamic_pointer(expressions, dst, tile_index, &mut accept)?; - accept.push( - Statement::Store { - pointer: tile_ptr, - value: values[i as usize].expect("loaded above"), - }, - Span::default(), + self.store_tile_value( + expressions, + &mut accept, + dst, + tile_ptr, + values[i as usize].expect("loaded above"), ); } Ok(accept) @@ -632,13 +679,14 @@ impl<'a> Lowerer<'a> { let storage_ptr = self.storage_dynamic_pointer(expressions, src, storage_index, &mut accept)?; let value = Self::emit_load(expressions, &mut accept, storage_ptr); - accept.push( - Statement::Store { - pointer: tile_ptr, - value, - }, - Span::default(), + let value = self.cast_tile_value( + expressions, + &mut accept, + value, + src.buffer.element, + dst.element, ); + self.store_tile_value(expressions, &mut accept, dst, tile_ptr, value); } Some(in_bounds) => { let mut in_bounds_body = Block::new(); @@ -655,13 +703,14 @@ impl<'a> Lowerer<'a> { &mut in_bounds_body, )?; let value = Self::emit_load(expressions, &mut in_bounds_body, storage_ptr); - in_bounds_body.push( - Statement::Store { - pointer: tile_ptr, - value, - }, - Span::default(), + let value = self.cast_tile_value( + expressions, + &mut in_bounds_body, + value, + src.buffer.element, + dst.element, ); + self.store_tile_value(expressions, &mut in_bounds_body, dst, tile_ptr, value); let zero_f32 = self.f32(expressions, 0.0); let mut out_of_bounds_body = Block::new(); @@ -672,13 +721,7 @@ impl<'a> Lowerer<'a> { ElementType::F32, dst.element, ); - out_of_bounds_body.push( - Statement::Store { - pointer: tile_ptr, - value: zero, - }, - Span::default(), - ); + self.store_tile_value(expressions, &mut out_of_bounds_body, dst, tile_ptr, zero); accept.push( Statement::If { condition: in_bounds, @@ -839,13 +882,7 @@ impl<'a> Lowerer<'a> { ElementType::F32, dst.element, ); - in_bounds_body.push( - Statement::Store { - pointer: ptr, - value, - }, - Span::default(), - ); + self.store_tile_value(expressions, &mut in_bounds_body, dst, ptr, value); } let zero_f32 = self.f32(expressions, 0.0); @@ -858,13 +895,7 @@ impl<'a> Lowerer<'a> { ); let mut out_of_bounds_body = Block::new(); for ptr in tile_ptrs { - out_of_bounds_body.push( - Statement::Store { - pointer: ptr, - value: zero, - }, - Span::default(), - ); + self.store_tile_value(expressions, &mut out_of_bounds_body, dst, ptr, zero); } accept.push( Statement::If { @@ -939,13 +970,7 @@ impl<'a> Lowerer<'a> { ElementType::F32, dst.element, ); - in_bounds_body.push( - Statement::Store { - pointer: tile_ptr, - value, - }, - Span::default(), - ); + self.store_tile_value(expressions, &mut in_bounds_body, dst, tile_ptr, value); let zero_f32 = self.f32(expressions, 0.0); let zero = self.cast_tile_value( expressions, @@ -955,13 +980,7 @@ impl<'a> Lowerer<'a> { dst.element, ); let mut out_of_bounds_body = Block::new(); - out_of_bounds_body.push( - Statement::Store { - pointer: tile_ptr, - value: zero, - }, - Span::default(), - ); + self.store_tile_value(expressions, &mut out_of_bounds_body, dst, tile_ptr, zero); accept.push( Statement::If { condition: in_bounds, @@ -974,6 +993,46 @@ impl<'a> Lowerer<'a> { }) } + /// Lower `Stmt::CoopStoreTile` to a `Statement::CooperativeStore` whose + /// destination is a workgroup tile — the staging step that lets per-lane + /// passes read fragment results. Never routed through per-lane stores. + fn lower_store_coop_tile( + &self, + expressions: &mut Arena, + body: &mut Block, + acc: &Local, + tile: &Tile, + row: &Expr, + col: &Expr, + ) -> Result<(), LowerError> { + // Flush any pending acc SSA so the Load below sees the current value. + self.flush_coop_acc_cache(expressions, body); + let acc_local = self.private_local(acc)?; + let stride_u = Self::row_major_tile_stride(self.tile_layout(tile))?; + let row_h = self.lower_expr(expressions, body, row)?; + let col_h = self.lower_expr(expressions, body, col)?; + let index = self.tile_matrix_index_inline(expressions, body, row_h, col_h, stride_u); + let ptr = self.tile_dynamic_pointer(expressions, tile, index, body)?; + let stride = self.u32(expressions, stride_u); + let acc_ptr = self.local_var(expressions, acc_local); + let acc_value = Self::emit_load(expressions, body, acc_ptr); + body.push( + Statement::CooperativeStore { + target: acc_value, + data: CooperativeData { + pointer: ptr, + stride, + // Accumulators are transposed internally; workgroup tiles + // are row-major, so the inverted flag is `false` exactly + // as in the row-major global-store case. + row_major: false, + }, + }, + Span::default(), + ); + Ok(()) + } + /// Lower `Stmt::CoopStore` to a `Statement::CooperativeStore`. Never routed /// through the per-lane `Store` path. fn lower_store_coop_acc( diff --git a/fusor-ml/tile-ir/src/lower/mod.rs b/fusor-ml/tile-ir/src/lower/mod.rs index ebde6fd8e..ac3ad866e 100644 --- a/fusor-ml/tile-ir/src/lower/mod.rs +++ b/fusor-ml/tile-ir/src/lower/mod.rs @@ -23,8 +23,18 @@ const WORKGROUP_ID_ARG: u32 = 1; const DEFAULT_WORKGROUP_INVOCATIONS: u32 = 256; const DEFAULT_WORKGROUP_SIZE: [u32; 3] = [16, 16, 1]; +/// Workgroup memory the lowered kernel allocates, in bytes — after +/// liveness-based tile sharing, i.e. what the GPU actually reserves. +pub(crate) fn workgroup_bytes(ir: &crate::KernelIr) -> u64 { + let info = crate::analysis::LivenessInfo::compute(ir); + arena::TileArena::assign(&info, ir.byte_arena).total_bytes() +} + pub(crate) fn lower_to_naga(ir: &KernelIr) -> Result { - Lowerer::new(ir)?.lower() + let info = crate::analysis::LivenessInfo::compute(ir); + let tile_arena = arena::TileArena::assign(&info, ir.byte_arena); + crate::analysis::verify_arena(&info, &tile_arena).map_err(LowerError::BarrierHazard)?; + Lowerer::new(ir, tile_arena)?.lower() } /// A validated Naga lowering result. @@ -36,6 +46,30 @@ pub struct NagaKernel { } impl NagaKernel { + /// Rebuild a kernel from a deserialized module, revalidating it with + /// every capability the lowerer can emit (a superset of the original + /// validation is safe: capabilities only gate what a module may use). + pub fn from_module(module: Module, subgroups: bool) -> Result { + let capabilities = naga::valid::Capabilities::SHADER_FLOAT16 + | naga::valid::Capabilities::SHADER_FLOAT16_IN_FLOAT32 + | naga::valid::Capabilities::SUBGROUP + | naga::valid::Capabilities::COOPERATIVE_MATRIX + | naga::valid::Capabilities::WORKGROUP_MEMORY_ALIAS; + let info = naga::valid::Validator::new(naga::valid::ValidationFlags::all(), capabilities) + .validate(&module) + .map_err(|error| LowerError::Validation(format!("{error:#?}")))?; + Ok(Self { + module, + info, + wgsl_extensions: WgslExtensions::new(subgroups), + }) + } + + /// Whether this kernel requires the subgroup WGSL extension. + pub fn subgroups(&self) -> bool { + self.wgsl_extensions.subgroups + } + /// The generated Naga module. pub fn module(&self) -> &Module { &self.module @@ -80,6 +114,8 @@ pub enum LowerError { UnsupportedMemoryLevel(MemoryLevel), /// The typed IR operation is outside the supported lowering subset. UnsupportedOperation(&'static str), + /// Tiles share workgroup bytes without a guaranteed separating barrier. + BarrierHazard(String), /// Naga rejected the generated module. Validation(String), } @@ -91,6 +127,7 @@ impl fmt::Display for LowerError { write!(f, "unsupported memory level {:?}", memory) } Self::UnsupportedOperation(op) => write!(f, "unsupported operation {op}"), + Self::BarrierHazard(error) => write!(f, "workgroup barrier hazard: {error}"), Self::Validation(error) => write!(f, "naga validation failed: {error}"), } } @@ -184,6 +221,9 @@ struct Lowerer<'a> { buffer_decls: Vec, tile_decls: Vec, local_decls: Vec, + /// Liveness-based sharing of workgroup allocations: barrier-separated + /// disjoint-lifetime tiles of one element type map onto one global. + tile_arena: arena::TileArena, } /// Snapshot of the per-iteration caches that are scoped to one loop iteration: @@ -213,6 +253,7 @@ fn local_key(local: &Local) -> *const () { } mod analysis; +pub(crate) mod arena; mod block; mod control; mod coop; diff --git a/fusor-ml/tile-ir/src/lower/setup.rs b/fusor-ml/tile-ir/src/lower/setup.rs index 5a0e640c0..a748d3fa0 100644 --- a/fusor-ml/tile-ir/src/lower/setup.rs +++ b/fusor-ml/tile-ir/src/lower/setup.rs @@ -1,7 +1,10 @@ use super::*; impl<'a> Lowerer<'a> { - pub(super) fn new(ir: &'a KernelIr) -> Result { + pub(super) fn new( + ir: &'a KernelIr, + tile_arena: super::arena::TileArena, + ) -> Result { let analysis = analysis::Analysis::run(ir); let mut caps = analysis.caps; // Cooperative-matrix lowering needs a subgroup id even if the kernel @@ -114,6 +117,7 @@ impl<'a> Lowerer<'a> { buffer_decls: analysis.buffers, tile_decls: analysis.tiles, local_decls: analysis.locals, + tile_arena, }) } @@ -231,6 +235,9 @@ impl<'a> Lowerer<'a> { if self.caps.uses_coop { capabilities |= naga::valid::Capabilities::COOPERATIVE_MATRIX; } + if matches!(self.tile_arena.mode, super::arena::ArenaMode::ByteArena) { + capabilities |= naga::valid::Capabilities::WORKGROUP_MEMORY_ALIAS; + } let info = naga::valid::Validator::new(naga::valid::ValidationFlags::all(), capabilities) .validate(&self.module) .map_err(|error| LowerError::Validation(format!("{error:#?}")))?; @@ -266,6 +273,7 @@ impl<'a> Lowerer<'a> { ty, init: None, memory_decorations: naga::MemoryDecorations::empty(), + workgroup_alias: None, }, Span::default(), ); @@ -277,13 +285,26 @@ impl<'a> Lowerer<'a> { } fn create_workgroup_globals(&mut self) -> Result<(), LowerError> { - let tiles = self.collect_tiles(); - for tile in &tiles { - if tile.layout.memory_level() != MemoryLevel::Workgroup { - continue; - } - let ty = self.tile_type(tile.element, &tile.layout)?; - let global = self.module.global_variables.append( + // One global per arena region: barrier-separated disjoint-lifetime + // tiles of one stride class share the allocation, sized for the + // widest occupant and typed with the region's canonical element. + // Every tile indexes from zero within its own logical extent, so + // sharing needs no address rewriting; a heterogeneous region + // bitcasts values at each access instead. + let mut region_globals: Vec>> = + vec![None; self.tile_arena.regions.len()]; + // ByteArena mode: one untyped-by-convention backing store, typed + // per-tile globals aliased into it at packed byte offsets. Created + // before any alias so the arena handle always precedes them. + let mut arena_global: Option> = None; + if matches!(self.tile_arena.mode, super::arena::ArenaMode::ByteArena) { + let quads = std::num::NonZeroU32::new(self.tile_arena.arena_bytes / 16) + .ok_or(LowerError::UnsupportedOperation("empty workgroup arena"))?; + let ty = self.array_type_with_size( + ElementType::vector(ScalarElement::U32, 4), + ArraySize::Constant(quads), + )?; + arena_global = Some(self.module.global_variables.append( GlobalVariable { name: None, space: AddressSpace::WorkGroup, @@ -291,12 +312,78 @@ impl<'a> Lowerer<'a> { ty, init: None, memory_decorations: naga::MemoryDecorations::empty(), + workgroup_alias: None, }, Span::default(), - ); - self.globals - .borrow_mut() - .insert(std::rc::Rc::as_ptr(tile) as *const (), global); + )); + } + let tiles = self.collect_tiles(); + for tile in &tiles { + if tile.layout.memory_level() != MemoryLevel::Workgroup { + continue; + } + let key = std::rc::Rc::as_ptr(tile) as *const (); + let global = match self.tile_arena.assignment.get(&key).copied() { + Some(super::arena::Placement::Region { index }) => { + if region_globals[index].is_none() { + let region = &self.tile_arena.regions[index]; + let elements = std::num::NonZeroU32::new(region.elements) + .ok_or(LowerError::UnsupportedOperation("empty workgroup tile"))?; + let ty = self.array_type_with_size( + region.canonical, + ArraySize::Constant(elements), + )?; + region_globals[index] = Some(self.module.global_variables.append( + GlobalVariable { + name: None, + space: AddressSpace::WorkGroup, + binding: None, + ty, + init: None, + memory_decorations: naga::MemoryDecorations::empty(), + workgroup_alias: None, + }, + Span::default(), + )); + } + region_globals[index].expect("created above") + } + Some(super::arena::Placement::Arena { byte_offset }) => { + let ty = self.tile_type(tile.element, &tile.layout)?; + self.module.global_variables.append( + GlobalVariable { + name: None, + space: AddressSpace::WorkGroup, + binding: None, + ty, + init: None, + memory_decorations: naga::MemoryDecorations::empty(), + workgroup_alias: Some(naga::WorkgroupAlias { + arena: arena_global.expect("created above for ByteArena mode"), + byte_offset, + }), + }, + Span::default(), + ) + } + // Declared but never touched: keep a private allocation. + None => { + let ty = self.tile_type(tile.element, &tile.layout)?; + self.module.global_variables.append( + GlobalVariable { + name: None, + space: AddressSpace::WorkGroup, + binding: None, + ty, + init: None, + memory_decorations: naga::MemoryDecorations::empty(), + workgroup_alias: None, + }, + Span::default(), + ) + } + }; + self.globals.borrow_mut().insert(key, global); } Ok(()) } @@ -443,36 +530,11 @@ impl<'a> Lowerer<'a> { } fn element_array_stride(element: ElementType) -> Result { - match element { - ElementType::F32 | ElementType::U32 => Ok(4), - ElementType::F16 => Ok(2), - ElementType::Vector { scalar, lanes } => Self::vector_array_stride(scalar, lanes), - ElementType::Bool => Err(LowerError::UnsupportedOperation( - "bool arrays are not supported", - )), - ElementType::CoopMatrix { .. } => Err(LowerError::UnsupportedOperation( - "cooperative-matrix arrays are not supported", - )), - } - } - - fn vector_array_stride(scalar: ScalarElement, lanes: u32) -> Result { - let scalar_size = match scalar { - ScalarElement::F32 | ScalarElement::U32 => 4, - ScalarElement::F16 => 2, - ScalarElement::Bool => { - return Err(LowerError::UnsupportedOperation( - "bool vector arrays are not supported", - )); - } - }; - match lanes { - 2 => Ok(2 * scalar_size), - 3 | 4 => Ok(4 * scalar_size), - _ => Err(LowerError::UnsupportedOperation( - "vectors must have 2, 3, or 4 lanes", - )), - } + element + .workgroup_array_stride() + .ok_or(LowerError::UnsupportedOperation( + "element cannot back a workgroup array", + )) } } diff --git a/fusor-ml/tile-ir/src/lower/tile_program/expr.rs b/fusor-ml/tile-ir/src/lower/tile_program/expr.rs index 40dec8827..533c9d7c9 100644 --- a/fusor-ml/tile-ir/src/lower/tile_program/expr.rs +++ b/fusor-ml/tile-ir/src/lower/tile_program/expr.rs @@ -45,7 +45,7 @@ impl<'a> Lowerer<'a> { ExprKind::LoadTile { tile, index } => { let index = self.lower_expr_lane(expressions, body, index, spill_depth)?; let ptr = self.tile_dynamic_pointer(expressions, tile, index, body)?; - Ok(Self::emit_load(expressions, body, ptr)) + Ok(self.load_tile_value(expressions, body, tile, ptr)) } ExprKind::LoadLocal(local) => { // Coop accumulators chain through the acc-value SSA memo: a live @@ -211,6 +211,17 @@ impl<'a> Lowerer<'a> { let handles = self.lower_dequantize(expressions, body, expr, spill_depth)?; Ok(handles[0]) } + ExprKind::VecComponent { vector, component } => { + let base = self.lower_expr_lane(expressions, body, vector, spill_depth + 1)?; + Ok(self.emit( + expressions, + body, + Expression::AccessIndex { + base, + index: *component, + }, + )) + } ExprKind::LaneOf { block, lane } => { let handles = self.lower_lane_of_block(expressions, body, block, spill_depth)?; handles diff --git a/fusor-ml/tile-ir/src/lower/tile_program/scalar.rs b/fusor-ml/tile-ir/src/lower/tile_program/scalar.rs index 50d3d20fa..a3071b416 100644 --- a/fusor-ml/tile-ir/src/lower/tile_program/scalar.rs +++ b/fusor-ml/tile-ir/src/lower/tile_program/scalar.rs @@ -137,13 +137,7 @@ impl<'a> Lowerer<'a> { Statement::ControlBarrier(Barrier::WORK_GROUP), Span::default(), ); - body.push( - Statement::Store { - pointer: lane_ptr, - value, - }, - Span::default(), - ); + self.store_tile_value(expressions, body, scratch_tile, lane_ptr, value); body.push( Statement::ControlBarrier(Barrier::WORK_GROUP), Span::default(), @@ -196,7 +190,7 @@ impl<'a> Lowerer<'a> { let result_ptr = self.tile_dynamic_pointer(expressions, scratch_tile, result_index, body)?; - Ok(Self::emit_load(expressions, body, result_ptr)) + Ok(self.load_tile_value(expressions, body, scratch_tile, result_ptr)) } fn lower_reduce_step( @@ -211,20 +205,14 @@ impl<'a> Lowerer<'a> { let rhs_index = self.add_literal_u32_emitted(expressions, lane, stride, &mut body); let lhs_ptr = self.tile_dynamic_pointer(expressions, scratch_tile, lane, &mut body)?; let rhs_ptr = self.tile_dynamic_pointer(expressions, scratch_tile, rhs_index, &mut body)?; - let lhs = Self::emit_load(expressions, &mut body, lhs_ptr); - let rhs = Self::emit_load(expressions, &mut body, rhs_ptr); + let lhs = self.load_tile_value(expressions, &mut body, scratch_tile, lhs_ptr); + let rhs = self.load_tile_value(expressions, &mut body, scratch_tile, rhs_ptr); let reduced = self.emit( expressions, &mut body, Self::tile_reduce_expression(op, lhs, rhs), ); - body.push( - Statement::Store { - pointer: lhs_ptr, - value: reduced, - }, - Span::default(), - ); + self.store_tile_value(expressions, &mut body, scratch_tile, lhs_ptr, reduced); Ok(body) } diff --git a/fusor-ml/tile-ir/src/lower/tile_program/types.rs b/fusor-ml/tile-ir/src/lower/tile_program/types.rs index 9e2c3e3b4..db9de9ceb 100644 --- a/fusor-ml/tile-ir/src/lower/tile_program/types.rs +++ b/fusor-ml/tile-ir/src/lower/tile_program/types.rs @@ -46,6 +46,55 @@ impl<'a> Lowerer<'a> { self.cast_as(expressions, body, value, scalar.kind, Some(scalar.width)) } + /// Element type a tile's backing array is emitted with: its region's + /// canonical element, which differs from the tile's own element only + /// when the region is shared across types. + pub(in crate::lower) fn tile_emitted_element(&self, tile: &Tile) -> ElementType { + match self.tile_arena.assignment.get(&super::super::tile_key(tile)) { + Some(super::super::arena::Placement::Region { index }) => { + self.tile_arena.regions[*index].canonical + } + _ => tile.element, + } + } + + /// Load one element through a tile pointer, bitcasting from the + /// region's canonical type back to the tile's element when they differ. + pub(in crate::lower) fn load_tile_value( + &self, + expressions: &mut Arena, + body: &mut Block, + tile: &Tile, + pointer: Handle, + ) -> Handle { + let value = Self::emit_load(expressions, body, pointer); + if self.tile_emitted_element(tile) == tile.element { + return value; + } + let scalar = Self::element_scalar(tile.element); + self.cast_as(expressions, body, value, scalar.kind, None) + } + + /// Bitcast a tile-element value to the region's canonical type (when + /// they differ) and store it through the tile pointer. + pub(in crate::lower) fn store_tile_value( + &self, + expressions: &mut Arena, + body: &mut Block, + tile: &Tile, + pointer: Handle, + value: Handle, + ) { + let emitted = self.tile_emitted_element(tile); + let value = if emitted == tile.element { + value + } else { + let scalar = Self::element_scalar(emitted); + self.cast_as(expressions, body, value, scalar.kind, None) + }; + body.push(Statement::Store { pointer, value }, Span::default()); + } + pub(in crate::lower) fn condition_value( &self, expressions: &mut Arena, diff --git a/fusor-ml/tile-ir/src/tests/arena.rs b/fusor-ml/tile-ir/src/tests/arena.rs new file mode 100644 index 000000000..f6ba3478e --- /dev/null +++ b/fusor-ml/tile-ir/src/tests/arena.rs @@ -0,0 +1,419 @@ +//! Workgroup-tile arena sharing: barrier soundness fixtures. +//! +//! Tile A (8x8 f32, 256 B) and tile B (4x8 f32, 128 B) touch on either side +//! of a barrier. Sharing collapses the footprint to 256 B; refusing keeps +//! 384 B. A barrier inside a loop that may break early, return, or run zero +//! iterations can be skipped at runtime, so it must not enable sharing for +//! tiles living outside that loop. + +use super::*; +use crate::tile; + +const SHARED: u64 = 256; +const UNSHARED: u64 = 256 + 128; + +/// Build the two-tile fixture with `body` between the touches. `body` runs +/// with the program block and must contain the only barrier. +fn two_tile_fixture( + between: impl FnOnce(&mut tile::TileBlock), +) -> KernelIr { + tile::build(|phase| { + let a = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + let b = phase.alloc_workgroup_tile(ScalarElement::F32, 4, 8); + phase.program_grid(32, [1, 1, 1], |program| { + let lane = program.lane(); + program.store_workgroup(&a, lane.clone(), 1.0f32); + between(program); + program.store_workgroup(&b, lane, 2.0f32); + }); + }) +} + +#[test] +fn shares_through_top_level_barrier() { + let ir = two_tile_fixture(|program| { + program.workgroup_barrier(); + }); + assert_eq!(ir.workgroup_bytes(), SHARED); +} + +#[test] +fn shares_through_static_loop_barrier() { + // A static-count, break-free loop executes its body on every pass, so an + // in-loop barrier is guaranteed and sharing through it is sound. + let ir = two_tile_fixture(|program| { + program.loop_range(4, |program, _| { + program.workgroup_barrier(); + }); + }); + assert_eq!(ir.workgroup_bytes(), SHARED); +} + +#[test] +fn no_share_through_unstructured_loop_barrier() { + // An unstructured loop's count is data-dependent: the barrier may never + // execute (or not on the exit path), so it cannot separate the tiles. + let ir = two_tile_fixture(|program| { + let lane = program.lane(); + program.loop_forever(|program| { + program.workgroup_barrier(); + program.break_if(lane.clone().lt(32u32)); + }); + }); + assert_eq!(ir.workgroup_bytes(), UNSHARED); +} + +#[test] +fn no_share_through_breaking_static_loop_barrier() { + // A break before the barrier can skip it on the final iteration. + let ir = two_tile_fixture(|program| { + let lane = program.lane(); + program.loop_range(4, |program, _| { + program.break_if(lane.clone().lt(32u32)); + program.workgroup_barrier(); + }); + }); + assert_eq!(ir.workgroup_bytes(), UNSHARED); +} + +#[test] +fn no_share_through_returning_static_loop_barrier() { + // A mid-body return exits before later barriers just like a break. + let ir = two_tile_fixture(|program| { + let lane = program.lane(); + program.loop_range(4, |program, _| { + program.if_then(lane.clone().lt(32u32), |program| program.return_()); + program.workgroup_barrier(); + }); + }); + assert_eq!(ir.workgroup_bytes(), UNSHARED); +} + +#[test] +fn accumulator_update_is_live_across_loop() { + // An accumulator update executes at the end of EVERY iteration — after + // any in-loop barrier — so a tile read only by the update is live across + // the whole loop and an in-loop barrier cannot separate it from a tile + // touched after the loop. Built as raw IR: no builder emits a counted + // accumulator loop without a break, but the IR admits it. + let workgroup = + |shape: [u32; 2]| Layout::contiguous(MemoryLevel::Workgroup, Shape::new(shape)); + let tile_a: Tile = std::rc::Rc::new(TileDecl { + element: ElementType::F32, + layout: workgroup([8, 8]), + }); + let tile_b: Tile = std::rc::Rc::new(TileDecl { + element: ElementType::F32, + layout: workgroup([4, 8]), + }); + let acc_local: Local = std::rc::Rc::new(LocalDecl { + element: ElementType::F32, + }); + let index_local: Local = std::rc::Rc::new(LocalDecl { + element: ElementType::U32, + }); + let lit_u32 = |value: u32| Expr::new(ExprKind::Literal(TileLiteral::U32(value)), ElementType::U32); + let lit_f32 = |value: f32| { + Expr::new( + ExprKind::Literal(TileLiteral::f32(value)), + ElementType::F32, + ) + }; + + let mut ir = KernelIr::default(); + ir.block = 32; + ir.body = vec![ + Stmt::Loop { + count: Some(lit_u32(4)), + index: Some(index_local), + accumulators: vec![Accumulator { + local: acc_local, + init: lit_f32(0.0), + update: Expr::new( + ExprKind::LoadTile { + tile: tile_a, + index: Box::new(lit_u32(0)), + }, + ElementType::F32, + ), + }], + body: vec![Stmt::Barrier], + }, + Stmt::StoreTile { + dst: tile_b, + index: Box::new(lit_u32(0)), + value: lit_f32(2.0), + }, + ]; + assert_eq!(ir.workgroup_bytes(), UNSHARED); +} + +#[test] +fn cross_type_tiles_share_one_region() { + // f32 and u32 have the same 4-byte stride: barrier-separated disjoint + // tiles share one region, emitted with the class-neutral u32 type and a + // value bitcast at each access. + let ir = tile::build(|phase| { + let a = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + let b = phase.alloc_workgroup_tile(ScalarElement::U32, 4, 8); + phase.program_grid(32, [1, 1, 1], |program| { + let lane = program.lane(); + program.store_workgroup(&a, lane.clone(), 1.0f32); + let read_back = program.load_workgroup(&a, lane.clone()); + program.store_workgroup(&a, lane.clone(), read_back); + program.workgroup_barrier(); + program.store_workgroup(&b, lane, 2u32); + }); + }); + assert_eq!(ir.workgroup_bytes(), SHARED); + let lowered = lower_or_fail(&ir, "cross-type region"); + let function = &lowered.module().entry_points[0].function; + let bitcasts = function + .expressions + .iter() + .filter(|(_, expr)| { + matches!( + expr, + naga::Expression::As { + convert: None, + .. + } + ) + }) + .count(); + // Both f32 stores, the f32 load, and the u32 store's canonical is u32 + // itself: three f32<->u32 casts. + assert_eq!(bitcasts, 3); +} + +#[test] +fn coop_consumed_tile_pins_its_region_type() { + // A tile consumed as a raw cooperative-matrix pointer cannot live in a + // widened region: the u32 tile must get its own allocation. + let ir = tile::build(|phase| { + let coop = crate::CoopMatrixToken::new_unchecked(); + let a = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + let b = phase.alloc_workgroup_tile(ScalarElement::U32, 4, 8); + let y = phase.storage_write(ScalarElement::F32.element(), Shape::new([8, 8])); + phase.program_grid(32, [1, 1, 1], |program| { + let acc = coop.alloc_coop_acc(program, ScalarElement::F32, 8, 8); + let a_frag = coop.coop_load_a(program, &a, 0u32, 0u32, ScalarElement::F32, 8, 8); + let b_frag = coop.coop_load_b(program, &a, 0u32, 0u32, ScalarElement::F32, 8, 8); + let c = coop.coop_zero(program, ScalarElement::F32, 8, 8); + coop.coop_store_local(program, &acc, coop.coop_mma(program, a_frag, b_frag, c)); + coop.coop_store(program, &acc, &y, 0u32, 0u32); + program.workgroup_barrier(); + let lane = program.lane(); + program.store_workgroup(&b, lane, 2u32); + }); + }); + assert_eq!(ir.workgroup_bytes(), UNSHARED); +} + +#[test] +fn mixed_stride_tiles_pack_into_byte_arena() { + // With the byte-arena backend proved, an f16 tile reuses the f32 tile's + // bytes after a barrier: footprint collapses to the f32 extent (16-byte + // aligned) instead of the sum. + let mut ir = tile::build(|phase| { + let a = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + let b = phase.alloc_workgroup_tile(ScalarElement::F16, 4, 8); + phase.program_grid(32, [1, 1, 1], |program| { + let lane = program.lane(); + program.store_workgroup(&a, lane.clone(), 1.0f32); + program.workgroup_barrier(); + program.store_workgroup(&b, lane, 2.0f32); + }); + }); + // Regions mode first: f16 cannot join the 4-byte class, so both + // allocations exist (256 + 64 bytes). + assert_eq!(ir.workgroup_bytes(), 256 + 64); + ir.byte_arena = true; + assert_eq!(ir.workgroup_bytes(), 256); + // The alias emission path validates end-to-end: one arena global plus + // typed aliased globals, accepted by the fork validator. + let lowered = lower_or_fail(&ir, "byte-arena emission"); + let aliased = lowered + .module() + .global_variables + .iter() + .filter(|(_, global)| global.workgroup_alias.is_some()) + .count(); + assert_eq!(aliased, 2); +} + +#[test] +fn byte_arena_keeps_concurrent_tiles_disjoint() { + // Without a separating barrier the tiles overlap in time: the packer + // pushes the f16 tile past the f32 extent. + let mut ir = tile::build(|phase| { + let a = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + let b = phase.alloc_workgroup_tile(ScalarElement::F16, 4, 8); + phase.program_grid(32, [1, 1, 1], |program| { + let lane = program.lane(); + program.store_workgroup(&a, lane.clone(), 1.0f32); + program.store_workgroup(&b, lane, 2.0f32); + }); + }); + ir.byte_arena = true; + assert_eq!(ir.workgroup_bytes(), 256 + 64); +} + +fn barrier_count(stmts: &[Stmt]) -> usize { + stmts + .iter() + .map(|stmt| match stmt { + Stmt::Barrier => 1, + Stmt::Loop { body, .. } => barrier_count(body), + Stmt::If { accept, reject, .. } => barrier_count(accept) + barrier_count(reject), + _ => 0, + }) + .sum() +} + +#[test] +fn elision_removes_duplicate_barrier() { + let ir = two_tile_fixture(|program| { + program.workgroup_barrier(); + program.workgroup_barrier(); + }); + assert_eq!(barrier_count(&ir.body), 1); + assert_eq!(ir.workgroup_bytes(), SHARED); +} + +#[test] +fn elision_removes_trailing_barrier() { + let ir = tile::build(|phase| { + let a = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + phase.program_grid(32, [1, 1, 1], |program| { + let lane = program.lane(); + program.store_workgroup(&a, lane, 1.0f32); + program.workgroup_barrier(); + }); + }); + assert_eq!(barrier_count(&ir.body), 0); +} + +#[test] +fn elision_keeps_wrap_around_separators() { + // The reverted-elision incident in kernel form: inside a loop, one + // barrier orders this iteration's write before its read, the other + // orders the read before the NEXT iteration's write. Neither backs the + // other up; both must survive. + let ir = tile::build(|phase| { + let t = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + phase.program_grid(32, [1, 1, 1], |program| { + let lane = program.lane(); + let scratch = program.private(ElementType::F32); + program.loop_range(4, |program, _| { + program.store_workgroup(&t, lane.clone(), 1.0f32); + program.workgroup_barrier(); + let value = program.load_workgroup(&t, lane.clone()); + program.store_local(&scratch, value); + program.workgroup_barrier(); + }); + }); + }); + assert_eq!(barrier_count(&ir.body), 2); +} + +/// Phased in-loop fixture: tiles touched in disjoint per-iteration phases, +/// with barriers controlled by the caller. A is 8x8 f32, B is 4x8 f32. +fn phased_loop_fixture( + barrier_between: bool, + barrier_at_end: bool, + unstructured: bool, +) -> KernelIr { + tile::build(|phase| { + let a = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + let b = phase.alloc_workgroup_tile(ScalarElement::F32, 4, 8); + phase.program_grid(32, [1, 1, 1], |program| { + let lane = program.lane(); + let scratch = program.private(ElementType::F32); + let body = |program: &mut tile::TileBlock| { + program.store_workgroup(&a, lane.clone(), 1.0f32); + let read = program.load_workgroup(&a, lane.clone()); + program.store_local(&scratch, read); + if barrier_between { + program.workgroup_barrier(); + } + program.store_workgroup(&b, lane.clone(), 2.0f32); + let read = program.load_workgroup(&b, lane.clone()); + program.store_local(&scratch, read); + if barrier_at_end { + program.workgroup_barrier(); + } + }; + if unstructured { + program.loop_forever(|program| { + body(program); + program.break_if(lane.clone().lt(32u32)); + }); + } else { + program.loop_range(4, |program, _| body(program)); + } + }); + }) +} + +#[test] +fn phased_in_loop_tiles_share() { + // Disjoint phases, a barrier between them, and a barrier covering the + // wrap back to the first phase: the tiles share one region. + let ir = phased_loop_fixture(true, true, false); + assert_eq!(ir.workgroup_bytes(), SHARED); + lower_or_fail(&ir, "phased in-loop sharing"); +} + +#[test] +fn phase_sharing_needs_the_forward_barrier() { + let ir = phased_loop_fixture(false, true, false); + assert_eq!(ir.workgroup_bytes(), UNSHARED); +} + +#[test] +fn phase_sharing_needs_the_wrap_barrier() { + let ir = phased_loop_fixture(true, false, false); + assert_eq!(ir.workgroup_bytes(), UNSHARED); +} + +#[test] +fn phase_sharing_survives_break_loops() { + // Taking the back edge means the full body executed, so in-loop + // barriers stay valid phase separators even in unstructured loops. + let ir = phased_loop_fixture(true, true, true); + assert_eq!(ir.workgroup_bytes(), SHARED); +} + +#[test] +fn phase_sharing_does_not_chain_transitively() { + // A->B and B->C phase separation does NOT cover the C->A wrap: with + // barriers only between the phases, C must not join A and B's region. + let ir = tile::build(|phase| { + let a = phase.alloc_workgroup_tile(ScalarElement::F32, 8, 8); + let b = phase.alloc_workgroup_tile(ScalarElement::F32, 4, 8); + let c = phase.alloc_workgroup_tile(ScalarElement::F32, 4, 4); + phase.program_grid(32, [1, 1, 1], |program| { + let lane = program.lane(); + let scratch = program.private(ElementType::F32); + program.loop_range(4, |program, _| { + let mut touch = |program: &mut tile::TileBlock, tile: &tile::WorkgroupTile| { + program.store_workgroup(tile, lane.clone(), 1.0f32); + let read = program.load_workgroup(tile, lane.clone()); + program.store_local(&scratch, read); + }; + touch(program, &a); + program.workgroup_barrier(); + touch(program, &b); + program.workgroup_barrier(); + touch(program, &c); + }); + }); + }); + // A and B share (forward barrier between phases, wrap covered by the + // second barrier for B and the first for... A's wrap: second barrier + // sits after B's phase, before the back edge). C shares with neither: + // its own phase has no trailing barrier, so the C->occupant wraps are + // uncovered. + assert_eq!(ir.workgroup_bytes(), 256 + 64); +} diff --git a/fusor-ml/tile-ir/src/tests/golden.rs b/fusor-ml/tile-ir/src/tests/golden.rs index 34414a0d9..8e7c13d20 100644 --- a/fusor-ml/tile-ir/src/tests/golden.rs +++ b/fusor-ml/tile-ir/src/tests/golden.rs @@ -2,7 +2,7 @@ //! builder DSL (no `fusor-tile-ir-kernels` high-level builders): the contiguous //! f16 workgroup-tile copy. //! -//! The full corpus (qgemv/qmatmul/coop/flash) lives in +//! The full corpus (qgemv/qmatmul/coop/attention) lives in //! `tile-ir-kernels/tests/golden.rs` because it needs the kernel-crate //! builders; this module proves the rewritten *builder* reaches the same module //! for the one kernel it can build against `tile-ir` alone. @@ -78,6 +78,10 @@ fn check_golden_structural(name: &str, ir: &KernelIr) { let lowered = lower_or_fail(ir, name); let serialized = format!("{:#?}", lowered.module()); let path = golden_dir().join(format!("{name}.txt")); + if std::env::var_os("FUSOR_UPDATE_GOLDEN").is_some() { + std::fs::write(&path, &serialized) + .unwrap_or_else(|error| panic!("writing {}: {error}", path.display())); + } let expected = std::fs::read_to_string(&path) .unwrap_or_else(|_| panic!("golden snapshot missing: {}", path.display())); assert_eq!( diff --git a/fusor-ml/tile-ir/src/tests/lowering.rs b/fusor-ml/tile-ir/src/tests/lowering.rs index afcf19d59..4534dafd4 100644 --- a/fusor-ml/tile-ir/src/tests/lowering.rs +++ b/fusor-ml/tile-ir/src/tests/lowering.rs @@ -221,6 +221,7 @@ fn if_branches_do_not_share_branch_local_expression_memos() { buffers: Vec::new(), grid: [1, 1, 1], block: 1, + byte_arena: false, body: vec![Stmt::If { condition: Expr::new( ExprKind::Literal(TileLiteral::Bool(true)), @@ -336,3 +337,73 @@ fn cooperative_load_store_layout_flags_use_transposed_internal_layout() { assert_eq!(loads, [false, false]); assert_eq!(stores, [true]); } + +#[test] +fn general_group_reduce_lowers_without_subgroup_intrinsics() { + // A general combine cannot use subgroup collectives (they are + // per-operator), so `group_reduce_with` stages through workgroup memory. + // It must therefore lower on devices with no subgroup support at all. + let ir = tile::build(|phase| { + let x = phase.storage_read(f32(), Shape::new([64])); + let y = phase.storage_write(f32(), Shape::new([64])); + phase.program_grid(64, [1, 1, 1], |program| { + let lane = program.lane(); + let mask = lane.clone().lt(64u32); + let value = program.load(x.at(lane.clone()), mask.clone(), 0.0); + // An arbitrary associative body, not one of the closed ops. + let combined = program.group_reduce_with(16, value, |program, a, b| { + let scaled = a.clone() * b.clone(); + program.bind(a + b + scaled) + }); + program.store(y.at(lane), combined, mask); + }); + }); + + let lowered = lower_or_fail(&ir, "general group reduce"); + assert_eq!( + lowered.wgsl_extension_prelude(), + "", + "a general combine must not require the subgroups extension" + ); +} + +#[test] +fn joint_carrier_group_reduce_stages_every_slot() { + // A two-slot carrier whose second slot reads the first on both sides — + // the shape that makes online softmax a carrier rather than two + // independent reductions. Each slot needs its own staging array, so the + // lowered program must hold two block-sized workgroup allocations. + let ir = tile::build(|phase| { + let x = phase.storage_read(f32(), Shape::new([64])); + let y = phase.storage_write(f32(), Shape::new([64])); + phase.program_grid(64, [1, 1, 1], |program| { + let lane = program.lane(); + let mask = lane.clone().lt(64u32); + let value = program.load(x.at(lane.clone()), mask.clone(), 0.0); + let one = program.bind(value.clone() * 0.0 + 1.0); + let combined = program.group_reduce_with_vec( + 16, + vec![value, one], + |program, acc, incoming| { + let mut acc = acc.into_iter(); + let (m, l) = (acc.next().unwrap(), acc.next().unwrap()); + let mut incoming = incoming.into_iter(); + let (m2, l2) = (incoming.next().unwrap(), incoming.next().unwrap()); + let joined = program.bind(m.clone().max(m2.clone())); + let scaled = l * (m - joined.clone()).exp() + l2 * (m2 - joined.clone()).exp(); + vec![joined, program.bind(scaled)] + }, + ); + let mut combined = combined.into_iter(); + let (m, l) = (combined.next().unwrap(), combined.next().unwrap()); + program.store(y.at(lane), m + l.log(), mask); + }); + }); + + let lowered = lower_or_fail(&ir, "joint carrier group reduce"); + assert_eq!( + lowered.wgsl_extension_prelude(), + "", + "a joint carrier must not require the subgroups extension" + ); +} diff --git a/fusor-ml/tile-ir/src/tests/mod.rs b/fusor-ml/tile-ir/src/tests/mod.rs index 58a1eb381..cb3bf221f 100644 --- a/fusor-ml/tile-ir/src/tests/mod.rs +++ b/fusor-ml/tile-ir/src/tests/mod.rs @@ -3,6 +3,7 @@ use super::*; // trees directly, so pull them from the internal module. use crate::ir::*; +mod arena; mod golden; mod layout; mod lowering; diff --git a/fusor-ml/tile-ir/src/tile/block.rs b/fusor-ml/tile-ir/src/tile/block.rs index 18a94210c..7a3cbdb84 100644 --- a/fusor-ml/tile-ir/src/tile/block.rs +++ b/fusor-ml/tile-ir/src/tile/block.rs @@ -274,6 +274,24 @@ impl TileBlock<'_> { ) } + /// Extract one component of a vector tile. + pub fn vector_component(&self, vector: Tile, component: u32) -> Tile { + let scalar = match vector.element() { + ElementType::Vector { scalar, lanes } => { + assert!(component < lanes, "vector component out of range"); + scalar + } + other => panic!("vector_component on non-vector element {other:?}"), + }; + Tile::new( + ExprKind::VecComponent { + vector: vector.into_expr().into(), + component, + }, + scalar.element(), + ) + } + /// Compose a `LANES`-wide vector by broadcasting one scalar `value` into /// every lane. pub fn vector_splat(&self, scalar: ScalarElement, value: Tile) -> Tile { @@ -312,6 +330,11 @@ impl TileBlock<'_> { self.push_stmt(Stmt::Barrier); } + /// Storage-memory barrier with workgroup execution scope. + pub fn storage_barrier(&mut self) { + self.push_stmt(Stmt::StorageBarrier); + } + /// Conditional block (no else). pub fn if_then(&mut self, condition: impl Into, accept: impl FnOnce(&mut Self)) { self.if_else(condition, accept, |_| {}); diff --git a/fusor-ml/tile-ir/src/tile/capability.rs b/fusor-ml/tile-ir/src/tile/capability.rs index b1562efc3..e46fda75c 100644 --- a/fusor-ml/tile-ir/src/tile/capability.rs +++ b/fusor-ml/tile-ir/src/tile/capability.rs @@ -54,6 +54,55 @@ impl SubgroupToken { pub fn subgroup_reduce_max(self, program: &TileBlock<'_>, value: Tile) -> Tile { self.subgroup_reduce(program, TileReduceOp::Max, value) } + + /// Whole-workgroup reduction accelerated with subgroup collectives: a + /// per-subgroup reduce, partials staged through workgroup memory, and a + /// serial fold — two barriers total versus one per tree level in + /// [`TileBlock::group_reduce`]. `subgroup_size` must be the device's + /// fixed subgroup width and divide the workgroup size; every lane + /// receives the combined value. + pub fn workgroup_reduce( + self, + program: &mut TileBlock<'_>, + op: TileReduceOp, + subgroup_size: u32, + value: Tile, + ) -> Tile { + program.workgroup_reduce_via_subgroups(op, subgroup_size, value) + } + + /// Reduction over subgroup-aligned lane groups of `group_size`, the + /// sub-block form of [`Self::workgroup_reduce`]: no barrier at all when a + /// group is one subgroup, two otherwise. `group_size` must be a multiple + /// of the device's fixed subgroup width and divide the workgroup size, and + /// the caller must derive group membership from `subgroup_id` (see + /// [`TileBlock::group_reduce_via_subgroups`]). + pub fn group_reduce( + self, + program: &mut TileBlock<'_>, + op: TileReduceOp, + subgroup_size: u32, + group_size: u32, + value: Tile, + ) -> Tile { + program.group_reduce_via_subgroups(op, subgroup_size, group_size, value) + } +} + +/// Capability token for byte-arena workgroup packing: the adapter supports +/// the workgroup-alias backend extension, so mixed-stride tiles may share +/// bytes at packed offsets. Same trust model as [`CoopMatrixToken`]: safe +/// construction lives in device code. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct ByteArenaToken { + _private: (), +} + +impl ByteArenaToken { + /// Construct a byte-arena token without checking device capabilities. + pub fn new_unchecked() -> Self { + Self { _private: () } + } } /// Capability token for tile-IR cooperative-matrix operations. @@ -136,6 +185,38 @@ impl CoopMatrixToken { program.coop_load_b(tile, row, col, scalar, rows, cols) } + /// Cooperatively load an A-role fragment of the tile's transpose: + /// fragment `(i, j)` reads tile element `(col + j, row + i)`. + #[allow(clippy::too_many_arguments)] + pub fn coop_load_a_transposed( + self, + program: &TileBlock<'_>, + tile: &WorkgroupTile, + row: impl Into, + col: impl Into, + scalar: ScalarElement, + rows: u32, + cols: u32, + ) -> Tile { + program.coop_load_a_transposed(tile, row, col, scalar, rows, cols) + } + + /// Cooperatively load a B-role fragment of the tile's transpose: + /// fragment `(i, j)` reads tile element `(col + j, row + i)`. + #[allow(clippy::too_many_arguments)] + pub fn coop_load_b_transposed( + self, + program: &TileBlock<'_>, + tile: &WorkgroupTile, + row: impl Into, + col: impl Into, + scalar: ScalarElement, + rows: u32, + cols: u32, + ) -> Tile { + program.coop_load_b_transposed(tile, row, col, scalar, rows, cols) + } + /// Cooperatively load a C-role fragment from a rank-1 storage vector. pub fn coop_load_c_broadcast( self, @@ -165,4 +246,18 @@ impl CoopMatrixToken { ) { program.coop_store(acc, dst, row, col); } + + /// Cooperatively store an accumulator into a workgroup tile at + /// `(row, col)` — the staging step between fragment math and per-lane + /// passes over the same values. + pub fn coop_store_tile( + self, + program: &mut TileBlock<'_>, + acc: &CoopAcc, + tile: &WorkgroupTile, + row: impl Into, + col: impl Into, + ) { + program.coop_store_tile(acc, tile, row, col); + } } diff --git a/fusor-ml/tile-ir/src/tile/coop.rs b/fusor-ml/tile-ir/src/tile/coop.rs index 32d61b222..275da9a3d 100644 --- a/fusor-ml/tile-ir/src/tile/coop.rs +++ b/fusor-ml/tile-ir/src/tile/coop.rs @@ -73,7 +73,21 @@ impl TileBlock<'_> { rows: u32, cols: u32, ) -> Tile { - self.coop_load_tile(CoopMatrixRole::A, tile, row, col, scalar, rows, cols) + self.coop_load_tile(CoopMatrixRole::A, tile, row, col, scalar, rows, cols, false) + } + + /// Cooperatively load an A-role fragment of the tile's transpose: + /// fragment `(i, j)` reads tile element `(col + j, row + i)`. + pub(crate) fn coop_load_a_transposed( + &self, + tile: &WorkgroupTile, + row: impl Into, + col: impl Into, + scalar: ScalarElement, + rows: u32, + cols: u32, + ) -> Tile { + self.coop_load_tile(CoopMatrixRole::A, tile, row, col, scalar, rows, cols, true) } /// Cooperatively load a B-role fragment from a region of a workgroup tile. @@ -86,9 +100,24 @@ impl TileBlock<'_> { rows: u32, cols: u32, ) -> Tile { - self.coop_load_tile(CoopMatrixRole::B, tile, row, col, scalar, rows, cols) + self.coop_load_tile(CoopMatrixRole::B, tile, row, col, scalar, rows, cols, false) + } + + /// Cooperatively load a B-role fragment of the tile's transpose: + /// fragment `(i, j)` reads tile element `(col + j, row + i)`. + pub(crate) fn coop_load_b_transposed( + &self, + tile: &WorkgroupTile, + row: impl Into, + col: impl Into, + scalar: ScalarElement, + rows: u32, + cols: u32, + ) -> Tile { + self.coop_load_tile(CoopMatrixRole::B, tile, row, col, scalar, rows, cols, true) } + #[allow(clippy::too_many_arguments)] #[allow(clippy::too_many_arguments)] fn coop_load_tile( &self, @@ -99,6 +128,7 @@ impl TileBlock<'_> { scalar: ScalarElement, rows: u32, cols: u32, + transposed: bool, ) -> Tile { assert!(rows == 8 || rows == 16, "coop rows must be 8 or 16"); assert!(cols == 8 || cols == 16, "coop cols must be 8 or 16"); @@ -112,6 +142,7 @@ impl TileBlock<'_> { tile: tile.decl().clone(), row: boxed_index(row), col: boxed_index(col), + transposed, }, }, ElementType::coop_matrix(scalar, role, rows, cols), @@ -177,4 +208,23 @@ impl TileBlock<'_> { }, }); } + + /// Cooperatively store an accumulator into a workgroup tile at + /// `(row, col)` — the staging step between fragment math and per-lane + /// passes over the same values. A distinct collective primitive — never + /// a per-lane store. + pub(crate) fn coop_store_tile( + &mut self, + acc: &CoopAcc, + tile: &WorkgroupTile, + row: impl Into, + col: impl Into, + ) { + self.push_stmt(Stmt::CoopStoreTile { + acc: acc.decl().clone(), + tile: tile.decl().clone(), + row: boxed_index(row), + col: boxed_index(col), + }); + } } diff --git a/fusor-ml/tile-ir/src/tile/grid.rs b/fusor-ml/tile-ir/src/tile/grid.rs index 519519627..2ea09ac58 100644 --- a/fusor-ml/tile-ir/src/tile/grid.rs +++ b/fusor-ml/tile-ir/src/tile/grid.rs @@ -25,5 +25,5 @@ use crate::ir::KernelIr; pub fn build(f: impl FnOnce(&mut Program)) -> KernelIr { let mut program = Program::new(); f(&mut program); - program.ir + program.into_ir() } diff --git a/fusor-ml/tile-ir/src/tile/mod.rs b/fusor-ml/tile-ir/src/tile/mod.rs index cedb6a990..3fae1b5c8 100644 --- a/fusor-ml/tile-ir/src/tile/mod.rs +++ b/fusor-ml/tile-ir/src/tile/mod.rs @@ -36,7 +36,7 @@ mod storage; mod value; pub use block::TileBlock; -pub use capability::{CoopMatrixToken, SubgroupToken}; +pub use capability::{ByteArenaToken, CoopMatrixToken, SubgroupToken}; pub use grid::build; pub use program::Program; pub use storage::{Storage, StorageIndex}; diff --git a/fusor-ml/tile-ir/src/tile/program.rs b/fusor-ml/tile-ir/src/tile/program.rs index 77a79c78c..aae7bac10 100644 --- a/fusor-ml/tile-ir/src/tile/program.rs +++ b/fusor-ml/tile-ir/src/tile/program.rs @@ -37,9 +37,12 @@ impl Program { } } - /// Consume the builder and return the constructed [`KernelIr`]. + /// Consume the builder and return the constructed [`KernelIr`], with + /// barriers that separate no hazardous access pair elided. pub(crate) fn into_ir(self) -> KernelIr { - self.ir + let mut ir = self.ir; + crate::analysis::elide_barriers(&mut ir); + ir } } @@ -169,6 +172,13 @@ impl Program { // ---- tile / local allocation ---------------------------------------- /// Allocate a rank-2 workgroup-scope tile of shape `[rows, cols]`. + /// Opt this kernel into byte-arena workgroup packing: mixed-stride + /// tiles share bytes at packed offsets via the workgroup-alias backend + /// extension the token proves. + pub fn enable_byte_arena(&mut self, _token: crate::tile::ByteArenaToken) { + self.ir.byte_arena = true; + } + pub fn alloc_workgroup_tile( &mut self, element: ScalarElement, @@ -193,6 +203,19 @@ impl Program { ) } + /// Allocate a rank-1 workgroup-scope scratch array of arbitrary element + /// type (vectors included — one load/store moves the whole element). + pub fn alloc_workgroup_array_elements( + &mut self, + element: ElementType, + len: u32, + ) -> WorkgroupTile { + self.alloc_tile( + element, + Layout::contiguous(MemoryLevel::Workgroup, Shape::new([len])), + ) + } + /// Allocate a rank-1 workgroup-scope scratch array. pub fn alloc_workgroup_array(&mut self, element: ScalarElement, len: u32) -> WorkgroupTile { self.alloc_tile( diff --git a/fusor-ml/tile-ir/src/tile/reduce.rs b/fusor-ml/tile-ir/src/tile/reduce.rs index e82f6206f..b0810ef3d 100644 --- a/fusor-ml/tile-ir/src/tile/reduce.rs +++ b/fusor-ml/tile-ir/src/tile/reduce.rs @@ -169,6 +169,184 @@ impl TileBlock<'_> { ) } + /// Whole-workgroup reduction built from subgroup collectives — the + /// `group_size == block` case of [`Self::group_reduce_via_subgroups`]. + pub(crate) fn workgroup_reduce_via_subgroups( + &mut self, + op: TileReduceOp, + subgroup_size: u32, + value: Tile, + ) -> Tile { + let block = self.block_size(); + self.group_reduce_via_subgroups(op, subgroup_size, block, value) + } + + /// Cross-lane reduction over lane groups of `group_size`, built from + /// subgroup collectives: one per-subgroup reduce, the per-subgroup + /// partials staged through a `num_subgroups`-sized workgroup array, and a + /// serial fold of the `group_size / subgroup_size` partials that belong to + /// the lane's own group. Two barriers total whatever the group size, + /// versus one per tree level in [`Self::group_reduce`] — and none at all + /// when a group is exactly one subgroup. The caller owns the device gating + /// (a fixed `subgroup_size`) via `SubgroupToken`. + /// + /// Groups are subgroup-aligned by construction: group `g` owns subgroup + /// ids `[g * s, (g + 1) * s)`. Callers packing several rows into one + /// workgroup must derive the row a lane serves from `subgroup_id` the same + /// way, because the mapping from `local_invocation_index` onto subgroups + /// is implementation defined. + pub(crate) fn group_reduce_via_subgroups( + &mut self, + op: TileReduceOp, + subgroup_size: u32, + group_size: u32, + value: Tile, + ) -> Tile { + let block = self.block_size(); + assert!( + subgroup_size > 0 + && group_size.is_multiple_of(subgroup_size) + && block.is_multiple_of(group_size), + "group_reduce_via_subgroups requires subgroup-aligned groups tiling the block" + ); + let element = value.element(); + let partial = self.subgroup_reduce(op, value); + let per_group = group_size / subgroup_size; + if per_group == 1 { + return self.bind(partial); + } + let partial = self.bind(partial); + let num_subgroups = block / subgroup_size; + let scratch = self.program.alloc_tile( + element, + Layout::contiguous(MemoryLevel::Workgroup, Shape::new([num_subgroups])), + ); + // Barrier before seeding the scratch: a previous reduction through + // the same call site (a reduce inside a loop) may still have lanes + // reading the prior partials. + self.workgroup_barrier(); + let subgroup_lane = self.subgroup_lane(); + let subgroup_id = self.subgroup_id(); + self.if_then(subgroup_lane.eq(0u32), |program| { + program.store_workgroup(&scratch, subgroup_id.clone(), partial); + }); + self.workgroup_barrier(); + // The whole block is one group: every lane folds the array from slot + // zero, so the base index stays a literal. + let base = (per_group != num_subgroups).then(|| { + let group_base = subgroup_id / per_group * per_group; + self.bind(group_base) + }); + let mut combined = match &base { + Some(base) => self.load_workgroup(&scratch, base.clone()), + None => self.load_workgroup(&scratch, 0u32), + }; + for index in 1..per_group { + let next = match &base { + Some(base) => self.load_workgroup(&scratch, base.clone() + index), + None => self.load_workgroup(&scratch, index), + }; + combined = combined.binary(op.binary(), next); + } + self.bind(combined) + } + + /// Cross-lane reduction over lane groups of `group_size` under a combine + /// the caller supplies, rather than one of the closed [`TileReduceOp`]s. + /// + /// Subgroup collectives are per-operator, so they cannot serve a general + /// monoid. This stages every lane's partial through a block-sized + /// workgroup array and has each lane fold its own group's slice, which + /// needs no intrinsic and leaves every lane holding the combined value — + /// the same broadcast contract as [`Self::group_reduce`]. + /// + /// `combine` must be associative over the staged values; it is invoked + /// `group_size - 1` times per lane in a fixed left-to-right order, so a + /// non-associative body silently produces an order-dependent result. + pub fn group_reduce_with(&mut self, group_size: u32, value: Tile, mut combine: F) -> Tile + where + F: FnMut(&mut Self, Tile, Tile) -> Tile, + { + let mut combined = self.group_reduce_with_vec( + group_size, + vec![value], + |program, mut acc, mut incoming| { + let acc = acc.pop().expect("one slot"); + let incoming = incoming.pop().expect("one slot"); + vec![combine(program, acc, incoming)] + }, + ); + combined.pop().expect("one slot") + } + + /// Cross-lane reduction of an `N`-slot carrier under a joint combine — + /// the tuple-valued [`Self::group_reduce_with`]. + /// + /// A joint carrier cannot be reduced slot by slot: every outgoing slot may + /// read every incoming one, which is exactly what makes online softmax a + /// carrier rather than three independent reductions (its normalizer is + /// rescaled by a factor derived from *both* sides' maxima). So each slot + /// gets its own block-sized staging array and the whole tuple folds in one + /// pass, leaving every lane holding the combined carrier. + /// + /// `combine` must be associative over the staged carriers; it is invoked + /// `group_size - 1` times per lane in a fixed left-to-right order. + pub fn group_reduce_with_vec( + &mut self, + group_size: u32, + values: Vec, + mut combine: F, + ) -> Vec + where + F: FnMut(&mut Self, Vec, Vec) -> Vec, + { + let block = self.block_size(); + assert!(!values.is_empty(), "a carrier needs at least one slot"); + assert!( + group_size > 0 && group_size <= block && block.is_multiple_of(group_size), + "group_reduce_with requires lane groups tiling the block" + ); + if group_size == 1 { + return values.into_iter().map(|value| self.bind(value)).collect(); + } + let scratch: Vec<_> = values + .iter() + .map(|value| { + self.program.alloc_tile( + value.element(), + Layout::contiguous(MemoryLevel::Workgroup, Shape::new([block])), + ) + }) + .collect(); + let lane = self.lane(); + // Barrier before seeding: an earlier reduction through this same call + // site (a reduce inside a loop) may still have lanes reading the + // previous round's partials. Mirrors `group_reduce_via_subgroups`. + self.workgroup_barrier(); + for (slot, value) in scratch.iter().zip(values) { + self.store_workgroup(slot, lane.clone(), value); + } + self.workgroup_barrier(); + let base = self.bind(lane / group_size * group_size); + let mut combined: Vec = scratch + .iter() + .map(|slot| self.load_workgroup(slot, base.clone())) + .collect(); + for index in 1..group_size { + let next: Vec = scratch + .iter() + .map(|slot| self.load_workgroup(slot, base.clone() + index)) + .collect(); + combined = combine(self, combined, next); + assert_eq!( + combined.len(), + scratch.len(), + "joint combine returned the wrong carrier width" + ); + } + combined.into_iter().map(|value| self.bind(value)).collect() + } + fn reduce(&mut self, op: TileReduceOp, value: Tile) -> Tile { let block = self.block_size(); self.group_reduce(op, block, value) diff --git a/fusor-ml/tile-ir/src/tile/value.rs b/fusor-ml/tile-ir/src/tile/value.rs index 47f1f5d42..998115c83 100644 --- a/fusor-ml/tile-ir/src/tile/value.rs +++ b/fusor-ml/tile-ir/src/tile/value.rs @@ -186,6 +186,10 @@ impl Tile { pub fn exp2(self) -> Self { self.unary(TileUnaryOp::Exp2) } + /// Natural logarithm. + pub fn log(self) -> Self { + self.unary(TileUnaryOp::Log) + } /// Hyperbolic tangent. pub fn tanh(self) -> Self { self.unary(TileUnaryOp::Tanh) diff --git a/fusor-ml/tile-ir/tests/golden_modules/f16_workgroup_tile_copy.txt b/fusor-ml/tile-ir/tests/golden_modules/f16_workgroup_tile_copy.txt index 897d0953f..899541bea 100644 --- a/fusor-ml/tile-ir/tests/golden_modules/f16_workgroup_tile_copy.txt +++ b/fusor-ml/tile-ir/tests/golden_modules/f16_workgroup_tile_copy.txt @@ -223,6 +223,7 @@ Module { memory_decorations: MemoryDecorations( 0x0, ), + workgroup_alias: None, }, [1]: GlobalVariable { name: None, @@ -242,6 +243,7 @@ Module { memory_decorations: MemoryDecorations( 0x0, ), + workgroup_alias: None, }, [2]: GlobalVariable { name: None, @@ -252,6 +254,7 @@ Module { memory_decorations: MemoryDecorations( 0x0, ), + workgroup_alias: None, }, }, global_expressions: {}, diff --git a/models/kalosm-llama/examples/profile_forward.rs b/models/kalosm-llama/examples/profile_forward.rs index 75e1ec539..7a4563f62 100644 --- a/models/kalosm-llama/examples/profile_forward.rs +++ b/models/kalosm-llama/examples/profile_forward.rs @@ -25,6 +25,16 @@ where println!( "llama_forward_profile tokens={tokens} elapsed={elapsed:?} per_token_ms={per_token_ms:.3}" ); + if tokens >= measured { + if let Ok(max_ms) = std::env::var("KALOSM_PROFILE_LLAMA_MAX_MS") { + if let Ok(max_ms) = max_ms.parse::() { + assert!( + per_token_ms <= max_ms, + "decode regression: {per_token_ms:.3} ms/token exceeds {max_ms:.3} ms/token" + ); + } + } + } } fn env_usize(name: &str, default: usize) -> usize { diff --git a/models/kalosm-llama/src/lib.rs b/models/kalosm-llama/src/lib.rs index 4d48ce4fa..d429e09a8 100644 --- a/models/kalosm-llama/src/lib.rs +++ b/models/kalosm-llama/src/lib.rs @@ -350,7 +350,6 @@ where pub struct LlamaBuilder { source: source::LlamaSource, device: Option, - flash_attn: bool, _marker: std::marker::PhantomData, } @@ -359,7 +358,6 @@ impl Default for LlamaBuilder { Self { source: Default::default(), device: None, - flash_attn: false, _marker: std::marker::PhantomData, } } @@ -377,12 +375,6 @@ impl LlamaBuilder { self } - /// Set whether to use Flash Attention. - pub fn with_flash_attn(mut self, use_flash_attn: bool) -> Self { - self.flash_attn = use_flash_attn; - self - } - /// Set the device to run the model with. (Defaults to an accelerator if available, otherwise the CPU) pub fn with_device(mut self, device: Device) -> Self { self.device = Some(device); diff --git a/models/kalosm-llama/src/raw/attention_layer.rs b/models/kalosm-llama/src/raw/attention_layer.rs index 104a3721b..a3617e5b6 100644 --- a/models/kalosm-llama/src/raw/attention_layer.rs +++ b/models/kalosm-llama/src/raw/attention_layer.rs @@ -316,7 +316,7 @@ where let query: Tensor<4, F> = query.cast(); if let Some(norm) = &self.attention_q_norm { - norm.forward_generic_4d(&query) + norm.forward_generic(&query) } else { query } @@ -332,7 +332,7 @@ where let key: Tensor<4, F> = key.cast(); if let Some(norm) = &self.attention_k_norm { - norm.forward_generic_4d(&key) + norm.forward_generic(&key) } else { key } @@ -375,7 +375,7 @@ where let query: Tensor<4, F> = query.cast(); if let Some(norm) = &self.attention_q_norm { - norm.forward_generic_4d(&query) + norm.forward_generic(&query) } else { query } @@ -395,7 +395,7 @@ where let key: Tensor<4, F> = key.cast(); if let Some(norm) = &self.attention_k_norm { - norm.forward_generic_4d(&key) + norm.forward_generic(&key) } else { key } @@ -629,7 +629,7 @@ where crate::raw::debug_check_nan_f32(&value_f32, layer_idx, "V_cache_view", start_pos); let scale = 1. / (head_dim as f64).sqrt(); - let attn_raw = query_f32.flash_attention( + let attn_raw = query_f32.attention( &key_f32, &value_f32, scale as f32, @@ -642,7 +642,7 @@ where (m.mask(), kind) }), ); - crate::raw::debug_check_nan_f32(&attn_raw, layer_idx, "flash_out", start_pos); + crate::raw::debug_check_nan_f32(&attn_raw, layer_idx, "attention_out", start_pos); let attn_output = attn_raw.transpose(1, 2); let attn_output = attn_output.reshape([b_sz, q_len, hidden_size]); let attn_output_f: Tensor<3, F> = attn_output.cast(); @@ -672,7 +672,7 @@ where f32: CastTo + CastTensor, { let scale = 1. / (head_dim as f64).sqrt(); - let attn_output = query_states.flash_attention( + let attn_output = query_states.attention( key_states, value_states, scale as f32, diff --git a/models/kalosm-llama/src/raw/vision/qwen.rs b/models/kalosm-llama/src/raw/vision/qwen.rs index 6eede3618..cde93ceec 100644 --- a/models/kalosm-llama/src/raw/vision/qwen.rs +++ b/models/kalosm-llama/src/raw/vision/qwen.rs @@ -414,7 +414,7 @@ where #[cfg(test)] #[tokio::test] async fn test_loading_qwen_vision() { - use fusor::ToVec2; + use fusor::ToVec; use kalosm_common::Cache; // Skip in CI @@ -469,7 +469,7 @@ async fn test_loading_qwen_vision() { let out = qwen_vision.rot_pos_emb(&vec![[2, 4, 4]]).unwrap(); let out_first_5_by_5 = out.slice([0..5, 0..5]).to_concrete(); let out_slice = out_first_5_by_5.as_slice().await.unwrap(); - let out_first_5_by_5: Vec> = out_slice.to_vec2(); + let out_first_5_by_5: Vec> = out_slice.to_vec(); println!("Rotary Pos Emb first 5x5: {:?}", out_first_5_by_5); let expected: [[f32; 5]; 5] = [ @@ -497,7 +497,7 @@ async fn test_loading_qwen_vision() { .unwrap(); let out_first_5_by_5 = out.slice([0..5, 0..5]).to_concrete(); let out_slice = out_first_5_by_5.as_slice().await.unwrap(); - println!("Qwen Vision: {:?}", out_slice.to_vec2()); + println!("Qwen Vision: {:?}", out_slice.to_vec()); // download image from https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg let image_bytes = @@ -515,7 +515,7 @@ async fn test_loading_qwen_vision() { let out = qwen_vision.forward_image(&pixels, grid).unwrap(); let out_first_5_by_5 = out.slice([0..5, 0..5]).to_concrete(); let out_slice = out_first_5_by_5.as_slice().await.unwrap(); - println!("Qwen Vision Image: {:?}", out_slice.to_vec2()); + println!("Qwen Vision Image: {:?}", out_slice.to_vec()); } fn generate_full_attention_blocks(block_count: usize, n_wa_pattern: u64) -> Vec { diff --git a/models/kalosm-llama/src/raw/vision/qwen_vision_block.rs b/models/kalosm-llama/src/raw/vision/qwen_vision_block.rs index e3e721417..42edd6aad 100644 --- a/models/kalosm-llama/src/raw/vision/qwen_vision_block.rs +++ b/models/kalosm-llama/src/raw/vision/qwen_vision_block.rs @@ -291,21 +291,21 @@ where } // The attention pattern is block-diagonal per `cu_seqlens`. The dense - // masked path runs the full M×M flash kernel even when most of the + // masked path runs the full M×M attention kernel even when most of the // mask is -inf — that's where ~95% of vision-encoder GPU time goes. // For window-attention layers we instead slice Q/K/V per window and - // run dense (unmasked) flash on each — same arithmetic the + // run dense (unmasked) attention on each — same arithmetic the // block-diagonal mask intended, but at window² scale (~64²) rather // than seq² (1944²). Full-attention layers (cu_seqlens=[0, seq]) - // skip the slicing and just run one regular flash call. - let t_flash = Instant::now(); + // skip the slicing and just run one regular attention call. + let t_attention = Instant::now(); let query_f32 = query_states; let key_f32 = key_states_f32; let value_f32 = value_states_f32; let is_full_attn = cu_seqlens.len() == 2 && cu_seqlens[0] == 0 && cu_seqlens[1] as usize == seq_len; let attn_out_4d = if is_full_attn { - query_f32.flash_attention( + query_f32.attention( &key_f32, &value_f32, 1.0 / (self.head_dim as f64).sqrt() as f32, @@ -345,7 +345,7 @@ where let k_run = key_f32.restride(run_specs).to_concrete(); let v_run = value_f32.restride(run_specs).to_concrete(); - let run_out = q_run.flash_attention(&k_run, &v_run, scale, None); + let run_out = q_run.attention(&k_run, &v_run, scale, None); let run_out = run_out .reshape([self.head_count, run_count * len, self.head_dim]) .unsqueeze(0) @@ -357,13 +357,13 @@ where fusor::cat(run_outputs, 2).to_concrete() }; - // After flash: transpose+reshape to [b_sz, seq, hidden_size] then proj. + // After attention: transpose+reshape to [b_sz, seq, hidden_size] then proj. let attn_output = attn_out_4d.transpose(1, 2); let attn_output = attn_output.reshape([bsz, seq_len, self.embed_dim]); let output: Tensor<3, F> = self.proj.forward_generic(&attn_output.cast()); if trace_attn { output.as_gpu().map(|g| g.materialize_sync()); - tracing::info!(" flash+proj: {:.2?}", t_flash.elapsed()); + tracing::info!(" attention+proj: {:.2?}", t_attention.elapsed()); } Ok(output) diff --git a/models/kalosm-llama/src/raw/vision/qwen_vision_embed.rs b/models/kalosm-llama/src/raw/vision/qwen_vision_embed.rs index c0d1d489a..1f0220eef 100644 --- a/models/kalosm-llama/src/raw/vision/qwen_vision_embed.rs +++ b/models/kalosm-llama/src/raw/vision/qwen_vision_embed.rs @@ -142,7 +142,7 @@ where #[cfg(test)] #[tokio::test] async fn test_vision_patch_embed() { - use fusor::ToVec2; + use fusor::ToVec; let embed_dim = 4; let in_channels = 3; @@ -239,7 +239,7 @@ async fn test_vision_patch_embed() { .unwrap() .cast::(); let output_slice = output.as_slice().await.unwrap(); - let output_vec = output_slice.to_vec2(); + let output_vec = output_slice.to_vec(); println!("Output: {output_vec:?}"); let expected_output = [[0.3058, 0.6866, -0.7391, -0.6952]]; assert_2d_vec_eq(output_vec, expected_output, 1e-2); diff --git a/models/rbert/src/language_model.rs b/models/rbert/src/language_model.rs index 7c0142e3b..e8baca8f6 100644 --- a/models/rbert/src/language_model.rs +++ b/models/rbert/src/language_model.rs @@ -8,7 +8,7 @@ use crate::BertBuilder; use crate::BertError; use crate::BertLoadingError; use crate::Pooling; -use fusor::ToVec2; +use fusor::ToVec; pub use kalosm_language_model::{ Embedder, EmbedderCacheExt, EmbedderExt, Embedding, EmbeddingInput, EmbeddingVariant, }; @@ -38,7 +38,7 @@ impl Bert { tensor: fusor::Tensor<2, f32>, ) -> Result { let slice = tensor.as_slice().await.map_err(BertError::Fusor)?; - let slice_data = slice.to_vec2(); + let slice_data = slice.to_vec(); Ok(Embedding::from( slice_data.into_iter().next().into_iter().next().unwrap(), )) diff --git a/models/rbert/src/raw/qwen/attention.rs b/models/rbert/src/raw/qwen/attention.rs index cd6fae354..52aba8426 100644 --- a/models/rbert/src/raw/qwen/attention.rs +++ b/models/rbert/src/raw/qwen/attention.rs @@ -74,10 +74,10 @@ impl QwenSelfAttention { // Apply optional Q/K normalization if let Some(ref q_norm) = self.q_norm { - query_states = q_norm.forward_4d(&query_states); + query_states = q_norm.forward(&query_states); } if let Some(ref k_norm) = self.k_norm { - key_states = k_norm.forward_4d(&key_states); + key_states = k_norm.forward(&key_states); } // Apply RoPE to Q and K @@ -87,9 +87,9 @@ impl QwenSelfAttention { let hidden_size = self.num_heads * self.head_dim; let scale = 1.0 / (self.head_dim as f32).sqrt(); - // Convert attention mask for flash attention if provided + // Convert the attention mask if provided. // The mask should be [b_sz, seq_len] where 1 = valid, 0 = pad - // Flash attention expects None for no mask, or a mask tensor + // Attention expects None for no mask, or a mask tensor // Note: We use a large negative value instead of NEG_INFINITY because // the GPU shader path does not support inf literals. -10000 is enough to effectively // zero out masked positions after softmax. @@ -105,7 +105,7 @@ impl QwenSelfAttention { ((ones - mask_f32) * MASK_NEG_VALUE).to_concrete() }); - let attn_output = query_states.flash_attention( + let attn_output = query_states.attention( &key_states, &value_states, scale, diff --git a/models/rbert/src/raw/self_attention.rs b/models/rbert/src/raw/self_attention.rs index 5bf1aef73..8ff084a53 100644 --- a/models/rbert/src/raw/self_attention.rs +++ b/models/rbert/src/raw/self_attention.rs @@ -84,7 +84,7 @@ impl BertSelfAttention { let attention_probs = attention_scores.softmax_last_dim::<3>(); attention_probs.mat_mul(&value_layer) } else { - query_layer.flash_attention(&key_layer, &value_layer, scale, None) + query_layer.attention(&key_layer, &value_layer, scale, None) }; let context_layer = context_layer.transpose(1, 2).to_concrete(); context_layer.flatten_last_n::<1, _>() diff --git a/models/rwhisper/src/quantized/mod.rs b/models/rwhisper/src/quantized/mod.rs index b5dbe5714..6c7df9707 100644 --- a/models/rwhisper/src/quantized/mod.rs +++ b/models/rwhisper/src/quantized/mod.rs @@ -4,7 +4,7 @@ use std::{num::NonZeroUsize, sync::Arc}; use fusor::{ cache::{AttentionMask, KvCache, MaskCache, TensorCache}, - layers::{Conv1d, Conv1dConfig, Embedding, LayerNorm, Linear}, + layers::{ConvNd, ConvNdConfig, Embedding, LayerNorm, Linear}, Device, Error, Result, Tensor, VarBuilder, }; use timestamps::extract_timestamps; @@ -14,13 +14,13 @@ use crate::config::Config; pub(crate) mod timestamps; fn conv1d( - config: Conv1dConfig, + config: ConvNdConfig<1>, device: &Device, vb: &mut VarBuilder, in_channels: usize, out_channels: usize, kernel_size: usize, -) -> Result> { +) -> Result> { let weight_q = vb.get("weight", device)?; let weight_shape = weight_q.shape(); @@ -51,7 +51,7 @@ fn conv1d( bias_2d.squeeze(1).to_concrete() } }; - Ok(Conv1d::new(weight, Some(bias), config)) + Ok(ConvNd::new(weight, Some(bias), config)) } struct MultiHeadAttentionCache { @@ -310,8 +310,8 @@ fn sinusoids(length: usize, channels: usize, device: &Device) -> Tensor<2, crate // https://github.com/openai/whisper/blob/f572f2161ba831bae131364c3bffdead7af6d210/whisper/model.py#L143 pub struct AudioEncoder { - conv1: Conv1d, - conv2: Conv1d, + conv1: ConvNd<1, 3, crate::WhisperDType>, + conv2: ConvNd<1, 3, crate::WhisperDType>, positional_embedding: Tensor<2, crate::WhisperDType>, blocks: Vec, ln_post: LayerNorm<1, crate::WhisperDType>, @@ -328,17 +328,15 @@ impl AudioEncoder { let n_state = cfg.d_model; let n_head = cfg.encoder_attention_heads; let n_ctx = cfg.max_source_positions; - let cfg1 = Conv1dConfig { - padding: 1, - stride: 1, + let cfg1 = ConvNdConfig { + padding: [1], + stride: [1], groups: 1, - dilation: 1, }; - let cfg2 = Conv1dConfig { - padding: 1, - stride: 2, + let cfg2 = ConvNdConfig { + padding: [1], + stride: [2], groups: 1, - dilation: 1, }; let n_mels = cfg.num_mel_bins; let conv1 = conv1d(cfg1, device, &mut vb.pp("conv1"), n_mels, n_state, 3)?; diff --git a/models/segment-anything-rs/src/lib.rs b/models/segment-anything-rs/src/lib.rs index 361748e1e..582f5e3ce 100644 --- a/models/segment-anything-rs/src/lib.rs +++ b/models/segment-anything-rs/src/lib.rs @@ -24,7 +24,7 @@ mod mask_generation; mod raw; -use fusor::{Concrete, Device, Tensor, ToVec1, VarBuilder}; +use fusor::{Concrete, Device, Tensor, ToVec, VarBuilder}; use image::{DynamicImage, GenericImage, GenericImageView, ImageBuffer, Rgba}; use kalosm_model_types::FileSource; use mask_generation::LowResMaskBatch; @@ -322,7 +322,7 @@ impl SegmentAnything { let mask_hwc = mask_t1.transpose(1, 2); // (H, W, 3); let mask_flat = mask_hwc.reshape([h * w * 3]); let mask_slice = mask_flat.as_slice().await?; - let mask_pixels: Vec = mask_slice.to_vec1().iter().map(|&v| v as u8).collect(); + let mask_pixels: Vec = mask_slice.to_vec().iter().map(|&v| v as u8).collect(); let mask_img: image::ImageBuffer, Vec> = image::ImageBuffer::from_raw(w as u32, h as u32, mask_pixels) @@ -427,12 +427,12 @@ impl SegmentAnything { let masks_flat = low_res_masks.reshape([total_mask_elems]); let masks_slice = masks_flat.as_slice().await?; - let masks_vec = masks_slice.to_vec1(); + let masks_vec = masks_slice.to_vec(); let total_iou_elems = batch * n_masks_per_point; let iou_flat = iou_preds.reshape([total_iou_elems]); let iou_slice = iou_flat.as_slice().await?; - let iou_vec = iou_slice.to_vec1(); + let iou_vec = iou_slice.to_vec(); mask_generation::collect_mask_candidates( LowResMaskBatch { diff --git a/models/segment-anything-rs/src/raw/mask_decoder.rs b/models/segment-anything-rs/src/raw/mask_decoder.rs index c54860d3c..45a63ea73 100644 --- a/models/segment-anything-rs/src/raw/mask_decoder.rs +++ b/models/segment-anything-rs/src/raw/mask_decoder.rs @@ -118,7 +118,7 @@ impl MlpMaskDecoder { fn forward(&self, xs: &Tensor<2, f32>) -> Tensor<2, f32> { let mut xs = xs.to_concrete(); for (i, layer) in self.layers.iter().enumerate() { - xs = layer.forward_2d(&xs); + xs = layer.forward(&xs); if i + 1 < self.layers.len() { xs = xs.relu(); } diff --git a/models/segment-anything-rs/src/raw/transformer.rs b/models/segment-anything-rs/src/raw/transformer.rs index d21cee7b5..8deb1006c 100644 --- a/models/segment-anything-rs/src/raw/transformer.rs +++ b/models/segment-anything-rs/src/raw/transformer.rs @@ -80,7 +80,7 @@ impl Attention { let c_per_head = q.shape()[3]; let scale = 1.0 / (c_per_head as f32).sqrt(); - let out = q.flash_attention(&k, &v, scale, None); + let out = q.attention(&k, &v, scale, None); let out = self.recombine_heads(&out); self.out_proj.forward(&out) }