hrx/tensor/
linear.rs

1use super::{DType, DeviceTensor, TensorDesc, TensorOps};
2use crate::{
3    Error, Result,
4    loom::Specialization,
5    model::{Command, Dispatch, ModelFragment, ModelSession},
6};
7
8impl TensorOps {
9    /// Single-row BF16 `x * weightᵀ`, accumulating in F32. Weights are contiguous
10    /// `[N, K]`; input is `[1, K]` with K a multiple of 128. Output may be BF16 or
11    /// F32. This bandwidth-oriented decode operation does not accept biases.
12    pub fn linear(
13        &self,
14        input: &DeviceTensor,
15        weight: &DeviceTensor,
16        output: DType,
17    ) -> Result<DeviceTensor> {
18        Ok(self
19            .linear_many(input, std::slice::from_ref(weight), output)?
20            .remove(0))
21    }
22
23    /// One dispatch for one, two or three single-row projections, e.g. Q/K/V or
24    /// gate/up. Weights and output widths can differ; K and output dtype agree.
25    /// Returned tensors retain a shared plan-slot lease. Three live results per
26    /// shape are supported; release one before requesting a fourth.
27    pub fn linear_many(
28        &self,
29        input: &DeviceTensor,
30        weights: &[DeviceTensor],
31        output: DType,
32    ) -> Result<Vec<DeviceTensor>> {
33        self.context.validate(input)?;
34        for weight in weights {
35            self.context.validate(weight)?;
36        }
37        let descriptions: Vec<_> = weights.iter().map(|w| w.desc().clone()).collect();
38        let key = (input.desc().clone(), descriptions.clone(), output);
39        let plan = self.linear.get_or_prepare(key, || {
40            self.linear_fragment(input.desc(), &descriptions, output)?
41                .prepare(3)
42        })?;
43        let mut bindings = vec![input.clone()];
44        bindings.extend_from_slice(weights);
45        Ok(plan.submit(&bindings)?.outputs().to_vec())
46    }
47
48    /// Record the same decode projections in a caller-owned inference graph.
49    /// Fragment inputs are `[input, weights...]`, outputs follow weight order.
50    pub fn linear_fragment(
51        &self,
52        input: &TensorDesc,
53        weights: &[TensorDesc],
54        output: DType,
55    ) -> Result<ModelFragment> {
56        let valid = input.dtype() == DType::BF16
57            && input.is_contiguous()
58            && input.shape().len() == 2
59            && input.shape()[0] == 1
60            && input.shape()[1] > 0
61            && input.shape()[1].is_multiple_of(128)
62            && (1..=3).contains(&weights.len())
63            && matches!(output, DType::BF16 | DType::F32)
64            && weights.iter().all(|w| {
65                w.dtype() == DType::BF16
66                    && w.is_contiguous()
67                    && w.shape().len() == 2
68                    && w.shape()[0] > 0
69                    && w.shape()[1] == input.shape()[1]
70                    && w.elements() <= 1073741824
71            });
72        if !valid {
73            return Err(Error::Message("linear requires BF16 [1,K], one to three contiguous BF16 [N,K] weights, K divisible by 128, and BF16/F32 output".into()));
74        }
75        let k = input.shape()[1];
76        let grid = weights
77            .iter()
78            .map(|w| w.shape()[0].div_ceil(8))
79            .max()
80            .unwrap();
81        let mut model = ModelSession::in_context(&self.context)?;
82        let x = model.allocate(input.bytes())?;
83        let mut inputs = vec![(x, input.clone())];
84        let mut outputs = Vec::new();
85        let mut bindings = vec![x.read()];
86        let mut source = "amdgpu.target<gfx11-generic> @target {subgroup_size = 32}\n".to_owned();
87        let mut spec = Specialization::new("hrx_linear");
88        for (key, value) in [("k", k), ("grid_x", grid), ("grid_y", weights.len())] {
89            source += &format!(
90                "config.decl @hrx.linear.{key} : %value: index where [range(%value, 1, 1073741824)]\n"
91            );
92            spec.set_config(format!("hrx.linear.{key}"), value.to_string());
93        }
94        for (i, w) in weights.iter().enumerate() {
95            for (key, value) in [
96                (format!("n{i}"), w.shape()[0]),
97                (format!("wn{i}"), w.elements()),
98            ] {
99                source += &format!(
100                    "config.decl @hrx.linear.{key} : %value: index where [range(%value, 1, 1073741824)]\n"
101                );
102                spec.set_config(format!("hrx.linear.{key}"), value.to_string());
103            }
104            let weight = model.allocate(w.bytes())?;
105            let desc = TensorDesc::new(output, vec![1, w.shape()[0]])?;
106            let out = model.allocate(desc.bytes())?;
107            inputs.push((weight, w.clone()));
108            outputs.push((out, desc));
109            bindings.extend([weight.read(), out.write()]);
110        }
111        source += r#"kernel.def target(@target) export("hrx_linear") @hrx_linear(%count: index) {
112 %one = index.constant 1 : index
113 %threads = index.constant 256 : index
114 %gx = config.get @hrx.linear.grid_x : index
115 %gy = config.get @hrx.linear.grid_y : index
116 kernel.launch.config workgroups(%gx, %gy, %one) workgroup_size(%threads, %one, %one) : index
117} launch(%count: index, %x: buffer, "#;
118        source += &(0..weights.len())
119            .map(|i| format!("%w{i}: buffer, %out{i}: buffer"))
120            .collect::<Vec<_>>()
121            .join(", ");
122        source += ") {\n %projection = kernel.workgroup.id<y> : index\n";
123        let (dtype, convert) = if output == DType::BF16 {
124            ("bf16", "scalar.fptrunc %sum1 : f32 to bf16")
125        } else {
126            ("f32", "scalar.addf %sum1, %zf : f32")
127        };
128        for i in 0..weights.len() {
129            source += &format!(
130                " %p{i} = index.constant {i} : index\n %is{i} = index.cmp eq, %projection, %p{i} : index\n scf.if %is{i} {{\n"
131            );
132            source += &include_str!("linear.loom")
133                .replace("@P@", &i.to_string())
134                .replace("@TYPE@", dtype)
135                .replace("@CONVERT@", convert);
136            source += "\n }\n";
137        }
138        source += " kernel.return\n}\n";
139        // Checked shapes bound every input/output view; each wave writes one row.
140        let kernel = unsafe { model.compile(&[(&source, spec)])? }[0];
141        unsafe {
142            model.freeze(&self.context)?.fragment(
143                &[Command::Dispatch(Dispatch::indices(
144                    kernel,
145                    [1],
146                    [grid as u32, weights.len() as u32, 1],
147                    bindings,
148                ))],
149                &inputs,
150                &outputs,
151            )
152        }
153    }
154}
hrx/tensor/linear.rs

hrx/tensor/
linear.rs
