Time Reducer

Time Reducer는 Packet Reducer의 [Lane, Packet] 출력을 시간 누산기로 Time에 걸쳐 OutTime으로 누산합니다.

인터페이스

.contract_time::<OutTime>()는 Time Reducer를 호출합니다. OutTime은 남는 Time 차원을 지정합니다(나머지는 합산되어 사라집니다).

impl<'l, const T: Tu, D: Scalar, Chip: M, Cluster: M, Slice: M, Lane: M, Time: M, Packet: M, B: Backend>
    ContractPacketTensor<'l, T, D, Chip, Cluster, Slice, Lane, Time, Packet, B>
{
    /// Accumulates per-cycle contractions over the `Time` dimension via the shared
    /// accumulator buffer, shrinking input `Time` to `OutTime`. The axes present in
    /// `Time` but absent from `OutTime` are reduce-added.
    #[primitive(ContractPacketTensor::contract_time)]
    pub fn contract_time<OutTime: M>(
        self,
    ) -> ContractTimeTensor<'l, T, D, Chip, Cluster, Slice, Lane, OutTime, Packet, B> {
        verify_contract_time(ContractTimeInput {
            in_time: Time::to_value(),
            out_time: OutTime::to_value(),
        });
        // Carry the deferred operands forward unreduced: the fused contraction at `contract_lane`
        // performs this Time reduction too. This stage only re-types the carrier to `OutTime`.
        ContractTimeTensor::new(self.device, self.inner, Time::to_value())
    }
}

예를 들어, 아래 커널은 2D 텐서를 B를 따라 리덕션합니다(A만 남습니다).

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 2048, B = 32];

/// Reduces along B; A survives.
fn reduce_b<'l, const T: Tu>(
    // Streaming operand: Slice = m![A / 8] (256 outer A chunks across slices).
    // Time = m![B / 16, A % 8]; Packet = m![B % 16].
    // B splits across Packet (B % 16) and Time (B / 16): each cycle produces a partial sum.
    input: CollectTensor<'l, T, bf16, m![1], m![1 # 2], m![A / 8], m![B / 16, A % 8], m![B % 16]>,
    // TRF operand: single-lane weight per slice.
    trf: &TrfTensor<bf16, m![1], m![1 # 2], m![A / 8], m![1], m![B]>,
    // Output: one f32 per (slice, A % 8) cell.
) -> ContractTensor<'l, T, f32, m![1], m![1 # 2], m![A / 8], m![A % 8], m![1 # 8]> {
    input
         // Outer: Lane = m![1], OutTime = m![B / 16, A % 8], OutPacket = m![B % 16].
         .contract_outer::<m![B / 16, A % 8], m![B % 16], _, _, _>(trf)
         // Packet Reducer: OutPacket = m![1]. Collapses B % 16 spatially.
         .contract_packet::<m![1]>()
         // Time Reducer: OutTime = m![A % 8]. Accumulator receives
         // Time::SIZE = (B / 16) × (A % 8) = 2 × 8 = 16 flits; B / 16 outer
         // chunks accumulate into 8 slots indexed by A % 8.
         .contract_time::<m![A % 8]>()
         // Lane Folder: Lane folds into OutPacket. Sequential mode (Lane = m![1]).
         .contract_lane::<m![A % 8], m![1 # 8]>(LaneMode::Sequential)
}

let mut device = Device::new(Topology { chips: 1, pes: 8 }).unwrap();

let a: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![A / 8], m![B / 16, A % 8], m![B % 16]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![A / 8], m![1], m![B]> = TrfTensor::zero();
let _o = reduce_b(a, &b);
}

아키텍처

Time Reducer는 Packet Reducer가 사이클마다 내보내는 [Lane, Packet] 출력을 받습니다. 하드웨어는 상류에서 Lane::SIZE ≤ 8(공간적으로 병렬인 레인)과 Packet::SIZE ≤ 32로 상한을 둡니다.

Time Reducer는 매 사이클 [Lane, Packet] 공간 그리드를 Time에 걸쳐 OutTime으로 접습니다. OutTime은 남는 차원의 상대적 순서를 보존한 Time의 부분집합이어야 합니다(verify_contract_time이 이를 강제합니다). Time에는 있지만 OutTime에는 없는 차원은 합산되어 사라지며, 그러한 차원 중 가장 바깥 차원이 플릿을 순회합니다.

InnerTime을 Time의 안쪽 비리덕션 차원(가장 바깥 리덕션 차원보다 안쪽에 있으면서 OutTime에 남는 차원)이라고 합시다. 위의 reduce_b에서 Time = m![B / 4, A % 8]이고 OutTime = m![A % 8]이므로, B / 4가 가장 바깥 리덕션 차원이며(Time::SIZE = 2 × 8 = 16개의 플릿을 순회합니다) InnerTime = m![A % 8]입니다.

누산에는 InnerTime 튜플 값마다 하나씩, InnerTime::SIZE개의 [Lane, Packet] 슬롯이 필요합니다. 튜플이 같은 플릿은 같은 슬롯에 누산됩니다. reduce_b의 경우 8개의 슬롯이 B / 4 = 2회의 반복에 걸쳐 누산되며, 플릿 15 이후 버퍼에는 최종 리덕션 결과가 담기고 버퍼는 이를 Lane Folder에 넘깁니다:

Time = m![B / 4, A % 8]
          ~~~~~  ~~~~~
        outer R  inner non-R (A % 8)

Flit sequence (B / 4 has values 0,1; A % 8 has values 0..7):

  flit #0:  B/4=0, A%8=0  ──→ ┌─────────────────┐
  flit #8:  B/4=1, A%8=0  ──→ │ slot 0 (A%8=0)  │  accumulates B for A%8=0
                              └─────────────────┘

  flit #1:  B/4=0, A%8=1  ──→ ┌─────────────────┐
  flit #9:  B/4=1, A%8=1  ──→ │ slot 1 (A%8=1)  │  accumulates B for A%8=1
                              └─────────────────┘
   ⋮

  flit #7:  B/4=0, A%8=7  ──→ ┌─────────────────┐
  flit #15: B/4=1, A%8=7  ──→ │ slot 7 (A%8=7)  │  accumulates B for A%8=7
                              └─────────────────┘

  8 non-reduce positions → 8 slots used

제약 조건

InnerTime::SIZE가 슬롯 용량(버퍼가 담을 수 있는 슬롯 수)을 넘지 않으면 매핑이 버퍼에 들어맞습니다.

슬롯 용량은 버퍼의 1,024개 셀과 하류 Lane Folder의 LaneMode에 따라 정해집니다. 각 슬롯은 LaneMode가 형태를 결정하는 [Lane, Packet] 청크이므로, 슬롯 수는 1,024를 청크당 셀 수로 나눈 값입니다.

Lane::SIZE = 1인 하위 단계 .contract_lane(LaneMode::Sequential) 아래의 reduce_b에서는 슬롯 용량이 32이므로 InnerTime::SIZE = 8이 여유 있게 들어갑니다. InnerTime::SIZE가 슬롯 용량을 초과한다면 Time을 재구성하거나(예: B를 더 분할) LaneMode를 전환하여 처리량을 일부 포기하는 대신 슬롯 여유 공간을 확보해야 합니다.

성능

처리량은 입력 측 기준으로 사이클당 한 패킷입니다. N개의 입력을 하나의 출력으로 리덕션하므로 실효 출력 속도는 입력의 1 / N입니다.

크기가 N인 Time 차원 리덕션의 지연 시간은 약 N 사이클입니다.