rlx_wgpu/backend/
set.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! `set` — extracted from the `backend` module for navigability (see `mod.rs`).
17
18#![allow(unused_imports)]
19
20use crate::buffer::{
21    Arena, ReadbackLayout, ReadbackStaging, TinyReadbackStaging, decode_mapped_readback_f32,
22    decode_tiny_mapped_f32, encode_readback_copies, plan_f32_uniform, read_f32_many_pooled,
23    schedule_readback_map, use_tiny_readback, wait_readback_map,
24};
25use crate::device::wgpu_device;
26use crate::kernels::{
27    ArgmaxParams, AttentionBwdParams, AttentionParams, BatchElementwiseRegionParams, BinaryParams,
28    Conv1dParams, Conv2dParams, Conv3dParams, CopyParams, CumsumBwdParams, CumsumParams,
29    DequantMatmulParams, ElementwiseRegionParams, ExpandParams, FmaParams, FusedResidualLnParams,
30    FusedResidualLnTeeParams, FusedResidualRmsNormParams, GatherAxisParams, GatherBwdParams,
31    GatherParams, GroupedMatmulParams, GruParams, Kernel, LayerNormBwdParams, LayerNormParams,
32    Mamba2Params, MatmulParams, MatmulQkvParams, NarrowConcatParams, Pool1dParams, Pool2dParams,
33    Pool3dParams, ReduceParams, RmsNormBwdParams, RnnParams, RopeBwdParams, RopeParams,
34    SampleParams, ScatterAddParams, SceParams, SelectiveScanParams, SoftmaxParams, TopKParams,
35    TransposeParams, UmapKnnParams, UnaryParams, WelchPeaksGpuParams, WhereParams, argmax_kernel,
36    attention_bwd_kernel, attention_kernel, batch_elementwise_region_kernel, binary_kernel,
37    cast_f32_to_f16_kernel, compare_kernel, concat_kernel, conv1d_kernel, conv2d_kernel,
38    conv3d_kernel, copy_kernel, cumsum_backward_kernel, cumsum_kernel, dequant_matmul_kernel,
39    elementwise_region_kernel, elementwise_region_spatial_kernel, expand_kernel, fma_kernel,
40    fused_residual_ln_kernel, fused_residual_ln_tee_kernel, fused_residual_rms_norm_kernel,
41    gather_axis_kernel, gather_backward_acc_kernel, gather_backward_zero_kernel, gather_kernel,
42    gather_split_kernel, grouped_matmul_kernel, gru_kernel,
43    layer_norm_backward_gamma_partial_kernel, layer_norm_backward_gamma_reduce_kernel,
44    layer_norm_backward_input_kernel, layernorm_kernel, mamba2_kernel,
45    matmul_coop_f16_vulkan_active_kernel, matmul_coop_f16_vulkan_kernel,
46    matmul_coop_f32_active_kernel, matmul_coop16_kernel, matmul_f16_compute_kernel,
47    matmul_f16w_kernel, matmul_kernel, matmul_qkv_coop_f16_vk_active_kernel,
48    matmul_qkv_coop_f16_vk_kernel, matmul_qkv_coop_f32_kernel, matmul_qkv_kernel,
49    matmul_wide_active_kernel, matmul_wide_kernel, narrow_kernel, pool1d_kernel, pool2d_kernel,
50    pool3d_kernel, reduce_kernel, rms_norm_backward_kernel, rms_norm_backward_param_kernel,
51    rnn_kernel, rope_backward_kernel, rope_kernel, sample_kernel, scatter_add_kernel,
52    selective_scan_kernel, softmax_cross_entropy_kernel, softmax_kernel, topk_kernel,
53    transpose_kernel, umap_knn_kernel, unary_f16_mirror_kernel, unary_kernel,
54    welch_peaks_gpu_kernel, where_kernel,
55};
56use rlx_ir::dynamic::{bind_graph, has_dynamic_dims, infer_bindings_from_f32_inputs, same_binding};
57use rlx_ir::op::{Activation, BinaryOp, CmpOp, MaskKind, ReduceOp};
58use rlx_ir::shape::DimBinding;
59use rlx_ir::{Graph, NodeId, Op};
60use std::collections::{HashMap, HashSet};
61use std::num::NonZeroU64;
62
63use super::*;
64
65impl WgpuExecutable {
66    /// Override RNG policy for in-graph random ops without recompiling.
67    pub fn set_rng(&mut self, rng: rlx_ir::RngOptions) {
68        *self.rng.write().expect("rng lock") = rng;
69    }
70
71    /// Hint the next `run` to process only the first `actual` rows
72    /// along the bucket axis (out of `upper`, the compile extent).
73    /// Honored when every Step is in the safe set. See PLAN L1.
74    pub fn set_active_extent(&mut self, extent: Option<(usize, usize)>) {
75        self.active_extent = extent;
76    }
77
78    pub fn set_param(&mut self, name: &str, data: &[f32]) {
79        const STASH_MAX_BYTES: usize = 16 * 1024 * 1024;
80        if data.len() * 4 <= STASH_MAX_BYTES {
81            self.stashed_params.insert(name.to_string(), data.to_vec());
82        }
83        if self.coop_f16_vk {
84            crate::coop_f16_vk::refresh_wide_b_flag(&mut self.coop_f16_vk_wide_b, name, data);
85        }
86        if self.unresolved.is_some() {
87            self.pending_params.insert(name.to_string(), data.to_vec());
88            return;
89        }
90        let dev = wgpu_device().expect("rlx-wgpu: device gone");
91        if let Some(&id) = self.param_offsets.get(name)
92            && self.arena.has(id)
93        {
94            self.arena.write_f32(&dev.queue, id, data);
95        }
96    }
97
98    /// Upload raw bytes for a Param. The bytes land tight-packed at
99    /// the param's slot offset — no f32 round-trip. Used for quantized
100    /// weights (int8 / int4) where the kernel reads the byte stream
101    /// via `bitcast<u32>` from the f32-typed arena.
102    pub fn set_param_bytes(&mut self, name: &str, data: &[u8]) {
103        if self.unresolved.is_some() {
104            self.pending_param_bytes
105                .insert(name.to_string(), data.to_vec());
106            return;
107        }
108        let dev = wgpu_device().expect("rlx-wgpu: device gone");
109        if let Some(&id) = self.param_offsets.get(name)
110            && self.arena.has(id)
111        {
112            dev.queue
113                .write_buffer(&self.arena.buffer, self.arena.offset(id) as u64, data);
114        }
115    }
116
117    pub fn set_gpu_handle_feed(&mut self, handle_name: &str, output_index: usize) {
118        self.gpu_handle_feeds
119            .insert(handle_name.to_string(), output_index);
120    }
121}
rlx_wgpu/backend/set.rs

rlx_wgpu/backend/
set.rs
