Vector Calculation Utilizing AMX

AMX를 활용한 벡터 계산

이 문서는 AMX(Advanced Matrix Extensions)를 활용한 벡터 계산에 대해 설명하며, 특히 BF16 (bfloat16) 데이터 타입을 활용한 연산에 중점을 둡니다. AMX의 동작 원리, 구현 세부사항, 이론적 최대 성능, 그리고 효율적인 계산을 위한 최적화 기법에 대해 다룹니다.

AMX란?

Intel® AMX (Advanced Matrix Extensions)는 Intel에서 도입한 매트릭스 연산 가속 기술로, 하드웨어에서 직접 대규모 행렬 연산을 효율적으로 수행할 수 있도록 설계된 x86 명령어 확장입니다. 타일(Tile)이라는 구조를 사용하여 데이터를 2차원 형태로 표현하며, 반복되는 연산의 오버헤드를 줄이는데 탁월한 성능을 제공합니다.
AMX는 인텔 제온 스케일러블 프로세서 코어에 내장된 전용 하드웨어 블록으로, 행렬 연산에 의존하는 딥러닝 학습 및 추론 작업을 최적화하고 가속화하는 데 도움을 줍니다.
AMX는 AI 작업을 별도의 가속기로 오프로드하는 대신 CPU에서 실행할 수 있도록 하여 성능을 크게 향상시킵니다.
이 아키텍처는 BF16(학습/추론) 및 int8(추론) 데이터 유형을 지원하며, 두 가지 주요 구성 요소로 이루어져 있습니다.
이러한 구성 요소들은 Intel® AMX가 각 코어에 더 많은 데이터를 저장하고 단일 연산으로 더 큰 행렬을 계산할 수 있도록 합니다. 또한 AMX는 완전히 확장 가능하고 스케일러블하도록 설계되었습니다.

AMX의 구성 요소

Intel AMX는 각 코어에 더 큰 데이터 조각을 저장한 다음 단일 연산으로 더 큰 행렬을 계산합니다. 첫 번째 구현은 TMUL(타일 행렬 곱셈 유닛)이라고 불리며, 타일에서 작동할 수 있는 융합 곱셈-덧셈 유닛의 그리드로 구성됩니다. 초기 형태에서는 최대 8개의 타일(TMM0 … TMM7이라고 명명된)을 구현하며, 이들은 16개의 행과 64바이트 크기의 배열로 구성되어 있습니다. 메모리에서 더 큰 이미지의 작은 부분을 나타내는 타일을 로드하고, 해당 타일에 대해 작업을 수행한 뒤, 이미지의 다음 부분을 나타내는 다음 타일로 반복합니다. 작업이 완료되면 결과 타일을 메모리에 저장합니다.
타일: 각 1킬로바이트 크기로, 대량의 데이터를 저장하는 2차원 레지스터 8개로 구성됩니다.
타일 행렬 곱셈(TILE MATRIX MULTIPLICATION, TMUL): TMUL은 타일에 연결된 가속 엔진으로, AI를 위한 행렬 곱셈 연산을 수행합니다.
TMUL 명령어 세트의 행렬 곱셈은 다음을 계산합니다: C[M][N] += A[M][K] * B[K][N].

타일 로드 및 저장과 가속기 명령은 다중 사이클 실행 유닛(TMUL)로 전송됩니다.
간단히 말해, Intel AMX는 각 코어에 더 큰 데이터 조각을 저장한 다음 단일 연산으로 더 큰 행렬을 계산합니다.

AMX의 주요 특징

int8 및 BF16 데이터 타입에 최적화되어 있음
- AMX가 딥러닝 모델을 가속화하려면 모델이 BF16 또는 INT8 형식이어야 합니다. 모델을 이러한 최적화된 형식으로 변환하려면 BF16의 경우 자동 혼합 정밀도(auto-mixed precision)를, INT8의 경우 양자화(quantization)를 사용하여 변환할 수 있습니다. 이는 프레임워크(e.g. PyTorch* 또는 TensorFlow*) 내에서 네이티브로 수행하거나, 추가 기능을 제공하는 Intel의 오픈소스 도구를 통해 가능합니다. BF16은 변환이 간단하며 일반적으로 정확도를 유지합니다. INT8은 더 효율적인 데이터 타입으로, Intel의 오픈소스 압축 도구를 사용하여 정확도를 유지할 수 있습니다.
- INT8은 FP32(주로 AI에서 사용되는 단일 정밀도 부동소수점 형식)의 정밀도가 필요하지 않을 때 추론에 사용되는 데이터 타입입니다. INT8 데이터 타입은 정밀도가 낮기 때문에, 계산 주기당 더 많은 INT8 연산을 처리할 수 있으며, 이는 속도와 효율성이 중요한 실시간 애플리케이션 및 행렬 곱셈 작업에 이상적입니다.
- BF16은 대부분의 학습 작업에서 충분한 정확도를 제공하는 데이터 타입입니다. 필요하다면 추론 시 더 높은 정확도를 제공할 수도 있습니다. BF16은 FP32로 얻을 수 있는 것과 거의 동일한 정확도로 머신러닝(ML) 모델 학습을 가능하게 하며, 계산 비용은 극히 일부에 불과합니다.
- Intel AMX는 BF16 및 INT8 데이터 타입만 지원하며, FP32 데이터 타입은 여전히 3세대 Intel Xeon 프로세서에 포함된 Intel® Advanced Vector Extensions 512 (Intel® AVX-512) 명령어를 통해 지원됩니다.

BF16이란?

*BF16 (bfloat16)**은 16비트 부동소수점 형식으로, 기계 학습 및 수치 계산 작업에서 널리 사용됩니다. IEEE 754 FP32 포맷과 유사하지만, 8비트를 지수로 사용하여 FP32와 동일한 범위를 지원합니다.

BF16의 장점

작은 메모리 크기: FP32 대비 메모리 사용량이 절감됩니다.
충분한 정밀도: 머신 러닝 및 벡터 계산에 적합.

이 기술은 대형 신경망을 학습시키는 데 사용되는 기술 중 하나로, 모델의 매개변수를 서로 다른 정밀도의 데이터 타입(주로 FP16과 FP32)에 저장하여 더 빠르게 실행하고 메모리 사용량을 줄일 수 있습니다.

현재 대부분의 모델은 단일 정밀도 부동소수점(FP32) 데이터 타입을 사용하며, 이는 32비트 메모리를 요구합니다. 그러나 두 가지 낮은 정밀도의 데이터 타입, float16과 bfloat16(BF16)은 각각 16비트 메모리만 요구합니다. BF16은 16비트의 컴퓨터 메모리를 차지하지만, 32비트 부동소수점 수의 동적 범위를 대략적으로 표현할 수 있는 부동소수점 형식입니다. BF16 형식은 다음과 같이 구성됩니다.

1비트 - 부호(sign)
8비트 - 지수(exponent)
7비트 - 가수(fraction)

Intel AMX는 추천 시스템, 자연어 처리, 이미지 감지와 같은 딥러닝 워크로드를 가속화합니다. 전통적인 머신 러닝 워크로드에서 사용되는 표형 데이터(tabular data)는 기존의 Intel AVX-512 명령어를 사용합니다. **많은 딥러닝 워크로드는 혼합 정밀도(mixed precision)를 사용하며, 4세대 Intel Xeon 프로세서는 Intel AMX와 Intel AVX-512 사이에서 가장 효율적인 명령어 세트를 사용하도록 원활히 전환할 수 있습니다.

행렬 데이터를 타일로 구성하여 하드웨어에서 효율적으로 처리.
행렬 곱셈 및 벡터 내적 같은 연산에서 큰 성능 향상.

AMX가 int8에서 동작하는 원리

int8 데이터에서 AMX는 타일 기반 연산을 통해 밀집 행렬 곱셈을 최적화합니다. 작동 방식은 다음과 같습니다:

타일 구성: 타일은 서브 매트릭스를 나타내도록 구성되며, 2D 그리드 형태로 데이터를 저장합니다.
데이터 로드: 입력 행렬(예: 가중치 및 활성화 값)을 타일에 로드합니다.
행렬 곱셈 수행: _tile_dpbusd() 같은 명령어를 통해 타일 간의 행과 열 간 내적을 계산합니다.
결과 저장: 계산된 결과는 타일 또는 메모리에 저장됩니다.

이 방식을 통해 AMX는 int8 데이터에서 병렬 연산을 극대화할 수 있습니다.

특히, 타일 2의 서명된 8비트 정수 쌍 4개를 인접 그룹으로 묶어 타일 3의 해당 서명된 8비트 정수와 곱셈을 수행하여 4개의 중간 32비트 결과를 생성합니다. 이 4개의 결과는 타일 1의 해당 32비트 정수에 더해지며, 최종 32비트 결과는 다시 타일 1에 저장됩니다.

타일의 바이트로 구성된 데이터를 소스/대상 누산기(accumulator)를 사용하여 점곱(dot-product)을 계산합니다. src0의 서명된 8비트 정수 쌍 4개를 인접 그룹으로 묶어 src1의 해당 서명된 8비트 정수와 곱셈을 수행하여 4개의 중간 32비트 결과를 생성합니다. 이 4개의 결과는 dst에 있는 해당 32비트 정수와 합산되며, 최종 32비트 결과는 다시 타일 dst에 저장됩니다.

타일의 형상(shape)은 __tile1024i 구조체에서 지정됩니다. 타일 레지스터는 컴파일러에 의해 할당됩니다.

DEFINE DPBD(c, x, y) {
	tmp1 := SignExtend32(x.byte[0]) * SignExtend32(y.byte[0])
	tmp2 := SignExtend32(x.byte[1]) * SignExtend32(y.byte[1])
	tmp3 := SignExtend32(x.byte[2]) * SignExtend32(y.byte[2])
	tmp4 := SignExtend32(x.byte[3]) * SignExtend32(y.byte[3])
	RETURN c + tmp1 + tmp2 + tmp3 + tmp4
}
FOR m := 0 TO dst.rows - 1
	tmp := dst.row[m]
	FOR k := 0 TO (src0.colsb / 4) - 1
		FOR n := 0 TO (dst.colsb / 4) - 1
			tmp.dword[n] := DPBD(tmp.dword[n], src0.row[m].dword[k], src1.row[k].dword[n])
		ENDFOR
	ENDFOR
	write_row_and_zero(dst, m, tmp, dst.colsb)
ENDFOR
zero_upper_rows(dst, dst.rows)
zero_tileconfig_start()

위 명령어는 타일의 바이트 데이터를 accumulator를 사용하여 dot-product 를 계산합니다. 구체적으로, 타일 2의 서명된 8비트 정수 쌍 4개를 인접 그룹으로 묶어 타일 3의 해당 서명된 8비트 정수와 곱셈을 수행하여 4개의 중간 32비트 결과를 생성합니다. 이 4개의 결과는 타일 1의 해당 32비트 정수에 더해지며, 최종 32비트 결과는 다시 타일 1에 저장됩니다.

AMX가 BF16에서 동작하는 원리

BF16 데이터에서 AMX는 float 연산에 최적화된 명령어를 사용하며, 작동 방식은 다음과 같습니다:

타일 구성: BF16 서브 매트릭스를 저장할 수 있도록 타일 크기를 설정합니다.
데이터 변환: 필요한 경우 FP32 데이터를 BF16으로 변환한 뒤 타일에 로드합니다.
행렬 곱셈 수행: _tile_dpbf16ps() 명령어를 통해 부동소수점 내적 연산을 수행합니다.
결과 저장: 최종 결과를 FP32 또는 BF16으로 변환해 저장합니다.

BF16 연산은 대규모 벡터 거리 계산이나 신경망 추론과 같은 작업에 특히 적합합니다.

BF16(16비트) 부동소수점 쌍을 src0와 src1 타일에서 dot-product으로 계산하고, 중간 결과인 단일 정밀도(32비트) 부동소수점 요소를 dst의 요소에 accumulate하여, 최종 32비트 결과를 다시 dst 타일에 저장합니다.

타일의 형상(shape)은 __tile1024i 구조체에서 지정되며, 타일의 레지스터는 컴파일러에 의해 할당됩니다.

FOR m := 0 TO dst.rows - 1
	tmp := dst.row[m]
	FOR k := 0 TO (src0.colsb / 4) - 1
		FOR n := 0 TO (dst.colsb / 4) - 1
			tmp.fp32[n] += FP32(src0.row[m].bf16[2*k+0]) * FP32(src1.row[k].bf16[2*n+0])
			tmp.fp32[n] += FP32(src0.row[m].bf16[2*k+1]) * FP32(src1.row[k].bf16[2*n+1])
		ENDFOR
	ENDFOR
	write_row_and_zero(dst, m, tmp, dst.colsb)
ENDFOR
zero_upper_rows(dst, dst.rows)
zero_tileconfig_start()

3세대 Intel Xeon Scalable 프로세서에서 AVX-512(Advanced Vector Extensions 512)에 비교하여, 4세대 Intel Xeon Scalable 프로세서에서 AMX의 성능이 훨씬 뛰어납니다.

AMX를 활용한 벡터 계산

Intel 컴파일러(icpx)를 사용해야 하는 이유

AMX를 효율적으로 사용하려면 **Intel® oneAPI DPC++/C++ Compiler (icpx)**를 사용해야 합니다. 이는 AMX 명령어 및 하드웨어 가속 기능을 제대로 지원합니다.

왜 icpx인가?

AMX 관련 명령어를 직접 지원.
AMX를 지원하는 Intel 프로세서에 최적화.
타일 구성 및 메모리 정렬 관리 도구 제공.

컴파일 예제:

icpx -march=native -o amx_program amx_vector_calculation.cpp

벡터 계산을 위한 AMX 데이터 정렬

벡터간 거리 계산을 cosine distance 를 이용한다고 할 때, inner product를 통해 그 대소를 비교할 수 있습니다. 이 때 AMX를 완전하게 활용하기 위해서는 데이터를 이에 맞춰 정렬해야합니다.

그 이유는 tile 안에 들어갈 수 있는 데이터가 한정적이고, AMX의 inner product 가 수행될 때 불필요한 연산을 줄여야 하기 때문입니다.

AMX는 INT8의 경우 4개씩, BF16의 경우 2개씩 set로 연산을 수행합니다. 그리고 A * B 의 경우 A는 row-major로 접근하고, B는 column-major 로 접근합니다.

따라서 BF16의 경우 다음과 같이 배열되어 있을 때 모든 연산이 필요한 연산이 됩니다.

v: query vector
d: data vector

matrix A

v1 v1 v1 v1 v1 v1 v1 v1
v2 v2 v2 v2 v2 v2 v2 v2
...

matrix B

d1 d1 d2 d2 d3 d3 d4 d4
d1 d1 ...
d1 d1
d1 d1

matrix result

v1*d1 v1*d2 v1*d3 v1*d4
v2*d1 ...
v3*d1
v4*d1

이 방법으로는 BF16 로 이루어진 32차원의 벡터를 쿼리벡터 16쌍, 그리고 데이터벡터 16쌍 을 한 instruction에 내적할 수 있습니다.