tenferro_gpu/webgpu/
exec_session.rs

1use std::any::TypeId;
2use tenferro_tensor::backend::{
3    BackendSession, BackendSessionHost, ElementwiseFusionPlan, GroupedGemmConfig, SessionCachedDot,
4    TensorAnalytic, TensorBuffer, TensorDeviceTransfer, TensorDot, TensorElementwise, TensorFusion,
5    TensorIndexing, TensorReduction, TensorStructural,
6};
7use tenferro_tensor::config::{
8    CompareDir, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
9};
10use tenferro_tensor::{
11    with_session_entry_guard, DotGeneralAccumulation, Tensor, TensorRead, TensorValue, TensorWrite,
12};
13
14use super::{WebGpuBackend, WebGpuRuntime, WebGpuRuntimeIdentity};
15
16/// Marker for the concrete erased WebGPU execution-session target.
17#[doc(hidden)]
18pub(super) struct WebGpuExecSessionMarker;
19
20/// Borrowed WebGPU execution capability.
21#[doc(hidden)]
22#[derive(Debug)]
23pub struct WebGpuExecSession<'a> {
24    backend: &'a mut WebGpuBackend,
25}
26
27impl WebGpuExecSession<'_> {
28    /// Borrow the provider runtime without exposing the owning backend.
29    #[doc(hidden)]
30    pub fn runtime(&self) -> &WebGpuRuntime {
31        self.backend.runtime()
32    }
33
34    /// Return the identity of the borrowed provider runtime.
35    #[doc(hidden)]
36    pub fn runtime_identity(&self) -> WebGpuRuntimeIdentity {
37        self.backend.runtime_identity()
38    }
39}
40
41/// Visit a WebGPU execution session through the erased backend-session surface.
42///
43/// The callback receives only the lifetime-bound session capability. The
44/// owning [`WebGpuBackend`] never crosses this boundary.
45#[doc(hidden)]
46pub fn with_webgpu_exec_session<B, R>(
47    session: &mut B,
48    f: impl for<'a> FnOnce(&'a mut WebGpuExecSession<'a>) -> R,
49) -> Option<R>
50where
51    B: BackendSession + ?Sized,
52{
53    if session.session_type_id() != std::any::TypeId::of::<WebGpuExecSessionMarker>() {
54        return None;
55    }
56    let data = unsafe { session.session_data_mut() };
57    // SAFETY: the exact marker check and BackendSession erased-pointer contract
58    // identify the value as WebGpuExecSession for this scoped visit.
59    Some(unsafe { f(&mut *(data.cast::<WebGpuExecSession<'static>>())) })
60}
61
62macro_rules! delegate {
63    ($trait:path {
64        $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) -> $ret:ty;)*
65    }) => {
66        impl $trait for WebGpuExecSession<'_> {
67            $(
68                fn $method(&mut self, $($arg: $arg_ty),*) -> $ret {
69                    self.backend.$method($($arg),*)
70                }
71            )*
72        }
73    };
74}
75
76delegate!(TensorElementwise {
77    fn add(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
78    fn sub(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
79    fn mul(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
80    fn neg(input: &Tensor) -> crate::Result<Tensor>;
81    fn conj(input: &Tensor) -> crate::Result<Tensor>;
82    fn div(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
83    fn abs(input: &Tensor) -> crate::Result<Tensor>;
84    fn sign(input: &Tensor) -> crate::Result<Tensor>;
85    fn maximum(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
86    fn minimum(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
87    fn compare(lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> crate::Result<Tensor>;
88    fn select(pred: &Tensor, on_true: &Tensor, on_false: &Tensor) -> crate::Result<Tensor>;
89    fn clamp(input: &Tensor, lower: &Tensor, upper: &Tensor) -> crate::Result<Tensor>;
90});
91
92delegate!(TensorAnalytic {
93    fn exp(input: &Tensor) -> crate::Result<Tensor>;
94    fn log(input: &Tensor) -> crate::Result<Tensor>;
95    fn sin(input: &Tensor) -> crate::Result<Tensor>;
96    fn cos(input: &Tensor) -> crate::Result<Tensor>;
97    fn tanh(input: &Tensor) -> crate::Result<Tensor>;
98    fn sqrt(input: &Tensor) -> crate::Result<Tensor>;
99    fn rsqrt(input: &Tensor) -> crate::Result<Tensor>;
100    fn pow(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
101    fn expm1(input: &Tensor) -> crate::Result<Tensor>;
102    fn log1p(input: &Tensor) -> crate::Result<Tensor>;
103});
104
105delegate!(TensorStructural {
106    fn to_contiguous_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
107    fn copy_read_into(src: TensorRead<'_>, dst: TensorWrite<'_>) -> crate::Result<()>;
108    fn transpose(input: &Tensor, perm: &[usize]) -> crate::Result<Tensor>;
109    fn reshape(input: &Tensor, shape: &[usize]) -> crate::Result<Tensor>;
110    fn broadcast_in_dim(input: &Tensor, shape: &[usize], dims: &[usize]) -> crate::Result<Tensor>;
111    fn cast(input: &Tensor, to: tenferro_tensor::DType) -> crate::Result<Tensor>;
112    fn extract_diagonal(input: &Tensor, axis_a: usize, axis_b: usize) -> crate::Result<Tensor>;
113    fn embed_diagonal(input: &Tensor, axis_a: usize, axis_b: usize) -> crate::Result<Tensor>;
114    fn tril(input: &Tensor, k: i64) -> crate::Result<Tensor>;
115    fn triu(input: &Tensor, k: i64) -> crate::Result<Tensor>;
116});
117
118delegate!(TensorReduction {
119    fn reduce_sum(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
120    fn reduce_prod(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
121    fn reduce_max(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
122    fn reduce_min(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
123});
124
125delegate!(TensorDot {
126    fn dot_general(lhs: &Tensor, rhs: &Tensor, config: &DotGeneralConfig) -> crate::Result<Tensor>;
127    fn dot_general_with_conj(
128        lhs: &Tensor,
129        rhs: &Tensor,
130        config: &DotGeneralConfig,
131        lhs_conj: bool,
132        rhs_conj: bool,
133    ) -> crate::Result<Tensor>;
134});
135
136delegate!(TensorIndexing {
137    fn gather(
138        operand: &Tensor,
139        start_indices: &Tensor,
140        config: &GatherConfig,
141    ) -> crate::Result<Tensor>;
142    fn scatter(
143        operand: &Tensor,
144        scatter_indices: &Tensor,
145        updates: &Tensor,
146        config: &ScatterConfig,
147    ) -> crate::Result<Tensor>;
148    fn slice(input: &Tensor, config: &SliceConfig) -> crate::Result<Tensor>;
149    fn dynamic_slice(
150        input: &Tensor,
151        starts: &Tensor,
152        slice_sizes: &[usize],
153    ) -> crate::Result<Tensor>;
154    fn dynamic_update_slice(
155        operand: &Tensor,
156        update: &Tensor,
157        starts: &Tensor,
158    ) -> crate::Result<Tensor>;
159    fn pad(input: &Tensor, config: &PadConfig) -> crate::Result<Tensor>;
160    fn concatenate(inputs: &[&Tensor], axis: usize) -> crate::Result<Tensor>;
161    fn reverse(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
162});
163
164delegate!(TensorFusion {
165    fn execute_elementwise_fusion(
166        inputs: &[&Tensor],
167        plan: &ElementwiseFusionPlan,
168    ) -> crate::Result<Option<Vec<Tensor>>>;
169    fn execute_broadcast_multiply(
170        lhs: TensorRead<'_>,
171        lhs_shape: &[usize],
172        lhs_dims: &[usize],
173        rhs: TensorRead<'_>,
174        rhs_shape: &[usize],
175        rhs_dims: &[usize],
176    ) -> crate::Result<Option<Tensor>>;
177    fn execute_broadcast_multiply_value(
178        lhs: TensorRead<'_>,
179        lhs_shape: &[usize],
180        lhs_dims: &[usize],
181        rhs: TensorRead<'_>,
182        rhs_shape: &[usize],
183        rhs_dims: &[usize],
184    ) -> crate::Result<Option<TensorValue>>;
185});
186
187delegate!(TensorBuffer {
188    fn reclaim_buffer(tensor: Tensor) -> ();
189});
190
191delegate!(TensorDeviceTransfer {
192    fn download_to_host(tensor: TensorRead<'_>) -> crate::Result<Tensor>;
193    fn upload_host_tensor(tensor: TensorRead<'_>) -> crate::Result<Tensor>;
194});
195
196macro_rules! delegate_cached {
197    ($(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) -> $ret:ty;)*) => {
198        impl SessionCachedDot for WebGpuExecSession<'_> {
199            $(
200                fn $method(&mut self, $($arg: $arg_ty),*) -> $ret {
201                    <WebGpuBackend as SessionCachedDot>::$method(self.backend, $($arg),*)
202                }
203            )*
204        }
205    };
206}
207
208delegate_cached! {
209    fn dot_general_cached(
210        cache_slot: Option<usize>,
211        lhs: &Tensor,
212        rhs: &Tensor,
213        config: &DotGeneralConfig,
214    ) -> crate::Result<Tensor>;
215    fn dot_general_read_cached(
216        cache_slot: Option<usize>,
217        lhs: TensorRead<'_>,
218        rhs: TensorRead<'_>,
219        config: &DotGeneralConfig,
220    ) -> crate::Result<Tensor>;
221    fn dot_general_with_conj_cached(
222        cache_slot: Option<usize>,
223        lhs: &Tensor,
224        rhs: &Tensor,
225        config: &DotGeneralConfig,
226        lhs_conj: bool,
227        rhs_conj: bool,
228    ) -> crate::Result<Tensor>;
229    fn dot_general_with_conj_read_cached(
230        cache_slot: Option<usize>,
231        lhs: TensorRead<'_>,
232        rhs: TensorRead<'_>,
233        config: &DotGeneralConfig,
234        lhs_conj: bool,
235        rhs_conj: bool,
236    ) -> crate::Result<Tensor>;
237    fn dot_general_read_into_accum_cached(
238        cache_slot: Option<usize>,
239        lhs: TensorRead<'_>,
240        rhs: TensorRead<'_>,
241        config: &DotGeneralConfig,
242        accumulation: DotGeneralAccumulation,
243        out: TensorWrite<'_>,
244    ) -> crate::Result<()>;
245    fn grouped_gemm_cached(
246        cache_slot: Option<usize>,
247        lhs: TensorRead<'_>,
248        rhs: TensorRead<'_>,
249        config: &GroupedGemmConfig<'_>,
250        out: TensorWrite<'_>,
251    ) -> crate::Result<()>;
252}
253
254impl BackendSession for WebGpuExecSession<'_> {
255    fn session_type_id(&self) -> TypeId {
256        TypeId::of::<WebGpuExecSessionMarker>()
257    }
258
259    unsafe fn session_data_mut(&mut self) -> *mut () {
260        self as *mut Self as *mut ()
261    }
262}
263
264impl BackendSessionHost for WebGpuBackend {
265    fn with_backend_session<R: Send>(
266        &mut self,
267        f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
268    ) -> R {
269        let mut session = WebGpuExecSession { backend: self };
270        // Nested entry is caught by the portable in-session guard in debug
271        // builds; the WebGPU runtime must never re-enter a session closure.
272        with_session_entry_guard(|| f(&mut session))
273    }
274}
tenferro_gpu/webgpu/exec_session.rs

tenferro_gpu/webgpu/
exec_session.rs
