Chapter 3: High-level Language-Specific Analysis and Transformation
원문: https://mlir.llvm.org/docs/Tutorials/Toy/Ch-3/
본 게시글에서는 MLIR의 Toy Tutorial - Chapter 3: High-level Language-Specific Analysis and Transformation 내용을 다루고 있습니다. 단순 해석글로, 이해를 쉽게 하기 위해 필자가 이해한 내용을 조금 덧붙이고 있습니다.
언어의 의미를 정확하게 표현하는 Dialect를 생성하는 것은 MLIR에서 고수준 언어에 대한 분석, 변환 및 최적화를 가능하게 합니다. 이러한 작업은 일반적으로 AST에서 수행됩니다.
컴파일러 변환은 로컬 변환과 글로벌 변환 두 가지 범주로 나눌 수 있습니다. 이번에는 Toy Dialect와 그 고수준 의미를 활용해 LLVM에서 로컬 패턴 매칭 변환을 수행하는 방법에 집중합니다. 이를 위해서 MLIR의 Generic DAG Rewriter을 사용합니다.
패턴 매칭 변환을 구현하는 방법에는 두 가지가 있습니다.
1. 명령형: C++ 패턴 매칭 및 변환
2. 선언형: 테이블 기반 선언적 재작성 규칙(DRR)을 사용한 패턴 매칭 및 변환
DRR을 사용하려면 연산이 ODS를 사용하여 정의되어야 합니다.
패턴 매칭? DRR?.. 지금 이해가 안 되더라도 끝까지 본 튜토리얼을 정독하면 이해 되실겁니다.
Optimize Transpose using C++ style pattern-match and rewrite
간단한 패턴을 사용해서 두 번의 transpose 연산이 서로 상쇄되는 경우를 최적화 해 보겠습니다.
이에 대한 Toy 예시는 다음과 같습니다.
def transpose_transpose(x) {
return transpose(transpose(x));
}위 코드는 아래와 같이 IR로 변환됩니다.
toy.func @transpose_transpose(%arg0: tensor<*xf64>) -> tensor<*xf64> {
%0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64>
%1 = toy.transpose(%0 : tensor<*xf64>) to tensor<*xf64>
toy.return %1 : tensor<*xf64>
}transpose_transpose 함수는 인수 x에 대해 두 번의 transpose 연산을 수행합니다. 이 두 개의 transpose 연산은 서로 상쇄되어 사실상 아무런 연산을 하지 않기에 이 코드는 최적화해서 단순히 x를 반환하는 코드로 만들 수 있습니다.
Clang을 예시로 들어보겠습니다. transpose 연산을 다루는 계산은 다음과 같은 반복문으로 표현됩니다.
#define N 100
#define M 100
void sink(void *);
void double_transpose(int A[N][M]) {
int B[M][N];
for(int i = 0; i < N; ++i) {
for(int j = 0; j < M; ++j) {
B[j][i] = A[i][j];
}
}
for(int i = 0; i < N; ++i) {
for(int j = 0; j < M; ++j) {
A[i][j] = B[j][i];
}
}
sink(A);
}C++ 스타일의 단순한 접근 방식으로, IR에서 트리 형태의 패턴을 매칭하고 이를 다른 연산 집합으로 바꾸는 방식이 있습니다. (=즉 코드에서 특정한 부분을 찾아내서 다른 연산으로 대체한다는 의미) 이는 MLIR의 Canonicalizer 패스를 사용하며 RewritePattern(= 특정한 부분을 다른 연산으로 대체하는 프로세스)을 구현하는 방식으로 적용할 수 있습니다.
/// Fold transpose(transpose(x)) -> x
struct SimplifyRedundantTranspose : public mlir::OpRewritePattern<TransposeOp> {
/// 이 패턴은 IR의 모든 `toy.transpose`를 매칭하도록 등록
/// "benefit"는 프레임워크에서 패턴들을 순서대로 처리하는 데 사용
SimplifyRedundantTranspose(mlir::MLIRContext *context)
: OpRewritePattern<TransposeOp>(context, /*benefit=*/1) {}
/// 이 메서드는 패턴을 매칭하고 이를 변환하려고 시도`rewriter` 인자는 변환 순서를 관리하는 도구로, IR을 변경할 때 사용
llvm::LogicalResult
matchAndRewrite(TransposeOp op,
mlir::PatternRewriter &rewriter) const override {
// 현재 transpose의 입력을 살펴보기
mlir::Value transposeInput = op.getOperand();
TransposeOp transposeInputOp = transposeInput.getDefiningOp<TransposeOp>();
// 입력이 다른 transpose에 의해 정의되지 않았다면 매칭하지 않음
if (!transposeInputOp)
return failure();
// 그렇지 않다면, 중복된 transpose가 있으므로 rewriter를 사용
rewriter.replaceOp(op, {transposeInputOp.getOperand()});
return success();
}
};위 구현은 ToyCombine.cpp에 존재하며 표준화 pass에서 변환을 적용하는 방식에 대해서 설명하고 있습니다.
표준화 pass는 연산들에 의해 정의된 변환들을 그리디하고 반복적인 방식으로 적용합니다. 앞서 작성한 새로운 변환을 적용하기 위해 hasCanonicalizer = 1로 설정하고 Canonicalization 프레임워크에 패턴을 등록합시다.
// Canonicalization 프레임워크에서 재작성할 패턴을 등록
void TransposeOp::getCanonicalizationPatterns(
RewritePatternSet &results, MLIRContext *context) {
results.add<SimplifyRedundantTranspose>(context);
}TransposeOp에 대한 최적화 패턴을 Canonicalization 프레임워크에 등록합니다. 이를 통해서 TransposeOp 연산에 대해 중복된 transpose를 최적화하도록 합니다.
다음으로 우리는 toyc.cpp에서 최적화 파이프라인을 추가해야 합니다. (toyc.cpp는 toy 언어에 대한 컴파일러 입니다. 깃허브에서 코드를 확인하실 수 있습니다.) MLIR에서는 LLVM과 같이 PassManager를 통해서 최적화를 실행합니다.
mlir::PassManager pm(module->getName());
pm.addNestedPass<mlir::toy::FuncOp>(mlir::createCanonicalizerPass());PassManager를 설정하고 createCanonicalizerPass()를 추가해서 표준화 패스를 실행하도록 합니다.
toyc-ch3을 실행해서 변환 패턴 적용을 실험할 수 있습니다.
toyc-ch3 test/Examples/Toy/Ch3/transpose_transpose.toy -emit=mlir -opt결과는 다음과 같이 출력됩니다.
이제 transpose 연산 없이 함수 인자값을 바로 변환하게 됩니다.
toy.func @transpose_transpose(%arg0: tensor<*xf64>) -> tensor<*xf64> {
%0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64>
toy.return %arg0 : tensor<*xf64>
}그러나 아직까지 transpose가 하나 남아있습니다.. 이를 해결하기 위해서 Transpose에 Pure라는 새로운 특성을 추가할 수 있습니다. 이는 Canonicalizer가 불필요한 연산을 제거할 수 있게 합니다.
def TransposeOp : Toy_Op<"transpose", [Pure]> {...}위와 같이 추가하고 다시 실행해 봅시다.
toyc-ch3 test/Examples/Toy/Ch3/transpose_transpose.toy -emit=mlir -opt그러면 아래와 같은 결과를 얻을 수 있습니다!
toy.func @transpose_transpose(%arg0: tensor<*xf64>) -> tensor<*xf64> {
toy.return %arg0 : tensor<*xf64>
}이제 1. 명령형: C++ 패턴 매칭 및 변환을 모두 살펴봤습니다.
Optimize Reshapes using DRR
여기부터 DDR을 사용해서 Reshape 연산의 중복 최적화를 수행하는 방법을 설명합니다. 구체적으로 Reshape(Reshape(x)) = Reshape(x) 형태의 중복된 연산을 최적화하는 규칙을 정의하는 예시입니다.
DDR은 DAG을 기반으로 한 선언적 재작성기 입니다. 이는 패턴 매칭 및 재작성 규칙을 테이블 기반 구문으로 제공하며 아래처럼 규칙을 선언적으로 작성할 수 있습니다.
class Pattern<
dag sourcePattern, list<dag> resultPatterns,
list<dag> additionalConstraints = [],
dag benefitsAdded = (addBenefit 0)>;SimplifyRedundantTranspose와 유사한 중복된 reshape 최적화는 DDR을 사용해서 더 간단하게 표현할 수 있습니다.
// Reshape(Reshape(x)) = Reshape(x)
def ReshapeReshapeOptPattern : Pat<(ReshapeOp(ReshapeOp $arg)),
(ReshapeOp $arg)>;DDR에서 자동으로 생성된 C++ 코드는 path/to/BUILD/tools/mlir/examples/toy/Ch3/ToyCombine.inc 경로에서 확인할 수 있습니다.
DDR은 변환이 인수 및 결과의 특정 속성에 따라 조건부로 수행될 때 인수 제약을 추가하는 방법도 제공합니다. 예를 들어서 입력과 출력의 형태가 동일할 때, 즉 reshape 연산이 중복된 경우 이를 제거하는 변환이 이에 해당합니다.
def TypesAreIdentical : Constraint<CPred<"$0.getType() == $1.getType()">>;
def RedundantReshapeOptPattern : Pat<
(ReshapeOp:$res $arg), (replaceWithValue $arg),
[(TypesAreIdentical $res, $arg)]>;인수 제약이 필요한 최적화 중 일부는 명령어 인수에 대해서 추가적인 변환이 필요할 수 있습니다. 이는 NativeCodeCall을 사용해서 더 복잡한 변환을 처리하는 방식입니다. 이는 C++ 헬퍼 함수 호출을 통해 변환을 처리하거나 인라인 C++을 사용해서 변환을 직접 처리할 수 있습니다. 이 예시로는 FoldConstantReshape가 있습니다. 상수 값을 Reshape할 때 상수를 바로 변환하고 Reshape 연산을 제거하는 방식으로 최적화하는 방식이 있습니다.
def ReshapeConstant : NativeCodeCall<"$0.reshape(($1.getType()).cast<ShapedType>())">;
def FoldConstantReshapeOptPattern : Pat<
(ReshapeOp:$res (ConstantOp $arg)),
(ConstantOp (ReshapeConstant $arg, $res))>;다음과 같은 trivial_reshape.toy 프로그램을 통해 이러한 Reshape 최적화를 시연해 보겠습니다.
def main() {
var a<2,1> = [1, 2];
var b<2,1> = a;
var c<2,1> = b;
print(c);
}module {
toy.func @main() {
%0 = toy.constant dense<[1.000000e+00, 2.000000e+00]> : tensor<2xf64>
%1 = toy.reshape(%0 : tensor<2xf64>) to tensor<2x1xf64>
%2 = toy.reshape(%1 : tensor<2x1xf64>) to tensor<2x1xf64>
%3 = toy.reshape(%2 : tensor<2x1xf64>) to tensor<2x1xf64>
toy.print %3 : tensor<2x1xf64>
toy.return
}
}아래 명령어를 실행해서 우리가 설정한 최적화 패턴을 적용한 결과를 확인해 봅시다.
toyc-ch3 test/Examples/Toy/Ch3/trivial_reshape.toy -emit=mlir -optmodule {
toy.func @main() {
%0 = toy.constant dense<[[1.000000e+00], [2.000000e+00]]> : tensor<2x1xf64>
toy.print %0 : tensor<2x1xf64>
toy.return
}
}위 코드를 보면 Reshape 연산이 남아있지 않은 것을 확인할 수 있습니다.
선언적 재작성 방법에 대한 자세한 내용은 DDR에서 확인할 수 있습니다. 다음 장에서는 인터페이스를 통해서 더 확장 가능하고 일반적인 솔루션을 사용하는 방법을 알아보겠습니다.
