Transpose Engine

Transpose Engine은 Chip, Cluster, Slice 차원은 그대로 두고 Time 차원과 Packet 차원을 맞바꿉니다.

인터페이스

CollectTensor와 VectorFinalTensor는 모두 .transpose()를 제공합니다. VectorFinalTensor 진입점은 Vector Engine 출력을 Transpose Engine에 직접 공급합니다.

// `D: MaterializableScalar` (see its doc) excludes i5/i9 stagings from transpose.
impl<
    'l,
    const T: Tu,
    P: CanApplyTranspose,
    D: MaterializableScalar,
    Chip: M,
    Cluster: M,
    Slice: M,
    Time: M,
    Packet: M,
    B: Backend,
> TuTensor<'l, T, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Performs the transpose operation.
    #[primitive(TuTensor::transpose)]
    pub fn transpose<OutTime: M, OutPacket: M>(
        self,
    ) -> TransposeTensor<'l, T, D, Chip, Cluster, Slice, OutTime, OutPacket, B> {
        verify_transpose::<D, Time, Packet, OutTime, OutPacket>();
        TransposeTensor::new(self.device, self.inner.transpose(false))
    }
}

커널 작성자는 OutTime과 OutPacket(출력 차원 레이아웃)을 선택하고, 컴파일러는 그 결과가 매개변수에 나열된 하드웨어 제약 조건을 만족하는지 검증합니다.

아래 예제는 너비 16인 각 행을 두 개의 입력 패킷(D = 2)에서 모아 구성한 8×16 i8 행렬을 전치합니다. 이 예제는 이 페이지의 나머지 부분에서 계속 사용하는 예제입니다.

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

fn basic_transpose<'l, const T: Tu>(
    input: CollectTensor<'l, T, i8, m![1], m![1 # 2], m![P], m![B, C, D], m![E # 32]>,
) -> TransposeTensor<'l, T, i8, m![1], m![1 # 2], m![P], m![B, D, E], m![C # 32]> {
    input.transpose()
}

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

let c: CollectTensor<'_, _, i8, m![1], m![1 # 2], m![P], m![B, C, D], m![E # 32]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = basic_transpose(c);
}

valid_size는 Transpose Engine이 32바이트 입력 버스에서 사이클마다 읽는 유효 요소의 수이며, 모든 입력 플릿은 bit-width × valid_size 형태로 도착하고, 32바이트 플릿에서 남는 바이트는 패딩으로 취급되어 Unpack 단계에서 버려집니다. 8비트 및 16비트 요소의 경우, 엔진은 이중 입력 패킹을 선택하여 각 플릿에서 기본값인 8개 대신 16개의 유효 요소를 읽을 수 있습니다. 4비트 또는 32비트 요소에서는 이중 패킹을 지원하지 않습니다. 데이터는 CollectTensor::transpose()(Fetch → [Switch →] Collect → [Cast →] Transpose 이후) 또는 VectorFinalTensor::transpose()(Vector Engine에서 직접)를 통해 Transpose Engine에 도달합니다. Contraction Engine은 32b × 8만 내보내는 반면, Vector Engine과 Fetch Engine은 아래 표의 어떤 조합이든 내보냅니다.

in_cols, in_rows, out_rows는 커널 작성자가 선택한 OutTime과 OutPacket에 따라 정해집니다. 네 가지 모두 요소 크기에 따라 제약을 받습니다:

요소 크기	단일 `valid_size`	이중 `valid_size`	최대 `in_rows`	유효한 `in_cols`
4비트	16	지원하지 않음	16	16, 32
8비트	8	16	8	8, 16, 32
16비트	8	16	4	8, 16, 32
32비트	8	지원되지 않음	2	8, 16, 32

이 문서 전체에서 사용하는 예제(i8이므로 valid_size = 8)에서 컴파일러는 다음 값을 도출합니다.

매개변수	값	비고
`in_cols`	16	수집된 패킷 `D = 2`개 × `valid_size = 8`
`in_rows`	8	`C::SIZE`
`out_rows`	16	`D·E`(= `in_cols`, 완전 활용)

                 in_cols                  in_rows # F
           ┌─────────────────┐         ┌──────────────────┐
           │ 12 13 14 15 ... │         │ 3  7  11 15  ... │
 in_rows   │ 8  9  10 11 ... │  ────►  │ 2  6  10 14  ... │  out_rows
           │ 4  5  6  7  ... │         │ 1  5  9  13  ... │
           │ 0  1  2  3  ... │         │ 0  4  8  12  ... │
           └─────────────────┘         └──────────────────┘
                data_in                      data_out

Unpack

32바이트 입력 패킷마다 유효한 요소가 valid_size개 들어 있습니다. Unpack 단계는 플릿의 나머지 부분을 패딩으로 보고 버립니다. 8비트 또는 16비트 패킷에 기본 8요소 접두부를 넘어서는 유효 요소가 있으면 컴파일러는 이중 패킹을 선택합니다. 이때도 선택된 valid_size는 유효 접두부 전체를 포함해야 하며, 나머지 레인은 모두 패딩이어야 합니다.

이 문서의 예제에서는 [C, D, E # 32] → [C, D, E]입니다.

예를 들어 i8 패킷 [B=16 # 32]는 이중 패킹을 사용하여 유효 요소 16개를 모두 한 사이클에 전달합니다. 이와 동등한 단일 패킹 레이아웃이라면 연속된 두 패킷에서 각각 요소를 8개씩 전달하게 됩니다. i4 패킷은 이중 형식을 요청할 수 없습니다. 지원되는 valid_size가 16 하나뿐이기 때문입니다.

Gather

입력 행렬의 한 행은 너비가 in_cols = packets_per_col × valid_size개 요소이며, 연속된 packets_per_col개 패킷, 즉 가장 안쪽의 시간 단계들로부터 조립됩니다. Gather 단계는 이 패킷들을 하나의 행으로 이어 붙이고, 그 위의 in_rows개 시간 단계가 쌓여 [in_rows × in_cols] 입력 행렬을 이룹니다. (Unpack과 Gather는 모두 엔진이 입력을 읽는 동안 일어나며, 별도로 버퍼링되는 패스가 아닙니다.)

실행 예제에서 가장 안쪽의 D = 2개 패킷이 각각 valid_size = 8을 제공하여 너비 in_cols = 16의 행을 이루고, 그 위의 C = 8개 시간 단계가 쌓여 [8 × 16] 입력 행렬을 이룹니다.

Transpose

행렬을 전치합니다: [in_rows × in_cols] → [in_cols × in_rows].

계속 사용하는 예제에서는 [C, D, E] → [D, E, C]입니다.

Trim

일부 입력 패킷이 valid_size보다 적은 유효 요소를 담고 있으면 전치된 행렬에 패딩 행이 생깁니다. Trim 단계는 이러한 행을 제거하여 [out_rows × in_rows]를 만들며, 이때 out_rows ≤ in_cols입니다.

계속 사용하는 예제에서는 [D, E, C] → [D, E, C]입니다(입력이 완전히 활용되므로 제거되는 행이 없습니다). Trim이 실제로 행을 버리는 경우는 작은 행렬 예제를 참고하십시오.

Align

전치된 행은 in_rows개 요소 너비이지만, DM 패킷은 32바이트여야 합니다. Align 단계는 각 행을 32바이트 플릿으로 패딩하여 [out_rows × (in_rows # F)] 형태를 만들며, 여기서 F는 D[F]가 32바이트가 되도록 선택됩니다.

계속 사용하는 예제에서는 [D, E, C] → [D, E, C # 32]입니다.

지연 시간

Note

공식은 먼저 성능을 읽어 보십시오.

계속 사용하는 예제에서는 in_cols = 16 ≤ 16이므로 더블 버퍼링이 선택됩니다. in_flits = 16, out_rows = 16, n = 2일 때 총 지연 시간은 16 + 1 × max(16, 16) + 16 = 48 사이클입니다.

예제

작은 행렬

이 예제는 out_rows < in_cols일 때 Trim 단계가 패딩 행을 버리는 과정을 보여 줍니다:

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

fn small_transpose<'l, const T: Tu>(
    input: CollectTensor<'l, T, i8, m![1], m![1 # 2], m![P], m![A], m![B # 32]>,
) -> TransposeTensor<'l, T, i8, m![1], m![1 # 2], m![P], m![B], m![A # 32]> {
    input.transpose()
}

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

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

매개변수:

매개변수	값	비고
`in_cols`	8	`B::SIZE = 2`, 8로 패딩
`in_rows`	4	`A::SIZE`
`out_rows`	2	`B::SIZE`

단계:

Unpack: [A, B # 32] → [A, B # 8].
Gather: [A, B # 8] → [A, B # 8] (packets_per_col = 1이므로 각 패킷이 이미 온전한 in_cols = 8 행입니다).
Transpose: [A, B # 8] → [B # 8, A].
Trim: [B # 8, A] → [B, A] (패딩된 행 6개가 잘려 나갑니다).
Align: [B, A] → [B, A # 32].

지연 시간: in_cols = 8 ≤ 16이므로 더블 버퍼링이 선택됩니다. in_flits = 4, out_rows = 2, n = 1일 때 총 사이클 수는 4 + 0 × max(4, 2) + 2 = 6입니다.

큰 열

이 예제는 in_cols > 16으로 만들어 싱글 버퍼링을 강제합니다. 이로 인해 입력과 출력이 겹칠 수 없으므로 총 사이클이 증가합니다:

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

fn large_col_transpose<'l, const T: Tu>(
    input: CollectTensor<'l, T, i8, m![1], m![1 # 2], m![P], m![B, C, D], m![E # 32]>,
) -> TransposeTensor<'l, T, i8, m![1], m![1 # 2], m![P], m![B, D, E], m![C # 32]> {
    input.transpose()
}

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

let c: CollectTensor<'_, _, i8, m![1], m![1 # 2], m![P], m![B, C, D], m![E # 32]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = large_col_transpose(c);
}

매개변수:

매개변수	값	비고
`in_cols`	32	`D::SIZE × E::SIZE`
`in_rows`	8	`C::SIZE`
`out_rows`	32	`D::SIZE × E::SIZE`

단계:

Unpack: [C, D, E # 32] → [C, D, E].
Gather: 가장 안쪽의 D = 4개 패킷이 각 in_cols = 32 행을 이루고, C = 8개의 시간 단계가 쌓여 [8 × 32] 입력 행렬이 됩니다.
Transpose: [C, D, E] → [D, E, C].
Trim: [D, E, C] → [D, E, C] (잘려 나가는 행이 없습니다).
Align: [D, E, C] → [D, E, C # 32].

지연 시간: in_cols = 32 > 16이므로 싱글 버퍼링이 선택됩니다. in_flits = 32, out_rows = 32, n = 2(B)일 때 총 사이클 수는 2 × (32 + 32) = 128입니다.

16비트 데이터 타입

이 예제는 bf16을 사용합니다. 더 넓은 요소로 인해 최대 in_rows가 절반(8 대신 4)으로 줄고, 32바이트 출력 플릿이 16개 요소로 줄어듭니다(i8의 경우 32개):

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

fn bf16_transpose<'l, const T: Tu>(
    input: CollectTensor<'l, T, bf16, m![1], m![1 # 2], m![P], m![C, D], m![E # 16]>,
) -> TransposeTensor<'l, T, bf16, m![1], m![1 # 2], m![P], m![C, E], m![D # 16]> {
    input.transpose()
}

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

let c: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![P], m![C, D], m![E # 16]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = bf16_transpose(c);
}

매개변수:

매개변수	값	비고
`in_cols`	8	`E::SIZE`
`in_rows`	4	`D::SIZE`
`out_rows`	8	`E::SIZE`

단계:

Unpack: [D, E # 16] → [D, E].
Gather: [D, E] → [D, E] (packets_per_col = 1, 각 패킷이 이미 in_cols = 8인 완전한 행입니다).
Transpose: [D, E] → [E, D].
Trim: [E, D] → [E, D] (잘라내는 행이 없습니다).
Align: [E, D] → [E, D # 16].

지연 시간: in_cols = 8 ≤ 16이므로 더블 버퍼링이 선택됩니다. in_flits = 4, out_rows = 8, n = 8(C)일 때 총 사이클 수는 4 + 7 × max(4, 8) + 8 = 68입니다.

4비트 데이터 타입

이 예제는 i4를 사용합니다. valid_size = 16으로 사이클당 요소 수가 두 배가 되고 최대 in_rows는 16(16 × 4비트 = 8바이트)으로 늘어나며, 32바이트 플릿은 64개 요소로 커집니다:

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

fn i4_transpose<'l, const T: Tu>(
    input: CollectTensor<'l, T, i4, m![1], m![1 # 2], m![P], m![B, C], m![E # 64]>,
) -> TransposeTensor<'l, T, i4, m![1], m![1 # 2], m![P], m![B, E], m![C # 64]> {
    input.transpose()
}

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

let c: CollectTensor<'_, _, i4, m![1], m![1 # 2], m![P], m![B, C], m![E # 64]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = i4_transpose(c);
}

매개변수:

매개변수	값	비고
`in_cols`	16	`E::SIZE`
`in_rows`	16	`C::SIZE`
`out_rows`	16	`E::SIZE`

단계:

Unpack: [C, E # 64] → [C, E].
Gather: [C, E] → [C, E] (packets_per_col = 1, 각 패킷이 이미 온전한 in_cols = 16 행입니다).
Transpose: [C, E] → [E, C].
Trim: [E, C] → [E, C] (잘라 내는 행이 없습니다).
Align: [E, C] → [E, C # 64].

지연 시간: in_cols = 16 ≤ 16이므로 더블 버퍼링이 선택됩니다. in_flits = 16, out_rows = 16, n = 4(B)일 때 총합은 16 + 3 × max(16, 16) + 16 = 80 사이클입니다.

32비트 데이터 타입

이 예제는 Contraction Engine이 출력하는 두 가지 32비트 형식(32b × 8, f32 또는 i32) 중 하나인 f32를 사용합니다. 더 넓은 요소로 인해 최대 in_rows가 2(2 × 4바이트 = 8바이트)로 줄고, 32바이트 플릿이 8개 요소로 줄어듭니다:

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

fn f32_transpose<'l, const T: Tu>(
    input: CollectTensor<'l, T, f32, m![1], m![1 # 2], m![P], m![B, D], m![E # 8]>,
) -> TransposeTensor<'l, T, f32, m![1], m![1 # 2], m![P], m![B, E], m![D # 8]> {
    input.transpose()
}

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

let c: CollectTensor<'_, _, f32, m![1], m![1 # 2], m![P], m![B, D], m![E # 8]> = CollectTensor::new(&mut device.main, Tensor::zero());
let _o = f32_transpose(c);
}

매개변수:

매개변수	값	비고
`in_cols`	8	`E::SIZE`
`in_rows`	2	`D::SIZE`
`out_rows`	8	`E::SIZE`

단계:

Unpack: [D, E # 8] → [D, E].
Gather: [D, E] → [D, E] (packets_per_col = 1; 각 패킷이 이미 완전한 in_cols = 8 행입니다).
Transpose: [D, E] → [E, D].
Trim: [E, D] → [E, D] (잘라내는 행이 없습니다).
Align: [E, D] → [E, D # 8].

지연 시간: in_cols = 8 ≤ 16이므로 더블 버퍼링이 선택됩니다. in_flits = 2, out_rows = 8, n = 4(B)일 때 합계는 2 + 3 × max(2, 8) + 8 = 34 사이클입니다.

입력 전용 단계(in_flits 사이클): 첫 번째 버퍼가 채워집니다.
중첩 단계((n - 1) × max(in_flits, out_rows) 사이클): 한 버퍼가 입력을 받는 동안 다른 버퍼가 동시에 출력을 생성하므로, 매 반복은 더 느린 쪽에 의해 결정됩니다.
출력 전용 단계(out_rows 사이클): 마지막 버퍼가 비워집니다.

Tensor Contraction Processor 프로그래밍

Transpose Engine

인터페이스

아키텍처

매개변수

Unpack

Gather

Transpose

Trim

Align

지연 시간

예제

작은 행렬

큰 열

16비트 데이터 타입

4비트 데이터 타입

32비트 데이터 타입

성능

버퍼링 모드

싱글 버퍼링 지연 시간

더블 버퍼링 지연 시간

Keyboard shortcuts

Tensor Contraction Processor 프로그래밍