MLIR 기반 컴파일러를 다루는 사람에게 linalg.generic의 ins, outs, 인덱싱 맵과 반복자 유형은 익혀야 할 개념이다. 알렉스 데니소프(AlexDenisov)는 이를 중첩된 scf.for 반복문으로 직접 풀어 쓰는 방법을 설명한다.
데니소프는 MLIR 문서와 설계 배경 등을 살펴본 뒤, linalg.generic을 익숙한 반복문 표기로 이해할 수 있다고 보고 손으로 변환하는 방법을 정리했다. 예시에서는 tensor 대신 다차원 메모리 버퍼인 memref를 사용하지만, 설명은 두 유형 모두에 적용된다고 말한다.
linalg.generic은 입력 N개와 출력 M개를 처리하는 K중 반복문을 간결하게 나타낸 연산이다. 입력과 출력마다 하나씩 총 N+M개의 인덱싱 맵이 있고, 반복문마다 반복자 유형이 하나씩 있다. 본문은 각 입력·출력에 대응하는 인자를 받아 작업을 수행한 뒤, linalg.yield로 출력값 M개를 돌려준다.
본문의 인자는 대응하는 입력·출력의 원소 유형을 따른다. 버퍼의 차원 수가 달라도 본문 인자의 유형은 원소 유형으로 정해지며, 원소는 벡터처럼 스칼라가 아닐 수도 있다. linalg.yield의 값은 각 출력의 원소 유형과 일치해야 한다.
반복자 유형은 parallel 또는 reduction이며, 반복문별로 지정된다. 데니소프의 설명에 따르면 parallel은 해당 반복문의 실행 순서가 정확성에 영향을 주지 않는다는 뜻이어서 재배치하거나 병렬화할 수 있다. 이 글에서는 반복자 유형의 세부 사항을 변환의 핵심에서 제외한다.
인덱싱 맵은 반복문의 개수와 각 메모리 접근에 사용할 인덱스를 알려준다. 모든 맵은 같은 입력 영역을 공유하며, 영역의 각 차원은 반복문의 유도 변수를 나타낸다. 예를 들어 행렬 곱셈에는 i, j, k 세 반복문이 있고, 각 반복에서 C[i, j] = A[i, k] * B[k, j] + C[i, j]처럼 입력과 출력을 읽고 쓴다.
맵의 결과 개수는 해당 버퍼의 차원 수와 같다. 반복문의 상한은 맵에 나타난 차원과 각 버퍼의 크기를 바탕으로 정한다. 예를 들어 A의 크기가 4x8, B가 8x6이고 C의 크기가 동적이라면, 맵의 차원 대응을 통해 C의 유효한 크기는 4x6으로 추론할 수 있다.
linalg.generic을 반복문으로 직접 옮길 때는 반복문 개수와 각 상한을 정하고, 대응하는 유도 변수를 둔 중첩 scf.for를 만든다. 이어 인덱싱 맵에 따라 memref.load로 본문 인자를 가져오고 본문 연산을 옮긴 뒤, 끝의 linalg.yield를 memref.store로 바꾼다. 동적 크기의 행렬 곱셈 예시에서는 상한을 A와 B의 차원에서 가져오고, A[%i, %k], B[%k, %j], C[%i, %j]를 읽어 처리한다.
글은 명명된 원소별 연산, linalg.pack/unpack과 linalg.index 같은 연산, 상수나 산술을 포함하는 더 복잡한 인덱싱 맵 등은 간략화를 위해 다루지 않는다. linalg.generic의 피연산자가 항상 memref나 tensor인 것도 아니며, 일부 세부 사항이 설명에서 빠졌다고 덧붙인다. 변환 예시는 Compiler Explorer에서 확인할 수 있다.
댓글 (0)
로그인하면 이 기사에 내 생각을 남길 수 있어요