안녕하세요, ENERZAi에서 Nadya를 개발하고 있는 김재우입니다. 이번 게시물에서는 Nadya에 적용된 기반 기술에 대해 조금 더 자세히 알아보겠습니다. 구체적으로는, Nadya가 어떻게 Garbage collector 없이 안전하게 메모리를 관리하고 Tensor의 메모리 안정성을 보장하는지 말씀드리려고 합니다. 추가적으로, 저희가 어떤 이유로 Nadya의 소유권 시스템을 개발했는지, 해당 시스템을 통해 기대하는 바는 무엇인지도 다룰 예정입니다.

기존 메모리 관리 모델의 성능적 한계

프로그래밍 언어는 메모리 관리를 어떻게 수행하느냐에 따라 두 그룹으로 나눌 수 있습니다. C#, Java, Python, JavaScript 등 대부분의 언어들은 Managed 언어에 해당하는데요. Managed 언어에서는 Garbage collector가 메모리 할당 및 해제를 수행합니다. Garbage collector는 실행 중에 작동하며, 프로그램에서 더 이상 사용되지 않는 모든 자원을 수집합니다. 이는 컴파일러의 관점에서 매우 편리하지만, Garbage collector 성능에 의한 병목이 발생할 수 있다는 단점이 있습니다.

반면, Unmanaged 프로그래밍 언어는 프로그래머가 메모리를 직접 할당하고 해제해야 하며, 대표적인 Unmanaged 언어로는 C와 C++가 있습니다. Unmanaged 언어는 Garbage collector를 필요로 하지 않으며, 메모리를 안전하게 관리하는 것이 전적으로 프로그래머의 책임이기 때문에 프로그래머에게 더 많은 노력과 주의가 요구됩니다.

대부분의 프로그래밍 언어는 위에서 언급한 두 그룹 중 하나에 속하지만, Rust는 ‘소유권(ownership)’이라는 개념을 도입하여 Garbage collector에 의존하지 않고 안전하게 메모리를 관리하는데 성공했습니다. Rust에서는 컴파일러가 프로그램 전체에서 메모리가 어떻게 사용되는지, 어떤 값이 어떤 메모리를 소유하고 있는지, 그리고 언제 해제해야 하는지를 정적으로 분석하는데요.

Nadya의 소유권 시스템은 Rust에서 영감을 받아 개발되었습니다. 사실 소유권 시스템의 도입은 빠른 추론 속도와 메모리 안전성을 모두 확보하기 위한 유일한 방법이었는데요. 프로그램 성능 및 메모리 사용 측면에서 Inconsistency를 초래할 가능성이 높은 Garbage collector를 사용할 수는 없었고, 프로그래머에게 메모리 안전성에 대한 부담을 주고 싶지 않았기 때문에 소유권 시스템을 도입하기로 결정하였습니다.

Nadya 소유권 시스템

Nadya는 Rust와 Descend의 기능들을 발전시켜 CPU와 GPU 모두에 대해 더 유연하고 안전한 모델을 제공합니다. Nadya v0.3.0부터 Nadya는 Garbage collector를 사용하지 않고 데이터를 관리하기 위해 자체 소유권 시스템을 도입했습니다. 이 시스템은 Rust에서 사용되는 유사한 소유권 모델을 사용하며, “Binding”이 데이터를 소유할 수 있고, 데이터를 Binding에서 차용하여 접근할 수 있습니다.

더 나아가기 전에, Nadya 소유권 시스템을 이해하기 위해 몇 가지 중요한 개념들을 소개해 드리겠습니다.

{
 let value = 10 // "value" is binding, and "10" is data
 ...
 // Lifetime of data "10" must end before here, since its owner binding "value" cannot escape the scope
}

1. Data (예시의 “10”)

   • 데이터는 저장을 위해 메모리를 필요로 하는 모든 정보를 의미합니다. 이 메모리는 stack, heap, 또는 프로세서 레지스터에 위치할 수 있습니다.

   • 데이터는 복사되거나 다른 메모리로 이동할 수 있습니다. “복사”될 경우 원본 데이터는 유지되며, 데이터가 “이동”될 경우 원본 데이터는 무효화됩니다.

2. Binding (예시의 “value”)

   • Binding은 데이터를 저장할 수 있는 공간이며, Binding의 Lifetime은 Binding이 정의된 Scope가 끝날 때까지만 지속됩니다. Binding은 데이터를 ‘소유’할 수 있습니다.

   • 데이터가 특정 Binding에 묶이게 되면, 해당 데이터의 Lifetime은 도중에 데이터가 Binding 밖으로 꺼내지지 않는 한, Binding과 동일합니다.

3. Lifetime (컴파일러가 “data”에 부여)

   • Lifetime은 데이터가 프로그램 내에서 유효한 기간을 지정하는 속성입니다.

   • 데이터가 더 이상 사용되지 않으면, Lifetime은 즉시 끝납니다.

   • 특정 데이터의 Lifetime이 무효화되면, 데이터는 파괴됩니다.

   • Nadya 컴파일러는 모든 데이터의 Lifetime을 결정하고 데이터를 언제 파괴해야 하는지 결정하는 Lifetime 추론까지 지원합니다.

그리고 Nadya 소유권 시스템 내에서 수행할 수 있는 작업은 아래와 같습니다.

1. 복사 (Copy)

   • 복사는 원본 데이터의 정확한 복제본을 생성합니다.

   • 복사된 데이터는 원본 데이터와 연결되지 않으며, 자체 수명을 가집니다.

   let a = 10
   let b = a // This is Copy!
   print(a) // binding "a" is still alive (owns value)

   
   

2. 이동 (Move)

   • “이동”은 “데이터”를 한 위치에서 다른 위치로 옮깁니다.

   • 새로운 Lifetime을 생성하지 않습니다.

   let a = 10
   let b = move(a) // This is Move!
   print(1) // Error! binding "a" does not have any data (it has been moved out)

   
   

3. 차용 (Borrow)

   • “차용”은 이미 Binding에 묶인 다른 데이터로부터 Loan을 생성합니다.

   • 차용된 데이터는 Binding에 묶인 데이터의 Stack 주소입니다.

   • 차용된 데이터의 Lifetime은 차용한 Binding의 Lifetime을 초과할 수 없습니다.

   let a = 10
   let borrowA = &a
   print(a) // prints 10
   print(borrowA) // prints 10. Points to 'a'

원활한 이해를 위해 몇 가지 예시를 더 살펴보겠습니다.

module core::Vector

let owner = core::Vector::ones<i32>() // (1) Creates empty vector
let borrower = &owner // (2) Immutably borrows from "owner". May access 'owner' but may not modify it

// (3) Prints inner value (Note : borrower.get(0) returns borrow type)
print(borrower.get(0))
// Note : lifetime of data stored in vector is same as vector instance itself

// Lifetime of borrower ends here since it does not have further use.

let mutableBorrow = &mut owner // Mutably borrows from "owner". May access 'owner'

1. owner Binding이 새롭게 생성된 벡터 데이터를 소유하게 됩니다.

2. borrower가 owner에 묶인 데이터를 영구적으로 차용합니다. 이로 인해 borrower는 owner의 차용된 데이터에 대한 참조를 보유하게 되며, 차용된 데이터는 곧 owner의 Stack 주소입니다.

   • borrower의 Lifetime은 owner에 Binding됩니다.

3. get 메서드는 데이터의 첫 번째 요소를 차용합니다.

   • 반환된 데이터의 Lifetime은 owner의 Lifetime보다 짧아야 합니다. 이 경우, 반환된 데이터의 Lifetime은 print 함수가 실행된 직후에 끝납니다.

Tensor에 대한 소유권 시스템

Rust에서 영감을 받은 Nadya의 소유권 시스템은 특히 Tensor를 다룰 때 그 진가를 발휘합니다. Nadya에서 Tensor는 컴파일러가 직접 최적화할 수 있는 원시 데이터 유형으로, Nadya의 성능을 높이는 데 중요한 역할을 합니다. Nadya는 **범위 분석(range analysis)**을 사용하여 데이터를 부분적으로 차용(partial borrow)할 수 있으며, 이 때 부분 차용된 Tensor는 Subview라고 합니다.

Subview를 사용하면 프로그래머는 Tensor를 다른 Tensor처럼 사용할 수 있습니다. 구체적으로는 인덱스 변환 작업을 통해 주어진 인덱스를 원래 Tensor의 인덱스로 변환하여 원본 데이터에 접근하는 것이 가능합니다.

Subview와 소유권의 결합은 Tensor 데이터에 대한 안전을 보장합니다. 소유된 Tensor의 일부 데이터는 다른 Binding에 의해 차용될 수 있으며, 이로 인해 가변성 및 차용 규칙이 위반되지 않도록 보장하여 병렬 프로그래밍에서 데이터 레이스(Data race) 및 동기화 장벽(Synchronization barrier)로 인한 문제를 최소화할 수 있습니다. Kopcke et al.에 의하면, 일부 연구자들은 여기에 착안하여 다양한 GPU 실행 리소스 간의 데이터 레이스를 방지하는데 성공했습니다. 구체적으로는 Rust의 Borrow checker를 GPU 프로그래밍에 확장 적용하여 안전성을 보장하고 동기화 작업을 효율화했는데요. 해당 모델에서는 데이터가 특정 실행 리소스에 의해 부분적으로 차용 및 소유될 수 있어, 서로 다른 실행 리소스 간의 데이터 레이스를 쉽게 감지할 수 있었던 것으로 보입니다.

Nadya에서는 이 개념을 일반 프로그래밍으로 확장하였습니다. Rust에서 도입된 개념과 Kopcke et al.에서 제안된 아이디어를 결합하여 안전성을 보장하면서 해당 연구를 General한 목적의 프로그래밍에도 적용할 수 있었는데요.

아래의 간단한 Nadya 코드를 통해 좀 더 구체적으로 설명드리겠습니다.

let myTensor = tensor((2, 3), 0) // (1) Initialize zero tensor of shape (2, 3)
// (2) Borrow whole range of the tensor
let borrowedTensorSimple = &myTensor
// (3) Borrow tensor with range (0:2:1, 1:3:1)
// Can be thought of tensor with shape (2, 2)
let borrowedTensorWitRange = &myTensor[(0:2:1, 1:3:1)]

위의 코드에서 Tensor는 (1)에서 생성되고, (2)와 (3)에서 차용됩니다. (2)에서는 Tensor의 모든 요소를 차용하고, (3)에서는 Tensor의 일부만 차용합니다. (2)와 (3) 모두 Tensor Subview를 생성하며, 이는 원본 Tensor의 Subset으로 인덱싱할 때 제약된 Shape의 Tensor로 간주됩니다.

(3)에서는 Tensor subview가 (0:2:1, 1:3:1)으로 인덱싱되므로 (2,2) 형태의 Tensor로 취급할 수 있습니다. (참고: a:b:c는 “a부터 b-1까지, c의 간격으로”를 의미합니다.)

Nadya 컴파일러는 Subview Tensor와 원본 Tensor 간의 Dependency를 정적으로 분석하려고 시도합니다. 이 과정이 성공하면, Nadya 컴파일러는 Subview 인덱싱 표현식을 원본 Tensor 인덱싱으로 변환하여 사전 최적화(Pre-optimize)할 수 있습니다. 만약 Nadya 컴파일러가 불확실한 실행 경로 등에 의해 정적 분석에 실패하면, 인덱싱 Logic은 런타임 단계에서 동적으로 계산됩니다.

위 코드의 경우, 컴파일러 차원에서 차용된 Tensor에서 소유한 Tensor까지의 전체 Use-def chain을 정적으로 확인할 수 있으므로, 컴파일러는 차용된 Tensor에 대한 접근을 최적화하여 원본 Tensor에 직접 접근할 수 있습니다.

하지만, 정적 분석 단계에서 차용된 Tensor의 소스를 확인하기 어려운 경우도 있습니다. 하나의 예시는 차용된 텐서를 함수에 전달하는 Case 입니다.

let myTensor = tensor((2, 3), 0)
let borrowedTensor = &myTensor[(0:2:1, 1:3:1)]

myFunc(borrowedTensor)


fun myFunc(input : &tensor<i32, 2>){
 // How can we find the location in original tensor outside the function?
 let borrowedTensor[(1,)] 
 ...
}

위 코드에서 borrowedTensor는 함수의 Argument로 전달됩니다. 이 때 컴파일러는 myFunc가 어디에 정의되어 있는지 파악할 수 없으며, myFunc는 다른 모듈이나 Translation unit에 위치할 수 있습니다. 따라서, 차용된 Tensor에서 원본 myTensor로 인덱스를 Resolve하는 과정은 동적으로 처리해야 합니다.

위와 같은 이유로, Nadya는 내부적으로 차용된 Tensor를 아래과 같이 표현합니다:

template<typename T, std::size_t rank>
struct Tensor{
 T* Data; // Pointer to the data where tensor is stored. Null if this tensor is borrowed
 int64_t Shape[rank]; // Shape of the tensor
 int64_t Stride[rank]; // Stride of the tensor
 int64_t Offset[rank]; // Index offsets
 Tensor<T>* ParantTensor; // Pointer to the stack address of parent tensor. Null if tensor is not borrow of other tensor
}

위와 같은 구조를 통해 Borrow checker는 차용된 Tensor가 사용되는 동안 원본 Tensor가 생존해 있음을 보장합니다. 따라서, Tensor는 프로그램의 어디에서든 사용될 수 있으며, 심지어 다른 Translation unit에서도 안전성을 저하시키지 않고 사용할 수 있습니다. 다만, 인덱스 Offset과 간접 참조를 동적으로 처리하는 과정에서 성능이 저하될 수 있기 때문에 Nadya 컴파일러는 가능한 인덱싱 규칙을 정적으로 분석하려고 합니다.

Nadya 소유권 시스템 개발 배경

Nadya의 소유권 시스템은 프로그래밍 경험을 단순하고 간결하게 유지하면서도 안전성을 보장하는 데 중요합니다. 게다가, 이 시스템은 컴파일러가 컴파일 과정에 대해 안전하게 추론할 수 있도록 도와줍니다.

Nadya의 프로그래밍 언어 및 컴파일러 산업에 대한 주요 기여 중 하나는 도메인 특화 언어에서만 가능했던 최적화를 일반 목적의 프로그래밍 언어에 통합한 것입니다. 이러한 통합은 언어의 의미 체계를 정의함으로써 이루어졌으며, 이 체계는 추론하기 쉽습니다.

Nadya가 Safe assumption에 집중하는 이유

최근 AI와 머신러닝 기술의 발전으로 인해, 컴파일러 엔지니어들은 AI 및 HPC 워크로드를 가속화하는 데 주력하고 있습니다. 예를 들어, Polygeist(William et al)는 C 코드를 MLIR로 파싱하고, MLIR을 SCoP 구조로 변환하며, CLoog 및 Pluto(Bondhugula et al)와 같은 기존 도구를 활용한 바 있는데요. 이는 Polyhedral 최적화를 통한 성능 향상 효과를 입증하고 있지만, 그 범위는 단순한 Control flow과 명확한 Load-store 의존성을 가진 Loop에 제한됩니다. 또한, 이러한 컴파일러에는 Higher-level 정보가 부재하기 때문에 Loop가 다른 형태로 변환될 수 있는지 여부를 판단할 수는 있지만, Loop의 목적을 이해하기는 어렵습니다. 이러한 문제를 해결하기 위해, Lower-level IR을 XLA 또는 MLIR Linalg dialect에서 사용하는 High-level abstraction으로 변환하려는 시도가 있었는데요. 예를 들어, mlirSynth(Brauckmann et al.)는 Lower-level MLIR dialect를 XLA 또는 MLIR Linalg dialect pass를 통해 최적화할 수 있는 Higher-level dialect로 변환시키는 코드 생성기를 도입했습니다. 이러한 접근법은 LLVM IR로 직접 Lowering하거나 Polyhedral 프레임워크를 사용할 때보다 더 나은 성능을 보여주었습니다.

이전 연구에서 확인할 수 있듯이, 고도로 최적화된 코드를 생성하기 위해서는 High-level context 정보를 확보하는 것이 중요합니다. 컴파일러가 High-level에서도 Safe assumption을 수행할 수 있을 때, 더 효과적으로 최적화를 수행하는 것이 가능합니다. 이러한 이유로 인해 XLA나 MLIR-Linalg pass와 같은 Domain-specific 컴파일러는 LLVM-opt나 특정 Polyhedral 컴파일러 등의 Lower-level 컴파일러보다 더 뛰어난 최적화 성능을 보유하고 있습니다. Nadya는 오류를 최소화하고 컴파일러가 Safe assumption을 수행할 수 있도록 지원하여 프로그래머가 더 안전하고 효과적으로 목표를 달성할 수 있도록 돕는 것을 목표로 하고 있습니다.