Transform 방언 튜토리얼에 앞서, MLIR의 Linalg 방언에서 구조화된 연산 개념과 구현을 개관한다. 벡터의 원소별 연산 확장, 리덕션과 컨트랙션, 메모리/텐서 상의 linalg.generic, 타일링과 루프 구체화, 생산자/소비자 퓨전과 재물질화, 그리고 명명된 Linalg 연산의 축약형을 예제 코드와 함께 설명한다.
Transform 방언 튜토리얼을 시작하기에 앞서, 구조화(Structured) 연산 개념과 그것의 Linalg 방언 내 구현을 간단히 살펴보자. 참고로 Transform 방언은 구조화 연산을 요구하지도, 그 반대도 아니다. 두 기술은 Transform 방언 초기 단계에 함께 발전했으며, 그 결과 구조화 연산을 대상으로 하는 변환들이 가장 성숙하고 튜토리얼에 가장 적합하다. 이미 이 개념에 익숙하다면 챕터 1로 건너뛰어도 좋다.
구조화된 코드 생성은 특정 변환들을 지원하기 위해 필요한 만큼, 그리고 가능한 한 오래 계산의 구조를 보존하려는 의도를 가진다. 이는 특정 변환을 지원하는 IR 추상화 설계까지 포함한다.
MLIR에서의 간단한 스칼라 부동소수점 덧셈 연산을 생각해 보자. 이는 대부분의 아키텍처에서 기계 명령어로 직접 매핑된다:
%2 = arith.addf %0, %1 : f32
이 연산은 1차원 벡터의 각 원소에 균일하게 적용되도록 손쉽게 확장할 수 있으며, 이는 벡터 머신의 명령어로도 자주 제공된다:
%2 = arith.addf %0, %1 : vector<8xf32>
현대의 소수 명령어 집합만이 2차원 이상의 벡터에 대한 명령어를 제공한다. 그러나 MLIR에서는 균일 원소별 적용을 임의 차수(rank)의 벡터로 투명하게 확장할 수 있다.
%2 = arith.addf %0, %1 : vector<8x4xf32>
%5 = arith.addf %3, %4 : vector<2x2x2x2x2x2x2xf32>
보듯이 MLIR의 벡터에 대한 산술 연산은 균일한 원소별 적용의 구조를 보존한다. 이 구조는 예를 들어 타깃에서 사용 가능한 더 낮은 차수의 연산으로 분해하거나, 곱셈과 덧셈을 하나로 융합하는 명령어가 있을 때 이를 활용하는 등의 방식으로 컴파일러가 이용할 수 있다(곱셈 100개 뒤에 덧셈 100개가 따르는 경우에는 융합이 훨씬 복잡해진다).
때로는 벡터의 원소들을 더해 스칼라를 얻어야 한다. 어떤 플랫폼은 이 연산을 위한 전용 명령어를 제공하고, 다른 플랫폼은 인접 원소의 덧셈과 셔플을 조합해 원하는 효과를 낼 수 있는 명령어들을 제공한다.
MLIR의 Vector 방언은 벡터 내부 리덕션을 명시적으로 나타내는 연산을 정의한다:
%1 = vector.reduction <add>, %0 : vector<8xf32> into f32
지원이 없을 때는 이러한 연산을 루프로 변환할 수 있다:
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c8 = arith.constant 8 : index
%init = arith.constant 0.0 : f32
%result = scf.for %i = %c0 to %c8 step %c1 iter_args(%partial = %init) -> (f32) {
%element = vector.extract %0[%i] : f32 into vector<8xf32>
%updated = arith.addf %partial, %element : f32
scf.yield %updated : f32
}
특수 명령어가 있더라도, 명령어 지연(latency)과 레지스터 압력에 따라 (언롤링된) 루프 형태를 사용하는 것이 바람직할 수 있다. 연산의 구조를 단일 리덕션으로 보존하면, 컴파일러는 벡터 내부 리덕션이 수행된다는 사실을 이해하여 구현 선택의 폭을 갖게 된다.
컨트랙션은 두 벡터의 원소를 곱한 뒤 합치는(더하는) 연산으로, 리덕션을 일반화한 것이다. 단순한 "더하기" 리덕션은 한 벡터의 원소가 곱셈의 항등원인 1.0으로 채워진 컨트랙션으로 생각할 수 있다. 컨트랙션은 컴파일러에 더 큰 유연성을 제공하며, MLIR에서는 전용 연산으로 표현된다:
// 덧셈을 위한 중립 초기값.
%init = arith.constant 0.0 : f32
// 곱셈의 항등원.
%ones = arith.constant dense<1.0> : vector<8xf32>
// 실제 컨트랙션.
%result = vector.contract {
indexing_maps = [affine_map<(i) -> (i)>,
affine_map<(i) -> (i)>,
affine_map<(i) -> ()>],
iterator_types = ["reduction"]
} %0, %ones, %init : vector<8xf32>, vector<8xf32> into f32
벡터 원소가 어떻게 인덱싱되는지를 나타내는 affine_map 표현에 주목하자. 이를 루프 형태의 의사코드로 쓰면 의미가 가장 분명해진다. 위 컨트랙션과 동등한 의사코드는 다음과 같다:
for i in 0 to 8:
init += p0[i] * ones[i]
여기서 %0과 %ones는 각각 해당 어파인 맵의 우변 (i) -> (i)가 가리키듯 루프 유도 변수 i를 사용하고, %init은 해당 어파인 맵의 우변 (i) -> ()에 반영되듯 사용하지 않는다.
균일 원소별 확장의 경우와 마찬가지로, MLIR의 벡터 컨트랙션은 1차원에 국한되지 않는다. 2차원 이상에서는 어떤 차원이 리덕션 대상이고 어떤 차원이 보존되는지 추가로 지정할 수 있다. 이는 각 차원이 리덕션 대상("reduction")인지 보존("parallel")되는지를 명시하는 iterator_types 속성으로 가능하다. 다음 3차원 컨트랙션은 행렬-행렬 곱셈을 인코딩한다:
%result = vector.contract {
indexing_maps = [affine_map<(i, j, k) -> (i, k)>,
affine_map<(i, j, k) -> (k, j)>,
affine_map<(i, j, k) -> (i, j)>],
iterator_types = ["parallel", "parallel", "reduction"]
} %lhs, %rhs, %init: vector<8x10xf32>, vector<10x16xf32> into vector<8x16xf32>
인덱싱 맵을 보면 루프 형태를 쉽게 떠올릴 수 있다:
for i in 0 to 8:
for j in 0 to 16:
for k in 0 to 10:
init[i, j] += lhs[i, k] * rhs[k, j]
이처럼 컨트랙션의 상위 수준 구조를 보존하면, 컴파일러가 행렬 곱셈과 점곱 같은 연산을 식별하기 훨씬 쉬워지고, 최신 명령어나 심지어 미리 생성된 마이크로커널을 활용하는 하위 수준 연산을 자유롭게 생성할 수 있게 된다.
지금까지는 가상 레지스터에 저장된 벡터에 대한 연산을 다뤘다. 메모리에서도 유사한 컨트랙션 추상화를 정의할 수 있다:
linalg.generic {
indexing_maps = [affine_map<(i, j, k) -> (i, k)>,
affine_map<(i, j, k) -> (k, j)>,
affine_map<(i, j, k) -> (i, j)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%lhs, %rhs : memref<8x10xf32>, memref<10x16xf32>)
outs(%init : memref<8x16xf32>) {
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32):
%0 = arith.mulf %lhs_one, %rhs_one : f32
%1 = arith.addf %init_one, %0 : f32
linalg.yield %1 : f32
}
좀 더 복잡해 보이므로 풀어보자. indexing_maps와 iterator_types는 앞서 본 벡터 컨트랙션과 정확히 동일하다. 이제 피연산자는 두 목록으로 나뉜다:
in 피연산자: 연산이 읽기만 하는 버퍼들;out 피연산자: 연산이 읽고 갱신하는 버퍼들.벡터는 MLIR에서 읽기 전용(SSA/함수형)이며, 벡터를 변경하는 연산은 실제로 새로운 벡터를 생성하므로, 벡터의 경우에는 이런 구분이 필요하지 않았다.
더 나아가, 이제 연산 안에는 컨트랙션에서 암묵적이었던 곱셈과 덧셈을 명시적으로 지정하는 리전(region)이 포함된다. 리전의 블록 인자는 버퍼에서 읽은 개별 원소에 해당한다. 처음 두 인자는 in 피연산자에, 마지막 인자는 out 피연산자에 대응한다. 리전에서 산출(yield)되는 값은 out 피연산자에 “쓰기”되며, 리전의 이후 실행에서 마지막 블록 인자로 제공된다. 다양한 튜플의 원소들에 대해 리전이 실행되는 순서는 지정되지 않으며, out 버퍼에 대한 쓰기는 연산의 끝에서 전체적으로 이루어진다는 점에 유의하라.
linalg.generic 연산의 리전에는 임의 개수의 연산을 포함할 수 있으므로, 리전 안에 더 많은 연산을 사슬처럼 연결하여 암묵적 루프의 “퓨전”을 표현할 수 있다. 예를 들어, 머신러닝에서 흔한 ReLU(Rectified Linear Unit) 계층은 relu(x) = max(0, x)로 정의되며, 비교-선택(compare-and-select) 이디엄을 사용해 하나의 linalg.generic 연산으로 표현할 수 있다. 이때 비교 결과를 위한 임시 버퍼가 필요 없고, 바깥쪽 연산도 반복하지 않는다:
linalg.generic {
indexing_maps [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
iterator_types = ["parallel"]
} ins(%in : memref<?xf32>) outs(%out : memref<?xf32>) {
^bb0(%in_one : f32, %out_one : f32):
%c0 = arith.constant 0.0 : f32
%0 = arith.cmpf ogt %in_one, %c0 : f32
%1 = arith.select %0, %in_one, %c0 : f32
linalg.yield %1 : f32
}
이러한 연산은 벡터 방언의 원시 연산에 각각 매핑되는 여러 연산들로 분해한 뒤 루프로 변환하거나 벡터 형태로 내릴 수 있다. 이와 같은 모델링은 다시 컴파일러가 코드 생성 전략을 선택하는 데 더 많은 자유도를 제공한다.
추상화 사다리를 한 단계 더 올라가 보자. MLIR은 다차원이면서도 규칙적인 데이터에 대해, 다차원 버퍼에서 필요해지는 별칭 분석(alias analysis)과 의존성 만족 같은 복잡한 문제를 해결하지 않고도 컴파일러가 쉽게 추론할 수 있게 해 주는 텐서 추상화를 제공한다. 텐서는 벡터 추상화와 매우 유사하다(주요 차이점으로는 무랭크 텐서의 존재, 텐서 레이아웃, 그리고 벡터가 텐서의 원소 타입으로는 가능하지만 다른 벡터의 원소 타입으로는 불가능하다는 점 등이 있다). 텐서는 읽기 전용이며, 텐서를 갱신하는 연산은 새로운 텐서를 생성한다.
앞서의 linalg.generic 연산은 버퍼 대신 텐서에 대해 동작하도록 끌어올릴 수 있다:
%result = linalg.generic {
indexing_maps = [affine_map<(i, j, k) -> (i, k)>,
affine_map<(i, j, k) -> (k, j)>,
affine_map<(i, j, k) -> (i, j)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%lhs, %rhs : tensor<8x10xf32>,tensor<10x16xf32>)
outs(%init :tensor<8x16xf32>) {
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32):
%0 = arith.mulf %lhs_one, %rhs_one : f32
%1 = arith.addf %init_one, %0 : f32
linalg.yield %1 : f32
} -> tensor<8x16xf32>
보다시피 이 연산의 대부분 구성 요소는 버퍼 버전과 동일하다. 이는 의도된 설계다. 피연산자 타입 외의 주요 차이는, 이제 연산이 out 버퍼를 갱신하는 대신 새로운 결과를 생성한다는 점이다. out 피연산자는 초기화 값으로만 사용된다.
만약 linalg.generic이 벡터에 대해서도 존재했다면, 구조는 완전히 동일했을 것이다.
이 추상화 수준에서는 일반적으로 고성능 코드 생성을 위해 요구되는 타일링 같은 고급 변환을 컴파일러가 쉽게 수행할 수 있다. 일반적으로 타일링은 반복 공간을 더 작은 부분(타일)으로 분할하는 것으로 볼 수 있으며, 예를 들어 각 부분이 필요로 하는 데이터가 특정 캐시 레벨에 들어맞도록 한다. 타일이 실행되는 순서는 원래의 데이터 의존성을 보존해야 한다.
linalg.generic 연산의 경우, 반복 공간은 암묵적이며 피연산자의 형태(shape)로 정의된다. 따라서 타일은 원래 데이터의 부분(슬라이스)에 대해 동일한 연산을 수행함으로써 표현할 수 있다. linalg.generic의 본문이 서로 다른 입력 원소 튜플들에 적용되는 순서는 지정되지 않으므로, 의존성 분석 없이도 타일은 어떤 순서로든 실행될 수 있다. 서로 다른 타일의 실행을 제어하기 위해, 타일링 구현은 루프를 생성한다. 즉, linalg.generic의 타일링은 지금까지 암묵적이던 루프를 구체화(materialize)하는 것으로 볼 수 있다.
예를 들어 앞서 제시한 행렬 곱셈을 타일 크기 (2, 8)로 타일링하면, 2x8 텐서에 대해 동일한 연산을 표현하는 linalg.generic 주위로 루프 중첩이 생긴다.
// 텐서 갱신이 아닌 텐서 삽입 의미론을 지원하는 특수한 "멀티-for" 루프.
// 이 루프가 8x16 텐서를 산출한다.
// 이터레이터의 반복 횟수는 원래 텐서 크기 8x16을 타일 크기 2x8로 나눠 4x2로 계산한다.
// 텐서 크기가 동적이면, 반복 횟수 계산은 IR로 생성되어 런타임에 계산된다.
%0 = scf.forall (%i, %j) in (4, 2)
shared_outs(%shared = %init) -> (tensor<8x16xf32>) {
// 루프 유도 변수를 타일 크기로 스케일한다.
%3 = affine.apply affine_map<(d0) -> (d0 * 2)>(%i)
%4 = affine.apply affine_map<(d0) -> (d0 * 8)>(%j)
// 입력과 출력의 슬라이스를 취한다. "i"와 "j" 차원만 슬라이스한다.
%lhs_slice = tensor.extract_slice %lhs[%3, 0] [2, 10] [1, 1]
: tensor<8x10xf32> to tensor<2x10xf32>
%rhs_slice = tensor.extract_slice %rhs[0, %4] [10, 8] [1, 1]
: tensor<10x16xf32> to tensor<10x8xf32>
%result_slice = tensor.extract_slice %shared[%3, %4] [2, 8] [1, 1]
: tensor<8x16xf32> to tensor<2x8xf32>
// 이전과 정확히 같은 연산이지만, 더 작은 데이터 슬라이스에 대해 동작한다.
%partial = linalg.generic {
indexing_maps = [affine_map<(i, j, k) -> (i, k)>,
affine_map<(i, j, k) -> (k, j)>,
affine_map<(i, j, k) -> (i, j)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%lhs_slice, %rhs_slice : tensor<2x10xf32>, tensor<10x8xf32>)
outs(%result_slice : tensor<2x8xf32>) -> tensor<2x8xf32> {
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32):
%0 = arith.mulf %lhs_one, %rhs_one : f32
%1 = arith.addf %init_one, %0 : f32
linalg.yield %1 : f32
} : tensor<2x8xf32>
// 텐서 삽입 의미론을 가진 루프의 종료자. 더 큰 텐서에 슬라이스를 삽입하며,
// 잠재적으로 병렬 수행될 수 있다.
scf.forall.in_parallel {
tensor.parallel_insert_slice %partial into %shared[%3, %4] [2, 8] [1, 1]
: tensor<2x8xf32> into tensor<8x16xf32>
}
}
타일링으로 루프를 구체화한 뒤에는 또 하나의 핵심 코드 생성 변환인 퓨전이 쉬워진다. 루프 퓨전과 달리, 구조화 연산 접근법은 연산들의(암묵적) 반복 공간이 일치하지 않아도 생산자/소비자 퓨전을 가능하게 한다. linalg.generic 같은 고수준의 텐서 기반 구조화 연산이 주어지면, use-def 체인을 따라가 다음을 식별할 수 있다:
indexing_map을 역변환하여 슬라이스를 통해 접근되는 원소 집합에 적용하면, 원래 전체 텐서를 정의하는 연산의 반복 공간 중 타일을 계산하는 데 필요한 부분을 구할 수 있다. 따라서 퓨전은 tensor.extract_slice 연산을 원 피연산자를 생성하던 linalg.generic의 타일로 대체하는 문제로 귀결된다.
행렬 곱셈 연산 뒤에, 결과 행렬의 각 원소를 자기 자신과 곱하는 후속 원소별 연산이 있다고 하자. 이 후속 원소별 연산의 반복 공간은 2차원으로, 행렬 곱셈의 3차원 반복 공간과 다르다. 그럼에도 불구하고, 후속 연산을 타일링한 다음 그 피연산자의 생산자인 matmul을 타일링으로 생성된 루프 안으로 퓨전할 수 있다. 타일링되지 않은 차원은 전체를 사용한다.
// 앞선 예와 동일한 루프.
%0 = scf.forall (%i, %j) in (4, 2)
shared_outs(%shared = %init)
-> (tensor<8x16xf32>, tensor<8x16xf32>) {
// 루프 유도 변수를 타일 크기로 스케일한다.
%1 = affine.apply affine_map<(d0) -> (d0 * 2)>(%i)
%2 = affine.apply affine_map<(d0) -> (d0 * 8)>(%j)
// 입력과 출력의 슬라이스를 취한다. "i"와 "j" 차원만 슬라이스한다.
%lhs_slice = tensor.extract_slice %lhs[%1, 0] [2, 10] [1, 1]
: tensor<8x10xf32> to tensor<2x10xf32>
%rhs_slice = tensor.extract_slice %rhs[0, %2] [10, 8] [1, 1]
: tensor<10x16xf32> to tensor<10x8xf32>
%result_slice = tensor.extract_slice %result[%1, %2] [2, 8] [1, 1]
: tensor<8x16xf32> to tensor<2x8xf32>
// 앞서의 matmul 슬라이스와 정확히 동일하다. 아래 제네릭 연산의 슬라이스 추출을 대체한다.
%partial = linalg.generic {
indexing_maps = [affine_map<(i, j, k) -> (i, k)>,
affine_map<(i, j, k) -> (k, j)>,
affine_map<(i, j, k) -> (i, j)>],
iterator_types = ["parallel", "parallel", "reduction"]
} ins(%lhs_slice, %rhs_slice : tensor<2x10xf32>, tensor<10x8xf32>)
outs(%result_slice : tensor<2x8xf32>) {
^bb0(%lhs_one: f32, %rhs_one: f32, %init_one: f32):
%5 = arith.mulf %lhs_one, %rhs_one : f32
%6 = arith.addf %init_one, %5 : f32
linalg.yield %6 : f32
} -> tensor<2x8xf32>
// 최종 결과의 슬라이스를 취한다. 위의 matmul 연산이 슬라이스를 제자리에서 계산하므로
// 피연산자의 슬라이스를 따로 취할 필요가 없다.
%shared_slice = tensor.extract_slice %shared[%1, %2] [2, 8] [1, 1]
: tensor<8x16xf32> to tensor<2x8xf32>
// 우리가 타일링한 원소별 연산.
%elemwise = linalg.generic {
indexing_maps = [affine_map<(i, j) -> (i, j)>,
affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"]
} ins(%partial : tensor<2x8xf32>)
outs(%shared_slice : tensor<2x8xf32>) {
^bb0(%in: f32, %out: f32):
%5 = arith.mulf %in, %in : f32
linalg.yield %5 : f32
} -> tensor<2x8xf32>
// 텐서 삽입 의미론을 가진 루프의 종료자. 더 큰 텐서에 슬라이스를 삽입하며,
// 잠재적으로 병렬 수행될 수 있다.
scf.forall.in_parallel {
tensor.parallel_insert_slice %elemwise into %shared[%1, %2] [2, 8] [1, 1]
: tensor<2x8xf32> into tensor<8x16xf32>
}
}
이 과정의 결과로, 피연산자 텐서의 일부 원소들이 루프의 각 반복에서 (재)계산될 수 있다. 이는 재물질화(rematerialization)라고도 하며, 중복 계산을 수행할지 아니면 그 결과를(느린) 메모리에 저장할지 사이의 트레이드오프를 표현한다.
Linalg는 행렬 곱셈, 점곱, 합성곱 등 일반적인 경우를 위한 사전 정의된 연산 집합을 제공한다. 이들 연산은 generic과 동등하지만, 접근 패턴과 본문을 일일이 적지 않아도 된다. 예를 들어, 행렬 곱셈은 다음처럼 간단하다:
%matmul = linalg.matmul ins(%lhs, %rhs: tensor<8x10xf32>, tensor<10x16xf32>)
outs(%init: tensor<8x10xf32xf32>) -> tensor<8x16xf32>