텐서(Tensor)의 위력

모든 훌륭한 추상화(Abstraction)는 문제를 해결하며, 이 포스트에서는 제가 지금까지 알고 있는 놀라운 수학적 추상화인 텐서(Tensor)에 관한 모든 것을 다룰 것입니다.

단순한 2층 MLP(Multi-Layer Perceptron)부터 GPT-5에 이르기까지, 모든 신경망(Neural networks)은 결국 동일한 것으로 귀결됩니다. 즉, 연산 그래프(Graph of operations)를 통해 흐르는 부동 소수점(Floating-point numbers)입니다. 이 포스트에서는 C 언어를 사용하여 처음부터 완전하고 가속화된 텐서 라이브러리를 구축합니다. 이 작업은 불행히도 아직 오픈 소스로 공개되지 않은 Bellard의 libnc에서 큰 영감을 받았습니다.

텐서(Tensor)는 숫자의 평탄한 배열(Flat array)일 뿐이며, 여기에 해당 숫자들을 다차원 객체로 어떻게 해석할지를 알려주는 메타데이터(Metadata)가 추가된 것입니다. 우리 모두는 2D 배열이 1D 배열과 행/열의 개수로 더 잘 표현될 수 있다는 것을 배웠습니다. 이것이 본질적으로 텐서(Tensor)가 하는 일입니다.

하지만 2차원을 넘어선다면, 일반화된 형상(Shape)과 같은 다른 메타데이터가 필요할 수 있습니다:

float data[32 * 3 * 28 * 28]; // 32개의 이미지, 3개의 채널, 28x28 픽셀
int shape[4] = {32, 3, 28, 28}; // 텐서의 형상 (shape)
int ndim = 4; // 차원의 수 (number of dimensions)

형상(Shape)이 있다면, 만약 우리가 텐서를 행 우선(Row-major) C-order 형식(오늘날 대부분의 텐서 라이브러리에서 기본값인 경우가 많음)으로 유지할 경우, 4D 텐서 내 위치 [n,c,h,w]에 있는 요소가 data[n*(3*28*28)+c*(28*28)+h*28+w] 오프셋(Offset)에 존재한다는 것을 파악할 수 있습니다.

이제 요소의 인덱스를 매번 이렇게 계산하는 것은 비효율적이므로, 형상(Shape)이 결정되면 *스트라이드(Strides)*를 한 번에 미리 계산할 수 있습니다. 스트라이드(Strides)는 주어진 차원에서 한 요소를 전진시키기 위해 몇 개의 요소를 건너뛰어야 하는지를 알려줍니다. 또한 형상(Shape)과 관련된 모든 필드를 하나로 그룹화할 수도 있습니다:

struct ut_shape {
int ndim; // 차원의 수 (number of dimensions)
int nelem; // 요소의 수 (number of elements)
...

형상(Shape)을 생성하고 다차원 인덱스로부터 평탄 인덱스(Flat index)를 얻기 위한 몇 가지 헬퍼(Helpers)를 추가할 수 있습니다:

ut_shape ut_shape_new(int ndims, int* dims) {
ut_shape s = {.ndims = ndims, .nelems = 1};
for (int i = 0; i < ndim; i++) {
...

텐서(Tensor)는 보통 동적으로 할당되므로, malloc/free 위에 래퍼(wrapper)를 제공하여 생성하고 파괴하는 방법을 마련해야 합니다.

때로는 다른 텐서와 동일한 데이터를 공유하는 텐서를 만들고 싶을 때가 있습니다. 예를 들어, 텐서의 일부에 대한 “뷰(view)”를 얻거나, 데이터를 복사하지 않고 텐서를 전치(transpose)하거나, 모양(shape)을 수정할 때(평탄화/flatten). 이런 경우 데이터 소유권(ownership)을 추적해야 합니다. 또한, 반복적인 학습 과정에서처럼 원래의 범위를 벗어나도 살아남아야 하는 텐서들을 “보존(retain)”하는 것도 유용합니다. 이를 위해 텐서 구조체에 참조 카운터 필드를 추가합니다:

struct ut_tensor {
struct ut_shape shape; // 텐서의 모양 (shape)
float *data; // 데이터 포인터
...
};

메모리 관리를 더욱 최적화하여 빈번한 malloc/free 호출을 피하기 위해 아레나 할당자(arena allocator) 또는 메모리 풀(memory pool)을 추가할 수 있지만, 현재 가진 것도 좋은 출발점입니다.

하지만 텐서만으로는 연산 없이는 거의 쓸모가 없습니다.

요소별 연산 (Elementwise)

텐서에 대한 가장 기본적인 연산은 요소별 연산(elementwise)입니다. 이는 모든 요소를 순회하며 각 요소에 함수를 적용하는 것(단항/unary) 또는 두 요소 쌍에 함수를 적용하는 것입니다. 우리는 이를 다음과 같이 구현할 수 있습니다:

static void ew_neg(float* out, const float* a, int n) {
for (int i = 0; i < n; i++) out[i] = -a[i];
}
...

이진 연산(binary operation)을 수행하기 전에 두 텐서의 모양이 같은지 assert() 하는 것이 좋을 것입니다. 또는, 이 시점에서 “브로드캐스팅(broadcasting)”을 지원하고 싶다고 결정할 수도 있습니다. 이는 더 작은 텐서를 더 큰 텐서의 모양에 맞게 확장하는 것을 의미합니다. 예를 들어, 모양이 [2, 3, 4]인 텐서와 모양이 [3, 4]인 텐서가 있습니다.

두 번째 텐서를 첫 번째 차원을 따라 "늘림(stretching)"으로써 함께 더할 수 있습니다. 이는 많은 텐서 라이브러리에서 흔히 볼 수 있는 기능이지만, 지금은 제외하기로 했습니다. 제가 목표로 하는 대부분의 모델은 기반이 되는 텐서들이 완벽하게 정렬되어 있습니다. 만약 정렬되어 있지 않다면, 텐서 모양을 명시적으로 맞추기 위해 데이터를 복제하거나, 나중에 브로드캐스팅 (broadcasting) 인덱스 계산을 추가할 수 있습니다.

이 시점에서 행렬 곱셈 (matrix multiplication), 컨볼루션 (convolutions) 등 더 많은 연산을 구현하고 싶은 유혹이 들 수도 있습니다. 하지만 바로 이 지점에서 우리는 질문을 던져야 합니다. 과연 모든 연산을 CPU에서 실행하기를 원하는가?

골든 프로세싱 유닛 (GPU)

현대적인 GPU 가격을 살펴보면, 숫자 배열을 처리하는 데 있어 GPU가 아마도 상당한 가치를 지니고 있다는 점이 명확합니다. 따라서 우리의 라이브러리를 CPU로만 제한하는 대신, 일부 연산을 GPU로 오프로딩 (offloading) 하는 것을 고려해 볼 수 있습니다.

여기서부터 상황이 흥미로워집니다. 이제 CPU와 GPU 모두에서 메모리를 관리하고, 데이터를 효율적으로 전송하며, GPU 커널 (kernel)을 작성하는 법을 배우는 등 훨씬 더 많은 것을 해야 하기 때문입니다. 설상가상으로, 단일한 "GPU 가속기"라는 것은 존재하지 않습니다. CUDA, OpenCL, WebGPU, Metal, Vulkan 등 각기 다른 API와 특성을 가진 수많은 벤더(vendor)들이 존재합니다.

저는 CPU 전용 텐서의 성능을 최대한 끌어올리기 위해 BLAS, LAPACK, OpenMP 등을 사용하여 여러 차례 시도해 보았지만, GPU 가속 프레임워크와 같은 수준의 성능 규모에는 도달할 수 없었습니다.

상황을 관리 가능한 수준으로 유지하기 위해, 저는 Metal로 시작하기로 결정했습니다. 한편으로는 Apple 기기로 제한된다는 단점이 있지만, 다른 한편으로는 상당히 현대적인 API이며 Apple Silicon의 "통합 메모리 (unified memory)" 덕분에 상황을 단순화해 줍니다.

Metal은 GPU 커널을 정의하기 위해 현대적인 C++와 유사한 자체 언어인 MSL (Metal Shading Language)을 사용하며, 런타임 (runtime)에 커널을 컴파일하므로 우리의 C 코드에서 문자열로 커널을 정의할 수 있습니다:

static const char *shaderSource =
"#include <metal_stdlib>\n"
"using namespace metal;\n"
...

이 커널(kernel)을 GPU에서 실행하려면 디바이스(device)를 생성하고, 커맨드 큐(command queue)를 구축하며(GPU는 비동기적입니다), 커널을 컴파일하고, 입력 및 출력을 위한 버퍼(buffer)를 생성한 다음, 마지막으로 커널을 GPU에서 실행하도록 디스패치(dispatch)해야 합니다. 간단한 GPU 가속 ReLU 연산은 다음과 같은 모습일 것입니다:

// 디바이스, 큐 준비 및 커널 컴파일
id<MTLDevice> device = MTLCreateSystemDefaultDevice();
id<MTLCommandQueue> queue = [device newCommandQueue];
...

우리는 이 코드의 첫 부분을 "Metal 컨텍스트 (Metal context)" 구조체로 감쌀 수 있으며, 이는 한 번 생성되어 라이브러리 전체의 수명 동안 유지됩니다. 그런 다음 버퍼를 할당하고, 읽고 쓰고, 커널을 디스패치하는 헬퍼(helper) 함수들을 추가할 수 있습니다. 공식 Metal API는 Objective-C 또는 Swift로 되어 있지만, 라이브러리를 순수하게 유지하기 위해(코드 전반에 걸쳐 objC_msgSend 사용) 저수준 ObjC 런타임 함수를 사용하여 C에서 이를 호출할 수 있습니다.

이제 텐서(tensor)에 대한 모든 연산은 두 번 구현되어야 합니다. 즉, 단순한(naïve) CPU 구현과 GPU 커널로서의 구현입니다. 텐서의 출처에 따라 두 구현 중 하나를 호출하게 됩니다. 이는 또한 텐서가 "디바이스 (device)" 필드를 보유해야 함을 의미하며, 텐서가 할당될 때 디바이스를 지정해야 한다는 것을 의미합니다. CPU 텐서는 일반 데이터 버퍼만 가지지만, Metal 텐서는 Metal 버퍼 객체에 대한 포인터도 가집니다.

또한 GPU와 CPU 버퍼가 서로 달라지는 시점을 추적하는 것도 중요해지고 있습니다. 쉬운 방법은 CPU/GPU 버퍼에 "더티 (dirty)" 플래그를 유지하고, 한 버퍼에서 다른 버퍼로 데이터를 복사하는 "싱크 (sync)" 함수를 제공하는 것입니다.

일부 연산은 GPU에서 수행할 수 있거나 수행해야 하지만, 어떤 경우에는 상황을 단순화하기 위해 먼저 CPU에서 생성한 다음, 처리를 수행하고 데이터를 다시 GPU로 전송/복사하기도 합니다.

결과적으로, 우리는 다음 사항들을 구현해야 합니다:

static const char *_mtl_src = "..."; // 모든 커널(kernel)을 위한 MSL 소스 코드
void *ut_mtl_init(); // 싱글톤(singleton)으로서 Metal 컨텍스트(context) 생성
void ut_mtl_dispatch(void *ctx, char *kernel, void **bufs, int nbufs, void *bytes, int blen, int n); // 버퍼(buffer)와 함께 커널(kernel) 디스패치(dispatch)
...

원소별 연산(element-wise ops)과 그에 상응하는 GPU 커널(kernel) 구현을 마쳤으므로, 이제 모든 신경망(neural networks)의 핵심인 행렬 곱셈(matrix multiplication)으로 넘어갈 수 있습니다.

행렬 곱셈 (Matmul)

즉석에서 지어낸 통계에 따르면, 어떤 신경망에서든 FLOPs(부동 소수점 연산 횟수)의 80%는 행렬 곱셈(matrix multiplications)이 차지한다고 합니다. 이를 정확하고 빠르게 수행하는 것은 모든 딥러닝(deep learning) 프레임워크에 있어 매우 중요합니다. 행렬 곱셈을 구현하는 방법은 많지만, 가장 효율적인 방법은 행렬의 크기, 하드웨어, 그리고 메모리 레이아웃(memory layout)에 따라 달라집니다.

먼저, 가장 단순한 경우인 두 개의 2D 행렬 A와 B를 곱하는 경우를 생각해 보겠습니다. 여기서 A의 형상(shape)은 [M, K]이고,

B의 형상은 [K, N]입니다.

결과물인 C의 형상은 [M, N]이 됩니다.

가장 단순한(naïve) 구현은 다음과 같을 수 있습니다:

void gemm(const float *A, const float *B, float *C,
int m, int n, int k, bool ta, bool tb) {
for (int i = 0; i < m; i++)
...

말할 필요도 없이 이는 매우 느릴 것입니다. 세 개의 중첩된 루프(nested loops)가 있고 캐시 최적화(cache optimisation)가 되어 있지 않기 때문입니다. Apple Silicon에서는 Accelerate 프레임워크의 cblas_sgemm을 사용하여 이를 직접 대체할 수 있습니다. 이는 AMX(행렬 보조 프로세서, matrix co-processor)를 사용하며, 32x32 이상의 모든 행렬에 대해 50~100배 더 빠릅니다:

cblas_sgemm(CblasRowMajor,
ta ? CblasTrans : CblasNoTrans,
tb ? CblasTrans : CblasNoTrans,
...

하지만 우리는 더 나아가 큰 행렬에 대해 GPU 행렬 곱셈(matmul)을 사용하는 Metal을 사용할 수 있습니다. Metal에는 MPSMatrix 클래스와 곱셈을 수행하기 위한 몇 가지 메서드(methods)가 포함되어 있습니다. 하지만 이는 2D 행렬에만 해당하므로, 2D 이상의 텐서(tensor)를 처리하려면 직접 커널(kernel)을 구현해야 합니다.

kernel void bmatmul(device
const float * a, device
const float * b_, device float * c,
...

이제 전체적인 ut_matmul 함수는 다음과 같은 모습일 것입니다:

ut_tensor *ut_matmul(ut_tensor *a, ut_tensor *b) {
// special case: 2D x 2D
if (a->shape.ndim == 2 && b->shape.ndim == 2) {
...

한 가지 빠진 것은 전치 행렬 (transposed matrices)을 위한 행렬 곱셈 (matmul)입니다. 앞서 우리는 데이터를 복사하지 않고도 (스트라이드 (strides)만 업데이트하여) 행렬을 "전치"하는 방법을 찾아냈지만, 곱셈을 위해서도 이를 특별한 방식으로 처리해야 합니다. 다행히 모든 구현체는 약간의 코드 변경만으로 전치를 지원하는 것으로 보입니다 (cblas_sgemm에서의 ta 및 tb 플래그, 그리고 Metal 커널에서의 몇 줄의 추가 코드).

수동 Autograd (Autograd by hand)

이 시점에서 우리는 텐서 (tensor)에 대한 산술 연산을 수행할 수 있으며, 새로운 연산을 추가하는 것도 대부분 간단해집니다. 이는 신경망 (neural networks)을 구축하기 위한 좋은 토대가 되지만, 여전히 한 가지 결정적인 기능이 빠져 있습니다. 바로 적절한 미분 (differentiation)입니다.

네트워크가 "학습"할 때, 네트워크는 손실 함수 (loss function)를 최소화하도록 매개변수 (parameters)를 조정합니다. 이를 위해 매개변수에 대한 손실 함수의 기울기 (gradients)를 계산해야 합니다. 여기서 보통 자동 미분 (automatic differentiation)이 등장합니다.

PyTorch나 Tensorflow와 같은 "거대" 프레임워크에서는 미분 엔진 (differentiation engine)이 순전파 (forward pass) 동안 계산 그래프 (computation graph)를 기록하고, 그 후 이를 역방향으로 따라가며 기울기를 계산합니다. 저는 더 단순한 방식인 수동 역전파 (manual backpropagation)를 구현해 보려고 합니다. 네트워크의 모든 레이어 (layer)는 "forward" 및 "backward" 함수를 가진 구조체 (struct)가 되며, 기울기는 각 학습 가능한 매개변수와 동일한 형상 (shape)을 가진 미리 할당된 텐서가 됩니다.

순전파 (forward pass)는 입력 텐서를 받아 출력 텐서를 생성합니다. 역전파 (backward pass)는 출력 텐서에 대한 손실의 기울기를 받아, 입력 텐서 및 매개변수에 대한 손실의 기울기를 계산합니다. Autograd만큼 사용자 친화적이지는 않겠지만, 대부분의 잘 알려진 네트워크 아키텍처 (architectures)의 경우 C에서 구현하기가 더 쉽습니다.

Linear 레이어부터 시작해 봅시다. 이는 NxM 크기의 가중치 행렬 (weights matrix, 여기서 N은 입력 특징 (input features)의 수이고 M은 출력 특징 (output features)의 수임)과 선택 사항인 편향 벡터 (bias vector)로 구성됩니다:

typedef struct ut_linear {
ut_tensor* weight; // [in, out]
ut_tensor* bias; // [out]
...

순전파 (Forward pass)는 매우 간단합니다. 행렬 곱셈 (matrix multiplication)을 수행한 뒤 편향 벡터 (bias vector)를 더하는 과정입니다 (나중에 속도를 조금 더 높이기 위해 이 두 연산을 하나로 합치는 (fusing) 것을 고려할 수도 있습니다):

ut_tensor *ut_linear_forward(ut_linear *l, ut_tensor *x) {
ut_tensor *out = ut_matmul(x, l->weight); // [B, in] @ [in, out]
if (l->bias)
...

실제 구현에서는 metal/non-metal 분기 때문에 순전파 (forward pass) 코드가 더 장황해지겠지만, 로직 자체는 단순하게 유지됩니다. 역전파 (backward pass)는 입력 (input) 및 파라미터 (parameters)에 대한 그래디언트 (gradients)를 계산해야 하므로 조금 더 어렵습니다.

행렬 곱셈 (matmul) 역전파 (backpropagation)에 대한 훌륭한 설명은 이 기사에서 찾아볼 수 있습니다. 수학적으로 부족한 저의 뇌에는 PyTorch가 이를 어떻게 구현했는지 읽는 것이 더 쉬웠습니다:

# forward(x: Tensor, w: Tensor): Tensor
y = x@w
# backward(dldy: Tensor): Tensor
...

만약 편향 (bias)을 포함한다면, y=x@w+b가 되며

dldb는

단순히 dldy입니다. 편향 (bias)은 덧셈 연산이므로, 그래디언트 (gradient)는 그대로 유지됩니다. 미니 배치 (mini-batch) 학습의 경우, 배치 차원 (batch dimension)을 따라 출력 그래디언트 (output gradients)를 합산해야 하며, 그것이 우리의 편향 그래디언트 (bias gradient)가 됩니다.

ut_tensor *ut_linear_backward(ut_linear *l,
ut_tensor *x,
ut_tensor *grad_out,
...