Chapter 2: Emitting Basic MLIR

moze·2025년 1월 1일

원문: https://mlir.llvm.org/docs/Tutorials/Toy/Ch-2/

본 게시글에서는 MLIR의 Toy Tutorial - Chapter 2: Emitting Basic MLIR 내용을 다루고 있습니다. 단순 해석글로, 이해를 쉽게 하기 위해 필자가 이해한 내용을 조금 덧붙이고 있습니다.

Chapter 1: Toy Language and AST에서 Toy 언어와 AST에 대해 감을 잡았으니 이제 MLIR이 Toy 언어를 어떻게 컴파일 하는지 알아봅시다.

Introduction: Multi-Level Intermediate Representation

LLVM은 미리 정의된 타입과 저수준 명령어 세트를 제공합니다. 특정 언어의 프론트엔드가 LLVM IR을 생성하기 전, 언어 특화된 타입 검사, 분석, 변환을 수행해야 합니다. C/C++ 보다 높은 수준에서 구성된 언어는 AST에서 LLVM IR을 생성하기 위해서 복잡한 Lowering 과정을 거쳐야 할 수 있습니다.

즉.. 다양한 프론트엔드에 대한 분석과 변환을 지원하기 위해서는 상당한 양의 인프라를 다시 구현해야 하는 상황이 발생합니다.

MLIR은 이러한 문제를 해결하기 위해서 "확장 가능성"이라는 핵심 키워드를 염두에 두고 설계되었습니다. MLIR에서는 미리 정의된 명령어나 타입이 거의 없다고 보면 됩니다. (사용자가 이를 편리하게 정의하도록 그 인프라를 제공)

Interfacing with MLIR

MLIR은 확장 가능한 인프라로 설계되었습니다. 쉽게 말해서 속성, 연산, 타입이 고정되어 있지 않아 확장이 쉽습니다.

MLIR은 Dialect라는 개념을 통해서 이러한 확장성을 지원합니다. Dialect는 고유한 네임스페이스 아래에서 추상화를 그룹화하는 메커니즘을 제공합니다. (지금 이 부분이 이해 안 되더라도, Toy Tutorial을 다 읽어보면 어느정도 감이 잡힐 것입니다.)

MLIR에서 연산은 추상화와 계산의 핵심 단위로 LLVM 명령어와 여러 면에서 유사합니다. 연산은 LLVM의 주요 IR 구조를 표현하는 데 사용됩니다.

예시로 Toy 언어의 transpose 연산에 대한 MLIR 어셈블리를 살펴보겠습니다. transepose 연산은 

%t_tensor = "toy.transpose"(%tensor) {inplace = true} : (tensor<2x3xf64>) -> tensor<3x2xf64> loc("example/file/path":12:1)

%t_tensor는 연산의 결과 이름으로, 충돌을 피하기 위해서 접두사가 추가된 기호를 포함하고 있습니다. 한 연산은 0개 이상의 결과를 반환할 수 있고 이들은 Static Single Assignment 값에 해당합니다. (transpose 연산에서는 1개의 결과를 내보냅니다.)

"toy.transpose"는 연산의 이름으로 Dialect의 네임스페이스가 접두사로 붙어 있습니다. 쉽게 말하자면, Toy Dialect의 transpose 연산이라고 볼 수 있습니다.

(%tensor)는 피연산자입니다.

{inplace = true}는 속성 사전이라는데.. 저도 잘 모르겠습니다 :( 링크를 참고해보면 좋겠습니다. 특별한 피연산자로 여기서는 inplace라는 이름의 부울 속성을 정의하며 값은 true로 되어 있습니다.

(tensor<2x3xf64>) -> tensor<3x2xf64>는 연산의 타입을 함수 형태로 표현한 것입니다. 괄호 안에 인자 타입을 나열하고 반환 값의 타입을 뒤에 표시합니다.

loc("example/file/path":12:1)는 소스 코드에서 연산이 유래한 위치입니다. 디버깅 목적으로 사용되는 부분입니다.

코드를 살펴보니 transpose() 함수 선언부만 MLIR로 변환해서 보여준거 같습니다.

위 내용을 토대로 MLIR 연산의 주요 구성 요소는 다음으로 추릴 수 있습니다.

  1. 연산의 이름
  2. 피연산자 값
  3. 속성 목록
  4. 결과 값의 타입 목록
  5. 디버깅을 위한 소스코드 위치
  6. (위 예시에 없음) 후속 블록
  7. (위 예시에 없음) 구조적 연산을 위한 region 목록 (블록의 정렬된 목록)

모든 MLIR 연산은 원 소스코드와 연결되어 있습니다. LLVM에서는 디버그 정보가 메타데이터로 처리되면서 생략될 수 있는데, MLIR은 소스코드와 긴밀히 연결되어 있으므로 생략을 걱정하지 않아도 됩니다. 소스코드 위치와 MLIR 연산은 연결되어 있습니다.

*mlir-opt 툴은 출력에 위치 정보를 포함하지 않습니다. -mlir-print-debuginfo 옵션을 사용하면 위치 정보를 포함하도록 할 수 있습니다.

Opaque API

MLIR은 속성, 연산, 타입과 같은 모든 IR 요소를 사용자가 정의할 수 있도록 설계되었습니다. 동시에 이러한 IR 요소는 앞에서 설명한 기본 개념들로 축약될 수 있습니다.

MLIR은 어떤 연산이라도 구문분석, 표현, 그리고 라운드 트립을 할 수 있습니다. 예를 들어서 Toy 연산을 .mlir 파일에 포함시키고 Toy 관련 Dialect를 등록하지 않은 상태에서 mlir-opt를 통해서 라운드 트립이 가능합니다.

*라운드 트립: 변환한 후 다시 원래 상태로 복원하는 것

등록되지 않은 속성, 연산, 타입의 경우에서 MLIR은 몇 가지 구조적 제약을 강제하지만 이 외의 것들은 Opaque하게 처리합니다. 즉.. MLIR이 다음에 대해서 알지 못한다는 겁니다.

  • 미등록 연산이 특정 데이터 타입에서 동작할 수 있는지
  • 몇 개의 피연산자를 받을 수 있는지
  • 몇 개의 결과값을 생성하는지

잘못된 IR을 정의해도 MLIR에서는 이를 필터링하지 않고 라운드 트립이 가능합니다.

성숙한 시스템에서는 이런 Opaque함이 권장되지 않습니다. 완전하지 않으니까요. 다르게 생각하면 초기 개발이나 부트스트래핑 과정에서 유용할 수 있습니다.

Defining a Toy Dialect

MLIR과 효과적으로 연동하기 위해서 Toy Dialect를 새로 정의해 보겠습니다. (지금까지는 Toy 언어를 다루었으니 혼란이 없으시길..) 이 Dialect는 Toy 언어의 구조를 모델링하며 고수준 분석 및 변환을 쉽게 수행할 수 있는 방법을 제공합니다.

C++를 통해서 Toy Dialect를 선언할 수 있습니다.

/// C++
/// Toy Dialect 정의
/// Dialect는 `mlir::Dialect`를 상속
/// 사용자 정의 속성(attributes), 연산(operations), 타입(types)을 등록
class ToyDialect : public mlir::Dialect {
public:
  explicit ToyDialect(mlir::MLIRContext *ctx);

  /// Dialect의 네임스페이스를 반환하는 유틸리티 접근자
  static llvm::StringRef getDialectNamespace() { return "toy"; }

  /// Toy Dialect 생성자에서 호출되는 초기화 함수
  /// 속성, 연산, 타입 등을 등록하는 데 사용
  void initialize();
};

그러나 MLIR에서는 더욱 쉬운 방법을 제공합니다. TableGen이라는 언어를 통해서 쉽게 Toy Dialect를 선언할 수 있습니다! 아래와 같이 진행됩니다.

// ODS 프레임워크에서 Toy Dialect를 정의합니다.
def Toy_Dialect : Dialect {
  // Dialect의 네임스페이스 `ToyDialect::getDialectNamespace`에서 제공한 문자열과 동일
  let name = "toy";

  // Dialect의 간단한 요약
  let summary = "A high-level dialect for analyzing and optimizing the "
                "Toy language";

  // Dialect에 대한 상세 설명
  let description = [{
    The Toy language is a tensor-based language that allows you to define
    functions, perform some math computation, and print results. This dialect
    provides a representation of the language that is amenable to analysis and
    optimization.
  }];

  // Dialect 클래스 정의가 위치한 C++ 네임스페이스
  let cppNamespace = "toy";
}

생각보다 선언하는 것은 간단합니다.

TableGen 명령어를 통해서 위 정의를 기반으로 결과물을 생성해낼 수 있습니다.

${build_root}/bin/mlir-tblgen -gen-dialect-decls 
${mlir_src_root}/examples/toy/Ch2/include/toy/Ops.td -I 
${mlir_src_root}/include/

이 명령들은 Dialect 선언을 생성합니다.

Dialect를 정의한 후에는 이를 MLIRContext에 로드할 수 있습니다.

context.loadDialect<ToyDialect>();

MLIRContext에는 Buitin Dialect 리스트가 존재합니다. 여기에 사용자가 만든 Dialect를 명시적으로 로드해 주어야 사용이 가능합니다.

자, 지금까지 Toy Dialect 선언을 MLIR에서 진행해 보았습니다.

Defining Toy Operations

이제 Toy Dialect를 기반으로 연산을 정의할 수 있습니다.
연산은 MLIR에서 실행 가능한 단위로 Dialect에 의미 정보를 제공하며 시스템이 이를 통해서 분석하고 최적화 작업을 수행할 수 있습니다.

다음은 toy.constant라는 Constant 연산을 정의하는 과정을 설명합니다.

toy.constant 연산의 예시는 다음과 같습니다.

%4 = "toy.constant"() {value = dense<1.0> : tensor<2x3xf64>} : () -> tensor<2x3xf64>

toy.constant는 다음과 같은 특징을 가집니다.

  1. 피연산자가 없음
  2. value라는 DenseElements 속성을 사용해서 상수 표현
  3. 하나의 RankedTensorType을 반환

그럼 본격적으로 연산을 정의해 봅시다.

먼저 Constantop라는 C++ 클래스를 작성해 봅시다.
C++에서 CRTP 패턴을 사용하는 mlir::Op 클래스를 상속해서 작성하면 됩니다. Traits를 활용해서 연산의 동작을 세부적으로 정의합니다.

class ConstantOp : public mlir::Op<
                     ConstantOp,                          // mlir::Op를 물려받는 클래스
                     mlir::OpTrait::ZeroOperands,         // 피연산자 없음
                     mlir::OpTrait::OneResult,            // 하나의 출력 반환
                     mlir::OpTraits::OneTypedResult<TensorType>::Impl> {
public:
  using Op::Op;  // 기본 생성자 상속

  /// Operation의 고유 이름을 반환 (Namespace + Operation 이름)
  static llvm::StringRef getOperationName() { return "toy.constant"; }

  /// 속성에서 상수 값을 반환하는 메서드
  mlir::DenseElementsAttr getValue();

  /// 추가 검증 로직: 결과 타입이 TensorType인지, 속성 `value`와 타입이 일치하는지 확인
  LogicalResult verifyInvariants();

  /// Operation을 생성하기 위한 인터페이스
  static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
                    mlir::Type result, mlir::DenseElementsAttr value);  // 속성과 반환 타입 기반 생성
  static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
                    mlir::DenseElementsAttr value);                     // 속성 타입 기반 생성
  static void build(mlir::OpBuilder &builder, mlir::OperationState &state,
                    double value);                                      // 스칼라 값으로 생성
};

toy.constant를 정의하고 나면 다음 이니셜라이저를 통해 Toy Dialect에 등록할 수 있습니다.

void ToyDialect::initialize() {
  addOperations<ConstantOp>();
}

지금까지 Toy Dialect를 선언하고 이에 constant라는 연산을 추가 등록해 보았습니다.

Operation vs Op: Using MLIR Operations

MLIR에서는 Operation과 Op가 다른 의미로 사용됩니다. 이 둘은 크게 보면 MLIR에서 연산을 정의하고 처리하는 데 같은 역할을 하는 것 처럼 보이지만 차이가 있습니다.

  • Operation
    모든 연산을 일반적으로 모델링하기 위한 클래스입니다.
    Opaque하다는 특징을 가집니다. 즉 특정 연산의 속성이나 유형에 대해 설명하지 않습니다.

    대신에 Operation은 연산 인스턴스에 대해 일반적인 API를 제공합니다. 어떤 연산이든 Operation 클래스를 통해서 접근할 수 있으며 이를 통해서 MLIR 내 연산을 범용적으로 다룰 수 있다는 특징이 있습니다.

  • Op
    특정 유형의 연산을 나타내기 위한 클래스입니다.
    예를 들어서.. ConstantOp는 입력이 없고 항상 동일한 값으로 출력되는 연산인데, 이러한 연산을 나타내기 위해 존재합니다. Op 클래스는 내부적으로 Operation 클래스를 참조하며 연산의 데이터는 Operation 클래스에 저장됩니다. Op 클래스는 특정 연산에 특화된 접근자 메서드와 타입 안정성을 제공합니다.

    Op 클래스를 정의하면 Operation 클래스에 대해 의미적으로 유용한 인터페이스를 정의할 수 있습니다. 예를 들어서 ConstantOp는 별도의 클래스 필드를 정의하지 않습니다. 모든 데이터를 내부 Operation에서 가져옵니다.

    MLIR에서는 Op 클래스의 인스턴스를 항상 by value 방식으로 전달합니다. (복사 전달하는 방식)

Using the Operation Definition Specification (ODS) Framework

MLIR의 Operation Definition Specification 프레임워크를 사용하면 MLIR에서 연산을 선언적으로 정의할 수 있습니다. TableGen을 사용하여 연산의 정보를 간결히 정의하고 이를 기반으로 컴파일 시 자동으로 C++ 템플릿으로 확장됩니다. ODS는 연산을 정의하는 간단하고 안정적인 방법입니다.

  • ODS로 ConstantOp 정의하기
    ConstantOp 연산을 ODS로 정의하려면 아래와 같은 단계를 따르면 됩니다.

1. Toy Dialect 기본 클래스 정의
ODS에서 연산은 Op 클래스를 상속받아서 정의됩니다. 연산의 정의를 단순화하기 위해 Toy Dialect의 연산을 위한 기본 클래스를 정의합니다.

class Toy_Op<string mnemonic, list<Trait> traits = []> :
  Op<Toy_Dialect, mnemonic, traits>;

Toy_Op 클래스는 Toy Dialect에서 사용할 연산의 기본 클래스를 정의합니다. 이를 통해서 모든 Toy Dialect의 연산이 공통적인 구조를 가질 수 있습니다.

2. 연산 정의
ConstantOp 연산은 앞서 정의한 Toy_Op를 상속받아 정의됩니다. mnemonic은 ConstantOp::getOperationName에 주어진 이름에서 Dialect 접두어 .toy를 제거한 것과 같습니다.

def ConstantOp : Toy_Op<"constant"> {
}

누락된 ZeroOperand와 OneResults는 나중에 정의할 인자와 결과에 따라 자동으로 추론됩니다.

이때 TableGen에 의해서 C++ 코드가 어떻게 생성되는지 궁금할 수 있는데, 이를 확인하기 위해서는 mlir-tblgen 명령을 gen-op-decls 또는 gen-op-defs 옵션과 함께 실행하면 됩니다.

${build_root}/bin/mlir-tblgen -gen-op-defs
${mlir_src_root}/examples/toy/Ch2/include/toy/Ops.td -I 
${mlir_src_root}/include/

3. 입력 및 출력 정의
구조를 정의했으므로 이제 연산의 입력과 출력을 만들어야 합니다. 연산의 입력 인자는 SSA 피연산자 값에 대한 속성이나 타입일 수 있습니다. 결과는 연산이 생성하는 값들이 타입 집합에 해당합니다.

def ConstantOp : Toy_Op<"constant"> {
 let arguments = (ins F64ElementsAttr:$value); // 64비트 요소 속성
 let results = (outs F64Tensor);              // 64비트 텐서 유형
}

argument -> ConstantOp는 F64ElementsAttr 속성인 $value를 입력으로 받습니다.
results -> ConstantOp는 F64Tensor 타입의 단일 값을 출력합니다.

4. 문서화 추가
연산의 요약 및 설명을 추가할 수 있습니다. 사용자가 연산을 이해하는 데 도움을 줍니다.

def ConstantOp : Toy_Op<"constant"> {
 let summary = "constant operation";
 let description = [{
   This operation converts a literal into an SSA value.
   For example:
   %0 = "toy.constant"() {value = dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf64>}
     : () -> tensor<2x2xf64>
 }];
 let arguments = (ins F64ElementsAttr:$value);
 let results = (outs F64Tensor);
}

5. 검증 추가
ODS는 기본적인 검증 로직을 자동으로 생성합니다. 추가적인 검증이 필요하다면 verifier 필드를 사용해서 C++ 코드로 정의할 수 있습니다.

def ConstantOp : Toy_Op<"constant"> {
  let hasVerifier = 1; // 추가 검증 활성화
}

6. 빌드 메서드 추가
build 메서드는 작업을 생성할 때 필요한 로직을 정의합니다. ODS는 기본적인 빌드 메서드를 자동으로 생성하지만 필요한 경우 사용자 정의 메서드를 추가할 수 있습니다.

def ConstantOp : Toy_Op<"constant"> {
 let builders = [
   OpBuilder<(ins "DenseElementsAttr":$value), [{
     build(builder, result, value.getType(), value);
   }]>,
   OpBuilder<(ins "double":$value)>
 ];
}

이쯤에서 Toy IR을 생성할 수 있습니다. 다음 코드를 Toy IR로 표현해 보겠습니다.

def multiply_transpose(a, b) {
  return transpose(a) * transpose(b);
}

def main() {
  var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
  var b<2, 3> = [1, 2, 3, 4, 5, 6];
  var c = multiply_transpose(a, b);
  var d = multiply_transpose(b, a);
  print(d);
}

위 코드는 다음과 같은 IR을 생성합니다.

module {
  "toy.func"() ({
  ^bb0(%arg0: tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":4:1), %arg1: tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":4:1)):
    %0 = "toy.transpose"(%arg0) : (tensor<*xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:10)
    %1 = "toy.transpose"(%arg1) : (tensor<*xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:25)
    %2 = "toy.mul"(%0, %1) : (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":5:25)
    "toy.return"(%2) : (tensor<*xf64>) -> () loc("test/Examples/Toy/Ch2/codegen.toy":5:3)
  }) {sym_name = "multiply_transpose", type = (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64>} : () -> () loc("test/Examples/Toy/Ch2/codegen.toy":4:1)
  "toy.func"() ({
    %0 = "toy.constant"() {value = dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>} : () -> tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":9:17)
    %1 = "toy.reshape"(%0) : (tensor<2x3xf64>) -> tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":9:3)
    %2 = "toy.constant"() {value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64>} : () -> tensor<6xf64> loc("test/Examples/Toy/Ch2/codegen.toy":10:17)
    %3 = "toy.reshape"(%2) : (tensor<6xf64>) -> tensor<2x3xf64> loc("test/Examples/Toy/Ch2/codegen.toy":10:3)
    %4 = "toy.generic_call"(%1, %3) {callee = @multiply_transpose} : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":11:11)
    %5 = "toy.generic_call"(%3, %1) {callee = @multiply_transpose} : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64> loc("test/Examples/Toy/Ch2/codegen.toy":12:11)
    "toy.print"(%5) : (tensor<*xf64>) -> () loc("test/Examples/Toy/Ch2/codegen.toy":13:3)
    "toy.return"() : () -> () loc("test/Examples/Toy/Ch2/codegen.toy":8:1)
  }) {sym_name = "main", type = () -> ()} : () -> () loc("test/Examples/Toy/Ch2/codegen.toy":8:1)
} loc(unknown)

보면 모든 Toy 연산이 일반적인 어셈블리 포맷을 사용해서 출력됩니다. 이 포맷은 toy.transpose를 분석할 때 본 것과 동일합니다. MLIR은 연산들이 자신의 사용자 정의 어셈블리 포맷을 정의할 수 있게 해줍니다. 사용자 정의 어셈블리 포맷을 정의하면 생성된 IR을 더 읽기 쉽게 바꿀 수 있습니다.

이 방법은 나중에 따로 다뤄보도록 하겠습니다.

Complete Toy Example

이제 우리가 만든 Toy IR을 생성할 수 있습니다! 아래 명령어로 위 예제를 실행해볼 수 있습니다.

toyc-ch2 test/Examples/Toy/Ch2/codegen.toy -emit=mlir -mlir-print-debuginfo

RoundTrip도 가능합니다.

toyc-ch2 test/Examples/Toy/Ch2/codegen.toy -emit=mlir -mlir-print-debuginfo 2> codegen.mlir

다음 명령어로 결과를 다시 실행해볼 수 있습니다.

toyc-ch2 codegen.mlir -emit=mlir

MLIR이 우리가 만든 Toy Dialect와 연산을 인식하게 되었습니다. 다음 장에서는 Dialect를 활용해서 Toy 언어를 위한 고수준 언어별 분석 및 변환을 구현해 보겠습니다.

profile
moze
이전 포스트

Chapter 1: Toy Language and AST

다음 포스트

Chapter 3: High-level Language-Specific Analysis and Transformation

0개의 댓글