logicaffeine_compile/spatial/
shim.rs

1//! What the spatial chips run whole (CUDA persistent, Tensix, the RDU, the fabric): their answer to the shared optimizer.
2//!
3//! A spatial chip lowers each op of the one IR to sections of the closed compute algebra — matrix tiles, row maps,
4//! reductions, recurrent and attention sections — places them on its sites and schedules the whole step as one program.
5//! It links one fused class whole: a region of pointwise ops, its members' sections joined into one, each value it makes
6//! handed from the member that makes it to the one that reads it and never stored. Every other fused op runs as its
7//! parts, each its own sections, and the spatial passes join sections after lowering. A class turns native here when
8//! the spatial link builds it as one section. Placements are refused: a section reads and writes each tensor through
9//! its own access map, never a part of another op's.
10
11use super::target::SpatialTarget;
12use crate::tensor_graph::graph::{Op, Sym, TensorGraph};
13use crate::tensor_graph::numerics::{holds_accumulator, Bound, Contract, Keep};
14use crate::tensor_graph::shim::{Shim, Support};
15use crate::tensor_graph::symbols::ClassId;
16
17/// A spatial chip, by the target its sections are lowered for.
18pub struct SpatialShim<'a> {
19    pub target: &'a SpatialTarget,
20}
21
22impl Shim for SpatialShim<'_> {
23    fn name(&self) -> String {
24        format!("spatial:{}", self.target.name)
25    }
26
27    fn supports(&self, class: ClassId) -> Support {
28        use ClassId::*;
29        use Support::{Decompose, Native, Refuse};
30        match class {
31            NormWeight | FlashDecode | FlashPrefill | GemmDequantInkernel | RouterTopk | MoeGateUpAct | MoeDownCombine | GdnStep => Native,
32            PersistentLayerMegakernel => Native,
33            AddNorm | NormThenAdd | NormScale | NormWeightOffset | NormQuant | AddNormQuant | DualNorm | GatedNorm | EmbedNorm => Decompose,
34            EmbedScale | GemvNormEpilogue | PrefillNormEpilogues | AttnResidualNorm | RecomputeFanout | QkNormRope => Decompose,
35            QkNormRopeKvAppend | RopeKvAppend | RopeQkOneLaunch | QkNormRopeGateSplit | QkvHeadEpilogue | HorizontalNormRopeInsert => Decompose,
36            CrossLayerLaunch | Mla | MlaPrologue | AttnOutputQuant | AttnOutputGate | RowWriteQuant | MlaUnabsorbInFold | GluActMul | GluActMulQuant => Decompose,
37            ElementwiseRegions => Native,
38            GemvGateUpGlu | ElementwiseChains | ReluSqr | MaskedExpertActivation | GemmBias | GemmScale => Decompose,
39            MergedQkv | MergedGateUp | AffineInputFold | RouterSharedGateMerge | MoeTokenSort | MoeGatherInOperand => Decompose,
40            MoeDeferredFinalize | SharedExpertFusion | GdnInputEpilogue | GdnDecodeWhole | GdnPostConvPrep | ConvBiasSilu => Decompose,
41            ConvParts | GdnGateProjections | GdnGating | FusedSampling | LastLayerRows | AllreduceNorm => Decompose,
42            ConvSiluL2norm | GdnOutputQuant | GemvQuant | ElementwiseQuant => Decompose,
43            StridedRead | RowWrite | JoinInPlace | StridedValues => Refuse("a section reads and writes each tensor through its own access map"),
44            Softcap => Native,
45            AttentionVariantsInline => Refuse("the attention sections attend a prefix of the cache, with no window or sink"),
46            MixedBatchAttention | SparseAttention => Refuse("the attention sections attend every cached row of one sequence's prefix"),
47            AttnGqaRowPacking | AttnLseMerge => Refuse("an attention section holds its scores whole, never in chunks"),
48            GroupedTopk => Refuse("the expert sections route by one top-k over every expert"),
49            Sampling => Refuse("a spatial step picks its token by arg-max"),
50            KvCacheQuant | HadamardRotation | CompensatedResidual => Refuse("a spatial chip holds its values in the precision the target's datapath reads"),
51            HostResidentKv | MoeDirectPairs => Refuse("a spatial program places every tensor on the chip's own memory"),
52            CudaGraphs => Refuse("a spatial program is placed and scheduled whole, never captured as launches"),
53        }
54    }
55
56    fn fits(&self, op: &Op, g: &TensorGraph) -> Result<(), String> {
57        holds_accumulator(op, g, self.target.matmul_acc_int_bits())?;
58        match &op.sym {
59            Sym::Prim(_) => Ok(()),
60            Sym::View { .. } => Err("a section reads a row of a table through its own access map, never a view of it".to_string()),
61            Sym::Quantize(_) => Err("a spatial chip reads no activation quantized".to_string()),
62            Sym::Fused(f) if f.class == ClassId::ElementwiseRegions => {
63                // The members join one into the next through each value they make, which the section then never stores:
64                // a value written out for a later reader, or read twice within, keeps its store.
65                let read_once = f.parts.iter().filter(|m| m.out != op.out).all(|m| f.parts.iter().flat_map(|p| &p.ins).filter(|t| **t == m.out).count() == 1);
66                if op.keeps.is_empty() && read_once {
67                    Ok(())
68                } else {
69                    Err("a region is one section where each value it makes is read once within it and none is written out".to_string())
70                }
71            }
72            Sym::Fused(_) => Err("a spatial chip runs every fused op as its parts".to_string()),
73        }
74    }
75
76    fn numerics(&self) -> Option<Contract> {
77        use crate::codegen_sva::fabric_target::Precision;
78        // A section reads its floats in the narrowest float the matrix datapath takes.
79        let keep = if self.target.matmul_supports(Precision::Bf16) {
80            Keep::Bf16
81        } else if self.target.matmul_supports(Precision::Fp16) {
82            Keep::F16
83        } else {
84            Keep::F32
85        };
86        let point = format!("spatial:{}", self.target.name);
87        // A joined section keeps a folded intermediate in its registers where the program apart packs it and reads it
88        // back: an exact program's joins are exact, an approximate one's within the tolerance its rewrites are held to.
89        let fused = Bound { of_block: 0.0, of_value: crate::spatial::equiv::APPROX_REL_TOL, floor: crate::spatial::equiv::APPROX_REL_TOL };
90        Some(Contract { point, activations: None, cache: keep, state: keep, attention: None, fused, accumulator_bits: self.target.matmul_acc_int_bits() })
91    }
92}
logicaffeine_compile/spatial/shim.rs

logicaffeine_compile/spatial/
shim.rs
