원문 내용 및 사진 출처는 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

• OpImplementation

  OpImplementation이 모여 OpSpecialization이 되고, OpSpecialization이 모여 OpStrategy가 됩니다.

  즉, OpImplementation이 기본 operator strategy 역할을 수행합니다. OpImplementation은 <compute/schedule function, implementation 이름, priority level>을 포함합니다.

  OpImplementation은 또 SpecializedCondition과 연관되어 있습니다. SpecializedCondition은 null 값이 될 수도 있고, 이때는 implementation이 generally applicable합니다. null이 아니라면, 특정 조건이 만족된 경우에만 applicable합니다. SpecializedCondition은 CNF(conjunctive normal form)로 표현된 Tensor Expression의 리스트로 구성되어, 텐서의 shape에 대해 다루고 있습니다.

• FTVMStrategy

  FTVMStrategy라는 strategy function은 workload가 주어졌을 때 어떤 compute/schedule function pair가 사용되어야 할지를 결정합니다. 각각의 Relay operator마다 FTVMStrategy가 register 되어 있어야 합니다. FTVMStrategy는 generic function으로, 다음과 같은 구조를 갖습니다.

  def OpStrategy(const Attrs& attrs, const Array<Tensor>& inputs, const Type& out_type, const Target& target) -> OpStrategy
  

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를 사용하면 됩니다.