Switch Engine

다른 모든 Tensor Unit 엔진은 각자의 DM 파티션에서 Slice 단위로 실행되지만, Switch Engine은 256개 Slice로 이루어진 링 네트워크를 통해 Slice 간에 데이터를 이동합니다. 즉, 한 Slice의 값을 그룹에 브로드캐스트하거나, Slice 간에 값을 교환하거나, 어느 Slice가 어떤 값을 가질지 치환합니다.

인터페이스

FetchTensor::switch()는 Chip, Cluster, Packet 및 기저 데이터 값을 보존하면서 SwitchTensor를 생성합니다. 선택한 구성을 반영하도록 Slice와 Time 매핑만 변경됩니다.

impl<'l, const T: Tu, P: CanApplySwitch, D: Scalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M, B: Backend>
    TuTensor<'l, T, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Applies switching network routing only. The packet passes through
    /// unchanged, no padding, no reshaping. Use `collect` afterwards to
    /// normalize the packet to flit-sized chunks.
    #[primitive(TuTensor::switch)]
    pub fn switch<OutSlice: M, OutTime: M>(
        self,
        config: SwitchConfig,
    ) -> SwitchTensor<'l, T, D, Chip, Cluster, OutSlice, OutTime, Packet, B> {
        verify_switch::<Slice, Time, OutSlice, OutTime>(&config);
        SwitchTensor::new(self.device, self.inner.transpose(true))
    }
}

커널 작성자는 OutSlice, OutTime, 그리고 구성과 그 매개변수를 선택하는 SwitchConfig 인자를 지정합니다. SwitchConfig는 일반적인 패턴을 위한 사전 정의된 변형(Broadcast01, Broadcast1, Transpose, InterTranspose, TransposedBroadcast1) 중 하나이거나, 보다 일반적인 패턴을 위한 CustomBroadcast입니다. 각 변형 이름의 숫자 접미사는 Slice에서 빠져나가는 slice 하위 차원을 나타냅니다. 예를 들어 Broadcast01은 slice0과 slice1을 모두 브로드캐스트하고, Broadcast1은 slice1만 브로드캐스트합니다.

컴파일러는 OutSlice와 OutTime이 구성에서 요구하는 차원 구조와 일치하는지 검증하며(아래의 각 구성별 섹션에서 입력/출력 다이어그램과 함께 요구되는 구조를 보여 줍니다), 일치하지 않으면 컴파일이 실패합니다. 또한 모든 구성은 InSlice::SIZE == OutSlice::SIZE를 요구합니다. 즉, switch는 전체 slice 수를 보존합니다.

정규 구성

각 정규 구성은 Chip의 256개 Slice를 각각 ring_size개의 Slice로 이루어진 병렬 서브링으로 분할합니다. ring_size는 구성의 매개변수(일반적으로 slice1 × slice0)에서 도출되며, 분할 단위와 사이클 비용을 함께 결정합니다. 링 토폴로지와 라우터별 결정 로직은 아래의 아키텍처를, 사이클 계수 분석은 성능을 참조하십시오.

정규 구성은 여섯 가지 일반적인 패턴을 다룹니다. 이 패턴들로 표현할 수 없는 임의의 Slice 하위 차원 순열에는 대신 CustomBroadcast가 필요합니다.

브로드캐스트 축(Broadcast01, Broadcast1, TransposedBroadcast1의 X 또는 Y)을 도입하는 구성에서는 해당 축이 새로운 축이어야 합니다. 즉, 입력 Slice나 입력 Time에 이미 나타나서는 안 됩니다.

Forwarding

Forwarding은 Slice 매핑과 Time 매핑을 변경하지 않습니다. 모든 라우터는 Slice 간 이동 없이 자기 Slice의 입력을 그대로 출력합니다.

SwitchConfig에는 Forwarding 변형이 없습니다. Slice 간 교환이 필요하지 않으면 .switch()를 건너뛰고 FetchTensor에서 .collect()를 직접 호출하십시오.

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

fn forwarding<'l, const T: Tu>(
    input: FetchTensor<'l, T, f32, m![1], m![1 # 2], m![A], m![B], m![C]>,
) -> CollectTensor<'l, T, f32, m![1], m![1 # 2], m![A], m![B, C / 8], m![C % 8]> {
    input.collect()
}

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

let f: FetchTensor<'_, _, f32, m![1], m![1 # 2], m![A], m![B], m![C]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o = forwarding(f);
}

Broadcast01

Broadcast01은 각 Slice의 데이터를 서브링 내 모든 Slice로 브로드캐스트하며, 안쪽 Slice 하위 차원 두 개(slice0과 slice1)를 모두 Time으로 옮깁니다.

입력 차원 구조(왼쪽에서 오른쪽으로, 가장 바깥쪽부터 가장 안쪽 순서)는 다음과 같습니다.

┌──────────────────────────┬───────────────┐
│          Slice           │      Time     │
├────────┬────────┬────────┼───────┬───────┤
│ slice2 │ slice1 │ slice0 │ time1 │ time0 │
└────────┴────────┴────────┴───────┴───────┘

스위칭 후 slice1과 slice0은 서브링 전체에 브로드캐스트하기 위해 Slice에서 Time으로 이동합니다. X와 Y로 표시된 새로운 브로드캐스트 차원 두 개가 비워진 Slice 위치를 채웁니다.

┌──────────────────────┬─────────────────────────────────┐
│        Slice         │              Time               │
├────────┬──────┬──────┼───────┬────────┬───────┬────────┤
│ slice2 │  X   │  Y   │ time1 │ slice1 │ time0 │ slice0 │
└────────┴──────┴──────┴───────┴────────┴───────┴────────┘

서브링은 ring_size = slice1 × slice0개의 Slice에 걸쳐 있으며, 고정된 slice2에서 (slice1, slice0) 조합마다 Slice가 하나씩 대응합니다. 모든 Slice는 자신의 패킷을 서브링을 따라 전송하고 ring_size개의 패킷을 모두 수신합니다. 따라서 각 출력 Slice는 브로드캐스트 그룹 전체의 데이터를 갖게 되며, 이 데이터는 출력 Slice의 새 X 축과 Y 축을 따라 인덱싱됩니다.

예시

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 256, B = 64, C = 63, D = 2, X = 2, Y = 2];

fn broadcast01<'l, const T: Tu>(
    input: FetchTensor<'l, T, f32, m![1], m![D], m![A], m![B], m![C # 64]>,
) -> SwitchTensor<'l, T, f32, m![1], m![D], m![A / 4, X, Y], m![B / 4, A / 2 % 2, B % 4, A % 2], m![C # 64]> {
    input.switch::<m![A / 4, X, Y], m![B / 4, A / 2 % 2, B % 4, A % 2]>(
        SwitchConfig::Broadcast01 {
            slice1: 2,
            slice0: 2,
            time0: 4
        }
    )
}

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

let f: FetchTensor<'_, _, f32, m![1], m![D], m![A], m![B], m![C # 64]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o= broadcast01(f);
}

slice1 = 2(브로드캐스트 X의 크기), slice0 = 2(브로드캐스트 Y의 크기), time0 = 4이면 컴파일러는 slice2 = 64, time1 = 16, ring_size = 4를 도출합니다(64개의 서브링이 256개의 Slice에 걸쳐 있습니다).

하위 차원은 slice2 = A / 4, slice1 = A / 2 % 2, slice0 = A % 2, time1 = B / 4, time0 = B % 4로 결정되며, 그 결과 OutSlice = m![A / 4, X, Y]와 OutTime = m![B / 4, A / 2 % 2, B % 4, A % 2]가 됩니다.

사이클 추정치: ring_size × Time::SIZE × flits_per_packet = 4 × 64 × 8 = 2048이며, 여기서 Time::SIZE = 64이고 flits_per_packet = sizeof(f32) × Packet::SIZE / 32 = 4 × 64 / 32 = 8입니다.

Broadcast1

입력 차원 구조(왼쪽에서 오른쪽으로, 가장 바깥쪽부터 가장 안쪽 순서)는 다음과 같습니다.

┌──────────────────────────┬────────┐
│           Slice          │  Time  │
├────────┬────────┬────────┼────────┤
│ slice2 │ slice1 │ slice0 │ time0  │
└────────┴────────┴────────┴────────┘

스위칭 후 slice1은 서브링 전체에 브로드캐스트하기 위해 Slice에서 Time으로 이동하고, slice0은 Slice에 그대로 남습니다. X로 표시된 새로운 브로드캐스트 차원이 slice1이 비운 Slice 위치를 채웁니다.

┌────────────────────────┬────────────────┐
│          Slice         │      Time      │
├────────┬──────┬────────┼───────┬────────┤
│ slice2 │  X   │ slice0 │ time0 │ slice1 │
└────────┴──────┴────────┴───────┴────────┘

서브링은 ring_size = slice1 × slice0개의 Slice에 걸쳐 있으며, 이는 Broadcast01의 서브링과 물리적 범위가 같습니다. 그러나 브로드캐스트는 slice1을 따라서만 일어납니다. 각 출력 Slice는 같은 slice0 위치에 있는 소스로부터 slice1개의 패킷을 수신하며, 이 패킷들은 가장 안쪽 Time을 따라 순서대로 배치됩니다. 출력 Slice의 새로운 X 축(크기 slice1)은 이렇게 수집된 데이터를 복제하며, slice0 자체는 원래 위치에 그대로 유지됩니다.

예시

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 256, B = 64, C = 63, X = 4];

fn broadcast1<'l, const T: Tu>(
    input: FetchTensor<'l, T, i8, m![1], m![1 # 2], m![A], m![B], m![C # 64]>,
) -> SwitchTensor<'l, T, i8, m![1], m![1 # 2], m![A / 32, X, A % 8], m![B, A / 8 % 4], m![C # 64]> {
    input.switch::<m![A / 32, X, A % 8], m![B, A / 8 % 4]>(
        SwitchConfig::Broadcast1 {
            slice1: 4,
            slice0: 8,
        }
    )
}

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

let f: FetchTensor<'_, _, i8, m![1], m![1 # 2], m![A], m![B], m![C # 64]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o = broadcast1(f);
}

slice1 = 4(브로드캐스트 X의 크기), slice0 = 8이면 컴파일러는 slice2 = 8과 ring_size = 32를 도출합니다(8개의 서브링이 256개의 Slice에 걸쳐 있습니다).

하위 차원은 slice2 = A / 32, slice1 = A / 8 % 4, slice0 = A % 8, time0 = B로 결정되며, 그 결과 OutSlice = m![A / 32, X, A % 8]와 OutTime = m![B, A / 8 % 4]가 됩니다.

사이클 추정치: ring_size × Time::SIZE × flits_per_packet = 32 × 64 × 2 = 4096이며, 여기서 Time::SIZE = 64이고 flits_per_packet = sizeof(i8) × Packet::SIZE / 32 = 1 × 64 / 32 = 2입니다.

Transpose

Transpose는 Slice의 가장 안쪽 부분에서 slice1과 slice0을 맞바꿉니다.

입력과 출력의 Slice 순서는 다음과 같습니다.

┌──────────────────────────┐         ┌──────────────────────────┐
│           Slice          │         │           Slice          │
├────────┬────────┬────────┤   ──►   ├────────┬────────┬────────┤
│ slice2 │ slice1 │ slice0 │         │ slice2 │ slice0 │ slice1 │
└────────┴────────┴────────┘         └────────┴────────┴────────┘

각 서브링은 slice0 × slice1개의 Slice에 걸쳐 있으며 데이터를 순환시킵니다. 그 결과 서브링의 모든 Slice는 교환 상대가 이전에 가지고 있던 값을 갖게 됩니다.

Transpose를 사용하려면 (정규화 후) 입력 Time과 출력 Time이 일치해야 합니다.

예시

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

fn transpose<'l, const T: Tu>(
    input: FetchTensor<'l, T, i8, m![1], m![1 # 2], m![A], m![B], m![C # 64]>,
) -> SwitchTensor<'l, T, i8, m![1], m![1 # 2], m![A / 64, A % 2, A / 2 % 32], m![B], m![C # 64]> {
    input.switch::<m![A / 64, A % 2, A / 2 % 32], m![B]>(SwitchConfig::Transpose {
        slice1: 32,
        slice0: 2,
    })
}

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

let f: FetchTensor<'_, _, i8, m![1], m![1 # 2], m![A], m![B], m![C # 64]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o = transpose(f);
}

slice1 = 32이고 slice0 = 2이면 컴파일러는 slice2 = 4와 ring_size = 64를 도출합니다(서브링 4개가 256개 Slice에 걸쳐 있습니다).

하위 차원은 slice2 = A / 64, slice1 = A / 2 % 32, slice0 = A % 2, time0 = B로 결정되며, 그 결과 OutSlice = m![A / 64, A % 2, A / 2 % 32]와 OutTime = m![B]가 됩니다(slice0과 slice1이 맞바뀌고 Time은 변하지 않습니다).

사이클 추정치: ring_size × Time::SIZE × flits_per_packet = 64 × 64 × 2 = 8192이며, 여기서 Time::SIZE = 64이고 flits_per_packet = sizeof(i8) × Packet::SIZE / 32 = 1 × 64 / 32 = 2입니다.

InterTranspose

InterTranspose는 Slice와 Time 사이에서 차원을 맞바꿉니다. slice1은 Time으로 이동하고 time1은 Slice로 이동합니다(일반 Transpose는 Slice 안에서만 이루어집니다).

Note

모든 대상 Slice가 이미 원본 영역의 사본을 가지고 있으면 축 리프팅으로 InterTranspose를 대체할 수 있습니다. 이 방식은 각 Slice에 서로 다른 읽기 오프셋을 할당하며 서브링 순회를 수행하지 않습니다. 데이터가 Slice 사이에서, 또는 Slice에서 Time으로 이동해야 할 때는 InterTranspose를 사용하십시오.

입력 차원 구조(왼쪽에서 오른쪽으로, 가장 바깥쪽에서 가장 안쪽 순)는 다음과 같습니다.

┌──────────────────────────┬───────────────────────┐
│           Slice          │         Time          │
├────────┬────────┬────────┼───────┬───────┬───────┤
│ slice2 │ slice1 │ slice0 │ time2 │ time1 │ time0 │
└────────┴────────┴────────┴───────┴───────┴───────┘

스위칭 후 slice1과 time1은 Slice/Time 경계를 넘어 위치를 맞바꿉니다.

┌─────────────────────────┬────────────────────────┐
│          Slice          │          Time          │
├────────┬───────┬────────┼───────┬───────┬────────┤
│ slice2 │ time1 │ slice0 │ time2 │ time0 │ slice1 │
└────────┴───────┴────────┴───────┴───────┴────────┘

각 서브링은 slice1 × slice0개의 Slice에 걸쳐 있으며 time1개의 시간 단계 동안 데이터를 순환시킵니다. 그 결과 각 Slice에서 이전에 slice1로 인덱싱되던 값은 time1로 인덱싱되고, 반대의 경우도 마찬가지입니다.

InterTranspose는 세 가지 크기 제약 조건을 적용합니다.

• InSlice는 256개의 Slice 전체에 걸쳐 있습니다: slice2 × slice1 × slice0 == 256.
• 맞바뀌는 차원들의 크기가 같습니다: time1.SIZE == slice1.
• time2 분해가 정수로 떨어지도록 InTime::SIZE는 slice1 × time0으로 나누어떨어져야 합니다.

예시

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

fn inter_transpose<'l, const T: Tu>(
    input: FetchTensor<'l, T, i8, m![1], m![1 # 2], m![A], m![B], m![C # 32]>,
) -> SwitchTensor<'l, T, i8, m![1], m![1 # 2], m![A / 32, B / 2 % 2, A % 16], m![B / 4, B % 2, A / 16 % 2], m![C # 32]> {
    input.switch::<m![A / 32, B / 2 % 2, A % 16], m![B / 4, B % 2, A / 16 % 2]>(
        SwitchConfig::InterTranspose {
            slice1: 2,
            slice0: 16,
            time0: 2,
        })
}

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

let f: FetchTensor<'_, _, i8, m![1], m![1 # 2], m![A], m![B], m![C # 32]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o = inter_transpose(f);
}

slice1 = 2, slice0 = 16, time0 = 2이면 컴파일러는 slice2 = 8, time2 = 2, ring_size = 32를 도출합니다(서브링 8개가 256개 Slice에 걸쳐 있습니다).

하위 차원은 slice2 = A / 32, slice1 = A / 16 % 2, slice0 = A % 16, time2 = B / 4, time1 = B / 2 % 2, time0 = B % 2로 결정되며, 그 결과 OutSlice = m![A / 32, B / 2 % 2, A % 16]와 OutTime = m![B / 4, B % 2, A / 16 % 2]가 됩니다(slice1과 time1이 Slice와 Time 사이에서 맞바뀝니다).

사이클 추정치: ring_size × Time::SIZE × flits_per_packet = 32 × 8 × 1 = 256이며, 여기서 Time::SIZE = 8이고 flits_per_packet = sizeof(i8) × Packet::SIZE / 32 = 1 × 32 / 32 = 1입니다.

TransposedBroadcast1

TransposedBroadcast1은 slice0을 Time의 가장 안쪽 위치로 브로드캐스트하고 slice1을 Slice의 가장 안쪽 위치로 옮깁니다. 이는 Transpose를 적용한 뒤 Broadcast1을 적용하는 것과 같습니다. 다음과 같은 구조의 입력 텐서는

┌──────────────────────────┬────────┐
│           Slice          │  Time  │
├────────┬────────┬────────┼────────┤
│ slice2 │ slice1 │ slice0 │ time0  │
└────────┴────────┴────────┴────────┘

아래 출력이 됩니다. 여기서 slice0은 서브링 전체에 브로드캐스트되도록 가장 안쪽 Time 위치로 이동하고, slice1은 가장 안쪽 Slice 위치로 옮겨지며, slice1이 비운 가운데 슬롯은 브로드캐스트 차원이 채웁니다:

┌────────────────────────┬────────────────┐
│          Slice         │      Time      │
├────────┬──────┬────────┼───────┬────────┤
│ slice2 │  Y   │ slice1 │ time0 │ slice0 │
└────────┴──────┴────────┴───────┴────────┘

각 서브링은 slice0 × slice1개의 Slice에 걸쳐 있으며 데이터를 순환시킵니다. 그 결과 모든 Slice는 교환 상대의 값을 갖게 되며, 이 값은 가장 안쪽 Time 하위 차원에서 slice0 위치 전체에 브로드캐스트됩니다.

예시

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

fn transposed_broadcast1<'l, const T: Tu>(
    input: FetchTensor<'l, T, i8, m![1], m![1 # 2], m![A], m![B], m![C # 32]>,
) -> SwitchTensor<'l, T, i8, m![1], m![1 # 2], m![A / 64, Y, A / 8 % 8], m![B, A % 8], m![C # 32]> {
    input.switch::<m![A / 64, Y, A / 8 % 8], m![B, A % 8]>(
        SwitchConfig::TransposedBroadcast1 {
            slice1: 8,
            slice0: 8,
        }
    )
}

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

let f: FetchTensor<'_, _, i8, m![1], m![1 # 2], m![A], m![B], m![C # 32]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o = transposed_broadcast1(f);
}

slice1 = 8이고 slice0 = 8(브로드캐스트 Y의 크기)이면 컴파일러는 slice2 = 4와 ring_size = 64를 도출합니다(서브링 4개가 256개 Slice에 걸쳐 있습니다).

하위 차원은 slice2 = A / 64, slice1 = A / 8 % 8, slice0 = A % 8, time0 = B로 결정되며, 그 결과 OutSlice = m![A / 64, Y, A / 8 % 8]와 OutTime = m![B, A % 8]가 됩니다.

사이클 추정치: ring_size × Time::SIZE × flits_per_packet = 64 × 16 × 1 = 1024이며, 여기서 Time::SIZE = 16이고 flits_per_packet = sizeof(i8) × Packet::SIZE / 32 = 1 × 32 / 32 = 1입니다.

아키텍처

앞서 소개한 구성을 실행하기 위해 Switch Engine은 Chip의 256개 Slice 전체를 하나의 물리적 링(Slice당 라우터 하나)으로 배치하며, 이 링은 각각 ring_size개의 Slice로 이루어진 256 / ring_size개의 병렬 서브링으로 분할됩니다(아래에 그러한 서브링 하나를 나타냈습니다). 정규 구성에서는 컴파일러가 구성의 매개변수(slice1, slice0, time0)로부터 ring_size를 도출하고, CustomBroadcast에서는 커널 작성자가 ring_size를 직접 설정합니다.

   ┌────────────────────┐    ┌────────────────────┐    ┌────────────────────┐    ┌────────────────────┐
   │      Router 0      │◀──▶│      Router 1      │◀──▶│        ...         │◀──▶│ Router ring_size-1 │
   └────────────────────┘    └────────────────────┘    └────────────────────┘    └────────────────────┘
             ▲                                                                                          ▲
             └──────────────────────────────────────────────────────────────────────────────────────────┘
                                            wrap-around (links are bidirectional)

성능

스위치 연산 하나에는 대략 ring_size × Time::SIZE × flits_per_packet 사이클이 소요됩니다. 모든 서브링은 병렬로 진행되므로, 이 링당 사이클 수가 곧 Chip 전체의 지연 시간이기도 합니다. 세 가지 요인은 다음과 같습니다:

• ring_size: 플릿 하나가 서브링 하나를 순회하는 데 걸리는 사이클 수입니다. ring_size가 클수록 링당 비용은 커지지만 링 하나가 더 많은 Slice에 도달하며, ring_size가 작을수록 Cluster가 더 많은 병렬 링으로 분할되고 링당 비용은 작아집니다.
• Time::SIZE: 입력 텐서의 시간 단계 수입니다. 서브링 순회는 시간 단계마다 한 번씩 반복됩니다.
• flits_per_packet: 패킷당 플릿 수로, D[Packet::SIZE]의 크기를 32바이트로 나눈 값과 같습니다. 순회는 또한 패킷의 플릿마다 한 번씩 반복됩니다.

사용자 정의 구성

사용자 정의 구성은 임의의 차원 순열이나 부분 차원 추출처럼 정규 구성으로 표현할 수 없는 이동 패턴을 처리합니다. 이러한 유연성에는 구성 오버헤드와 이 절의 끝에 나열된 제약 조건이 따릅니다.

SwitchConfig::CustomBroadcast 변형은 필드 하나를 가집니다:

/// Routes data across slices using a custom snoop bitmap.
/// The bitmap is computed by the compiler from the input shape and
/// topology parameters.
CustomBroadcast {
    /// Ring group size for the custom routing.
    ring_size: usize,
},

정규 구성은 snoop 비트맵을 위한 내장 생성기를 제공하는 반면, CustomBroadcast에서는 컴파일러가 커널 작성자의 입력/출력 매핑과 ring_size로부터 비트맵을 직접 합성합니다.

지원되는 변환 패턴

사용자 정의 비트맵은 정규 구성으로 처리할 수 없는 두 가지 패턴을 다룹니다:

• 브로드캐스트를 동반한 자유 전치: Transpose나 TransposedBroadcast1의 고정된 형태를 넘어서는, 분할 차원의 임의 순열과 브로드캐스트입니다.
• 부분 차원 추출: 브로드캐스트 중에 한 차원의 값 중 일부만 Time으로 이동합니다. 반면 Broadcast01 같은 정규 구성은 항상 차원 전체를 이동시킵니다.

아래 예제에서 이러한 패턴을 보여 줍니다.

구성 오버헤드

사용자 정의 스누프 비트맵을 기록하면 구성 데이터가 Switch Engine의 특수 기능 레지스터(Special Function Register, SFR)로 스트리밍되며, 이 SFR 기록이 진행되는 동안 DMA Engine과 서브 컨텍스트가 모두 점유됩니다. 비트맵이 로드되는 동안 DMA 컨텍스트와 sub 컨텍스트는 다른 연산을 실행할 수 없습니다. 따라서 이 비용은 고정 사이클 정지가 아니라 스케줄링 병렬성의 감소로 나타납니다.

예제 1: 임의 순열

이 예시는 가장 안쪽의 Slice 하위 차원 네 개(A / 4, A % 4, B / 4, B % 4)를 [3, 2, 1, 0] 순서로 뒤집습니다. 이는 어떤 정규 구성으로도 표현할 수 없는 패턴입니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 16, B = 16, C = 8, D = 8, E = 8];

fn arbitrary_permutation<'l, const T: Tu>(
    input: FetchTensor<'l, T, f32, m![1], m![1 # 2], m![A, B], m![C], m![D, E]>,
) -> SwitchTensor<'l, T, f32, m![1], m![1 # 2], m![B % 4, B / 4, A % 4, A / 4], m![C], m![D, E]> {
    input.switch::<m![B % 4, B / 4, A % 4, A / 4], m![C]>(
        SwitchConfig::CustomBroadcast { ring_size: 256 }
    )
}

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

let f: FetchTensor<'_, _, f32, m![1], m![1 # 2], m![A, B], m![C], m![D, E]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o = arbitrary_permutation(f);
}

출력 Slice = m![B % 4, B / 4, A % 4, A / 4]는 입력 Slice 형태 [0, 1, 2, 3]을 [3, 2, 1, 0]으로 치환합니다. 이 순열은 어떤 정규 구성으로도 다룰 수 없지만 사용자 정의 비트맵으로는 다룰 수 있습니다.

사이클 공식에 대입하면 cycles ≈ ring_size × Time::SIZE × flits_per_packet = 256 × 8 × 8 = 16384이며, 각 값은 다음과 같습니다:

• ring_size = 256
• Time::SIZE = C::SIZE = 8
• flits_per_packet = sizeof(f32) × Packet::SIZE / 32 = 4 × 64 / 32 = 8 (Packet = m![D, E], D::SIZE × E::SIZE = 64)

최대값인 ring_size = 256이 필요한 이유는 이 순열이 반복 구조 없이 모든 Slice에 걸친 의존성을 만들기 때문입니다. 따라서 입력 Slice와 출력 Slice는 링 인덱스상에서 임의로 멀리 떨어져 있을 수 있으며, 더 작은 서브링으로는 그러한 쌍 중 적어도 하나를 포괄하지 못합니다.

예제 2: 다차원 브로드캐스트

예제 1의 순수 순열과 달리, 이 예제는 연속하지 않는 두 차원(A % 2와 B % 2)을 Slice에서 Time으로 옮기고 원래 위치에서는 브로드캐스트합니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 16, B = 16, C = 8, D = 8, E = 8, X = 2, Y = 2];

fn multi_axis_broadcast<'l, const T: Tu>(
    input: FetchTensor<'l, T, f32, m![1], m![1 # 2], m![A, B], m![C], m![D, E]>,
) -> SwitchTensor<'l, T, f32, m![1], m![1 # 2], m![A / 2, X, B / 2, Y], m![C, A % 2, B % 2], m![D, E]> {
    input.switch::<m![A / 2, X, B / 2, Y], m![C, A % 2, B % 2]>(
        SwitchConfig::CustomBroadcast { ring_size: 32 }
    )
}

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

let f: FetchTensor<'_, _, f32, m![1], m![1 # 2], m![A, B], m![C], m![D, E]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o = multi_axis_broadcast(f);
}

출력은 A % 2와 B % 2를 Slice에서 Time으로 옮기고, 브로드캐스트 차원 X와 Y를 통해 원래 위치에서 브로드캐스트합니다. Broadcast01도 비슷한 형태를 지원하지만 브로드캐스트 차원(slice0, slice1)이 입력 Slice에서 연속해야 하므로, 연속하지 않는 차원이 Slice에서 Time으로 이동하는 경우는 표현할 수 없습니다. 대신 사용자 정의 비트맵으로 이를 표현합니다.

사이클 공식 cycles ≈ ring_size × Time::SIZE × flits_per_packet = 32 × 8 × 8 = 2048에 대입하면 다음과 같습니다.

• ring_size = 32(가장 바깥쪽 A / 2 분할은 서브링 간 교환이 필요 없으므로 각 서브링 안에서 가장 안쪽 32개 Slice끼리만 통신합니다)
• Time::SIZE = 8
• flits_per_packet = sizeof(f32) × Packet::SIZE / 32 = 4 × 64 / 32 = 8(Packet = m![D, E], D::SIZE × E::SIZE = 64)

예제 3: 부분 축 추출(슬라이싱)

이동하는 차원의 모든 값을 포함하는 예제 1, 2와 달리, 여기서는 C # 4의 유효한 값 3개만 Slice에서 Time으로 이동하고 패딩 셀은 그대로 남습니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 16, B = 4, C = 3, D = 8, E = 8, X = 4];

fn partial_axis_extraction<'l, const T: Tu>(
    input: FetchTensor<'l, T, f32, m![1], m![1 # 2], m![A, B, C # 4], m![D], m![E]>,
) -> SwitchTensor<'l, T, f32, m![1], m![1 # 2], m![A, B, X], m![D, C], m![E]> {
    input.switch::<m![A, B, X], m![D, C]>(
        SwitchConfig::CustomBroadcast { ring_size: 4 }
    )
}

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

let f: FetchTensor<'_, _, f32, m![1], m![1 # 2], m![A, B, C # 4], m![D], m![E]> = FetchTensor::new(&mut device.main, Tensor::zero());
let _o = partial_axis_extraction(f);
}

출력은 C # 4를 Slice에서 Time으로 옮기고 비워진 Slice 위치에 브로드캐스트 축 X를 배치합니다. 순수한 패딩 셀인 네 번째 값(C # 4 = 3)은 버려지므로 유효한 값 3개인 C만 추출됩니다. Broadcast1도 비슷한 형태를 지원하지만 항상 차원 전체를 옮기므로 부분 집합을 표현할 수 없습니다.

이 부분 추출이 허용되는 이유는 버려지는 값이 패딩이기 때문입니다. C # 4는 4로 패딩된 슬롯에 C = 3개의 유효한 요소를 담고 있으므로, 이를 C로 슬라이싱하여 줄여도(C # 4 → C) 실제 값은 전혀 버려지지 않습니다. 반면 유효한 범위를 슬라이싱하면 살아 있는 입력이 버려지며, 이는 버려지는 값은 패딩이어야 함 제약 조건에 의해 금지됩니다. 대신 이처럼 패딩만 제거하는 추출은 사용자 정의 비트맵으로 표현합니다.

사이클 공식에 대입하면 cycles ≈ ring_size × Time::SIZE × flits_per_packet = 4 × 24 × 1 = 96입니다.

• ring_size = 4(바깥쪽 A, B 분할은 서브링 간 교환이 필요 없으므로, 각 서브링 안에서 가장 안쪽의 Slice 4개(C # 4 그룹 하나)만 통신합니다)
• Time::SIZE = D::SIZE × C::SIZE = 8 × 3 = 24
• flits_per_packet = sizeof(f32) × Packet::SIZE / 32 = 4 × 8 / 32 = 1(Packet = m![E], E::SIZE = 8)

비트맵을 보면 패딩만 걸러 내는 추출 방식이 그대로 드러납니다. bitmap[0] = {0, 1, 2}는 출력 Slice 0이 자신의 C # 4 그룹에 있는 유효한 입력 Slice 3개에서 데이터를 받고, 패딩 셀인 인덱스 3은 건너뛴다는 뜻입니다. {0, 1, 2, 3}을 읽으면 그 패딩을 실제 데이터처럼 가져오게 됩니다.

제약 조건

커스텀 구성에는 이러한 유연성을 제한하는 제약 조건 일곱 가지가 있습니다.

브로드캐스트 축은 새 축이어야 합니다

정규 구성 소개 부분과 같은 규칙입니다. 출력 Slice에 새로 추가하는 각 브로드캐스트 축은 입력 Slice나 입력 Time에 나타나면 안 됩니다.

예를 들어 axes![A = 256, B = 64, C = 32]이고 입력이 Slice = m![A], Time = m![B], Packet = m![C # 32]일 때, 출력 Slice = m![A / 4, B / 32, A % 4]는 B가 이미 입력 Time에 나타나므로 이 제약 조건을 위반합니다.

각 브로드캐스트 축은 정확히 한 번만 사용해야 합니다

각 브로드캐스트 축은 출력 Slice에 정확히 한 번만 나타나야 합니다. 같은 축을 출력의 두 위치에 반복하면 라우팅 비트맵에서 그 의미가 정의되지 않습니다.

예를 들어 출력 Slice = m![A / 4, X, X](X는 두 번 사용된 새 축)는 이 제약 조건을 위반합니다.

브로드캐스트 축은 패딩될 수 없습니다

출력 Slice의 브로드캐스트 축에는 패딩이 있어서는 안 됩니다(Axis # N 형태 불가). 브로드캐스트 축에 패딩이 있으면 패딩된 위치의 라우팅 목적지가 정의되지 않은 상태로 남습니다.

예를 들어 출력 Slice = m![A / 4, X # 4, Y]는 X가 패딩이 있는 브로드캐스트 축이므로 이 제약 조건을 위반합니다.

순서 보존

Slice에서 Time으로 이동하는 축은 입력 Slice 차원에서의 상대 순서를 유지해야 하며, 그렇지 않으면 SwitchConfig::CustomBroadcast의 검증기가 커널 컴파일 시점에 패닉을 일으킵니다. 각 라우터의 버퍼는 최소한(패킷 하나)이므로 로컬로 출력할지 전달할지를 즉시 결정해야 하며, 여러 패킷을 버퍼에 담아 순서를 바꿀 여지가 없습니다.

예를 들어 axes![A = 16, B = 16, C = 8, D = 8, E = 8]이고 dtype = i8일 때 입력 Slice = m![A, B], Time = m![C], Packet = m![D, E]를 출력 Slice = m![A, B / 4, 4], Time = m![C, B % 2, B / 2], Packet = m![D, E]로 매핑하면 이 제약 조건을 위반합니다. 여기서 B % 2와 B / 2는 입력 Slice에서의 배치와 비교해 순서가 뒤바뀌어 있습니다. 출력 Time = m![C, B / 2, B % 2]는 B / 2, B % 2가 입력 순서와 일치하므로 유효합니다.

가장 안쪽 시간 위치

Slice에서 Time으로 이동하는 축은 출력 Time의 가장 안쪽 위치를 차지해야 합니다. 파이프라인에서 다른 Slice의 데이터는 각 패킷 안에서 마지막에 도착하므로, Slice에서 Time으로 이동하는 하위 차원은 자연스럽게 가장 안쪽 시간 차원에 놓입니다. 이들을 다른 위치에 두려면 전체 시간 시퀀스를 버퍼링하고 재정렬해야 하는데, 하드웨어는 이를 수행할 수 없습니다.

예를 들어 위와 같은 axes!와 dtype에서 입력 Slice = m![A, B], Time = m![C], Packet = m![D, E]를 출력 Slice = m![A / 2, 2, B / 2, 2], Time = m![A % 2, C, B % 2], Packet = m![D, E]로 매핑하면 이 제약 조건을 위반합니다. 여기서 A % 2와 B % 2는 상대적인 순서를 올바르게 유지하지만, 그 사이에 C가 끼어 있습니다.

Note

Broadcast01은 time0 파라미터를 통해 이 제약 조건을 우회하지만, 사용자 정의 구성에는 그러한 메커니즘이 없으므로 이 제약 조건을 엄격히 따라야 합니다.

버려지는 값은 패딩이어야 함

어떤 차원에서 Slice에서 Time으로 이동하는 값의 개수가 그 차원의 전체 크기보다 적을 때(예시 3과 같은 부분 추출), 남겨지는 값은 모두 패딩이어야 합니다. 유효한 값을 버리면 실제 입력이 아무런 경고 없이 폐기되므로, 검증기는 커널 컴파일 시점에 이를 거부합니다.

예를 들어 axes![A = 16, B = 4, C = 3, D = 8, E = 8, X = 4]이고 입력이 Slice = m![A, B, C # 4]일 때, C # 4 → C로 추출하는 것은 버려지는 네 번째 값이 패딩 셀이므로 허용됩니다. 반면 완전히 유효한 축을 슬라이싱하는 것, 예를 들어 B = 4 → B = 3은 버려지는 값이 실제 데이터를 담고 있으므로 이 제약 조건을 위반합니다.

링 크기

ring_size 파라미터는 2의 거듭제곱이어야 합니다. 또한 컴파일러는 입력/출력 매핑(가장 바깥쪽의 non-direct-cast 경계)에서 예상 ring_size를 도출하며, 이와 일치하지 않는 사용자 지정 값은 거부합니다.