1use cubecl::{prelude::*, unexpanded};
4
5use cubecl_std::tensor::layout::CoordsDyn;
6
7use cubecl::zspace::SmallVec;
8
9use crate::stage::pipeline::payload::base::{StageOperand, StageSpec, stage_one};
10use crate::*;
11
12#[derive(CubeType, Clone)]
14#[expand(derive(Clone))]
15pub struct Tile {
16 pub(crate) kind: TileKind,
17 pub(crate) factor: Factor,
20 #[cube(comptime)]
22 pub(crate) place: Placement,
23 #[cube(comptime)]
26 pub(crate) element: ElemType,
27}
28
29#[cube]
30impl Tile {
31 pub(crate) fn new<T: Numeric>(kind: TileKind, #[comptime] place: Placement) -> Tile {
33 Tile {
34 kind,
35 factor: Factor::none(),
36 place,
37 element: elem_type_of::<T>(),
38 }
39 }
40
41 pub fn space(&self) -> comptime_type!(Space) {
43 comptime!(self.place.space.clone())
44 }
45
46 pub fn with_levels(&self, #[comptime] depth: usize, #[comptime] levels: Vec<Level>) -> Tile {
48 Tile {
49 kind: self.kind.clone(),
50 factor: self.factor.clone(),
51 place: comptime!(Placement::new(self.place.space.clone(), depth, levels)),
52 element: comptime!(self.element),
53 }
54 }
55
56 pub(crate) fn mem(&self, #[comptime] site: &str) -> &Memory {
58 match &self.kind {
59 TileKind::Memory(g) => g,
60 TileKind::PlaneTile(_) | TileKind::PlanePartition(_) => {
61 panic!("Tile::{site}: a plane tile has no memory view")
62 }
63 TileKind::TmaGmem(_) => panic!("Tile::{site}: a tma source has no element view"),
64 TileKind::Procedural(_) | TileKind::Lines(_) => {
65 panic!("Tile::{site}: a procedural tile and the plane's units have no memory view")
66 }
67 }
68 }
69
70 pub(crate) fn recipe(&self) -> &Procedural {
72 match &self.kind {
73 TileKind::Procedural(recipe) => recipe,
74 TileKind::Memory(_)
75 | TileKind::PlaneTile(_)
76 | TileKind::PlanePartition(_)
77 | TileKind::TmaGmem(_)
78 | TileKind::Lines(_) => panic!("Tile::recipe: only a procedural tile has a recipe"),
79 }
80 }
81
82 pub(crate) fn mem_mut(&mut self, #[comptime] site: &str) -> &mut Memory {
83 match &mut self.kind {
84 TileKind::Memory(g) => g,
85 TileKind::PlaneTile(_) | TileKind::PlanePartition(_) => {
86 panic!("Tile::{site}: a plane tile has no memory view")
87 }
88 TileKind::TmaGmem(_) => panic!("Tile::{site}: a tma source is not writable"),
89 TileKind::Procedural(_) | TileKind::Lines(_) => {
90 panic!("Tile::{site}: a procedural tile and the plane's units are not writable")
91 }
92 }
93 }
94
95 pub fn delivery(&self) -> comptime_type!(Delivery) {
97 match &self.kind {
98 TileKind::Memory(d) => comptime!(d.access.delivery),
99 TileKind::TmaGmem(_) => comptime!(Delivery::Tma),
100 TileKind::PlaneTile(_) | TileKind::PlanePartition(_) => {
101 panic!("Tile::delivery: a resident fragment is not a stage source")
102 }
103 TileKind::Procedural(_) | TileKind::Lines(_) => comptime!(Delivery::Procedural),
104 }
105 }
106
107 pub(crate) fn fragments_bounce(&self, #[comptime] in_place: bool) -> comptime_type!(bool) {
111 match &self.kind {
112 TileKind::Memory(m) => {
113 let addressed = m.store.addressed();
114 comptime!(
115 FragmentDrain::of(&m.access, addressed, in_place) == FragmentDrain::Bounce
116 )
117 }
118 TileKind::PlaneTile(_)
119 | TileKind::PlanePartition(_)
120 | TileKind::TmaGmem(_)
121 | TileKind::Procedural(_)
122 | TileKind::Lines(_) => comptime!(false),
123 }
124 }
125
126 pub(crate) fn runtime_map(&self) -> RuntimeMap {
127 match &self.kind {
128 TileKind::Memory(g) => g.map.clone(),
129 TileKind::PlaneTile(_)
130 | TileKind::PlanePartition(_)
131 | TileKind::TmaGmem(_)
132 | TileKind::Procedural(_)
133 | TileKind::Lines(_) => RuntimeMap::integral(comptime!(self.place.space.rank())),
134 }
135 }
136
137 pub(crate) fn units(&self) -> comptime_type!(usize) {
139 let fill = self.fill_units();
140 comptime!(fill.count)
141 }
142
143 pub(crate) fn fill_units(&self) -> comptime_type!(FillUnits) {
146 match &self.kind {
147 TileKind::Memory(d) => comptime!(d.access.fill),
148 TileKind::TmaGmem(t) => comptime!(FillUnits::cube(t.units)),
149 TileKind::Procedural(_)
150 | TileKind::Lines(_)
151 | TileKind::PlaneTile(_)
152 | TileKind::PlanePartition(_) => comptime!(FillUnits::cube(0)),
153 }
154 }
155
156 pub fn vector_size(&self) -> comptime_type!(usize) {
159 match &self.kind {
160 TileKind::Memory(d) => {
161 let load = d.vector_tile(&comptime!(self.place.space.clone()));
162 comptime!(load.along_one_axis("a reader asking Tile::vector_size").1)
163 }
164 TileKind::Lines(c) => c.line(),
165 TileKind::TmaGmem(t) => comptime!(t.packing.served(1)),
167 TileKind::PlaneTile(_) | TileKind::PlanePartition(_) | TileKind::Procedural(_) => {
168 comptime!(1usize)
169 }
170 }
171 }
172
173 pub(crate) fn vector_tile(&self) -> comptime_type!(VectorTile) {
176 let space = comptime!(self.place.space.clone());
177 match &self.kind {
178 TileKind::Memory(d) => d.vector_tile(&space),
179 TileKind::Lines(_)
180 | TileKind::PlaneTile(_)
181 | TileKind::PlanePartition(_)
182 | TileKind::TmaGmem(_)
183 | TileKind::Procedural(_) => {
184 let width = self.vector_size();
185 comptime!(VectorTile::run(space.axis_at(space.rank() - 1), width))
186 }
187 }
188 }
189
190 pub(crate) fn request_init_from(
193 &mut self,
194 #[comptime] init_from: InitFrom,
195 ) -> comptime_type!(InitFrom) {
196 match &mut self.kind {
197 TileKind::Memory(d) => {
198 let init_from = comptime!(match d.access.write {
199 Write::Accumulate | Write::Exclusive(_) => InitFrom::Identity,
200 Write::Replace => init_from,
201 });
202 d.set_init_from(comptime!(init_from));
203 comptime!(init_from)
204 }
205 TileKind::PlaneTile(_)
206 | TileKind::PlanePartition(_)
207 | TileKind::TmaGmem(_)
208 | TileKind::Procedural(_)
209 | TileKind::Lines(_) => {
210 comptime!(InitFrom::Cell)
211 }
212 }
213 }
214
215 pub(crate) fn stage_element(&self) -> comptime_type!(StageElement) {
218 let decodes = self.decodes();
219 let staged = match &self.kind {
220 TileKind::Memory(d) => d.stage_element(),
221 TileKind::TmaGmem(t) => comptime!(match t.packing {
222 Packing::Plain => StageElement::Served,
223 Packing::Packed { .. } => StageElement::Stored,
224 }),
225 TileKind::PlaneTile(_)
226 | TileKind::PlanePartition(_)
227 | TileKind::Procedural(_)
228 | TileKind::Lines(_) => {
229 comptime!(StageElement::Served)
230 }
231 };
232 comptime!(match decodes {
233 true => StageElement::Decoded,
234 false => staged,
235 })
236 }
237
238 pub(crate) fn read_element(&self) -> comptime_type!(ElemType) {
242 let packing = self.packing();
243 comptime!(match packing {
244 Packing::Plain => self.element,
245 Packing::Packed { .. } => f32::elem_type_native(),
246 })
247 }
248
249 pub(crate) fn staged_element(
252 &self,
253 #[comptime] decoded: Option<ElemType>,
254 ) -> comptime_type!(ElemType) {
255 let staged = self.stage_element();
256 let read = self.read_element();
257 comptime!(match staged {
258 StageElement::Decoded => decoded.unwrap_or(read),
259 StageElement::Served | StageElement::Stored => self.element,
260 })
261 }
262
263 pub(crate) fn packing(&self) -> comptime_type!(Packing) {
265 match &self.kind {
266 TileKind::Memory(d) => d.packing(),
267 TileKind::Lines(c) => c.packing(),
268 TileKind::TmaGmem(t) => comptime!(t.packing),
269 TileKind::PlaneTile(_) | TileKind::PlanePartition(_) | TileKind::Procedural(_) => {
270 comptime!(Packing::Plain)
271 }
272 }
273 }
274
275 pub fn gathered(&self) -> comptime_type!(bool) {
277 let projection = self.projection();
278 comptime!(projection.composition() == Composition::Overlapping)
279 }
280
281 pub(crate) fn is_procedural(&self) -> comptime_type!(bool) {
283 match &self.kind {
284 TileKind::Procedural(_) => comptime!(true),
285 TileKind::Memory(_)
286 | TileKind::PlaneTile(_)
287 | TileKind::PlanePartition(_)
288 | TileKind::TmaGmem(_)
289 | TileKind::Lines(_) => comptime!(false),
290 }
291 }
292
293 pub(crate) fn is_memory(&self) -> comptime_type!(bool) {
295 match &self.kind {
296 TileKind::Memory(_) => comptime!(true),
297 TileKind::Procedural(_)
298 | TileKind::PlaneTile(_)
299 | TileKind::PlanePartition(_)
300 | TileKind::TmaGmem(_)
301 | TileKind::Lines(_) => comptime!(false),
302 }
303 }
304
305 pub fn is_shared(&self) -> comptime_type!(bool) {
307 match &self.kind {
308 TileKind::Memory(m) => comptime!(m.address == AddressSpace::Shared),
309 TileKind::Procedural(_)
310 | TileKind::PlaneTile(_)
311 | TileKind::PlanePartition(_)
312 | TileKind::TmaGmem(_)
313 | TileKind::Lines(_) => comptime!(false),
314 }
315 }
316
317 pub(crate) fn factors(&self) -> comptime_type!(Option<usize>) {
319 match &self.kind {
320 TileKind::Procedural(data) => data.factorization(),
321 TileKind::Memory(_)
322 | TileKind::PlaneTile(_)
323 | TileKind::PlanePartition(_)
324 | TileKind::TmaGmem(_)
325 | TileKind::Lines(_) => comptime!(None),
326 }
327 }
328
329 pub(crate) fn window_signed(&self) -> comptime_type!(bool) {
330 match &self.kind {
331 TileKind::Memory(d) => comptime!(d.window.signed),
332 _ => comptime!(false),
333 }
334 }
335
336 pub(crate) fn window_boundaries(
338 &self,
339 ) -> comptime_type!(SmallVec<[Option<Boundary>; Space::MAX_RANK]>) {
340 match &self.kind {
341 TileKind::Memory(d) => comptime!(d.window.boundaries.clone()),
342 _ => comptime!(SmallVec::new()),
343 }
344 }
345
346 pub(crate) fn separable_boundaries(
348 &self,
349 ) -> comptime_type!(SmallVec<[Option<Boundary>; Space::MAX_RANK]>) {
350 match &self.kind {
351 TileKind::Memory(d) =>
352 {
353 #[comptime]
354 match &d.source_window {
355 ComptimeOption::Some(source) => comptime!(source.boundaries.clone()),
356 ComptimeOption::None => comptime!(match d.address {
357 AddressSpace::Global => d.window.boundaries.clone(),
358 AddressSpace::Shared => SmallVec::new(),
359 }),
360 }
361 }
362 TileKind::Procedural(_)
363 | TileKind::PlaneTile(_)
364 | TileKind::PlanePartition(_)
365 | TileKind::TmaGmem(_)
366 | TileKind::Lines(_) => comptime!(SmallVec::new()),
367 }
368 }
369
370 pub(crate) fn factor_normalization(&self) -> comptime_type!(Option<Normalization>) {
372 match &self.kind {
373 TileKind::Procedural(data) => comptime!(data.normalization.clone()),
374 TileKind::Memory(_)
375 | TileKind::PlaneTile(_)
376 | TileKind::PlanePartition(_)
377 | TileKind::TmaGmem(_)
378 | TileKind::Lines(_) => comptime!(None),
379 }
380 }
381
382 pub(crate) fn separable_factor<T: Numeric>(
384 &self,
385 pos: CoordsDyn,
386 #[comptime] factor: usize,
387 ) -> T {
388 match &self.kind {
389 TileKind::Procedural(data) => {
390 data.read_factor_at::<T>(&pos, factor, comptime!(self.place.space.clone()))
391 }
392 TileKind::Memory(_)
393 | TileKind::PlaneTile(_)
394 | TileKind::PlanePartition(_)
395 | TileKind::TmaGmem(_)
396 | TileKind::Lines(_) => {
397 panic!(
398 "Tile::separable_factor: a tile read from a buffer states no factorization, \
399 so `factors` answered `None` and there is no factor {factor} to evaluate"
400 )
401 }
402 }
403 }
404
405 pub(crate) fn witnesses(&self, #[comptime] axis: Axis) -> comptime_type!(bool) {
407 let bounded = self.bounded();
408 let projection = self.projection();
409 comptime!(
410 bounded
411 && self.place.space.contains(axis)
412 && self.place.space.is_dynamic(axis)
413 && Witness::new(&projection, axis).is_some()
414 )
415 }
416
417 pub(crate) fn projection(&self) -> comptime_type!(Projection) {
419 match &self.kind {
420 TileKind::Memory(g) => comptime!(g.projection.clone()),
421 TileKind::Lines(c) => c.projection(),
422 TileKind::PlaneTile(_)
423 | TileKind::PlanePartition(_)
424 | TileKind::TmaGmem(_)
425 | TileKind::Procedural(_) => {
426 comptime!(Projection::direct_over(&self.place.space))
427 }
428 }
429 }
430
431 pub(crate) fn nested_like(self, like: &Tile) -> Tile {
433 Tile {
434 kind: self.kind,
435 factor: self.factor,
436 place: comptime!(Placement::new(
437 self.place.space.clone(),
438 like.place.depth,
439 like.place.levels.clone()
440 )),
441 element: comptime!(self.element),
442 }
443 }
444
445 pub fn within(&self, #[comptime] axis: Axis, from: usize, until: usize) -> Tile {
447 let tile_kind = match &self.kind {
448 TileKind::Memory(g) => TileKind::new_Memory(g.within(axis, from, until)),
449 TileKind::TmaGmem(_)
450 | TileKind::Procedural(_)
451 | TileKind::Lines(_)
452 | TileKind::PlaneTile(_)
453 | TileKind::PlanePartition(_) => panic!(
454 "Tile::within: only a memory operand carries a window to place; a fragment, a \
455 tensor-map source and a recipe have none"
456 ),
457 };
458 Tile {
459 kind: tile_kind,
460 factor: self.factor.clone(),
461 place: comptime!(self.place.clone()),
462 element: comptime!(self.element),
463 }
464 }
465
466 pub fn read_page_entry(&self, region: &Region) -> PageEntry {
470 let cut = self.at(region);
471 match &cut.kind {
472 TileKind::Memory(g) => g.layout.pinned_entry(),
473 TileKind::TmaGmem(_)
474 | TileKind::Procedural(_)
475 | TileKind::Lines(_)
476 | TileKind::PlaneTile(_)
477 | TileKind::PlanePartition(_) => PageEntry::unlisted(),
478 }
479 }
480
481 pub fn with_page_entry(&self, entry: &PageEntry) -> Tile {
487 let kind = match &self.kind {
488 TileKind::Memory(g) => TileKind::new_Memory(g.with_page_entry(entry)),
489 TileKind::TmaGmem(_)
490 | TileKind::Procedural(_)
491 | TileKind::Lines(_)
492 | TileKind::PlaneTile(_)
493 | TileKind::PlanePartition(_) => self.kind.clone(),
494 };
495 Tile {
496 kind,
497 factor: self.factor.clone(),
498 place: comptime!(self.place.clone()),
499 element: comptime!(self.element),
500 }
501 }
502
503 pub fn at(&self, region: &Region) -> Tile {
505 let skip = comptime!({
506 let path = ®ion.path;
507 assert!(
508 path.base() <= self.place.depth && self.place.depth < path.depth(),
509 "Tile::at: this tile sits {} levels down its nest, and the region's path runs \
510 from depth {} through {} levels; state the loops the tile is missing above it",
511 self.place.depth,
512 path.base(),
513 path.len()
514 );
515 self.place.depth - path.base()
516 });
517 self.at_from(region, skip)
518 }
519
520 fn at_from(&self, region: &Region, #[comptime] i: usize) -> Tile {
521 let sub = self.at_step(®ion.step(i));
522 if comptime!(i + 1 == region.path.len()) {
523 sub
524 } else {
525 sub.at_from(region, comptime!(i + 1))
526 }
527 }
528
529 pub(crate) fn at_step(&self, step: &Step) -> Tile {
531 let tile_kind = match &self.kind {
532 TileKind::Memory(g) => {
533 TileKind::new_Memory(g.at(step, comptime!(self.place.space.clone())))
534 }
535 TileKind::TmaGmem(t) => {
536 TileKind::new_TmaGmem(t.at(step, comptime!(self.place.space.clone())))
537 }
538 TileKind::Procedural(p) => {
539 TileKind::new_Procedural(p.at(step, comptime!(self.place.space.clone())))
540 }
541 TileKind::Lines(c) => {
542 TileKind::new_Lines(c.at(step, comptime!(self.place.space.clone())))
543 }
544 TileKind::PlaneTile(t) => {
546 comptime!(assert!(
547 !MatrixGrid::new(&step.level, &self.place.space).cuts(),
548 "Tile::at: a level that cuts tiles cannot select into a single plane \
549 tile (it needs a partition, or a memory output)"
550 ));
551 TileKind::new_PlaneTile(t.clone())
552 }
553 TileKind::PlanePartition(p) => p.at_step(step, comptime!(self.place.space.clone())),
554 };
555 let windowed = Tile {
556 kind: tile_kind,
557 factor: self.factor.at(step),
558 place: comptime!(self.place.at_step(&step.level)),
559 element: comptime!(self.element),
560 };
561 windowed.indexed_at(step)
562 }
563
564 fn bounded(&self) -> comptime_type!(bool) {
566 match &self.kind {
567 TileKind::Memory(_) | TileKind::TmaGmem(_) => comptime!(true),
568 TileKind::PlaneTile(_)
569 | TileKind::PlanePartition(_)
570 | TileKind::Procedural(_)
571 | TileKind::Lines(_) => {
572 comptime!(false)
573 }
574 }
575 }
576
577 pub(crate) fn runtime_extent(&self, #[comptime] axis: Axis) -> usize {
580 let projection = self.projection();
581 let p = comptime!(
582 Witness::new(&projection, axis)
583 .unwrap_or_else(|| panic!(
584 "Tile::runtime_extent: no bound of this operand is {axis:?}'s own extent (it \
585 gathers over it, or splits it across storage fragments); ask an operand that \
586 witnesses it"
587 ))
588 .dim()
589 );
590 let raw = match &self.kind {
591 TileKind::Memory(g) => g.walked_bound(p).cast::<usize>(),
592 TileKind::TmaGmem(t) => t.bound[p].cast::<usize>(),
593 TileKind::PlaneTile(_) | TileKind::PlanePartition(_) => {
594 panic!("Tile::runtime_extent: a plane tile has no extent")
595 }
596 TileKind::Procedural(_) | TileKind::Lines(_) => {
597 panic!(
598 "Tile::runtime_extent: a procedural tile and the plane's units have no extent"
599 )
600 }
601 };
602 let last = comptime!(projection.physical_rank() - 1);
604 let w = self.vector_size();
605 comptime!(if p == last { w } else { 1usize }) * raw
606 }
607
608 pub(crate) fn runtime_space(&self) -> Space {
610 witnessed_space(comptime!(self.place.space.clone()), self, self, self)
611 }
612
613 pub fn walk(&self) -> Walk {
615 let space = self.runtime_space();
616 Region::rooted(
617 &space,
618 comptime!(self.place.levels.clone()),
619 comptime!(self.place.depth),
620 )
621 .walk()
622 }
623
624 pub fn over(&self, #[comptime] level: &Level) -> Walk {
626 let space = self.runtime_space();
627 Walk::of(
628 &space,
629 comptime!(level.clone()),
630 Region::rooted(
631 &space,
632 comptime!(self.place.levels.clone()),
633 comptime!(self.place.depth),
634 ),
635 )
636 }
637
638 pub fn zero(&mut self) {
640 let define!(Served) = self.element;
641 match &mut self.kind {
642 TileKind::Memory(d) => d.zero::<Served>(),
643 TileKind::PlaneTile(t) => t.zero::<Served>(),
644 TileKind::PlanePartition(p) => p.zero::<Served>(),
645 TileKind::TmaGmem(_) => panic!("Tile::zero: a tma source is not writable"),
646 TileKind::Procedural(_) | TileKind::Lines(_) => {
647 panic!("Tile::zero: a procedural tile and the plane's units are not writable")
648 }
649 }
650 }
651
652 pub(crate) fn init_identity<S: Numeric>(&mut self, #[comptime] monoid: Monoid) {
654 match comptime!(monoid) {
655 Monoid::Sum => self.zero(),
656 Monoid::Prod | Monoid::Max | Monoid::Min => {
657 self.init::<S>(Monoid::identity::<S>(monoid))
658 }
659 }
660 }
661
662 pub fn scale_by<E: Numeric>(&mut self, factor: E) {
664 let define!(Served) = self.element;
665 let factor = Served::cast_from(factor);
666 match &mut self.kind {
667 TileKind::PlaneTile(t) => t.scale::<Served>(factor),
668 TileKind::PlanePartition(p) => p.scale::<Served>(factor),
669 TileKind::Memory(_) => panic!(
670 "Tile::scale_by: a memory tile is scaled by the cube, not by one unit (Tile::mul)"
671 ),
672 TileKind::TmaGmem(_) | TileKind::Procedural(_) | TileKind::Lines(_) => {
673 panic!("Tile::scale_by: not writable")
674 }
675 }
676 }
677
678 pub(crate) fn fragment_grid(&self) -> comptime_type!(((usize, usize), usize, usize)) {
680 let edges = comptime!(MatrixAxes::edges(&self.place.space));
681 let rows = comptime!(self.place.space.extent_at(edges.row_split));
682 let cols = comptime!(self.place.space.extent_at(edges.col_split));
683 match &self.kind {
684 TileKind::PlanePartition(p) => comptime!(((p.m_tiles, p.n_tiles), p.rows, p.cols)),
686 TileKind::PlaneTile(_) => comptime!(((1, 1), rows, cols)),
687 TileKind::Memory(_)
688 | TileKind::TmaGmem(_)
689 | TileKind::Procedural(_)
690 | TileKind::Lines(_) => {
691 panic!("Tile::fragment_grid: only a plane-resident accumulator holds fragments")
692 }
693 }
694 }
695
696 pub fn init<E: Numeric>(&mut self, val: E) {
698 let define!(Served) = self.element;
699 let val = Served::cast_from(val);
700 match &mut self.kind {
701 TileKind::Memory(d) => d.init::<Served>(val),
702 TileKind::PlaneTile(t) => t.init::<Served>(val),
703 TileKind::PlanePartition(p) => p.init::<Served>(val),
704 TileKind::TmaGmem(_) => panic!("Tile::init: a tma source is not writable"),
705 TileKind::Procedural(_) | TileKind::Lines(_) => {
706 panic!("Tile::init: a procedural tile and the plane's units are not writable")
707 }
708 }
709 }
710
711 pub fn stage(&self, #[comptime] level: Level, #[comptime] storage: StageStorage) -> Tile {
713 let define!(Staged) = self.staged_element(comptime!(None));
714 Memory::stage::<Staged>(
715 &self.shaped_for(comptime!(level.clone())),
716 level,
717 storage,
718 comptime!(None),
719 comptime!(StageOwner::CUBE),
720 )
721 }
722
723 pub fn stage_for(
731 &self,
732 walk: &Walk,
733 #[comptime] storage: StageStorage,
734 #[comptime] decoded: Option<ElemType>,
735 ) -> Tile {
736 let owner = walk.stage_owner();
737 stage_one(
738 self,
739 comptime!(StageSpec {
740 level: walk.level.clone(),
741 depth: walk.depth(),
742 storage,
743 width: None,
744 decoded,
745 owner,
746 }),
747 )
748 }
749
750 pub fn scratch<T: Numeric>(
758 walk: &Walk,
759 #[comptime] axes: Vec<Axis>,
760 #[comptime] storage: StageStorage,
761 ) -> Tile {
762 let owner = walk.stage_owner();
763 let place = comptime!(Placement::new(
764 walk.level.child(&walk.space).subspace(&axes),
765 walk.depth(),
766 walk.parent.path.root_levels(),
767 ));
768 let tile =
769 Memory::smem_owned::<T>(place.space.clone(), 1usize, storage, 0usize, 0usize, owner);
770 Tile::new::<T>(tile.kind, place)
771 }
772
773 pub fn shared<T: Numeric>(#[comptime] space: Space, #[comptime] storage: StageStorage) -> Tile {
775 Memory::smem::<T>(space, comptime!(1usize), storage, comptime!(0usize))
776 }
777
778 pub fn smem<T: Numeric>(
781 #[comptime] space: Space,
782 #[comptime] vector_size: usize,
783 #[comptime] storage: StageStorage,
784 #[comptime] units: usize,
785 ) -> Tile {
786 Memory::smem::<T>(space, vector_size, storage, units)
787 }
788
789 pub fn smem_box(&self) -> Tile {
795 let define!(Served) = self.element;
796 let units = self.units();
797 Memory::smem_owned::<Served>(
798 comptime!(self.place.space.clone()),
799 1usize,
800 comptime!(StageStorage::Strided),
801 units,
802 comptime!(Delivery::TMA_STAGE_ALIGNMENT),
803 comptime!(StageOwner::CUBE),
804 )
805 .nested_like(self)
806 }
807
808 pub fn copy_from(&mut self, src: &Tile) {
819 let decodes = src.decodes();
820 comptime!(assert!(
821 !decodes,
822 "Tile::copy_from: this source carries scales or a table that a copy as stored cannot \
823 keep; write it with them applied (`Tile::dequantize_from`)"
824 ));
825 let (beside, kept) = (src.scales_beside(), self.scales_beside());
826 comptime!(assert!(
827 !beside || kept,
828 "Tile::copy_from: this source's scales are kept beside its values, which only a stage \
829 that keeps them too takes as they are stored"
830 ));
831 let define!(Served) = self.element;
832 let define!(Source) = src.element;
833 self.copy_unscaled_from::<Served, Source>(src);
834 }
835
836 pub fn dequantize_from(&mut self, src: &Tile) {
843 match &src.kind {
844 TileKind::Memory(_) => {
845 let into_memory = self.is_memory();
846 comptime!(assert!(
847 into_memory,
848 "Tile::dequantize_from: values are decoded into memory; to load them into \
849 fragments, contract them: the fragment lands them"
850 ));
851 let define!(Served) = self.element;
852 let define!(Source) = src.read_element();
853 self.copy_scaled_from::<Served, Source>(src);
854 }
855 TileKind::PlaneTile(_) | TileKind::PlanePartition(_) => {
856 let dest = self.clone();
857 src.drained_into(&dest);
858 }
859 TileKind::TmaGmem(_) | TileKind::Procedural(_) | TileKind::Lines(_) => panic!(
860 "Tile::dequantize_from: values are decoded from memory, and a sum from its \
861 fragments; a tensor-map source lands bytes as they lie, so stage it first"
862 ),
863 }
864 }
865
866 pub fn quantize_from(&mut self, src: &Tile) {
876 let (from, into) = (src.is_memory(), self.is_memory());
877 comptime!(assert!(
878 from && into,
879 "Tile::quantize_from: values are encoded from memory into memory"
880 ));
881 let define!(Source) = src.element;
882 self.copy_quantized_from::<Source>(src);
883 }
884
885 fn copy_unscaled_from<S: Numeric, R: Numeric>(&mut self, src: &Tile) {
886 let space = comptime!(self.place.space.clone());
888 match &src.kind {
889 TileKind::PlanePartition(s) => match &mut self.kind {
890 TileKind::Memory(d) => s.fragment().store_window::<S>(d, space),
891 TileKind::PlaneTile(_)
892 | TileKind::PlanePartition(_)
893 | TileKind::TmaGmem(_)
894 | TileKind::Procedural(_)
895 | TileKind::Lines(_) => {
896 panic!("Tile::copy_from: a fragment stores into memory")
897 }
898 },
899 TileKind::Memory(_)
900 | TileKind::PlaneTile(_)
901 | TileKind::TmaGmem(_)
902 | TileKind::Procedural(_)
903 | TileKind::Lines(_) => match (&mut self.kind, &src.kind) {
904 (TileKind::Lines(d), TileKind::Memory(_)) => d.load::<R>(src),
905 (TileKind::PlanePartition(d), TileKind::Memory(_)) => {
907 d.fill_from::<S>(src, DeepLoad::Steps)
908 }
909 (TileKind::PlaneTile(d), TileKind::Memory(_)) => {
910 d.load_window::<S>(src, DeepLoad::Steps)
911 }
912 (TileKind::Memory(d), TileKind::PlaneTile(s)) => s.store_window::<S>(d, space),
913 (TileKind::Memory(d), TileKind::TmaGmem(s)) => s.load_into::<S>(d),
914 (TileKind::TmaGmem(d), TileKind::Memory(s)) => d.store_from::<S>(s),
915 (TileKind::Memory(d), TileKind::Memory(s)) => d.load_from::<S>(s, space),
916 (TileKind::Memory(d), TileKind::Procedural(s)) => d.fill_procedural::<S>(s, space),
917 (TileKind::PlaneTile(_), TileKind::PlaneTile(_)) => {
918 panic!("Tile::copy_from: plane tile to plane tile cast not wired")
919 }
920 _ => panic!("Tile::copy_from: unsupported kind pairing"),
921 },
922 }
923 }
924
925 pub fn add_cast_from(&mut self, src: &Tile) {
927 let define!(Served) = self.element;
928 let define!(Source) = src.element;
929 match (&mut self.kind, &src.kind) {
930 (TileKind::PlaneTile(d), TileKind::PlaneTile(s)) => {
931 d.add_cast_from::<Served, Source>(s)
932 }
933 (TileKind::PlanePartition(d), TileKind::PlanePartition(s)) => {
934 let mut fragment = d.fragment();
936 fragment.add_cast_from::<Served, Source>(&s.fragment())
937 }
938 _ => panic!(
939 "Tile::add_cast_from: a resident block promotes into a resident block; memory \
940 is reached through `copy_cast_from`"
941 ),
942 }
943 }
944
945 pub fn copy_cast_from(&mut self, src: &Tile) {
948 let define!(Served) = self.element;
949 let define!(Source) = src.element;
950 let space = comptime!(self.place.space.clone());
951 let holder = comptime!(self.place.holder());
952 match (&mut self.kind, &src.kind) {
953 (TileKind::Memory(d), TileKind::PlaneTile(s)) => {
954 s.store_cast_window::<Source, Served>(d, space)
955 }
956 (TileKind::Memory(d), TileKind::PlanePartition(s)) => {
957 s.fragment().store_cast_window::<Source, Served>(d, space)
958 }
959 (TileKind::Memory(d), TileKind::Memory(s)) => d
962 .filled_by(holder)
963 .fill_cast_from::<Served, Source>(s, space),
964 _ => panic!(
965 "Tile::copy_cast_from: a fragment or a memory window stores into memory; nothing \
966 else casts"
967 ),
968 }
969 }
970
971 pub(crate) fn spill_to_scratch<S: Numeric>(&self) {
973 match &self.kind {
974 TileKind::PlaneTile(t) => t.spill_to_scratch::<S>(),
975 TileKind::PlanePartition(p) => p.fragment().spill_to_scratch::<S>(),
976 TileKind::Memory(_)
977 | TileKind::TmaGmem(_)
978 | TileKind::Procedural(_)
979 | TileKind::Lines(_) => {
980 panic!("Tile::spill_to_scratch: a plane-resident tile spills; nothing else does")
981 }
982 }
983 }
984
985 pub(crate) fn add_from_scratch<S: Numeric, R: Numeric>(&mut self, src: &Tile) {
987 let space = comptime!(self.place.space.clone());
988 match (&mut self.kind, &src.kind) {
989 (TileKind::Memory(d), TileKind::PlaneTile(s)) => s.add_from_scratch::<R, S>(d, space),
990 (TileKind::Memory(d), TileKind::PlanePartition(s)) => {
991 s.fragment().add_from_scratch::<R, S>(d, space)
992 }
993 _ => panic!("Tile::add_from_scratch: a spilled plane tile adds into memory"),
994 }
995 }
996
997 pub(crate) fn separable_physical_tap_in_bounds(
999 &self,
1000 pos: &CoordsDyn,
1001 #[comptime] axis: Axis,
1002 ) -> bool {
1003 match &self.kind {
1004 TileKind::Memory(data) => {
1006 let projection = comptime!(data.projection.clone());
1007 if comptime!(!projection.addresses(axis)) {
1008 true.runtime()
1009 } else {
1010 #[comptime]
1011 match &data.source_window {
1012 ComptimeOption::Some(source) => source.axes_in_bounds(
1013 &data.window.origin,
1014 pos,
1015 comptime!(projection.carriers(axis).to_vec()),
1016 ),
1017 ComptimeOption::None => match comptime!(data.address) {
1018 AddressSpace::Global => data
1019 .window
1020 .axes_in_bounds(pos, comptime!(projection.carriers(axis).to_vec())),
1021 AddressSpace::Shared => true.runtime(),
1022 },
1023 }
1024 }
1025 }
1026 TileKind::Procedural(data) => data.axis_in_bounds(pos, axis),
1027 TileKind::PlaneTile(_)
1028 | TileKind::PlanePartition(_)
1029 | TileKind::TmaGmem(_)
1030 | TileKind::Lines(_) => {
1031 panic!(
1032 "Tile::separable_physical_tap_in_bounds: a separable gather needs an addressable rhs"
1033 )
1034 }
1035 }
1036 }
1037}
1038
1039impl Tile {
1040 pub(crate) fn factor_dependencies(
1041 &self,
1042 _factors: Option<usize>,
1043 _row: Axis,
1044 _col: Axis,
1045 ) -> comptime_type!(Option<Vec<(bool, bool)>>) {
1046 unexpanded!()
1047 }
1048}
1049
1050impl TileExpand {
1051 pub fn space(&self) -> Space {
1053 self.place.space.clone()
1054 }
1055
1056 pub(crate) fn __expand_factor_dependencies_method(
1057 &self,
1058 scope: &Scope,
1059 factors: Option<usize>,
1060 row: Axis,
1061 col: Axis,
1062 ) -> Option<Vec<(bool, bool)>> {
1063 factors.map(|factors| match &self.kind {
1064 TileKindExpand::Procedural(data) => (0..factors)
1065 .map(|f| {
1066 (
1067 data.__expand_factor_reads_method(scope, f, row),
1068 data.__expand_factor_reads_method(scope, f, col),
1069 )
1070 })
1071 .collect(),
1072 TileKindExpand::Memory(_)
1073 | TileKindExpand::PlaneTile(_)
1074 | TileKindExpand::PlanePartition(_)
1075 | TileKindExpand::TmaGmem(_)
1076 | TileKindExpand::Lines(_) => (0..factors).map(|_| (true, true)).collect(),
1077 })
1078 }
1079}
1080
1081#[cube]
1084pub(crate) fn witnessed_space(#[comptime] space: Space, a: &Tile, b: &Tile, c: &Tile) -> Space {
1085 let mut sizes = Sequence::<usize>::new();
1086 if comptime!(!space.is_static()) {
1087 #[unroll]
1088 for p in 0..comptime!(space.rank()) {
1089 let axis = comptime!(space.axis_at(p));
1090 let size = match comptime!(space.extent_raw(axis)) {
1092 Extent::Static(n) => comptime!(n).runtime(),
1093 Extent::Dynamic => {
1094 let by_a = a.witnesses(axis);
1095 let by_b = b.witnesses(axis);
1096 let by_c = c.witnesses(axis);
1097 if comptime!(by_a) {
1098 a.runtime_extent(axis)
1099 } else if comptime!(by_b) {
1100 b.runtime_extent(axis)
1101 } else if comptime!(by_c) {
1102 c.runtime_extent(axis)
1103 } else {
1104 panic!(
1105 "witnessed_space: {axis:?} is Dynamic and no operand states its size; \
1106 every operand spanning it gathers over it, holds it Static, or is a \
1107 fragment. Keep it Static in the kernel space, or give the operation \
1108 an operand that maps it identically"
1109 )
1110 }
1111 }
1112 };
1113 sizes.push(size);
1114 }
1115 }
1116 Space::with_sizes(space, sizes)
1117}
1118
1119impl IntoIterator for Tile {
1121 type Item = Region;
1122 type IntoIter = std::vec::IntoIter<Region>;
1123
1124 fn into_iter(self) -> Self::IntoIter {
1125 unexpanded!()
1126 }
1127}
1128
1129impl IntoIterator for &Tile {
1130 type Item = Region;
1131 type IntoIter = std::vec::IntoIter<Region>;
1132
1133 fn into_iter(self) -> Self::IntoIter {
1134 unexpanded!()
1135 }
1136}
1137
1138impl Iterable for TileExpand {
1140 type Item = RegionExpand;
1141
1142 fn expand(self, scope: &Scope, body: &mut dyn FnMut(&Scope, RegionExpand)) {
1143 let walk = self.__expand_walk_method(scope);
1144 if walk.const_len() == Some(1) {
1145 walk.expand_unroll(scope, body)
1146 } else {
1147 walk.expand(scope, body)
1148 }
1149 }
1150
1151 fn expand_unroll(self, scope: &Scope, body: &mut dyn FnMut(&Scope, RegionExpand)) {
1152 self.__expand_walk_method(scope).expand_unroll(scope, body)
1153 }
1154}
1155
1156impl Iterable for &TileExpand {
1158 type Item = RegionExpand;
1159
1160 fn expand(self, scope: &Scope, body: &mut dyn FnMut(&Scope, RegionExpand)) {
1161 self.clone().expand(scope, body)
1162 }
1163
1164 fn expand_unroll(self, scope: &Scope, body: &mut dyn FnMut(&Scope, RegionExpand)) {
1165 self.clone().expand_unroll(scope, body)
1166 }
1167}
1168
1169impl Tile {
1170 pub fn mul(&self, _scale: &Tile) -> Tile {
1173 unexpanded!()
1174 }
1175
1176 pub(crate) fn reader(
1178 &self,
1179 _axes: MatrixAxes,
1180 _matrix: usize,
1181 _side: Side,
1182 _out: Space,
1183 _acc_axes: MatrixAxes,
1184 ) -> FactorReader {
1185 unexpanded!()
1186 }
1187
1188 pub fn factor(&self, _space: Space, _instruction: Instruction) -> (Tile, Factor) {
1200 unexpanded!()
1201 }
1202
1203 pub fn scaled(&self, _factor: &Factor) -> Tile {
1205 unexpanded!()
1206 }
1207
1208 pub(crate) fn placed(&self) -> Tile {
1213 unexpanded!()
1214 }
1215
1216 #[allow(dead_code)] pub(crate) fn lookup(&self, _table: &Tile) -> Tile {
1221 unexpanded!()
1222 }
1223
1224 pub(crate) fn refuse_factor(&self, _site: &str) {
1225 unexpanded!()
1226 }
1227
1228 pub(crate) fn quantized(&self) -> comptime_type!(bool) {
1231 unexpanded!()
1232 }
1233
1234 pub(crate) fn factor_levels(&self) -> comptime_type!(usize) {
1236 unexpanded!()
1237 }
1238
1239 pub(crate) fn unscaled(&self) -> Tile {
1241 unexpanded!()
1242 }
1243
1244 pub(crate) fn innermost_factor(&self) -> Factor {
1247 unexpanded!()
1248 }
1249
1250 pub(crate) fn decodes(&self) -> comptime_type!(bool) {
1257 unexpanded!()
1258 }
1259
1260 pub(crate) fn scales_beside(&self) -> comptime_type!(bool) {
1262 unexpanded!()
1263 }
1264
1265 pub(crate) fn with_scales_staged(&self, _operand: &Tile, _spec: StageSpec) -> Tile {
1268 unexpanded!()
1269 }
1270
1271 pub(crate) fn fill_scales_from(&self, _src: &Tile, _meeting: &Meeting) {
1274 unexpanded!()
1275 }
1276
1277 pub(crate) fn free_scales(&self) {
1279 unexpanded!()
1280 }
1281
1282 pub(crate) fn scale_operands(&self) -> comptime_type!(Vec<StageOperand>) {
1285 unexpanded!()
1286 }
1287}
1288
1289impl TileExpand {
1290 pub fn __expand_mul_method(&self, scope: &Scope, scale: &TileExpand) -> TileExpand {
1291 let values = self.place.space.clone();
1292 let projection = scale.clone().__expand_projection_method(scope);
1293 let varied: Vec<Axis> = scale
1296 .place
1297 .space
1298 .axes()
1299 .filter(|&axis| projection.addresses(axis))
1300 .collect();
1301 assert!(
1302 varied.iter().all(|&axis| values.contains(axis)),
1303 "Tile::mul: the scales vary along {varied:?} where the values span {:?}; a scale is \
1304 looked up at the value's coordinates, so every axis the scales vary along is one of \
1305 the values'",
1306 values.axes().collect::<Vec<_>>()
1307 );
1308 let mut out = self.clone();
1309 out.factor = out.factor.and(FactorExpand::of(scale));
1310 out
1311 }
1312
1313 pub(crate) fn __expand_reader_method(
1314 &self,
1315 scope: &Scope,
1316 axes: MatrixAxes,
1317 matrix: NativeExpand<usize>,
1318 side: Side,
1319 out: Space,
1320 acc_axes: MatrixAxes,
1321 ) -> FactorReaderExpand {
1322 self.refuse_codebook("a leaf's read");
1323 let values = self.place.space.clone();
1324 let vector_size = self.clone().__expand_vector_tile_method(scope).extents()[0].1;
1327 let factor = self.factor.clone();
1328 if let Some(inner) = factor.innermost_tile() {
1329 let projection = inner.clone().__expand_projection_method(scope);
1330 check_scales_omit_rather_than_divide(&projection);
1331 check_scales_ride(side, &inner.place.space, &out, acc_axes);
1332 let innermost = values.axis_at(values.rank() - 1);
1333 assert!(
1334 vector_size == 1 || !projection.addresses(innermost),
1335 "Tile::mul: a line runs {vector_size} values along {innermost:?}, which its \
1336 scales address, so one line lies under several scales; serve the factor one \
1337 value a line, or omit {innermost:?} from the scales"
1338 );
1339 }
1340 let lift = self.__expand_packing_method(scope).lift();
1343 let coarse = match lift == 1.0 {
1344 true => factor.coarse(scope),
1345 false => {
1346 let lift: NativeExpand<f32> =
1347 cubecl::ir::ExpandValue::constant((lift as u64).into(), f32::elem_type(scope))
1348 .into();
1349 MulExpand::__expand_mul_method(factor.coarse(scope), scope, lift)
1350 }
1351 };
1352 FactorReaderExpand {
1353 coarse,
1354 inner: factor.innermost(),
1355 scaled: factor.scaled() || lift != 1.0,
1356 values,
1357 axes,
1358 vector_size,
1359 matrix,
1360 }
1361 }
1362
1363 pub fn __expand_factor_method(
1364 &self,
1365 scope: &Scope,
1366 space: Space,
1367 instruction: Instruction,
1368 ) -> (TileExpand, FactorExpand) {
1369 let gives_up = !self.looks_up()
1370 && (self.factor.staged_beside(scope) || instruction.scales_its_sum_in_place());
1371 if !gives_up {
1372 return (self.clone(), FactorExpand::default());
1373 }
1374 let (inner, outer) = self.factor.split_on(scope, &space);
1375 let mut values = self.clone();
1376 values.factor = inner;
1377 (values, outer)
1378 }
1379
1380 pub fn __expand_scaled_method(&self, _scope: &Scope, factor: &FactorExpand) -> TileExpand {
1381 let mut out = self.clone();
1382 out.factor = out.factor.and(factor.clone());
1383 out
1384 }
1385
1386 pub(crate) fn __expand_placed_method(&self, _scope: &Scope) -> TileExpand {
1387 let mut out = self.clone();
1388 if let TileKindExpand::Memory(memory) = &mut out.kind {
1389 memory.store.packing = memory.store.packing.placed();
1390 }
1391 out
1392 }
1393
1394 fn looks_up(&self) -> bool {
1396 match &self.kind {
1397 TileKindExpand::Memory(memory) => memory.codebook.present(),
1398 _ => false,
1399 }
1400 }
1401
1402 pub(crate) fn __expand_quantized_method(&self, _scope: &Scope) -> bool {
1403 self.looks_up() || self.factor.scaled()
1404 }
1405
1406 pub(crate) fn __expand_factor_levels_method(&self, _scope: &Scope) -> usize {
1407 self.factor.depth()
1408 }
1409
1410 pub(crate) fn __expand_unscaled_method(&self, _scope: &Scope) -> TileExpand {
1411 let mut out = self.clone();
1412 out.factor = FactorExpand::default();
1413 out
1414 }
1415
1416 pub(crate) fn __expand_innermost_factor_method(&self, _scope: &Scope) -> FactorExpand {
1417 self.factor.innermost()
1418 }
1419
1420 pub(crate) fn __expand_decodes_method(&self, scope: &Scope) -> bool {
1421 self.__expand_quantized_method(scope) && !self.__expand_scales_beside_method(scope)
1422 }
1423
1424 pub(crate) fn __expand_scales_beside_method(&self, scope: &Scope) -> bool {
1425 !self.looks_up() && self.factor.staged_beside(scope)
1426 }
1427
1428 pub(crate) fn __expand_with_scales_staged_method(
1429 &self,
1430 scope: &Scope,
1431 operand: &TileExpand,
1432 spec: StageSpec,
1433 ) -> TileExpand {
1434 let mut stage = self.clone();
1435 if operand.__expand_scales_beside_method(scope) {
1436 stage.factor = operand.factor.staged(scope, &spec);
1437 }
1438 stage
1439 }
1440
1441 pub(crate) fn __expand_fill_scales_from_method(
1442 &self,
1443 scope: &Scope,
1444 src: &TileExpand,
1445 meeting: &MeetingExpand,
1446 ) {
1447 if self.factor.staged_beside(scope) {
1448 self.factor.fill_from(scope, &src.factor, meeting);
1449 }
1450 }
1451
1452 pub(crate) fn __expand_scale_operands_method(&self, scope: &Scope) -> Vec<StageOperand> {
1453 let keys = self.pending_keys();
1455 match self.__expand_scales_beside_method(scope) {
1456 true => self
1457 .factor
1458 .stage_operands(scope)
1459 .into_iter()
1460 .map(|scale| StageOperand {
1461 keys: keys.clone(),
1462 ..scale
1463 })
1464 .collect(),
1465 false => Vec::new(),
1466 }
1467 }
1468
1469 pub(crate) fn __expand_free_scales_method(&self, scope: &Scope) {
1470 if self.factor.staged_beside(scope) {
1471 self.factor.free_stage(scope);
1472 }
1473 }
1474
1475 pub fn __expand_lookup_method(&self, _scope: &Scope, table: &TileExpand) -> TileExpand {
1476 assert!(
1477 table.place.space.rank() == 1,
1478 "Tile::lookup: a table is a tile of one axis, its positions; this one spans {:?}",
1479 table.place.space.axes().collect::<Vec<_>>()
1480 );
1481 let mut out = self.clone();
1482 match &mut out.kind {
1483 TileKindExpand::Memory(memory) => {
1484 let packing = memory.store.packing;
1485 let Packing::Packed {
1486 field: Field::Index { bits },
1487 } = packing
1488 else {
1489 panic!(
1490 "Tile::lookup: the values index a table through an index field \
1491 (`Field::Index`), and these are stored as {packing:?}"
1492 )
1493 };
1494 assert!(
1495 table.place.space.extent_at(0) >= 1 << bits,
1496 "Tile::lookup: a {bits}-bit index names {} positions; the table holds {}",
1497 1usize << bits,
1498 table.place.space.extent_at(0)
1499 );
1500 assert!(
1501 !memory.codebook.present(),
1502 "Tile::lookup: these values already index a table"
1503 );
1504 memory.codebook = CodebookExpand::of(table);
1505 }
1506 TileKindExpand::PlaneTile(_)
1507 | TileKindExpand::PlanePartition(_)
1508 | TileKindExpand::TmaGmem(_)
1509 | TileKindExpand::Procedural(_)
1510 | TileKindExpand::Lines(_) => {
1511 panic!("Tile::lookup: a table is indexed by values stored in memory")
1512 }
1513 }
1514 out
1515 }
1516
1517 pub(crate) fn __expand_refuse_factor_method(&self, _scope: &Scope, site: &str) {
1518 self.refuse_codebook(site);
1519 if let TileKindExpand::Memory(_) = &self.kind {
1520 assert!(
1521 !self.factor.scaled(),
1522 "{site}: this leaf takes its operands from registers, where scales have nowhere to land; contract through a fragment or in memory"
1523 );
1524 }
1525 }
1526
1527 fn refuse_codebook(&self, site: &str) {
1529 if let TileKindExpand::Memory(memory) = &self.kind {
1530 assert!(
1531 !memory.codebook.present(),
1532 "{site}: these values are indices into a codebook (`TileArgLaunch::codebook`), \
1533 which only a copy or a decoded read decodes; copy them into a stage first \
1534 (`stage.copy_from(&w)`), or read them `decoded`"
1535 );
1536 }
1537 }
1538}
1539
1540#[cube]
1541impl Tile {
1542 #[allow(dead_code)] pub(crate) fn scale_at<S: Numeric>(&self, coords: &Coords<u32>, #[comptime] guard: Guard) -> S {
1546 match &self.kind {
1547 TileKind::Lines(lines) => lines.read::<S>(coords),
1548 TileKind::Memory(_) => self.value_in_line::<S>(coords, guard),
1549 TileKind::PlaneTile(_)
1550 | TileKind::PlanePartition(_)
1551 | TileKind::TmaGmem(_)
1552 | TileKind::Procedural(_) => {
1553 panic!("Tile::scale_at: a scale is read from memory or from the plane's units")
1554 }
1555 }
1556 }
1557
1558 #[allow(dead_code)] pub(crate) fn scale_for<S: Numeric>(
1563 &self,
1564 coords: &Coords<u32>,
1565 #[comptime] values: Space,
1566 #[comptime] guard: Guard,
1567 ) -> S {
1568 let own = self.scale_coords(coords, values);
1569 self.scale_at::<S>(&own, guard)
1570 }
1571
1572 pub(crate) fn scale_coords(
1575 &self,
1576 coords: &Coords<u32>,
1577 #[comptime] values: Space,
1578 ) -> Coords<u32> {
1579 let rank = comptime!(self.place.space.rank());
1580 let projection = self.projection();
1581 let mut own = Coords::<u32>::new();
1582 #[unroll]
1583 for p in 0..rank {
1584 let axis = comptime!(self.place.space.axis_at(p));
1585 if comptime!(values.contains(axis)) {
1586 own.push(coords.at(comptime!(values.position(axis))));
1587 } else {
1588 comptime!(assert!(
1589 !projection.addresses(axis),
1590 "Tile::scale_for: these scales vary along {axis:?}, which the values they \
1591 scale do not span"
1592 ));
1593 own.push(0u32);
1594 }
1595 }
1596 own
1597 }
1598
1599 #[allow(dead_code)] fn value_in_line<S: Numeric>(&self, coords: &Coords<u32>, #[comptime] guard: Guard) -> S {
1604 let define!(Served) = self.element;
1605 let space = comptime!(self.place.space.clone());
1606 let rank = comptime!(space.rank());
1607 let load = self.vector_tile();
1608 let width = comptime!(load.values());
1609 let packing = self.packing();
1610 let mut position = 0u32.runtime();
1611 #[unroll]
1612 for p in 0..rank {
1613 let axis = comptime!(space.axis_at(p));
1614 let extent = comptime!(load.extent_along(axis) as u32);
1615 if comptime!(extent > 1) {
1616 let step = comptime!(load.step_along(axis) as u32);
1617 position += coords.at(p).remainder(extent).times(step);
1618 }
1619 }
1620 match comptime!(packing) {
1621 Packing::Plain => {
1622 let size!(W) = width;
1623 let line = self
1624 .nd_packed::<Served, W>(guard)
1625 .read(load.index(coords, &space));
1626 let value = if comptime!(width > 1) {
1627 line.extract_dynamic(position.cast::<usize>())
1628 } else {
1629 line.extract(0usize)
1630 };
1631 S::cast_from(value)
1632 }
1633 Packing::Packed { field } => {
1634 let (bits, per_word) =
1635 comptime!((field.size_bits() as u32, field.per_word() as u32));
1636 let words_wide = comptime!(packing.physical(width));
1637 let size!(WP) = words_wide;
1638 let words = self.nd_words::<WP>(guard).read(load.index(coords, &space));
1639 let word = if comptime!(words_wide > 1) {
1640 words.extract_dynamic((position / per_word).cast::<usize>())
1641 } else {
1642 words.extract(0usize)
1643 };
1644 let shifted = word >> ((position % per_word) * bits);
1645 unpack_line::<S, Const<1>, Const<1>>(Vector::new(shifted), field).extract(0usize)
1646 }
1647 }
1648 }
1649}
1650
1651#[cube]
1652impl Tile {
1653 pub fn cells<E: Float>(&self, #[comptime] count: usize) -> Array<E> {
1657 self.assert_cells("Tile::cells", count);
1658 match &self.kind {
1659 TileKind::Memory(window) => window.cells(count),
1660 TileKind::PlanePartition(_)
1661 | TileKind::PlaneTile(_)
1662 | TileKind::TmaGmem(_)
1663 | TileKind::Procedural(_)
1664 | TileKind::Lines(_) => Tile::refuse_cells(),
1665 }
1666 }
1667
1668 pub fn set_cells<E: Float>(&mut self, values: &Array<E>, #[comptime] count: usize) {
1671 self.assert_cells("Tile::set_cells", count);
1672 let writers = comptime!(self.place.holder());
1673 match &mut self.kind {
1674 TileKind::Memory(window) => window.set_cells(values, count, writers),
1675 TileKind::PlanePartition(_)
1676 | TileKind::PlaneTile(_)
1677 | TileKind::TmaGmem(_)
1678 | TileKind::Procedural(_)
1679 | TileKind::Lines(_) => Tile::refuse_cells(),
1680 }
1681 }
1682
1683 fn assert_cells(&self, #[comptime] who: &str, #[comptime] count: usize) {
1685 comptime!(assert!(
1686 matches!(
1687 self.place.holder(),
1688 ComputeScope::Unit | ComputeScope::Plane
1689 ) && self.place.space.cells() == count,
1690 "{who}: a window of {count} cells a unit or a plane holds, not {:?} held by {:?}",
1691 self.place.space,
1692 self.place.holder()
1693 ));
1694 }
1695
1696 fn refuse_cells() -> ! {
1697 panic!("Tile::cells: a window of memory holds the cells every unit of its holder reads")
1698 }
1699}