2장: 간단한 새로운 변환 연산 추가

새 변환을 추가하기 위한 설정¶

새로운 Transform 연산을 정의하기 전에, 그 구현을 어디에 둘지 결정해야 한다. MLIR은 업스트림 기여를 권장하지만, 변환이 업스트림에 존재하지 않는 트리 외부(out-of-tree) 다이얼렉트에 특화된 경우처럼 메인 Transform 다이얼렉트를 수정하는 것이 항상 가능하거나 바람직한 것은 아니다.

Transform 다이얼렉트는 다이얼렉트 확장(dialect extension) 메커니즘을 사용하여 다이얼렉트 자체를 수정하지 않고도 추가 연산을 주입할 수 있도록 한다. 다이얼렉트 확장은 컨텍스트에 등록되며 다이얼렉트가 로드될 때 함께 로드된다. 확장 정의는 다음과 같이 간단하다.

// In MyExtension.cpp.
#include "mlir/Dialect/Transform/IR/TransformDialect.h"

// 새로운 Transform 다이얼렉트 확장을 정의한다. 이는 확장을 식별하기 위해
// CRTP 관용구를 사용한다.
class MyExtension : public ::mlir::transform::TransformDialectExtension<MyExtension> {
public:
  // 확장은 기본 생성자를 상속해야 한다.
  using Base::Base;

  // 이 함수는 다이얼렉트 정의의 `initialize`와 유사하게 확장을 초기화한다.
  // 여기서 개별 연산과 의존 다이얼렉트를 나열한다.
  void init();
};

void MyExtension::init() {
  // 다이얼렉트와 마찬가지로, 확장은 의존 다이얼렉트를 선언할 수 있다.
  // 이 다이얼렉트는 확장과 함께, 따라서 Transform 다이얼렉트와 함께 로드된다.
  // Transform 연산에서 사용하는 속성(attribute)이나 타입(type)을 포함하는
  // 다이얼렉트만 의존 대상으로 선언하라. 변환 과정에서 생성되는 다이얼렉트는
  // 의존 대상으로 선언하지 말 것.
  //
  // declareDependentDialect<MyDialect>();

  // 변환을 적용하면, 이전에 로드되지 않은 다이얼렉트의 새 연산이 생성될 수 있다.
  // 일반적으로 패스는 이러한 새 연산을 포함하는 다이얼렉트에 의존한다고 선언해야 한다.
  // 확장 자체가 의존하는 다이얼렉트와 혼동을 피하기 위해 Transform 다이얼렉트는
  // 다음을 구분한다:
  //   - 의존 다이얼렉트: Transform 연산이 사용하는 다이얼렉트;
  //   - 생성 다이얼렉트: 변환을 적용하면, 원래 페이로드 IR에 없더라도 생성될 수 있는
  //     엔티티(속성, 연산, 타입)를 포함하는 다이얼렉트.
  // 다음 장에서는 함수 호출과 구조적 제어 흐름 연산을 생성하는 연산을 추가할 것이므로
  // 해당 다이얼렉트를 생성 다이얼렉트로 선언하자.
  declareGeneratedDialect<::mlir::scf::SCFDialect>();
  declareGeneratedDialect<::mlir::func::FuncDialect>();

  // 마지막으로, 추가 Transform 연산을 다이얼렉트에 등록한다.
  registerTransformOps<
    // TODO: 연산 클래스를 나열한다.
  >();
}

연산 자체는 ODS를 사용하여, 다이얼렉트의 일반 연산을 정의하는 것과 정확히 동일한 방식으로 정의할 수 있다.

// In MyExtension.td
#ifndef MY_EXTENSION
#define MY_EXTENSION

include "mlir/Dialect/Transform/IR/TransformDialect.td"
include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
include "mlir/IR/OpBase.td"
include "mlir/Interfaces/SideEffectInterfaces.td"

def MyOp : Op<Transform_Dialect, "transform.my.op", [
    // TODO: 여기에 인터페이스와 트레이트를 추가한다.
   ]> {
  let summary = "my transform op";
  // TODO: 연산 속성을 정의한다.
}

#endif // MY_EXTENSION

다이얼렉트와 마찬가지로, 이 연산들의 헤더와 구현을 생성하기 위해 TableGen을 사용해야 한다. CMake에 다음과 같이 지시할 수 있다.

# In CMakeLists.txt next to MyExtension.td.

# TableGen에 MyExtension.td를 입력으로 사용하도록 지시한다.
set(LLVM_TARGET_DEFINITIONS MyExtension.td)

# ODS로부터 연산 선언과 정의 생성을 요청한다.
mlir_tablegen(MyExtension.h.inc -gen-op-decls)
mlir_tablegen(MyExtension.cpp.inc -gen-op-defs)

# 컴파일 전에 생성이 이루어지도록 의존할 수 있는 CMake 타깃을 추가한다.
add_public_tablegen_target(MyExtensionIncGen)

# 문서 생성도 잊지 말자. 이는 Dialects/ 아래에 MyExtension.md를 생성한다.
add_mlir_doc(MyExtension MyExtension Dialects/ -gen-op-doc)

# In CMakeLists.txt next to MyExtension.cpp
add_mlir_library(
  # 라이브러리 이름은 MyExtension.
  MyExtension

  # 다음 소스 파일들로 빌드한다.
  MyExtension.cpp

  # 이를 컴파일하기 전에 ODS 선언과 정의가 생성되도록 보장한다.
  DEPENDS
  MyExtensionIncGen

  # Transform 다이얼렉트와 모든 생성 다이얼렉트를 링크한다.
  LINK_LIBS PUBLIC
  MLIRTransformDialect
  MLIRFuncDialect
  MLIRSCFDialect
)

이렇게 하면, 각각 Transform 연산의 선언과 정의에 포함되어야 하는 MyExtension.h.inc와 MyExtension.cpp.inc 두 파일이 생성된다.

// In MyExtension.h.
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"

#define GET_OP_CLASSES
#include "MyExtension.h.inc"

// In MyExtension.cpp.

#include "MyExtension.h"

#define GET_OP_CLASSES
#include "MyExtension.cpp.inc"

// …
void MyExtension::init() {
  // …

  // 마지막으로, 추가 Transform 연산을 다이얼렉트에 등록한다.
  // ODS로부터 생성된 모든 연산을 나열한다. 이 호출은 다이얼렉트 인터프리터가
  // 요구하는 Transform 및 메모리 효과 인터페이스를 연산이 구현하는지 추가 검사를 수행하며,
  // 구현하지 않은 경우 assert한다.
  registerTransformOps<
#define GET_OP_LIST
#include "MyExtension.cpp.inc"
  >();
}

Transform 연산 정의¶

이제 함수 호출을 리라이트하는 새로운 Transform 연산을 정의할 준비가 되었다. 이는 다이얼렉트의 일반 연산을 정의하는 것과 동일하다. Transform 다이얼렉트는 연산이 TransformOpInterface와 함께 피연산자가 소비되는지 혹은 읽기 전용인지 표시하기 위한 MemoryEffectsOpInterface를 구현할 것을 요구한다. 우리의 연산은 다음과 같은 방식으로 정의할 수 있다.

// In MyExtension.td.

// 새로운 연산을 정의한다. 관례에 따라, 이름 앞에 다이얼렉트 확장 이름인 "my."를 붙인다.
// 전체 연산 이름은 여기에 "transform." 접두사가 더해진다.
def ChangeCallTargetOp : Op<Transform_Dialect, "my.change_call_target",
    // 이 연산이 필요한 TransformOpInterface와 MemoryEffectsOpInterface를
    // 구현함을 나타낸다.
    [DeclareOpInterfaceMethods<TransformOpInterface>,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
  // 간단한 설명과 전체 설명을 제공한다. 후자는 피연산자에 미치는 영향과
  // 연산이 다양한 실패 모드를 어떻게 처리하는지 설명하는 것을 권장한다.
  let summary = "Changes the callee of a call operation to the specified one";
  let description = [{
    For each `func.call` payload operation associated with the handle, changes
    its callee to be the symbol whose name is provided as an attribute to this operation.

    Generates a silenceable failure if the operand is associated with payload operations that are not `func.call`. Only reads the operand.
  }];

  // 인자는 페이로드 연산에 대한 핸들과 새로운 피호출자(callee)를 지정하는 속성을 포함한다.
  // 핸들은 TransformHandleTypeInterface를 구현해야 한다.
  // 심볼이 transform IR에 존재하지 않을 수도 있으므로 검증이 실패할 수 있어 문자열 속성을 사용한다.
  let arguments = (ins
    TransformHandleTypeInterface:$call,
    StrAttr:$new_target);

  // 결과는 비어있다. 이 변환은 새로운 페이로드를 생성하지 않기 때문이다.
  let results = (outs);

  // 읽기 쉬운 어셈블리 구문을 제공한다.
  let assemblyFormat = "$call `,` $new_target attr-dict `:` type($call)";
}

Transform 연산 정의를 마무리하려면 인터페이스 메서드를 구현해야 한다. TransformOpInterface는 현재 실제 변환을 수행하는 하나의 메서드 apply만을 요구한다. 이 메서드의 본문은 Transform 다이얼렉트 구성 요소 조작으로 한정하고, 실제 변환은 독립 함수로 구현하여 코드의 다른 곳에서도 사용할 수 있도록 하는 것이 좋다. 리라이트 패턴과 유사하게, 모든 IR 수정은 제공된 rewriter를 통해 수행해야 한다.

// In MyExtension.cpp

// 우리의 Transform 다이얼렉트 연산 구현.
// 이 연산은 다음 셋 중 하나의 삼진 결과를 반환한다:
// - 성공: 변환이 성공했을 때;
// - 명백한 실패(definite failure): 변환이 후속 변환이 불가능하거나 바람직하지 않은 방식으로
//   실패했을 때. 보통 페이로드 IR을 잘못된 상태로 남겨둘 수 있다. 이 경우 반환 직전에
//   진단(diagnostic)을 즉시 방출하는 것이 기대된다;
// - 소거 가능한 실패(silenceable failure): 변환이 실패했지만 후속 변환은 여전히 적용 가능할 때.
//   일반적으로 이는 변환의 전제조건이 충족되지 않아 페이로드 IR이 수정되지 않았음을 의미한다.
//   소거 가능한 실패는 사용자가 방출할 수 있는 Diagnostic을 추가로 담는다.
::mlir::DiagnosedSilenceableFailure mlir::transform::ChangeCallTargetOp::apply(
    // IR을 수정할 때 사용해야 하는 rewriter.
    ::mlir::transform::TransformRewriter &rewriter,
    // 이 Transform 연산이 정의하는 transform IR 값과 연관될 페이로드 IR 엔티티 목록.
    // 이 경우 결과가 없으므로 비어 있을 수 있다.
    ::mlir::transform::TransformResults &results,
    // Transform 적용 상태. 이 객체는 transform IR 값과 페이로드 IR 엔티티 간의 현재
    // 연관을 조회하는 데 사용할 수 있다. 또한 사용자 정의 상태를 전달할 수도 있다.
    ::mlir::transform::TransformState &state) {

  // 먼저, 피연산자 핸들과 연관된 페이로드 연산 목록을 얻는다.
  auto payload = state.getPayloadOps(getCall());

  // 그런 다음 피연산자 목록을 순회하며 실제 IR 변경 함수를 호출한다.
  // 이곳에서 전제조건 검사도 수행한다.
  for (Operation *payloadOp : payload) {
    auto call = dyn_cast<::mlir::func::CallOp>(payloadOp);
    if (!call) {
      DiagnosedSilenceableFailure diag = emitSilenceableError()
          << "only applies to func.call payloads";
      diag.attachNote(payloadOp->getLoc()) << "offending payload";
      return diag;
    }

    updateCallee(call, getNewTarget());
  }

  // 모든 것이 잘 진행되었으면 성공을 반환한다.
  return DiagnosedSilenceableFailure::success();
}

MemoryEffectsOpInterface 구현은 이 연산이 피연산자(소비 또는 읽기 전용)와 페이로드 IR(변경 또는 읽기 전용)에 대해 갖는 효과를 지정해야 한다. Transform 다이얼렉트 검증기는 부작용(side effect)의 존재를 확인하며, 없을 경우 디버그 빌드에서 assert한다.

// In MyExtension.cpp

void ChangeCallTargetOp::getEffects(
    ::llvm::SmallVectorImpl<::mlir::MemoryEffects::EffectInstance> &effects) {
  // 이 연산이 `call` 핸들을 읽기만 함을 나타낸다. 연관된 연산은 삭제되지 않고
  // 제자리에서 수정되므로 그 참조는 유효하게 남는다.
  onlyReadsHandle(getCall(), effects);

  // 이 연산이 페이로드를 수정함을 나타낸다.
  modifiesPayload(effects);
}

등록 및 사용¶

Transform 연산을 정의하기에는 이 정도면 충분하다. 남은 것은 프로젝트의 main에서 호출할 수 있는 확장 등록 훅을 제공하는 일이다.

// In TransformDialect.cpp (TransformDialect.h에 선언도 잊지 말자);

void registerMyExtension(::mlir::DialectRegistry &registry) {
  registry.addExtensions<MyExtension>();
}

확장을 등록하면 Transform 다이얼렉트 인터프리터에서 우리의 새 연산을 사용할 수 있다. 업스트림 테스트용 패스는 그대로 사용할 수 있다. 실제로 mlir/test/Examples/transform/Ch2/sequence.mlir에 존재하며, 여기에는 microkernel 구현이 들어 있다.

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(
      %arg0: !transform.any_op,
      %arg1: !transform.op<"linalg.matmul">,
      %arg2: !transform.op<"linalg.elementwise">) {
    // %arg2 핸들은 두 개의 elementwise 연산 모두와 연관되어 있으므로,
    // 두 개의 핸들로 분할해 두 번째 elementwise 연산만을 대상으로 할 수 있도록 한다.
    %add, %max = transform.split_handle %arg2
        : (!transform.op<"linalg.elementwise">)
        -> (!transform.any_op, !transform.any_op)

    // 실제 타일링 변환은 타일 크기를 속성으로 받는다.
    // 타일링 중에 생성된 루프에 대한 핸들을 생성한다.
    %tiled, %loop = transform.structured.tile_using_forall %max
                    tile_sizes [8, 32]
        : (!transform.any_op) -> (!transform.any_op, !transform.any_op)

    // 이제 다른 연산들을 루프 안으로 퓨즈할 수 있다. 여기서는 연산을 하나씩 퓨즈한다.
    // 이때 퓨즈되는 연산이 루프 안에서 사용되는 값을 정의해야 하므로, 이러한 퓨즈의
    // 순서가 중요하다. 또한 "transform.merge_handles"를 사용해 모든 연산에 대한
    // 단일 핸들을 얻고, 해당 경우에는 `fuse_into_containing_op`가 순서를 알아서
    // 처리하도록 할 수도 있다.
    %add_fused, %loop2 = transform.structured.fuse_into_containing_op %add into %loop
        : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
    %matmul_fused, %loop3 = transform.structured.fuse_into_containing_op %arg1
                    into %loop2
        : (!transform.op<"linalg.matmul">, !transform.any_op)
       -> (!transform.any_op, !transform.any_op)

    // 원하는 크기를 얻기 위해 다시 타일링한다. 이번에는 "add" 연산을 타일링하고
    // matmul을 루프로 퓨즈하지만, "max" 연산에는 영향이 없다. 이는 transform
    // 다이얼렉트를 통해 정밀 타기팅을 할 수 있음을 보여준다. 그렇지 않으면 같은 종류를 가진
    // "add"와 "max"를 구분하기 어렵다.
    %tiled_second, %loop_second = transform.structured.tile_using_forall %add_fused
                        tile_sizes [4, 4]
        : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
    %matmul_fused_2, %loop_second_2 = transform.structured.fuse_into_containing_op %matmul_fused
                      into %loop_second
        : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)

    // 현재 아웃라이닝은 루프와 같은 리전 보유 연산에만 구현되어 있으므로,
    // 크기 1로 타일링하여 아웃라인할 바깥 루프를 구체화(materialize)한다.
    %_0, %loop_third = transform.structured.tile_using_forall %tiled_second tile_sizes [1]
        : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
    %_1, %outline_target = transform.structured.fuse_into_containing_op %matmul_fused_2 into %loop_third
        : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
    %func, %call = transform.loop.outline %outline_target
                   {func_name = "outlined"}
        : (!transform.any_op) -> (!transform.any_op, !transform.any_op)

    // 호출 대상 리라이트.
    transform.my.change_call_target %call, "microkernel" : !transform.any_op

    transform.yield
  }
}

부록: 자동 생성 문서¶

transform.my.change_call_target (transform::ChangeCallTargetOp)¶

호출 연산의 callee를 지정된 대상으로 변경한다

문법:

operation ::= `transform.my.change_call_target` $call `,` $new_target attr-dict `:` type($call)

핸들과 연관된 각 func.call 페이로드 연산에 대해, 이 연산에 속성으로 제공된 이름을 가진 심볼로 그 callee를 변경한다.

피연산자가 func.call이 아닌 페이로드 연산과 연관되어 있으면 소거 가능한 실패를 발생시킨다. 피연산자는 읽기만 한다.

Interfaces: MemoryEffectOpInterface, TransformOpInterface

속성:¶

피연산자:¶