IR 구조 이해 - MLIR

IR 구조 이해

MLIR 언어 레퍼런스는 고수준 구조를 설명합니다. 이 문서는 예제를 통해 해당 구조를 보여주고, 동시에 이를 조작하는 데 사용되는 C++ API를 소개합니다.

우리는 어떤 MLIR 입력이든 순회하면서 IR 내부의 엔티티를 출력하는 패스를 구현할 것입니다. 패스(또는 일반적으로 거의 모든 IR 조각)는 항상 하나의 operation을 루트로 합니다. 대부분의 경우 최상위 operation은 ModuleOp이며, MLIR PassManager는 실제로 최상위 ModuleOp에 대한 operation으로 제한됩니다. 따라서 패스는 하나의 operation에서 시작하며, 우리의 순회도 다음과 같이 시작합니다:

void runOnOperation() override {
    Operation *op = getOperation();
    resetIndent();
    printOperation(op);
  }

IR 중첩 구조 순회¶

IR은 재귀적으로 중첩됩니다. 하나의 Operation은 하나 이상의 중첩된 Region을 가질 수 있고, 각 Region은 Block의 목록이며, 각 Block은 다시 Operation의 목록을 감쌉니다. 우리의 순회는 이 구조를 따라 printOperation(), printRegion(), printBlock()의 세 가지 메서드로 진행됩니다.

첫 번째 메서드는 operation의 속성을 살펴본 다음, 중첩된 region을 순회하며 각각 출력합니다:

void printOperation(Operation *op) {
    // operation 자체와 일부 속성을 출력
    printIndent() << "visiting op: '" << op->getName() << "' with "
                  << op->getNumOperands() << " operands and "
                  << op->getNumResults() << " results\n";
    // operation의 attribute 출력
    if (!op->getAttrs().empty()) {
      printIndent() << op->getAttrs().size() << " attributes:\n";
      for (NamedAttribute attr : op->getAttrs())
        printIndent() << " - '" << attr.getName() << "' : '"
                      << attr.getValue() << "'\n";
    }

    // operation에 연결된 각 region으로 재귀 진입
    printIndent() << " " << op->getNumRegions() << " nested regions:\n";
    auto indent = pushIndent();
    for (Region &region : op->getRegions())
      printRegion(region);
  }

Region은 Block 목록 외에는 아무것도 보유하지 않습니다:

void printRegion(Region &region) {
    // region 자체는 block 목록 외에는 아무것도 보유하지 않음
    printIndent() << "Region with " << region.getBlocks().size()
                  << " blocks:\n";
    auto indent = pushIndent();
    for (Block &block : region.getBlocks())
      printBlock(block);
  }

마지막으로, Block은 인자 목록을 가지며 Operation 목록을 보유합니다:

void printBlock(Block &block) {
    // 블록의 고유 속성(주로: 인자 목록)을 출력
    printIndent()
        << "Block with " << block.getNumArguments() << " arguments, "
        << block.getNumSuccessors()
        << " successors, and "
        // 참고: 이 `.size()`는 연결 리스트를 순회하므로 O(n)입니다.
        << block.getOperations().size() << " operations\n";

    // 블록의 주된 역할은 Operation 목록을 보유하는 것: 각 operation을 재귀적으로 출력
    auto indent = pushIndent();
    for (Operation &op : block.getOperations())
      printOperation(&op);
  }

해당 패스의 코드는 레포지토리 여기에서 확인할 수 있으며, mlir-opt -test-print-nesting으로 실행해 볼 수 있습니다.

예시¶

앞서 소개한 패스는 다음 IR에 mlir-opt -test-print-nesting -allow-unregistered-dialect llvm-project/mlir/test/IR/print-ir-nesting.mlir로 적용할 수 있습니다:

"builtin.module"() ( {
  %results:4 = "dialect.op1"() {"attribute name" = 42 : i32} : () -> (i1, i16, i32, i64)
  "dialect.op2"() ( {
    "dialect.innerop1"(%results#0, %results#1) : (i1, i16) -> ()
  },  {
    "dialect.innerop2"() : () -> ()
    "dialect.innerop3"(%results#0, %results#2, %results#3)[^bb1, ^bb2] : (i1, i32, i64) -> ()
  ^bb1(%1: i32):  // pred: ^bb0
    "dialect.innerop4"() : () -> ()
    "dialect.innerop5"() : () -> ()
  ^bb2(%2: i64):  // pred: ^bb0
    "dialect.innerop6"() : () -> ()
    "dialect.innerop7"() : () -> ()
  }) {"other attribute" = 42 : i64} : () -> ()
}) : () -> ()

그리고 다음과 같은 출력이 생성됩니다:

visiting op: 'builtin.module' with 0 operands and 0 results
 1 nested regions:
  Region with 1 blocks:
    Block with 0 arguments, 0 successors, and 2 operations
      visiting op: 'dialect.op1' with 0 operands and 4 results
      1 attributes:
       - 'attribute name' : '42 : i32'
       0 nested regions:
      visiting op: 'dialect.op2' with 0 operands and 0 results
      1 attributes:
       - 'other attribute' : '42 : i64'
       2 nested regions:
        Region with 1 blocks:
          Block with 0 arguments, 0 successors, and 1 operations
            visiting op: 'dialect.innerop1' with 2 operands and 0 results
             0 nested regions:
        Region with 3 blocks:
          Block with 0 arguments, 2 successors, and 2 operations
            visiting op: 'dialect.innerop2' with 0 operands and 0 results
             0 nested regions:
            visiting op: 'dialect.innerop3' with 3 operands and 0 results
             0 nested regions:
          Block with 1 arguments, 0 successors, and 2 operations
            visiting op: 'dialect.innerop4' with 0 operands and 0 results
             0 nested regions:
            visiting op: 'dialect.innerop5' with 0 operands and 0 results
             0 nested regions:
          Block with 1 arguments, 0 successors, and 2 operations
            visiting op: 'dialect.innerop6' with 0 operands and 0 results
             0 nested regions:
            visiting op: 'dialect.innerop7' with 0 operands and 0 results
             0 nested regions:

기타 IR 순회 방법¶

많은 경우 IR의 재귀적 구조를 직접 풀어내어 다루는 것은 번거롭습니다. 이럴 때는 다른 헬퍼들을 사용할 수 있습니다.

필터링된 이터레이터: getOps<OpTy>()¶

예를 들어 Block 클래스는 필터링된 이터레이터를 제공하는 편리한 템플릿 메서드 getOps<OpTy>()를 노출합니다. 예시는 다음과 같습니다:

auto varOps = entryBlock.getOps<spirv::GlobalVariableOp>();
  for (spirv::GlobalVariableOp gvOp : varOps) {
     // 블록에 있는 각 GlobalVariable operation을 처리
     ...
  }

마찬가지로, Region 클래스도 동일한 getOps 메서드를 제공하며, region에 있는 모든 블록을 대상으로 순회합니다.

워커(walkers)¶

getOps<OpTy>()는 단일 블록(또는 단일 region) 바로 아래에 나열된 일부 Operation을 순회할 때 유용하지만, IR을 중첩된 방식으로 순회하는 것이 더 흥미로운 경우가 자주 있습니다. 이를 위해 MLIR은 Operation, Block, Region에 walk() 헬퍼를 제공합니다. 이 헬퍼는 하나의 인자, 즉 제공된 엔티티(그리고 그 엔티티 자체)에 재귀적으로 중첩된 모든 operation에 대해 호출될 콜백 함수를 받습니다.

// 함수 내부에 중첩된 모든 region과 block을 재귀적으로 순회하고
  // 모든 operation에 대해 포스트오더로 콜백을 적용
  getFunction().walk([&](mlir::Operation *op) {
    // Operation `op`를 처리
  });

제공된 콜백은 특정 종류의 Operation만 필터링하도록 특수화할 수 있습니다. 예를 들어, 다음은 함수 내부에 중첩된 LinalgOp에 대해서만 콜백을 적용합니다:

getFunction().walk([](LinalgOp linalgOp) {
    // LinalgOp `linalgOp` 처리
  });

마지막으로, 콜백은 WalkResult::interrupt() 값을 반환하여 순회를 중단할 수도 있습니다. 예를 들어, 다음 순회는 함수 내부에 중첩된 모든 AllocOp를 찾되, 그중 하나라도 기준을 만족하지 않으면 순회를 중단합니다:

WalkResult result = getFunction().walk([&](AllocOp allocOp) {
    if (!isValid(allocOp))
      return WalkResult::interrupt();
    return WalkResult::advance();
  });
  if (result.wasInterrupted())
    // 하나의 alloc이라도 일치하지 않음
    ...

def-use 체인 순회¶

IR에서 또 다른 관계는 Value와 그 사용자들 간의 연결입니다. 언어 레퍼런스에 정의된 대로, 각 Value는 BlockArgument이거나 정확히 하나의 Operation 결과입니다(하나의 Operation은 여러 결과를 가질 수 있으며, 각 결과는 별도의 Value입니다). Value의 사용자들은 Operation이며, 그들의 피연산자(operand)를 통해 참조합니다: 각 Operation의 피연산자는 하나의 Value를 참조합니다.

다음은 Operation의 피연산자를 검사하고 이에 대한 정보를 출력하는 코드 예시입니다:

// 각 피연산자의 생성자(producer)에 대한 정보를 출력
  for (Value operand : op->getOperands()) {
    if (Operation *producer = operand.getDefiningOp()) {
      llvm::outs() << "  - Operand produced by operation '"
                   << producer->getName() << "'\n";
    } else {
      // 정의하는 op이 없다면, 해당 Value는 반드시 Block 인자입니다.
      auto blockArg = cast<BlockArgument>(operand);
      llvm::outs() << "  - Operand produced by Block argument, number "
                   << blockArg.getArgNumber() << "\n";
    }
  }

유사하게, 다음 코드 예시는 Operation이 생성한 결과 Value들을 순회하고, 각 결과의 사용자들을 순회하며 그 정보들을 출력합니다:

// 각 결과의 사용자에 대한 정보를 출력
  llvm::outs() << "Has " << op->getNumResults() << " results:\n";
  for (auto indexedResult : llvm::enumerate(op->getResults())) {
    Value result = indexedResult.value();
    llvm::outs() << "  - Result " << indexedResult.index();
    if (result.use_empty()) {
      llvm::outs() << " has no uses\n";
      continue;
    }
    if (result.hasOneUse())
      llvm::outs() << " has a single use: ";
    else
      llvm::outs() << " has " << result.getNumUses() << " uses:\n";
    for (Operation *userOp : result.getUsers()) {
      llvm::outs() << "    - " << userOp->getName() << "\n";
    }
  }

이 패스를 위한 예시 코드는 레포지토리 여기에서 확인할 수 있으며, mlir-opt -test-print-defuse로 실행해 볼 수 있습니다.

Value들과 그 사용(use)들의 연결은 다음과 같이 볼 수 있습니다:

Value의 사용들(OpOperand 또는 BlockOperand) 역시 이중 연결 리스트로 체이닝되어 있습니다. 이는 한 Value의 모든 사용을 새로운 값으로 교체(“RAUW”)할 때 특히 유용합니다: