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, GroupedGemmConfig, SessionCachedDot,
7    TensorAnalytic, TensorBuffer, TensorDeviceTransfer, TensorDot, TensorElementwise, TensorFusion,
8    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(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
409    fn sub(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
410    fn mul(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
411    fn neg(input: &Tensor) -> crate::Result<Tensor>;
412    fn conj(input: &Tensor) -> crate::Result<Tensor>;
413    fn div(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
414    fn rem(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
415    fn abs(input: &Tensor) -> crate::Result<Tensor>;
416    fn sign(input: &Tensor) -> crate::Result<Tensor>;
417    fn maximum(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
418    fn minimum(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
419    fn compare(lhs: &Tensor, rhs: &Tensor, dir: &CompareDir) -> crate::Result<Tensor>;
420    fn select(pred: &Tensor, on_true: &Tensor, on_false: &Tensor) -> crate::Result<Tensor>;
421    fn clamp(input: &Tensor, lower: &Tensor, upper: &Tensor) -> crate::Result<Tensor>;
422});
423
424delegate!(TensorAnalytic {
425    fn exp(input: &Tensor) -> crate::Result<Tensor>;
426    fn log(input: &Tensor) -> crate::Result<Tensor>;
427    fn sin(input: &Tensor) -> crate::Result<Tensor>;
428    fn cos(input: &Tensor) -> crate::Result<Tensor>;
429    fn tanh(input: &Tensor) -> crate::Result<Tensor>;
430    fn sqrt(input: &Tensor) -> crate::Result<Tensor>;
431    fn rsqrt(input: &Tensor) -> crate::Result<Tensor>;
432    fn pow(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor>;
433    fn expm1(input: &Tensor) -> crate::Result<Tensor>;
434    fn log1p(input: &Tensor) -> crate::Result<Tensor>;
435});
436
437delegate!(TensorStructural {
438    fn to_contiguous_read(input: TensorRead<'_>) -> crate::Result<Tensor>;
439    fn copy_read_into(src: TensorRead<'_>, dst: TensorWrite<'_>) -> crate::Result<()>;
440    fn transpose(input: &Tensor, perm: &[usize]) -> crate::Result<Tensor>;
441    fn reshape(input: &Tensor, shape: &[usize]) -> crate::Result<Tensor>;
442    fn broadcast_in_dim(input: &Tensor, shape: &[usize], dims: &[usize]) -> crate::Result<Tensor>;
443    fn cast(input: &Tensor, to: tenferro_tensor::DType) -> crate::Result<Tensor>;
444    fn extract_diagonal(input: &Tensor, axis_a: usize, axis_b: usize) -> crate::Result<Tensor>;
445    fn embed_diagonal(input: &Tensor, axis_a: usize, axis_b: usize) -> crate::Result<Tensor>;
446    fn tril(input: &Tensor, k: i64) -> crate::Result<Tensor>;
447    fn triu(input: &Tensor, k: i64) -> crate::Result<Tensor>;
448});
449
450delegate!(TensorReduction {
451    fn reduce_sum(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
452    fn reduce_sum_squares_read(input: TensorRead<'_>, axes: &[usize]) -> crate::Result<Tensor>;
453    fn reduce_prod(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
454    fn reduce_max(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
455    fn reduce_min(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
456});
457
458delegate!(TensorDot {
459    fn dot_general(lhs: &Tensor, rhs: &Tensor, config: &DotGeneralConfig) -> crate::Result<Tensor>;
460    fn dot_general_with_conj(
461        lhs: &Tensor,
462        rhs: &Tensor,
463        config: &DotGeneralConfig,
464        lhs_conj: bool,
465        rhs_conj: bool,
466    ) -> crate::Result<Tensor>;
467    fn dot_general_read_into_accum(
468        lhs: TensorRead<'_>,
469        rhs: TensorRead<'_>,
470        config: &DotGeneralConfig,
471        accumulation: DotGeneralAccumulation,
472        out: TensorWrite<'_>,
473    ) -> crate::Result<()>;
474});
475
476delegate!(TensorIndexing {
477    fn gather(
478        operand: &Tensor,
479        start_indices: &Tensor,
480        config: &GatherConfig,
481    ) -> crate::Result<Tensor>;
482    fn scatter(
483        operand: &Tensor,
484        scatter_indices: &Tensor,
485        updates: &Tensor,
486        config: &ScatterConfig,
487    ) -> crate::Result<Tensor>;
488    fn slice(input: &Tensor, config: &SliceConfig) -> crate::Result<Tensor>;
489    fn dynamic_slice(
490        input: &Tensor,
491        starts: &Tensor,
492        slice_sizes: &[usize],
493    ) -> crate::Result<Tensor>;
494    fn dynamic_update_slice(
495        operand: &Tensor,
496        update: &Tensor,
497        starts: &Tensor,
498    ) -> crate::Result<Tensor>;
499    fn pad(input: &Tensor, config: &PadConfig) -> crate::Result<Tensor>;
500    fn concatenate(inputs: &[&Tensor], axis: usize) -> crate::Result<Tensor>;
501    fn reverse(input: &Tensor, axes: &[usize]) -> crate::Result<Tensor>;
502});
503
504delegate!(TensorFusion {
505    fn execute_elementwise_fusion(
506        inputs: &[&Tensor],
507        plan: &ElementwiseFusionPlan,
508    ) -> crate::Result<Option<Vec<Tensor>>>;
509    fn execute_broadcast_multiply(
510        lhs: TensorRead<'_>,
511        lhs_shape: &[usize],
512        lhs_dims: &[usize],
513        rhs: TensorRead<'_>,
514        rhs_shape: &[usize],
515        rhs_dims: &[usize],
516    ) -> crate::Result<Option<Tensor>>;
517    fn execute_broadcast_multiply_value(
518        lhs: TensorRead<'_>,
519        lhs_shape: &[usize],
520        lhs_dims: &[usize],
521        rhs: TensorRead<'_>,
522        rhs_shape: &[usize],
523        rhs_dims: &[usize],
524    ) -> crate::Result<Option<TensorValue>>;
525});
526
527delegate!(TensorBuffer {
528    fn reclaim_buffer(tensor: Tensor) -> ();
529});
530
531delegate!(TensorDeviceTransfer {
532    fn download_to_host(tensor: TensorRead<'_>) -> crate::Result<Tensor>;
533    fn upload_host_tensor(tensor: TensorRead<'_>) -> crate::Result<Tensor>;
534});
535
536macro_rules! delegate_cached {
537    ($(fn $method:ident($($arg:ident: $arg_ty:ty),* $(,)?) -> $ret:ty;)*) => {
538        impl SessionCachedDot for CudaExecSession<'_> {
539            $(
540                fn $method(&mut self, $($arg: $arg_ty),*) -> $ret {
541                    <CudaBackend as SessionCachedDot>::$method(self.backend, $($arg),*)
542                }
543            )*
544        }
545    };
546}
547
548delegate_cached! {
549    fn dot_general_cached(
550        cache_slot: Option<usize>,
551        lhs: &Tensor,
552        rhs: &Tensor,
553        config: &DotGeneralConfig,
554    ) -> crate::Result<Tensor>;
555    fn dot_general_read_cached(
556        cache_slot: Option<usize>,
557        lhs: TensorRead<'_>,
558        rhs: TensorRead<'_>,
559        config: &DotGeneralConfig,
560    ) -> crate::Result<Tensor>;
561    fn dot_general_with_conj_cached(
562        cache_slot: Option<usize>,
563        lhs: &Tensor,
564        rhs: &Tensor,
565        config: &DotGeneralConfig,
566        lhs_conj: bool,
567        rhs_conj: bool,
568    ) -> crate::Result<Tensor>;
569    fn dot_general_with_conj_read_cached(
570        cache_slot: Option<usize>,
571        lhs: TensorRead<'_>,
572        rhs: TensorRead<'_>,
573        config: &DotGeneralConfig,
574        lhs_conj: bool,
575        rhs_conj: bool,
576    ) -> crate::Result<Tensor>;
577    fn dot_general_read_into_accum_cached(
578        cache_slot: Option<usize>,
579        lhs: TensorRead<'_>,
580        rhs: TensorRead<'_>,
581        config: &DotGeneralConfig,
582        accumulation: DotGeneralAccumulation,
583        out: TensorWrite<'_>,
584    ) -> crate::Result<()>;
585    fn grouped_gemm_cached(
586        cache_slot: Option<usize>,
587        lhs: TensorRead<'_>,
588        rhs: TensorRead<'_>,
589        config: &GroupedGemmConfig<'_>,
590        out: TensorWrite<'_>,
591    ) -> crate::Result<()>;
592}
593
594impl BackendSession for CudaExecSession<'_> {
595    fn session_type_id(&self) -> TypeId {
596        TypeId::of::<CudaExecSessionMarker>()
597    }
598
599    unsafe fn session_data_mut(&mut self) -> *mut () {
600        self as *mut Self as *mut ()
601    }
602}
603
604impl BackendSessionHost for CudaBackend {
605    fn with_backend_session<R: Send>(
606        &mut self,
607        f: impl FnOnce(&mut dyn BackendSession) -> R + Send,
608    ) -> R {
609        let mut session = CudaExecSession {
610            backend: self,
611            _not_send_sync: PhantomData,
612        };
613        // Nested entry is caught by the portable in-session guard in debug
614        // builds; the CUDA runtime must never re-enter a session closure.
615        with_session_entry_guard(|| f(&mut session))
616    }
617}
tenferro_gpu/cubecl/exec_session.rs

tenferro_gpu/cubecl/
exec_session.rs
