텐서 의미론

텐서는 HBM, Chip 내부 DM 또는 파이프라인 스트림에 존재하며, 연산은 이러한 텐서를 변환합니다. 이 장에서는 텐서의 수학적 의미를 정의합니다. 즉, 텐서 변수가 수학적 텐서를 보유한다는 것이 무엇을 의미하는지, 그리고 연산이 수학적 함수를 명세한다는 것이 무엇을 의미하는지를 정의합니다. 이러한 정의를 통해 vISA 프로그램을 텐서 수준에서 추론할 수 있습니다. 어떤 매핑이나 메모리 계층 단계를 사용하든 관계없이, 함수의 출력이 올바른 수학적 텐서를 보유하면 그 함수는 올바릅니다.

텐서 보유 의미론

각 요소가 각 차원의 매핑이 생성한 부분 인덱스를 합하여 만든 텐서 인덱스에서의 \(T\) 값을 저장할 때, 텐서 변수는 수학적 텐서 \(T\)를 보유합니다.

HostTensor<D, E>는 가장 단순한 경우로, 단일 매핑 E가 버퍼 위치와 텐서 인덱스 사이의 대응 관계를 완전히 결정합니다. 예를 들어 A = 8, B = 512인 HostTensor<bf16, m![A, B]>는 4,096개의 bf16 요소를 A-major, B-minor 순서로 저장합니다. 이 텐서는 다음 조건을 만족할 때 텐서 \(T\)를 보유합니다.

• E::map(i) = Cell::Index(ti)인 모든 버퍼 인덱스 i에 대해,
• i번째 요소가 ti에서의 \(T\) 값을 저장합니다.

HbmTensor<D, Chip, Element>는 단일 매핑을 둘로 나누어 이를 확장합니다. Chip은 Chip 인덱스를 부분 텐서 인덱스로 매핑하고, Element는 Chip별 요소 인덱스를 나머지 부분 인덱스로 매핑하며, 각각이 서로 겹치지 않는 축 부분집합을 담당하므로 둘의 합이 전체 텐서 인덱스를 복원합니다. 이 텐서는 다음 조건을 만족할 때 \(T\)를 보유합니다.

• Chip::map(i) = Cell::Index(ti)이고 Element::map(j) = Cell::Index(tj)인 모든 Chip 인덱스 i와 요소 인덱스 j에 대해,
• i번째 Chip의 j번째 요소는 인덱스 ti + tj에서의 \(T\) 값을 저장합니다.

다른 모든 텐서 타입은 같은 규칙을 더 많은 차원에 적용합니다. 즉, 각 요소는 모든 매핑 매개변수가 반환한 부분 인덱스의 합에서의 \(T\) 값을 저장합니다.

선형 결합 의미론

선형 결합 표현식 $(e1:n1, ..., ed:nd)는 지정된 스트라이드로 여러 차원을 결합합니다. 그 크기는 size_S($(e1:n1, ..., ed:nd)) = 1 + sum_k((size_S(ek) - 1) * nk)입니다. 매핑 S, $(e1:n1, ..., ed:nd) |- si ~ ti는 모든 k에 대해 S, ek |- sik ~ tik, si = sum_k(sik * nk), ti = sum_k(tik * nk)를 만족하는 si1...sid, ti1...tid가 존재할 때 유효합니다.

선형 결합은 외합(outer sum)을 표현할 수 있습니다. e1 * e2는 $(e1 : size_S(e2), e2 : 1)와 동등합니다. 축 순서 변경이 중요한 경우에는 외합이 선호됩니다. e1 * e2를 e2 * e1로 바꿀 때 스트라이드를 수동으로 갱신할 필요가 없기 때문입니다.

슬라이딩 연산은 서로 겹치는 데이터 블록에 접근합니다. \(\{N=5, F=3\}\) 형태의 텐서를 나타내는 9개 요소의 버퍼를 생각해 보십시오. 여기서 각 행은 한 번에 한 요소씩 슬라이딩하는 3개 요소 Slice입니다. \(N, F\) 위치의 텐서 요소는 버퍼 인덱스 \(N + 2F\)에 매핑됩니다.

$$ \begin{array}{c|ccc} & F=0 & F=1 & F=2 \\ \hline N=0 & 0 & 2 & 4 \\ N=1 & 1 & 3 & 5 \\ N=2 & 2 & 4 & 6 \\ N=3 & 3 & 5 & 7 \\ N=4 & 4 & 6 & 8 \\ \end{array} $$

이 슬라이딩 패턴에서는 하나의 공간 인덱스가 여러 텐서 인덱스에 매핑될 수 있습니다. 예를 들어 공간 인덱스 4는 {4_N}, {2_N, 1_F}, {2_F}에 동시에 매핑되며, 이는 (S, e).maps(si, ti)가 일대일 관계가 아님을 보여 줍니다. 이 선형 결합은 N에 스트라이드 1을, F에 스트라이드 2를 사용하므로 1 + (5-1)*1 + (3-1)*2 = 9가 됩니다.

함수 명세

함수를 명세한다는 것은 그 출력이 입력에 대해 무엇을 보유하는지 선언하는 것을 의미합니다. 예를 들어 함수 elementwise_add는 다음과 같은 의미에서 수학적 연산 \(f(T_1, T_2) = T_1 + T_2\)를 명세합니다.

• 모든 텐서 \(T_1\)와 \(T_2\)에 대해,
• lhs가 \(T_1\)를 담고 rhs가 \(T_2\)를 담으면,
• 반환값은 \(T_1 + T_2\)를 담습니다.

extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 512];

// This signature specifies the mathematical contract; the engine pipeline
// implementation is shown in the Computing Tensors chapter.
fn elementwise_add(
    lhs: &HbmTensor<bf16, m![A], m![B]>,
    rhs: &HbmTensor<bf16, m![A], m![B]>,
) -> HbmTensor<bf16, m![A], m![B]> {
    // The implementation is intentionally omitted; this example states only
    // the function's tensor-level contract.
}

수학적 텐서 이동은 \(f(T) = T\)를 명세합니다. 즉, 출력은 표현 방식과 관계없이 입력과 같은 수학적 텐서를 담습니다. .to_dm()은 수학적 텐서 이동입니다. 예를 들어 .to_dm() 메서드는 다음과 같은 의미에서 \(f(T) = T\)를 명세합니다.

• hbm이 \(T\)를 담으면,
• 반환값은 \(T\)를 담습니다.

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

fn hbm_to_dm(
    device: &mut Device,
    hbm: &HbmTensor<bf16, m![A], m![B]>,
) -> DmTensor<bf16, m![A], m![1], m![B / 2], m![B % 2]> {
    hbm.to_dm(&mut device.tdma)
}
}