1use alloc::sync::Arc;
34use core::{cell::UnsafeCell, marker::PhantomData};
35use cubecl::prelude::*;
36use cubecl_core::{
37 self as cubecl,
38 frontend::select,
39 ir::{ExpandValue, VectorSize},
40 unexpanded,
41};
42
43use crate::tensor::r#virtual::{VirtualTensor, VirtualTensorExpand};
44use crate::tensor::{
45 ViewOperations, ViewOperationsExpand, ViewOperationsMut, ViewOperationsMutExpand,
46 layout::Coords1d,
47};
48use cubecl_core::prelude::barrier::Barrier;
49
50#[derive(Clone, Copy)]
56pub struct WriteOnly;
57
58pub trait ErasedIo: Clone + Copy + Send + Sync + 'static {}
60
61pub trait ErasedIoRead: ErasedIo {}
63
64pub trait ErasedIoWrite: ErasedIo {}
66
67impl ErasedIo for ReadOnly {}
68impl ErasedIo for ReadWrite {}
69impl ErasedIo for WriteOnly {}
70
71impl ErasedIoRead for ReadOnly {}
72impl ErasedIoRead for ReadWrite {}
73
74impl ErasedIoWrite for ReadWrite {}
75impl ErasedIoWrite for WriteOnly {}
76
77pub struct ErasedTensor<E: Numeric, IO = ReadOnly> {
79 _e: PhantomData<E>,
80 _p: PhantomData<IO>,
81}
82
83pub struct ErasedTensorExpand<E: Numeric, IO> {
85 state: Arc<UnsafeCell<dyn ErasedTensorOperationsExpand<E>>>,
86 _p: PhantomData<IO>,
87}
88
89impl<E: Numeric, IO> Clone for ErasedTensor<E, IO> {
90 fn clone(&self) -> Self {
91 *self
92 }
93}
94
95impl<E: Numeric, IO> Copy for ErasedTensor<E, IO> {}
96
97impl<E: Numeric, IO> Clone for ErasedTensorExpand<E, IO> {
98 fn clone(&self) -> Self {
99 Self {
100 state: self.state.clone(),
101 _p: PhantomData,
102 }
103 }
104}
105
106pub trait ErasedTensorOperationsExpand<E: Numeric> {
126 fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize;
128
129 fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize>;
132
133 fn __expand_read_line_method(
136 &self,
137 _scope: &Scope,
138 _index: NativeExpand<usize>,
139 ) -> ExpandValue {
140 unimplemented!("ErasedTensor: this backing does not serve reads")
141 }
142
143 fn __expand_write_line_method(
146 &mut self,
147 _scope: &Scope,
148 _index: NativeExpand<usize>,
149 _value: ExpandValue,
150 ) {
151 unimplemented!("ErasedTensor: this backing does not serve writes")
152 }
153}
154
155pub trait ReadsLines<E: Numeric>: ErasedTensorOperationsExpand<E> {}
158
159pub trait WritesLines<E: Numeric>: ErasedTensorOperationsExpand<E> {}
162
163pub trait ErasedBacking<E: Numeric, IO>: ErasedTensorOperationsExpand<E> {}
170
171impl<E: Numeric, T: ReadsLines<E>> ErasedBacking<E, ReadOnly> for T {}
172impl<E: Numeric, T: WritesLines<E>> ErasedBacking<E, WriteOnly> for T {}
173impl<E: Numeric, T: ReadsLines<E> + WritesLines<E>> ErasedBacking<E, ReadWrite> for T {}
174
175impl<E: Numeric, IO> ErasedTensorExpand<E, IO> {
176 pub fn new<S: ErasedBacking<E, IO> + 'static>(backing: S) -> Self {
181 Self {
182 state: Arc::new(UnsafeCell::new(backing)),
183 _p: PhantomData,
184 }
185 }
186
187 fn state_read(&self) -> &dyn ErasedTensorOperationsExpand<E> {
188 unsafe { &*self.state.get() }
191 }
192
193 #[allow(clippy::mut_from_ref)]
194 fn state_write(&self) -> &mut dyn ErasedTensorOperationsExpand<E> {
195 unsafe { &mut *self.state.get() }
198 }
199
200 fn check_width<N: Size>(&self, scope: &Scope, op: &str) {
207 let served = self.state_read().__expand_vector_size_method(scope);
208 let asked = <N as Size>::__expand_value(scope);
209 assert_eq!(
210 served, asked,
211 "ErasedTensor::{op}: the tensor takes {served}-wide lines and the {op} is {asked}-wide"
212 );
213 }
214}
215
216#[cube]
217impl<E: Numeric, IO: ErasedIoRead> ErasedTensor<E, IO> {
218 #[allow(unused)]
224 pub fn read<N: Size>(&self, index: usize) -> Vector<E, N> {
225 intrinsic!(|scope| {
226 self.check_width::<N>(scope, "read");
227 self.state_read()
228 .__expand_read_line_method(scope, index)
229 .into()
230 })
231 }
232}
233
234#[cube]
235impl<E: Numeric, IO: ErasedIoWrite> ErasedTensor<E, IO> {
236 #[allow(unused)]
242 pub fn write<N: Size>(&mut self, index: usize, value: Vector<E, N>) {
243 intrinsic!(|scope| {
244 self.check_width::<N>(scope, "write");
245 self.state_write()
246 .__expand_write_line_method(scope, index, value.into())
247 })
248 }
249}
250
251#[cube]
252impl<E: Numeric, IO: ErasedIo> ErasedTensor<E, IO> {
253 #[allow(unused, clippy::len_without_is_empty)]
255 pub fn len(&self) -> usize {
256 intrinsic!(|scope| self.state_read().__expand_lines_method(scope))
257 }
258}
259
260impl<E: Numeric, IO: ErasedIo> Vectorized for ErasedTensor<E, IO> {}
261impl<E: Numeric, IO: ErasedIo> VectorizedExpand for ErasedTensorExpand<E, IO> {
262 fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
263 self.state_read().__expand_vector_size_method(scope)
264 }
265}
266
267impl<E: Numeric, N: Size, IO: Clone> ErasedTensorOperationsExpand<E>
275 for VirtualTensorExpand<E, N, IO>
276{
277 fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
278 VectorizedExpand::__expand_vector_size_method(self, scope)
279 }
280
281 fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
282 self.clone().__expand_len_method(scope)
283 }
284
285 fn __expand_read_line_method(&self, scope: &Scope, index: NativeExpand<usize>) -> ExpandValue {
286 Self::__expand_read_method(self, scope, index).expand
287 }
288
289 fn __expand_write_line_method(
290 &mut self,
291 scope: &Scope,
292 index: NativeExpand<usize>,
293 value: ExpandValue,
294 ) {
295 self.state_write()
296 .__expand_write_method(scope, index, value.into())
297 }
298}
299
300impl<E: Numeric, N: Size, IO: Clone> ReadsLines<E> for VirtualTensorExpand<E, N, IO> {}
301impl<E: Numeric, N: Size> WritesLines<E> for VirtualTensorExpand<E, N, ReadWrite> {}
302
303impl<E: Numeric, N: Size> ErasedTensorOperationsExpand<E> for TensorExpand<Vector<E, N>> {
308 fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
309 VectorizedExpand::__expand_vector_size_method(self, scope)
310 }
311
312 fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
313 self.__expand_len_method(scope)
314 }
315
316 fn __expand_read_line_method(&self, scope: &Scope, index: NativeExpand<usize>) -> ExpandValue {
317 unsafe {
318 self.__expand_get_unchecked_method(scope, index)
319 .__expand_deref_method(scope)
320 .expand
321 }
322 }
323
324 fn __expand_write_line_method(
325 &mut self,
326 scope: &Scope,
327 index: NativeExpand<usize>,
328 value: ExpandValue,
329 ) {
330 unsafe {
331 self.__expand_get_unchecked_mut_method(scope, index)
332 .__expand_assign_method(scope, value.into())
333 };
334 }
335}
336
337impl<E: Numeric, N: Size> ReadsLines<E> for TensorExpand<Vector<E, N>> {}
338impl<E: Numeric, N: Size> WritesLines<E> for TensorExpand<Vector<E, N>> {}
339
340pub struct ErasedView<V, N: Size> {
347 view: V,
348 _n: PhantomData<N>,
349}
350
351impl<E: Numeric, N: Size, V> ErasedTensorOperationsExpand<E> for ErasedView<V, N>
352where
353 V: ViewOperationsExpand<Vector<E, N>, Coords1d>,
354{
355 fn __expand_vector_size_method(&self, _scope: &Scope) -> VectorSize {
356 <N as Size>::__expand_value(_scope)
357 }
358
359 fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
360 self.view.__expand_shape_method(scope)
361 }
362
363 fn __expand_read_line_method(&self, scope: &Scope, index: NativeExpand<usize>) -> ExpandValue {
364 self.view.__expand_read_method(scope, index).expand
365 }
366}
367
368impl<E: Numeric, N: Size, V> ReadsLines<E> for ErasedView<V, N> where
369 V: ViewOperationsExpand<Vector<E, N>, Coords1d>
370{
371}
372
373pub struct ErasedViewMut<V, N: Size> {
375 view: V,
376 _n: PhantomData<N>,
377}
378
379impl<E: Numeric, N: Size, V> ErasedTensorOperationsExpand<E> for ErasedViewMut<V, N>
380where
381 V: ViewOperationsMutExpand<Vector<E, N>, Coords1d>,
382{
383 fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
384 <N as Size>::__expand_value(scope)
385 }
386
387 fn __expand_lines_method(&self, scope: &Scope) -> NativeExpand<usize> {
388 <V as ViewOperationsExpand<Vector<E, N>, Coords1d>>::__expand_shape_method(
389 &self.view, scope,
390 )
391 }
392
393 fn __expand_read_line_method(&self, scope: &Scope, index: NativeExpand<usize>) -> ExpandValue {
394 <V as ViewOperationsExpand<Vector<E, N>, Coords1d>>::__expand_read_method(
395 &self.view, scope, index,
396 )
397 .expand
398 }
399
400 fn __expand_write_line_method(
401 &mut self,
402 scope: &Scope,
403 index: NativeExpand<usize>,
404 value: ExpandValue,
405 ) {
406 self.view.__expand_write_method(scope, index, value.into())
407 }
408}
409
410impl<E: Numeric, N: Size, V> ReadsLines<E> for ErasedViewMut<V, N> where
411 V: ViewOperationsMutExpand<Vector<E, N>, Coords1d>
412{
413}
414impl<E: Numeric, N: Size, V> WritesLines<E> for ErasedViewMut<V, N> where
415 V: ViewOperationsMutExpand<Vector<E, N>, Coords1d>
416{
417}
418
419impl<E: Numeric, IO: ErasedIo> ErasedTensor<E, IO> {
422 pub fn of_view<V: CubeType, N: Size>(_view: V) -> Self {
428 unexpanded!()
429 }
430
431 pub fn __expand_of_view<V: CubeType, N: Size>(
433 _scope: &Scope,
434 view: V::ExpandType,
435 ) -> ErasedTensorExpand<E, IO>
436 where
437 V::ExpandType: ViewOperationsExpand<Vector<E, N>, Coords1d> + 'static,
438 ErasedView<V::ExpandType, N>: ErasedBacking<E, IO>,
439 {
440 ErasedTensorExpand::new(ErasedView::<V::ExpandType, N> {
441 view,
442 _n: PhantomData,
443 })
444 }
445
446 pub fn of_view_mut<V: CubeType, N: Size>(_view: V) -> Self {
448 unexpanded!()
449 }
450
451 pub fn __expand_of_view_mut<V: CubeType, N: Size>(
453 _scope: &Scope,
454 view: V::ExpandType,
455 ) -> ErasedTensorExpand<E, IO>
456 where
457 V::ExpandType: ViewOperationsMutExpand<Vector<E, N>, Coords1d> + 'static,
458 ErasedViewMut<V::ExpandType, N>: ErasedBacking<E, IO>,
459 {
460 ErasedTensorExpand::new(ErasedViewMut::<V::ExpandType, N> {
461 view,
462 _n: PhantomData,
463 })
464 }
465
466 pub fn of_tensor<N: Size>(_tensor: &Tensor<Vector<E, N>>) -> Self {
468 unexpanded!()
469 }
470
471 pub fn __expand_of_tensor<N: Size>(
473 _scope: &Scope,
474 tensor: &TensorExpand<Vector<E, N>>,
475 ) -> ErasedTensorExpand<E, IO>
476 where
477 TensorExpand<Vector<E, N>>: ErasedBacking<E, IO>,
478 {
479 ErasedTensorExpand::new(ExpandTypeClone::clone_unchecked(tensor))
480 }
481
482 pub fn of_tensor_mut<N: Size>(_tensor: &mut Tensor<Vector<E, N>>) -> Self {
489 unexpanded!()
490 }
491
492 pub fn __expand_of_tensor_mut<N: Size>(
494 _scope: &Scope,
495 tensor: &mut TensorExpand<Vector<E, N>>,
496 ) -> ErasedTensorExpand<E, IO>
497 where
498 TensorExpand<Vector<E, N>>: ErasedBacking<E, IO>,
499 {
500 ErasedTensorExpand::new(ExpandTypeClone::clone_unchecked(tensor))
501 }
502
503 pub fn of_virtual<N: Size, IO2: Clone>(_tensor: VirtualTensor<E, N, IO2>) -> Self {
505 unexpanded!()
506 }
507
508 pub fn __expand_of_virtual<N: Size, IO2: Clone + 'static>(
510 _scope: &Scope,
511 tensor: VirtualTensorExpand<E, N, IO2>,
512 ) -> ErasedTensorExpand<E, IO>
513 where
514 VirtualTensorExpand<E, N, IO2>: ErasedBacking<E, IO>,
515 {
516 ErasedTensorExpand::new(tensor)
517 }
518}
519
520impl<E: Numeric, N: Size> From<VirtualTensorExpand<E, N, ReadWrite>>
521 for ErasedTensorExpand<E, ReadWrite>
522{
523 fn from(tensor: VirtualTensorExpand<E, N, ReadWrite>) -> Self {
524 ErasedTensorExpand::new(tensor)
525 }
526}
527
528impl<E: Numeric, N: Size> From<VirtualTensorExpand<E, N, ReadOnly>>
529 for ErasedTensorExpand<E, ReadOnly>
530{
531 fn from(tensor: VirtualTensorExpand<E, N, ReadOnly>) -> Self {
532 ErasedTensorExpand::new(tensor)
533 }
534}
535
536mod __cube_type {
539 use super::*;
540
541 impl<E: Numeric, IO: ErasedIo> CubeType for ErasedTensor<E, IO> {
542 type ExpandType = ErasedTensorExpand<E, IO>;
543 }
544
545 impl<E: Numeric, IO> IntoExpand for ErasedTensorExpand<E, IO> {
546 type Expand = ErasedTensorExpand<E, IO>;
547
548 fn into_expand(self, _: &Scope) -> Self::Expand {
549 self
550 }
551 }
552
553 impl<E: Numeric, IO> ExpandTypeClone for ErasedTensorExpand<E, IO> {
554 fn clone_unchecked(&self) -> Self {
555 self.clone()
556 }
557 }
558
559 impl<E: Numeric, IO> IntoMut for ErasedTensorExpand<E, IO> {
560 fn into_mut(self, _scope: &Scope) -> Self {
561 self
562 }
563 }
564
565 impl<E: Numeric, IO> CubeDebug for ErasedTensorExpand<E, IO> {}
566
567 impl<E: Numeric, IO> AsRefExpand for ErasedTensorExpand<E, IO> {
568 fn __expand_ref_method(&self, _: &Scope) -> &Self {
569 self
570 }
571 }
572
573 impl<E: Numeric, IO> AsMutExpand for ErasedTensorExpand<E, IO> {
574 fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
575 self
576 }
577 }
578}
579
580impl<E: Numeric, N: Size, IO: ErasedIo> ViewOperations<Vector<E, N>, Coords1d>
596 for ErasedTensor<E, IO>
597{
598}
599
600impl<E: Numeric, N: Size, IO: ErasedIo> ViewOperationsExpand<Vector<E, N>, Coords1d>
601 for ErasedTensorExpand<E, IO>
602{
603 fn __expand_read_method(
604 &self,
605 scope: &Scope,
606 pos: NativeExpand<usize>,
607 ) -> <Vector<E, N> as CubeType>::ExpandType {
608 self.check_width::<N>(scope, "read");
609 self.state_read()
610 .__expand_read_line_method(scope, pos)
611 .into()
612 }
613
614 fn __expand_read_checked_method(
615 &self,
616 scope: &Scope,
617 pos: NativeExpand<usize>,
618 ) -> <Vector<E, N> as CubeType>::ExpandType {
619 let zero = <Vector<E, N>>::__expand_cast_from(scope, 0.into());
620 <Self as ViewOperationsExpand<Vector<E, N>, Coords1d>>::__expand_read_masked_method(
621 self, scope, pos, zero,
622 )
623 }
624
625 fn __expand_read_masked_method(
626 &self,
627 scope: &Scope,
628 pos: NativeExpand<usize>,
629 mask_value: <Vector<E, N> as CubeType>::ExpandType,
630 ) -> <Vector<E, N> as CubeType>::ExpandType {
631 let in_bounds =
632 <Self as ViewOperationsExpand<Vector<E, N>, Coords1d>>::__expand_is_in_bounds_method(
633 self, scope, pos,
634 );
635 let keep = usize::__expand_cast_from(scope, in_bounds);
638 let pos = pos.__expand_mul_method(scope, keep);
639 let value = <Self as ViewOperationsExpand<Vector<E, N>, Coords1d>>::__expand_read_method(
640 self, scope, pos,
641 );
642 select::expand::<Vector<E, N>>(scope, in_bounds, value, mask_value)
643 }
644
645 fn __expand_read_unchecked_method(
646 &self,
647 scope: &Scope,
648 pos: NativeExpand<usize>,
649 ) -> <Vector<E, N> as CubeType>::ExpandType {
650 <Self as ViewOperationsExpand<Vector<E, N>, Coords1d>>::__expand_read_method(
651 self, scope, pos,
652 )
653 }
654
655 fn __expand_as_linear_slice_method(
656 &self,
657 _scope: &Scope,
658 _pos: NativeExpand<usize>,
659 _end: NativeExpand<usize>,
660 ) -> &SliceExpand<Vector<E, N>> {
661 unimplemented!("ErasedTensor: no slice yet, see the module docs")
662 }
663
664 fn __expand_shape_method(&self, scope: &Scope) -> NativeExpand<usize> {
665 self.state_read().__expand_lines_method(scope)
666 }
667
668 fn __expand_is_in_bounds_method(
669 &self,
670 scope: &Scope,
671 pos: NativeExpand<usize>,
672 ) -> NativeExpand<bool> {
673 let lines = self.state_read().__expand_lines_method(scope);
674 pos.__expand_lt_method(scope, &lines)
675 }
676
677 fn __expand_tensor_map_load_method(
678 &self,
679 _scope: &Scope,
680 _barrier: &NativeExpand<Barrier>,
681 _shared_memory: &mut SliceExpand<Vector<E, N>>,
682 _pos: NativeExpand<usize>,
683 ) {
684 unimplemented!("ErasedTensor: not a tensor map")
685 }
686}
687
688impl<E: Numeric, N: Size, IO: ErasedIoWrite> ViewOperationsMut<Vector<E, N>, Coords1d>
689 for ErasedTensor<E, IO>
690{
691}
692
693impl<E: Numeric, N: Size, IO: ErasedIoWrite> ViewOperationsMutExpand<Vector<E, N>, Coords1d>
694 for ErasedTensorExpand<E, IO>
695{
696 fn __expand_write_method(
697 &self,
698 scope: &Scope,
699 pos: NativeExpand<usize>,
700 value: <Vector<E, N> as CubeType>::ExpandType,
701 ) {
702 self.check_width::<N>(scope, "write");
703 self.state_write()
704 .__expand_write_line_method(scope, pos, value.into())
705 }
706
707 fn __expand_write_checked_method(
708 &self,
709 scope: &Scope,
710 pos: NativeExpand<usize>,
711 value: <Vector<E, N> as CubeType>::ExpandType,
712 ) {
713 let in_bounds =
714 <Self as ViewOperationsExpand<Vector<E, N>, Coords1d>>::__expand_is_in_bounds_method(
715 self, scope, pos,
716 );
717 if_expand(scope, in_bounds, |scope| {
718 <Self as ViewOperationsMutExpand<Vector<E, N>, Coords1d>>::__expand_write_method(
719 self, scope, pos, value,
720 )
721 })
722 }
723
724 fn __expand_as_linear_slice_mut_method(
725 &self,
726 _scope: &Scope,
727 _pos: NativeExpand<usize>,
728 _end: NativeExpand<usize>,
729 ) -> &mut SliceExpand<Vector<E, N>> {
730 unimplemented!("ErasedTensor: no slice yet, see the module docs")
731 }
732
733 fn __expand_tensor_map_store_method(
734 &self,
735 _scope: &Scope,
736 _shared_memory: &SliceExpand<Vector<E, N>>,
737 _pos: NativeExpand<usize>,
738 ) {
739 unimplemented!("ErasedTensor: not a tensor map")
740 }
741}