원문 내용 및 사진 출처는 https://tvm.apache.org/docs/dev/relay_op_strategy.html 에서 확인할 수 있습니다.

여기서는 relay lowering strategy에 대해 자세히 다루어보겠습니다. Relay operator를 TOPI 라이브러리로 lowering 하려면, 각각의 compute/schedule function이 relay operator에 미리 정의가 되어있어야 합니다. 하지만 compute/schedule function은 대개 특정 타깃이 정해져 있고, 또 같은 타깃에 대해 다양한 알고리즘으로 나타낼 수 있습니다.

Operator Strategy Design

Write a Strategy Function

OpStrategy class는 strategy에 implementation을 추가하는 API 하나만 가지고 있습니다:

def add_implementation(self, compute, schedule, name="default", plevel=10)

본문을 보면 strategy에 하나의 implementation을 추가하는 예시, 여러개 implementation을 추가하는 예시, third-party library implementation으로 추가하는 예시, 특정 shape에서만 구동되는 implementation을 추가하는 예시를 들고 있습니다.

Register Strategy Function to An Operator

위에서는 strategy function을 작성하는 방법에 대해 배웠고, 여기서는 그 strategy function을 operator에 적용하는 방법에 대해 다룹니다. operator에 strategy function을 등록하려면 다음과 같이 하면 됩니다:

register_strategy("topk", strategy.topk_strategy)

하지만, operator에 대해 일일이 추가하는 것이 귀찮기 때문에, 다른 두 가지 간단한 방법들을 추가로 소개합니다:

  1. injective, broadcast, reduction pattern을 갖는 operator의 경우

    각각 register_injective_schedule, register_broadcast_schedule, register_reduce_schedule을 사용하면 됩니다. 이러한 패턴을 갖는 schedule function의 경우에는 각각의 target에서 미리 등록을 해 두었고, operator에 바로 적용이 가능합니다.

    register_broadcast_schedule("add")
    
  2. 그 외 operator의 경우

    모든 target에 대해 같은 compute function을 갖는다면, register_schedule API를 사용하면 됩니다.

    어떤 schedule function을 사용할지 지정만 하면 되기 때문에, FTVMSchedule 함수를 작성하는 것이 더 쉬울 수도 있습니다. 본문을 보면 pooling을 예로 든 예시가 나옵니다.

Register Strategies for a new target

새로운 target에 대해 strategy를 추가하는 방법은 크게 두 가지가 있습니다.

  1. 새로운 target file을 python/tvm/relay/op/strategy 디렉토리에 추가하기

    새로운 target에서 사용되기 위해 만들어진 op의 strategy를 customize하고, 나머지는 기존의 generic strategy를 사용하면 됩니다.