Tensor Program Abstraction

TensorIR에 대해 알아보기 전에 primitive tensor function에 대해 먼저 알아보자 primitive tensor function은 computational operation의 single “unit”에 대응하는 function이다.

• 예를 들어 convolution 또는 convolution+relu(fused) operation

primitive tensor function의 추상화는 tensor computation을 위한 multi-dimensional buffers, loop nests을 포함한다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

from tvm.script import tir as T

@T.prim_func
def main(
    A: T.Buffer((128,), "float32"),
    B: T.Buffer((128,), "float32"),
    C: T.Buffer((128,), "float32"),
) -> None:
    for i in range(128):
        with T.block("C"):
            vi = T.axis.spatial(128, i)
            C[vi] = A[vi] + B[vi]

Key Elements of Tensor Programs

위의 코드는 primitive tensor function의 예제로 다음의 Key Elements를 가진다.

1. multi-dimensional buffer (input 2개, output 1개)
2. 계산을 용이하게 하는 단일 loop nest i
3. 단일 compute statement (element-wise sum)

Extra Structure in TensorIR

(TVM의 TensorIR에서 annotation(추가적인 구조)가 필요한지를 설명하는 내용)
어떤 프로그램들은 loop’s sequence에 의존하기 때문에 마음대로 루프를 재배열하거나 병렬화할 수는 없지만 우리가 주로 다루는 대부분의 primitive tensor function은 루프 반복 간의 독립성이라는 좋은 성질을 가지고 있다. 이전 예제에서는 block / iteration annotations을 들 수 있다.

• block annotation : with T.block("C")은 블록이 스케줄링을 위해 지정된 기본 계산 단위임을 의미, block()은 single computation statement, multiple computation statements, opaque intrinsics(Tensor Core instructions) 을 포함할 수 있다.
• iteration annotation : T.axis.spatial은 vi가 에 매핑되고 모든 iterations이 독립적임을 의미한다. 블록과 축의 annotation는 프로그램 실행 자체에는 필수는 아니지만, 프로그램을 최적화하거나 변환할 때 매우 중요하다. (vi와 관련된 루프를 parallelize, reorder 할 수 있음을 나타낸다.)

Understand TensorIR Abstraction

tensor program abstraction의 주요 목적은 hardware acceleration options 및 loop를 묘사하는 것(threading, 특수 hardware instructions 적, memory access)

예를 들어, 128*128 Matrix_Multiplication + Relu 를 python numpy code와 TensorIR로 나타내면 아래와 같다 (TVMScript 이용). 다음 각 sub-chapter에서 아래 코드에 대한 element에 대하여 설명한다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    Y = np.empty((128, 128), dtype="float32")
    for i in range(128):
        for j in range(128):
            for k in range(128):
                if k == 0:
                    Y[i, j] = 0
                Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
    for i in range(128):
        for j in range(128):
            C[i, j] = max(Y[i, j], 0)


@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def mm_relu(A: T.Buffer((128, 128), "float32"),
                B: T.Buffer((128, 128), "float32"),
                C: T.Buffer((128, 128), "float32")):
        Y = T.alloc_buffer((128, 128), dtype="float32")
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                vk = T.axis.reduce(128, k)
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

Function Parameters and Buffers

TensorIR에서 A,B,C는 float32, (128,128) 형태의 T.Buffer를 취하며 이러한 추가 정보는 MLC(Machine Learning Compile)가 shape 와 data type에 특화된 코드를 생성하는 데 도움이 됩니다.

1
2
3
4
5
6
7
8

# TensorIR
def mm_relu(A: T.Buffer((128, 128), "float32"),
            B: T.Buffer((128, 128), "float32"),
            C: T.Buffer((128, 128), "float32")):
    ...
# NumPy
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    ...

상기 내용과 비슷하게 중간 할당에서도 버퍼를 사용함

1
2
3
4

# TensorIR
Y = T.alloc_buffer((128, 128), dtype="float32")
# NumPy
Y = np.empty((128, 128), dtype="float32")

Loop Iterations

Loop Iterations은 python 구문과 바로 대응 되기에 for i in range(num)를 바로 사용하거나, T.grid를 사용하여 코드를 축약할 수 있다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# TensorIR with `T.grid`
for i, j, k in T.grid(128, 128, 128):
    ...
# TensorIR with `range`
for i in range(128):
    for j in range(128):
        for k in range(128):
            ...
# NumPy
for i in range(128):
    for j in range(128):
        for k in range(128):
            ...

Computational Block

computational statements에서 중요 차이점은 T.block이라는 추가적인 요소를 포함한다. T.block은 TensorIR내부에서 기본적인 computation unit을 나타며 numpy코드보다 더 많은 정보(block axes 세트인 (vi, vj, vk)와 이를 둘러 싼 computations)를 포함한다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

# TensorIR
with T.block("Y"):
    vi = T.axis.spatial(128, i)
    vj = T.axis.spatial(128, j)
    vk = T.axis.reduce(128, k)
    with T.init():
        Y[vi, vj] = T.float32(0)
    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
# NumPy
vi, vj, vk = i, j, k
if vk == 0:
    Y[vi, vj] = 0
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]

T.axis 코드(vi = T.axis.spatial(128, i) 3줄) block axes의 key properties를 나타낸다. 문법은 다음과 같다.

• [block_axis] = T.axis.[axis_type]([axis_range], [mapped_value]) 이 코드는 하기의 세부 정보를 전달한다.
• vi, vj, vk의 i, j, k와의 binding
• vi, vj, vk의 range 선언( = 128)
• iterators 의 속성 : spatial(독립적인 계산), reduce(축소/누적 계산) 또한 위의 코드는 T.axis.remap을 이용하여 하기와 같이 표현 가능하다.
• vi, vj, vk = T.axis.remap(“SSR”, [i, j, k])

역자 주

• spatial axis : 이 블록이 출력으로 쓰는 인덱스로 병렬화 가능
• reduce axis : 이 블록이 누적에 사용하는 인덱스로 병렬화가 가능할 수도 있다.(race condition 조심)
• Reduction 연산 : 다차원 배열(tensor)에서 하나의 축을 따라 값을 누적(sum(x, axis=1), matmul(A, B), mean(x), softmax(x), conv2d) 하는 연산

Block Axis Properties

T.axis는 해당 axis가 현재 수행 중인 computation과 어떤 관계가 있는지를 나타낸다. 이 블록은 세 개의 축(vi, vj, vk)을 가지고 있으며 A[vi, vk], B[vk, vj]를 읽어 Y[vi, vj]를 업데이트한다. vi, vj의 값이 고정되면 (Y[vi, vj]라는 한 점을 정하면) 이 점의 계산은 오직 vk를 반복하면서 이루어지며 다른 위치인 Y[vi’, vj’]와는 독립적이다.(안전하게 병렬 처리도 가능함) 그래서 vi, vj는 spatial axis (공간 좌표) vk는 reduce axis (누적을 위한 축) 라고 부른다. (챗GPT의 해석 : TVM은 블록 내에서 어떤 축이 “누적 대상"이고, 어떤 축이 “출력 대상"인지를 명시적으로 구분하며 구분을 통해 병렬화 가능성 판단, 메모리 스케줄링, 코드 자동 최적화 등을 가능하게 함.)

Why Extra Information in Block

블록 내부의 축 정보(축의 범위와 속성 등)가 포함되기 때문에, 블록은 외부 루프와 독립적인 “자기 완결(self-contained)” 구조가 된다 즉 vi = T.axis.spatial(128, i) 처럼 축의 역할과 범위를 명시적으로 선언함으로써 루프에 의존하지 않고, 자신만으로도 무슨 계산을 어디서 수행해야 하는지 알 수 있는 완전한 구조가 되는 것. 또한 블록의 축 정보는 외부 루프가 올바른지 검사하는 기준도 제공해준다.

TensorIR Creation

이번 챕터에서는 TensorIR function을 코딩하는 방법을 기술한다.(이 챕터에서의 기술내용은 Relax모델을 컴파일하는 유저에게 필 수적인 것은 아니다.)

Create TensorIR using TVMScript

앞에서 사용한 mm_relu TVM Script를 다시 설명함(본 포스트에서는 스킵함, 필요하면 TVM DOC를 볼 것)

TensorIR Function with Dynamic Shapes

TVMScript가 Python 인터프리터에 의해 실행되지 않지만 Python과의 상호 작용은 제한적으로 가능하다. Python 변수를 사용하여 TensorIR의 shape과 Data type 변수 지정 가능하다. 또한 이를 통해 runtime dynamic shape inference가 가능하다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

@I.ir_module
class DynamicShapeModule:
    @T.prim_func
    def mm_relu(a: T.handle, b: T.handle, c: T.handle):
        # Dynamic shape definition
        M, N, K = T.int32(), T.int32(), T.int32()

        # Bind the input buffers with the dynamic shapes
        A = T.match_buffer(a, [M, K], dtype)
        B = T.match_buffer(b, [K, N], dtype)
        C = T.match_buffer(c, [M, N], dtype)
        Y = T.alloc_buffer((M, N), dtype)
        for i, j, k in T.grid(M, N, K):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    Y[vi, vj] = T.cast(T.float32(0), dtype)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(M, N):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype)

def evaluate_dynamic_shape(lib: tvm.runtime.Module, m: int, n: int, k: int):
    A = tvm.nd.array(np.random.uniform(size=(m, k)).astype("float32"))
    B = tvm.nd.array(np.random.uniform(size=(k, n)).astype("float32"))
    C = tvm.nd.array(np.zeros((m, n), dtype="float32"))
    lib(A, B, C)
    return C.numpy()


# Compile lib only once
dyn_shape_lib = tvm.compile(DynamicShapeModule, target="llvm")
# Able to handle different shapes
print(evaluate_dynamic_shape(dyn_shape_lib, m=4, n=4, k=4))
print(evaluate_dynamic_shape(dyn_shape_lib, m=64, n=64, k=128))

Create TensorIR using Tensor Expression

TensorIR을 더 간결하게 표현할 수 있는 Tensor Expression 표현 방식이 있다. 레거시 인듯 해서 우선은 넘어가기로 하겠다.

Transformation

primitive tensor functions의 transformation은 컴파일 flow의 주요 과정이다. 이제까지 사용한 mm_relu를 평가하는 함수를 만들어 TensorIR transformation 과정을 진행한다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

import numpy as np

a_np = np.random.uniform(size=(128, 128)).astype("float32")
b_np = np.random.uniform(size=(128, 128)).astype("float32")
c_np = a_np @ b_np

a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
c_nd = tvm.nd.array(np.zeros((128, 128), dtype="float32"))


def evaluate(mod: tvm.IRModule):
   lib = tvm.tir.build(mod, target="llvm")
   # check correctness
   lib(a_nd, b_nd, c_nd)
   np.testing.assert_allclose(c_nd.numpy(), c_np, rtol=1e-5)
   # evaluate performance
   f_timer = lib.time_evaluator("main", tvm.cpu())
   print(f_timer(a_nd, b_nd, c_nd))


evaluate(MyModule)

Initialization Schedule

Schedule helper class를 사용하여 transformation 프로세스를 초기화 할 수 있다.

1

sch = tvm.tir.Schedule(MyModule)

Loop Tiling

Loop Split

다음 명령을 수행하여 block Y와 이와 관련된 루프의 reference를 얻는다.

1
2

block_Y = sch.get_block("Y")
i, j, k = sch.get_loops(block_Y)

그 다음 j를 두개의 loop로 분리(inner loop의 길이가 8, 즉 128회를 16*8회로 나눔)하는 transformation을 수행할 수 있다. (이 프로세스는 절차?적 이기 때문에 위의 내용을 두번 수행하면 j가 존재하지 않는다는 에러가 출력됨, 역자주, 아마 첫번째 수행에서 j가 j_0, j_1로 분리되므로 j는 사라진다는 뜻일 듯)

1

j0, j1 = sch.split(j, factors=[None, 8])

sch.mod.show() 로 transformation의 결과를 확인 할 수 있다. 원문을 보면 j에 관련된 루프가 j_0, j_1로 분리된 결과를 확인할 수 있다.

Loop Reorder

루프 j_0, j_1를 reordering 할 수 있다.

1
2
3

sch.reorder(j0, k, j1)
sch.mod.show()
evaluate(sch.mod)

원문을 보면 j에 관련된 루프가 j_0, k, j_1로 reorder 된 것을 확인할 수 있다.

역자주

TensorIR에서 Transformation은 고성능 코드 생성을 위한 스케줄 최적화 작업이다.

• Loop Tiling은 큰 Loop를 작은 Loop(블록) 단위로 나누어 중첩 루프 구조로 바꾸는 최적화이다. 큰 배열이 L1/L2 캐시에 모드 들어가지 않을 때 이를 나누어 타일 단위로 계산하면 작은 덩어리의 데이터만 반복 사용하게 되어 캐시 히트율이 올라간다. 또한 벡터화나 병렬화가 가능해지며 loop fusion / compute_at 등의 고급 스케줄링을 하기 전 준비 단계로 쓰인다.
• Loop Reordering도 메모리 접근의 지역성(Locality) 향상시켜 CPU캐시 히트율 향상시킨다. 또한 벡터화의 가능성을 증진 시킨다.

Leverage Localities

block C를 loop Y 내부로 재배치 하는 reverse_compute_at이라는 primitive를 사용할 수 있다.

1
2
3

block_C = sch.get_block("C")
sch.reverse_compute_at(block_C, j0)
sch.mod.show()

본문에 T.block("C") 블록이 j_1 하위로 이동한 결과를 확인 할 수 있다.

역자주

• compute_at, reverse_compute_at은 블록의 계산 위치를 다른 루프 또는 블록 내부로 이동시키는 스케줄링
  • compute_at : 앞선 블록 안에서 같이 계산, reverse_compute_at : 나중 블록을 앞 루프 안으로 옮김
• 연산을 필요한 시점에만 수행하여 메모리 접근 줄이고 연산 지역성(locality) 향상
• 버퍼 공유 또는 중간 결과를 바로 사용하는 구조에서 타일링 이후, 해당 블록을 타일 단위로 옮길 때 효과적

Rewrite Reduction

이제까지의 reduction initialization / update 는 single block body에서 수행되었으며 이러한 형태는 outer loops i, j의 initialization / update가 동기화를 필요. loop transformation 후 Y elements의 초기화를 reduction update로 부터 분리할 수 있다.

1
2
3

sch.decompose_reduction(block_Y, k)
sch.mod.show()
evaluate(sch.mod)

역자 주

decompose_reduction은 TVM TensorIR에서 reduction 연산(누적 연산)을 초기화 단계와 갱신 단계로 분리해 주는 스케줄링 기법 T.init()이 loop body에 있으면 병렬화/재배열 어려워 진다. 또한 루프 제어를 단순화 시킬 수 있으며 메모리 접근이 최적화 된다. reorder, compute_at, unroll 등을 적용하기도 쉬워진다.

Trace the Transformation

TensorIR schedule은 절차적 언어로 단계별로 실행된다. TVM은 schedule 또는 schedule 이력을 추적할 수 있다.

• sch.trace.show() : schedule 이력 출력
• sch.mod.show() : schedule 출력