tenferro_gpu/cubecl/
exec_session.rs

1use cubecl::prelude::{CubeElement, CubePrimitive};
2use std::any::TypeId;
3use std::marker::PhantomData;
4use std::rc::Rc;
5use tenferro_tensor::backend::{
6    BackendSession, BackendSessionHost, ElementwiseFusionPlan, ElementwiseReadOp,
7    GroupedGemmConfig, SessionCachedDot, TensorAnalytic, TensorBuffer, TensorDeviceTransfer,
8    TensorDot, TensorElementwise, TensorFusion, TensorIndexing, TensorReduction, TensorStructural,
9};
10use tenferro_tensor::config::{
11    CompareDir, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
12};
13use tenferro_tensor::{
14    with_session_entry_guard, TensorRank, TensorScalar, TensorViewCanonicalization, TypedTensorView,
15};
16use tenferro_tensor::{
17    DotGeneralAccumulation, Tensor, TensorRead, TensorValue, TensorWrite, TypedTensor,
18};
19
20use super::identity::GpuExtensionCapability;
21use super::runtime::RawContextRestore;
22use super::{
23    raw, session_cubecl, CudaBackend, CudaDeviceInfo, CudaExtensionCache, CudaRuntime,
24    CudaRuntimeIdentity,
25};
26
27/// Best-effort exit flush for a `with_cubecl` session.
28///
29/// Flushes once eagerly (returned to the caller as an error if it fails) and
30/// once more on `Drop` so a panic/unwind path still drains pending CubeCL
31/// work.
32struct CubeclExitFlush<'a> {
33    op: &'static str,
34    client: &'a cubecl::client::ComputeClient<cubecl_cuda::CudaRuntime>,
35    flushed: bool,
36}
37
38impl<'a> CubeclExitFlush<'a> {
39    fn new(
40        op: &'static str,
41        client: &'a cubecl::client::ComputeClient<cubecl_cuda::CudaRuntime>,
42    ) -> Self {
43        Self {
44            op,
45            client,
46            flushed: false,
47        }
48    }
49
50    /// Flush now and return the typed result.
51    fn flush_now(&mut self) -> crate::Result<()> {
52        self.client
53            .flush()
54            .map_err(|err| crate::Error::backend_source(self.op, err))?;
55        self.flushed = true;
56        Ok(())
57    }
58}
59
60impl Drop for CubeclExitFlush<'_> {
61    fn drop(&mut self) {
62        if !self.flushed {
63            let _ = self.client.flush();
64        }
65    }
66}
67
68/// Marker for the concrete erased CUDA execution-session target.
69#[doc(hidden)]
70pub(super) struct CudaExecSessionMarker;
71
72/// Borrowed CUDA execution capability.
73///
74/// This is the single public execution-authority boundary for CUDA kernel
75/// extensions (issue #1597). External operation crates obtain it through
76/// [`with_cuda_exec_session`] and then borrow backend/device-scoped extension
77/// sessions via [`CudaExecSession::with_cubecl`] and
78/// [`CudaExecSession::with_raw`].
79///
80/// The session is not constructible by users and is `!Send + !Sync`: it
81/// carries thread-local execution capability. Success of an enrolled operation
82/// means the work was enqueued; only [`CudaExecSession::synchronize`] is a
83/// host barrier.
84#[derive(Debug)]
85pub struct CudaExecSession<'a> {
86    backend: &'a mut CudaBackend,
87    _not_send_sync: PhantomData<Rc<()>>,
88}
89
90impl CudaExecSession<'_> {
91    /// Borrow the provider runtime without exposing the backend.
92    pub fn runtime(&self) -> &CudaRuntime {
93        self.backend.runtime()
94    }
95
96    /// Return the identity of the borrowed provider runtime.
97    pub fn runtime_identity(&self) -> CudaRuntimeIdentity {
98        self.backend.runtime_identity()
99    }
100
101    /// Report whether this session supports a GPU extension capability.
102    ///
103    /// # Examples
104    ///
105    /// ```
106    /// use tenferro_gpu::cuda::{CudaExecSession, GpuExtensionCapability};
107    ///
108    /// // Method-call check only: `CudaExecSession` is not user-constructible, so
109    /// // the example asserts the method is callable from an external crate.
110    /// fn check(session: &CudaExecSession<'_>, capability: GpuExtensionCapability) -> bool {
111    ///     session.supports(capability)
112    /// }
113    /// let _ = check;
114    /// ```
115    pub fn supports(&self, capability: GpuExtensionCapability) -> bool {
116        self.backend.runtime().supports_extension(capability)
117    }
118
119    /// Borrow immutable metadata for the session's device.
120    ///
121    /// # Examples
122    ///
123    /// ```
124    /// use tenferro_gpu::cuda::CudaExecSession;
125    ///
126    /// // Method-call check only: `CudaExecSession` is not user-constructible, so
127    /// // the example asserts the method is callable from an external crate.
128    /// fn check(session: &CudaExecSession<'_>) {
129    ///     let _ = session.device_info();
130    /// }
131    /// let _ = check;
132    /// ```
133    pub fn device_info(&self) -> &CudaDeviceInfo {
134        self.backend.runtime().device_info()
135    }
136
137    /// Return the allocation ownership domain of this session.
138    ///
139    /// # Examples
140    ///
141    /// ```
142    /// use tenferro_gpu::cuda::CudaExecSession;
143    ///
144    /// // Method-call check only: `CudaExecSession` is not user-constructible, so
145    /// // the example asserts the method is callable from an external crate.
146    /// fn check(session: &CudaExecSession<'_>) -> tenferro_tensor::AllocationDomainId {
147    ///     session.allocation_domain()
148    /// }
149    /// let _ = check;
150    /// ```
151    pub fn allocation_domain(&self) -> tenferro_tensor::AllocationDomainId {
152        self.backend.runtime().allocation_domain()
153    }
154
155    /// Validate that a dense GPU tensor is resident on this exact session:
156    /// CubeCL-backed, same allocation domain, and placed on this runtime's
157    /// CUDA device. Rejects host tensors, foreign-backend buffers, and
158    /// foreign-runtime/device tensors without an implicit transfer.
159    ///
160    /// This is the credentialed public-seam residency guard for extension
161    /// crates that receive a session but must validate inputs before entering
162    /// a `with_raw`/`with_cubecl` sub-session.
163    ///
164    /// # Errors
165    ///
166    /// Returns [`crate::Error::RuntimeState`] when the tensor is not resident
167    /// on this exact session runtime/device.
168    ///
169    /// # Examples
170    ///
171    /// ```
172    /// use tenferro_gpu::cuda::CudaExecSession;
173    ///
174    /// // Method-call check only: `CudaExecSession` is not user-constructible.
175    /// fn check(session: &CudaExecSession<'_>, tensor: &tenferro_tensor::Tensor) -> tenferro_tensor::Result<()> {
176    ///     session.ensure_gpu_resident(tensor, "test.ensure_gpu_resident")
177    /// }
178    /// let _ = check;
179    /// ```
180    pub fn ensure_gpu_resident(&self, input: &Tensor, op: &'static str) -> crate::Result<()> {
181        match input {
182            Tensor::F32(t) => super::dispatch::ensure_resident_on_runtime(self.runtime(), t, op),
183            Tensor::F64(t) => super::dispatch::ensure_resident_on_runtime(self.runtime(), t, op),
184            Tensor::I32(t) => super::dispatch::ensure_resident_on_runtime(self.runtime(), t, op),
185            Tensor::I64(t) => super::dispatch::ensure_resident_on_runtime(self.runtime(), t, op),
186            Tensor::Bool(t) => super::dispatch::ensure_resident_on_runtime(self.runtime(), t, op),
187            Tensor::C32(t) => super::dispatch::ensure_resident_on_runtime(self.runtime(), t, op),
188            Tensor::C64(t) => super::dispatch::ensure_resident_on_runtime(self.runtime(), t, op),
189        }
190    }
191
192    /// Block the host until work enqueued on the session's stream completes.
193    ///
194    /// This is the only host barrier on the success path; ordinary successful
195    /// session operations only enqueue.
196    ///
197    /// # Errors
198    ///
199    /// Returns [`crate::Error::BackendSource`] when CUDA stream
200    /// synchronization fails.
201    pub fn synchronize(&mut self) -> crate::Result<()> {
202        self.backend.runtime().synchronize()
203    }
204
205    /// Borrow the type-safe raw CUDA extension session for one operation.
206    ///
207    /// The enter/exit protocol is fully contained in this call: a definite
208    /// CubeCL stream is captured on the current thread, pending CubeCL work is
209    /// flushed, the calling thread's previous device/context is saved, the
210    /// tenferro primary context is activated, the callback runs, and the
211    /// previous device/context is best-effort restored on return, `Err`, or
212    /// unwind (restoration failures are logged to stderr). The success path
213    /// does not synchronize.
214    ///
215    /// # Errors
216    ///
217    /// Returns [`crate::Error::BackendSource`] when CubeCL cannot expose or
218    /// flush the stream, or when the CUDA context cannot be entered. Context
219    /// restoration on exit is best-effort: a failure to restore the caller's
220    /// previous device/context is logged to stderr rather than propagated, so
221    /// a callback result is never replaced by a restore error.
222    ///
223    /// # Examples
224    ///
225    /// ```
226    /// use tenferro_gpu::cuda::CudaExecSession;
227    ///
228    /// // Method-call check only: `CudaExecSession` is not user-constructible.
229    /// fn check(session: &mut CudaExecSession<'_>) -> tenferro_tensor::Result<()> {
230    ///     session.with_raw("test.raw", |raw| {
231    ///         let _ = raw.stream();
232    ///         Ok(SessionOutcome::Done)
233    ///     })?;
234    ///     Ok(())
235    /// }
236    /// enum SessionOutcome { Done }
237    /// let _ = check;
238    /// ```
239    pub fn with_raw<R>(
240        &mut self,
241        op: &'static str,
242        f: impl for<'s> FnOnce(&mut raw::Session<'s>) -> crate::Result<R>,
243    ) -> crate::Result<R> {
244        let runtime = self.backend.runtime().clone();
245        let cache = self.backend.cuda_extension_cache();
246        // 1. Capture the definite CubeCL stream on this thread.
247        let stream = runtime.raw_cuda_stream()?;
248        // 2. Flush pending CubeCL work so raw library calls observe it.
249        runtime.flush_cubecl(op)?;
250        // 3-4. Save previous context, activate the tenferro primary context.
251        let device_ordinal = i32::try_from(runtime.device_ordinal())
252            .map_err(|source| crate::Error::backend_source(op, source))?;
253        let _guard = RawContextRestore::enter(op, device_ordinal, runtime.primary_context())?;
254        // 5. Build the unique raw session and run the callback.
255        // SAFETY: `_guard` keeps the primary context current for the whole
256        // `Session<'s>` borrow; `stream` is the captured CubeCL stream bound to
257        // the current thread.
258        let mut session = unsafe { raw::Session::new(runtime, cache, stream) };
259        f(&mut session)
260    }
261
262    /// Borrow the public tenferro-wide CubeCL session for one operation.
263    ///
264    /// The session exposes the exact tenferro CubeCL client bound to this
265    /// runtime. Pending CubeCL work is flushed before entering and again on
266    /// exit (including `Err` and unwind) so a later raw-session or host read
267    /// observes the enqueued work. The success path does not synchronize.
268    ///
269    /// # Examples
270    ///
271    /// ```
272    /// use tenferro_gpu::cuda::CudaExecSession;
273    ///
274    /// fn check(session: &mut CudaExecSession<'_>) {
275    ///     let _ = session.with_cubecl("test.cubecl", |_cubecl| Ok(()));
276    /// }
277    /// let _ = check;
278    /// ```
279    ///
280    /// # Errors
281    ///
282    /// Returns the callback's error, or [`crate::Error::BackendSource`] when
283    /// pending CubeCL work cannot be flushed on entry or exit.
284    pub fn with_cubecl<R>(
285        &mut self,
286        op: &'static str,
287        f: impl for<'s> FnOnce(&session_cubecl::Session<'s>) -> crate::Result<R>,
288    ) -> crate::Result<R> {
289        let runtime = self.backend.runtime().clone();
290        runtime.flush_cubecl(op)?;
291        let session = unsafe { session_cubecl::Session::new(runtime) };
292        // Best-effort exit flush on every path via Drop.
293        let mut _flush_guard = CubeclExitFlush::new(op, session.client());
294        let result = f(&session);
295        let flush_result = _flush_guard.flush_now();
296        match result {
297            Ok(value) => {
298                flush_result?;
299                Ok(value)
300            }
301            Err(err) => {
302                let _ = flush_result;
303                Err(err)
304            }
305        }
306    }
307
308    #[doc(hidden)]
309    pub fn tril_typed<T>(&self, input: &TypedTensor<T>, k: i64) -> crate::Result<TypedTensor<T>>
310    where
311        T: CubeElement + TensorScalar + CubePrimitive + Clone,
312    {
313        self.backend.tril_typed(input, k)
314    }
315
316    #[doc(hidden)]
317    pub fn slice_typed<T>(
318        &self,
319        input: &TypedTensor<T>,
320        config: &SliceConfig,
321    ) -> crate::Result<TypedTensor<T>>
322    where
323        T: CubeElement + TensorScalar + CubePrimitive + Clone,
324    {
325        self.backend.slice_typed(input, config)
326    }
327
328    /// Borrow the CUDA extension cache owned by the provider runtime.
329    #[doc(hidden)]
330    pub fn cuda_extension_cache(&self) -> &CudaExtensionCache {
331        self.backend.cuda_extension_cache()
332    }
333
334    #[doc(hidden)]
335    pub fn triu_typed<T>(&self, input: &TypedTensor<T>, k: i64) -> crate::Result<TypedTensor<T>>
336    where
337        T: CubeElement + TensorScalar + CubePrimitive + Clone,
338    {
339        self.backend.triu_typed(input, k)
340    }
341
342    #[doc(hidden)]
343    pub fn to_contiguous<T, R>(
344        &mut self,
345        view: &TypedTensorView<'_, T, R>,
346    ) -> crate::Result<TypedTensor<T, R>>
347    where
348        T: TensorScalar,
349        R: TensorRank,
350        CudaBackend: TensorViewCanonicalization<T, R>,
351    {
352        self.backend.to_contiguous(view)
353    }
354}
355
356/// Visit a CUDA execution session through the erased backend-session surface.
357///
358/// This is the public entry point that borrows CUDA execution authority for
359/// the duration of the callback (issue #1597). The callback cannot return a
360/// borrow of the reconstructed session, so the authority cannot escape the
361/// scope.
362///
363/// Returns `None` when `session` is not a CUDA execution session.
364///
365/// # Examples
366///
367/// ```
368/// use tenferro_gpu::cuda::{with_cuda_exec_session, CudaExecSession};
369///
370/// // Call-check only: the visitor borrows CUDA execution authority for the
371/// // duration of the callback.
372/// fn check(session: &mut dyn tenferro_tensor::backend::BackendSession) {
373///     let _ = with_cuda_exec_session(session, |_session| 0usize);
374/// }
375/// let _ = check;
376/// ```
377pub fn with_cuda_exec_session<B, R>(
378    session: &mut B,
379    f: impl for<'a> FnOnce(&'a mut CudaExecSession<'a>) -> R,
380) -> Option<R>
381where
382    B: BackendSession + ?Sized,
383{
384    if session.session_type_id() != std::any::TypeId::of::<CudaExecSessionMarker>() {
385        return None;
386    }
387    let data = unsafe { session.session_data_mut() };
388    // SAFETY: the exact marker check and the BackendSession erased-pointer
389    // contract identify the value as CudaExecSession for this scoped visit.
390    Some(unsafe { f(&mut *(data.cast::<CudaExecSession<'static>>())) })
391}
392
393macro_rules! delegate {
394    ($trait:path {
395        $(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) -> $ret:ty;)*
396    }) => {
397        impl $trait for CudaExecSession<'_> {
398            $(
399                fn $method(&mut self, $($arg: $arg_ty),*) -> $ret {
400                    self.backend.$method($($arg),*)
401                }
402            )*
403        }
404    };
405}
406
407delegate!(TensorElementwise {
408    fn add_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
409    fn sub_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
410    fn mul_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
411    fn neg_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
412    fn conj_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
413    fn div_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
414    fn rem_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
415    fn abs_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
416    fn sign_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
417    fn maximum_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
418    fn minimum_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
419    fn compare_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>, dir: &CompareDir) -> crate::Result<Tensor>;
420    fn select_read(pred: TensorRead<'_>, on_true: TensorRead<'_>, on_false: TensorRead<'_>) -> crate::Result<Tensor>;
421    fn clamp_read(input: TensorRead<'_>, lower: TensorRead<'_>, upper: TensorRead<'_>) -> crate::Result<Tensor>;
422    fn elementwise_read_into(op: ElementwiseReadOp, inputs: &[TensorRead<'_>], out: TensorWrite<'_>) -> crate::Result<()>;
423    fn add(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
424    fn sub(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
425    fn mul(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
426    fn neg(input: &Tensor) -> crate::Result<Tensor>;
427    fn conj(input: &Tensor) -> crate::Result<Tensor>;
428    fn div(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
429    fn rem(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
430    fn abs(input: &Tensor) -> crate::Result<Tensor>;
431    fn sign(input: &Tensor) -> crate::Result<Tensor>;
432    fn maximum(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
433    fn minimum(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
434    fn compare(lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> crate::Result<Tensor>;
435    fn select(pred: &Tensor, on_true: &Tensor, on_false: &Tensor) -> crate::Result<Tensor>;
436    fn clamp(input: &Tensor, lower: &Tensor, upper: &Tensor) -> crate::Result<Tensor>;
437});
438
439delegate!(TensorAnalytic {
440    fn exp_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
441    fn log_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
442    fn sin_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
443    fn cos_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
444    fn tanh_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
445    fn sqrt_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
446    fn rsqrt_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
447    fn pow_read(lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor>;
448    fn expm1_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
449    fn log1p_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
450    fn exp(input: &Tensor) -> crate::Result<Tensor>;
451    fn log(input: &Tensor) -> crate::Result<Tensor>;
452    fn sin(input: &Tensor) -> crate::Result<Tensor>;
453    fn cos(input: &Tensor) -> crate::Result<Tensor>;
454    fn tanh(input: &Tensor) -> crate::Result<Tensor>;
455    fn sqrt(input: &Tensor) -> crate::Result<Tensor>;
456    fn rsqrt(input: &Tensor) -> crate::Result<Tensor>;
457    fn pow(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
458    fn expm1(input: &Tensor) -> crate::Result<Tensor>;
459    fn log1p(input: &Tensor) -> crate::Result<Tensor>;
460});
461
462delegate!(TensorStructural {
463    fn transpose_read(input: TensorRead<'_>, perm: &[usize]) -> crate::Result<Tensor>;
464    fn reshape_read(input: TensorRead<'_>, shape: &[usize]) -> crate::Result<Tensor>;
465    fn broadcast_in_dim_read(input: TensorRead<'_>, shape: &[usize], dims: &[usize]) -> crate::Result<Tensor>;
466    fn to_contiguous_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
467    fn copy_read_into(src: TensorRead<'_>, dst: TensorWrite<'_>) -> crate::Result<()>;
468    fn transpose(input: &Tensor, perm: &[usize]) -> crate::Result<Tensor>;
469    fn reshape(input: &Tensor, shape: &[usize]) -> crate::Result<Tensor>;
470    fn broadcast_in_dim(input: &Tensor, shape: &[usize], dims: &[usize]) -> crate::Result<Tensor>;
471    fn cast(input: &Tensor, to: tenferro_tensor::DType) -> crate::Result<Tensor>;
472    fn extract_diagonal(input: &Tensor, axis_a: usize, axis_b: usize) -> crate::Result<Tensor>;
473    fn embed_diagonal(input: &Tensor, axis_a: usize, axis_b: usize) -> crate::Result<Tensor>;
474    fn tril(input: &Tensor, k: i64) -> crate::Result<Tensor>;
475    fn triu(input: &Tensor, k: i64) -> crate::Result<Tensor>;
476});
477
478delegate!(TensorReduction {
479    fn reduce_sum_read(input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
480    fn reduce_prod_read(input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
481    fn reduce_max_read(input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
482    fn reduce_min_read(input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
483    fn reduce_sum(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
484    fn reduce_sum_squares_read(input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
485    fn reduce_prod(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
486    fn reduce_max(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
487    fn reduce_min(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
488});
489
490delegate!(TensorDot {
491    fn dot_general(lhs: &Tensor, rhs: &Tensor, config: &DotGeneralConfig) -> crate::Result<Tensor>;
492    fn dot_general_with_conj(
493        lhs: &Tensor,
494        rhs: &Tensor,
495        config: &DotGeneralConfig,
496        lhs_conj: bool,
497        rhs_conj: bool,
498    ) -> crate::Result<Tensor>;
499    fn dot_general_read_into_accum(
500        lhs: TensorRead<'_>,
501        rhs: TensorRead<'_>,
502        config: &DotGeneralConfig,
503        accumulation: DotGeneralAccumulation,
504        out: TensorWrite<'_>,
505    ) -> crate::Result<()>;
506});
507
508delegate!(TensorIndexing {
509    fn gather(
510        operand: &Tensor,
511        start_indices: &Tensor,
512        config: &GatherConfig,
513    ) -> crate::Result<Tensor>;
514    fn scatter(
515        operand: &Tensor,
516        scatter_indices: &Tensor,
517        updates: &Tensor,
518        config: &ScatterConfig,
519    ) -> crate::Result<Tensor>;
520    fn slice(input: &Tensor, config: &SliceConfig) -> crate::Result<Tensor>;
521    fn dynamic_slice(
522        input: &Tensor,
523        starts: &Tensor,
524        slice_sizes: &[usize],
525    ) -> crate::Result<Tensor>;
526    fn dynamic_update_slice(
527        operand: &Tensor,
528        update: &Tensor,
529        starts: &Tensor,
530    ) -> crate::Result<Tensor>;
531    fn pad(input: &Tensor, config: &PadConfig) -> crate::Result<Tensor>;
532    fn concatenate(inputs: &[&Tensor], axis: usize) -> crate::Result<Tensor>;
533    fn reverse(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
534});
535
536delegate!(TensorFusion {
537    fn execute_elementwise_fusion(
538        inputs: &[&Tensor],
539        plan: &ElementwiseFusionPlan,
540    ) -> crate::Result<Option<Vec<Tensor>>>;
541    fn execute_broadcast_multiply(
542        lhs: TensorRead<'_>,
543        lhs_shape: &[usize],
544        lhs_dims: &[usize],
545        rhs: TensorRead<'_>,
546        rhs_shape: &[usize],
547        rhs_dims: &[usize],
548    ) -> crate::Result<Option<Tensor>>;
549    fn execute_broadcast_multiply_value(
550        lhs: TensorRead<'_>,
551        lhs_shape: &[usize],
552        lhs_dims: &[usize],
553        rhs: TensorRead<'_>,
554        rhs_shape: &[usize],
555        rhs_dims: &[usize],
556    ) -> crate::Result<Option<TensorValue>>;
557});
558
559delegate!(TensorBuffer {
560    fn reclaim_buffer(tensor: Tensor) -> ();
561});
562
563delegate!(TensorDeviceTransfer {
564    fn download_to_host(tensor: TensorRead<'_>) -> crate::Result<Tensor>;
565    fn upload_host_tensor(tensor: TensorRead<'_>) -> crate::Result<Tensor>;
566});
567
568macro_rules! delegate_cached {
569    ($(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) -> $ret:ty;)*) => {
570        impl SessionCachedDot for CudaExecSession<'_> {
571            $(
572                fn $method(&mut self, $($arg: $arg_ty),*) -> $ret {
573                    <CudaBackend as SessionCachedDot>::$method(self.backend, $($arg),*)
574                }
575            )*
576        }
577    };
578}
579
580delegate_cached! {
581    fn dot_general_cached(
582        cache_slot: Option<usize>,
583        lhs: &Tensor,
584        rhs: &Tensor,
585        config: &DotGeneralConfig,
586    ) -> crate::Result<Tensor>;
587    fn dot_general_read_cached(
588        cache_slot: Option<usize>,
589        lhs: TensorRead<'_>,
590        rhs: TensorRead<'_>,
591        config: &DotGeneralConfig,
592    ) -> crate::Result<Tensor>;
593    fn dot_general_with_conj_cached(
594        cache_slot: Option<usize>,
595        lhs: &Tensor,
596        rhs: &Tensor,
597        config: &DotGeneralConfig,
598        lhs_conj: bool,
599        rhs_conj: bool,
600    ) -> crate::Result<Tensor>;
601    fn dot_general_with_conj_read_cached(
602        cache_slot: Option<usize>,
603        lhs: TensorRead<'_>,
604        rhs: TensorRead<'_>,
605        config: &DotGeneralConfig,
606        lhs_conj: bool,
607        rhs_conj: bool,
608    ) -> crate::Result<Tensor>;
609    fn dot_general_read_into_accum_cached(
610        cache_slot: Option<usize>,
611        lhs: TensorRead<'_>,
612        rhs: TensorRead<'_>,
613        config: &DotGeneralConfig,
614        accumulation: DotGeneralAccumulation,
615        out: TensorWrite<'_>,
616    ) -> crate::Result<()>;
617    fn grouped_gemm_cached(
618        cache_slot: Option<usize>,
619        lhs: TensorRead<'_>,
620        rhs: TensorRead<'_>,
621        config: &GroupedGemmConfig<'_>,
622        out: TensorWrite<'_>,
623    ) -> crate::Result<()>;
624}
625
626impl BackendSession for CudaExecSession<'_> {
627    fn vdot_read(&mut self, lhs: TensorRead<'_>, rhs: TensorRead<'_>) -> crate::Result<Tensor> {
628        BackendSession::vdot_read(self.backend, lhs, rhs)
629    }
630
631    fn norm_squared_read(&mut self, input: TensorRead<'_>) -> crate::Result<Tensor> {
632        BackendSession::norm_squared_read(self.backend, input)
633    }
634
635    fn axpby_read_into_accum(
636        &mut self,
637        alpha: tenferro_tensor::ContractionScalar,
638        x: TensorRead<'_>,
639        beta: tenferro_tensor::ContractionScalar,
640        y: TensorWrite<'_>,
641    ) -> crate::Result<()> {
642        BackendSession::axpby_read_into_accum(self.backend, alpha, x, beta, y)
643    }
644
645    fn session_type_id(&self) -> TypeId {
646        TypeId::of::<CudaExecSessionMarker>()
647    }
648
649    unsafe fn session_data_mut(&mut self) -> *mut () {
650        self as *mut Self as *mut ()
651    }
652}
653
654impl BackendSessionHost for CudaBackend {
655    fn with_backend_session<R: Send>(
656        &mut self,
657        f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
658    ) -> R {
659        let mut session = CudaExecSession {
660            backend: self,
661            _not_send_sync: PhantomData,
662        };
663        // Nested entry is caught by the portable in-session guard in debug
664        // builds; the CUDA runtime must never re-enter a session closure.
665        with_session_entry_guard(|| f(&mut session))
666    }
667}
tenferro_gpu/cubecl/exec_session.rs

tenferro_gpu/cubecl/
exec_session.rs
