Introduction
텐서 계산은 고차원 배열을 계산하는 걸 말한다. GPU에서는 cuDNN이, Intel CPU에서는 MKL이, FPGA에서는 FBLAS가 이를 도와준다. 이들은 모두 사람이 직접 만든 라이브러리다. 그렇기 때문에 알고리즘이 개발되는 것에 비해 느리다는 단점이 있다. 그리고 하드웨어에따라 라이브러리가 각기 다른 것이 있다시피 하드웨어에따라 최적화해줄 수 있는 부분이 다양해지기 때문에 모든 부분을 고려해주기는 힘들다. 그래서 본논문에서는 이를 AutoML을 이용해 다양한 schedule primitives를 fine-tuning해주는 프레임워크인 FlexTensor를 개발했다. Dense tensors에 대해서만 고려하였고, sparse tensor에 대해서는 고려하지 않았다. 경량/가속을 위해 pruning과 quantization이 사용되는데 이 부분은 아쉽다.
FlexTensor
FlexTensor는 프론트와 백엔드로 나뉘어 최적화된 스케쥴을 생성한다. 프론트에서는 통계적인 정보와 구조적인 정보를 이용해 schedule space를 생성한다. 백엔드에서는 휴리스틱한 방법과 머신러닝을 이용해 schedule space를 탐색하여 최적의 schedule을 만들어낸다. 텐서 계산을 할 때는 이를 computation graph의 형태로 저장하여 계산한다. computation graph는 node와 edge로 구성되어 있는데, node는 nested loop를 의미하고 입력과 출력으로 텐서를 갖는다. edge는 텐서 포맷으로 저장된 데이터를 의미한다. node는 spatial loop와 reduce loop로 나뉜다. spatial loop는 data dependency가 없어 병렬화에 적합한 loop를 뜻하고, reduce loop는 data dependency 때문에 병렬화가 어려워 순차적으로 실행해야하는 loop를 뜻한다.
우선 프론트에서는 static analysis를 한다. 이는 computation graph를 인풋으로 받았을 때, 그래프에 대한 구조적인 정보와 통계적인 정보를 분석하는 것을 말한다. 여기서 통계적인 정보는 computation graph의 node에 해당한다. 여기에 포함되는 정보로는 spatial/reduce loop의 개수 및 trip counts, loop order가 있다. 구조적인 정보는 computation graph의 edge에 해당한다. 여기에 포함되는 정보로는 mini-graph내의 edge 개수, 각 node의 인풋/아웃풋 텐서 개수, 각 node에서 consumer node의 개수가 있다.
이렇게 얻은 정보와 schedule primitives와 파라미터간 서로 다른 조합들을 열거하면 schedule space를 생성할 수 있다. 이때 열거 과정은 split, reorder, fuse와 같은 primitives에 대해 먼저 시도하여 같은 개수의 파라미터를 가진 configuration point를 구성한다. 각 point는 벡터로 인코딩돼있고 특정 primitive나 파라미터 선택을 나타낸다. 이렇게 schedule space를 구성하면 space자체가 너무 크기 때문에 효율적인 탐색이 불가능하다. 본 논문에서는 좋은 성능을 보이지 않을 것 같은 space내의 point를 세 가지 방법으로 프루닝하여, 백엔드에서 탐색 전 schedule space의 크기를 줄였다.
먼저, 몇몇 primitive들 간의 조합은 재귀적이다. 그렇기 때문에 영영 끝나지 않을 수도 있어 primitive간 조합의 개수를 제한했다. 두 번째로 split factor가 parameter space에서 큰 부분을 차지하기 때문에 parameter space 자체를 제한했다. 세 번째로 각기 다른 하드웨어에는 중요한 조합은 미리 결정했다. 이전 연구 결과들을 참고하여 CPU에서는 fusion이후에 outer-most loop를 병렬화, inner-most loop를 벡터화하는 것, GPU에서는 outer loop를 스레드 블록에 Inner loop를 스레드에 묶는 거, FPGA에서는 3단계 파이프라이닝을 하는 것과 같은 중요한 조합은 미리 결정하여 고정해두었다.
그리고 structural similarity를 활용하기 위해 space를 재배열했는데, 이는 프루닝 이후에 시작했다. 왜냐하면 pruned point가 탐색 탄계에서는 다시 나타나지 않기 때무이다. 여튼, 위와 같이 schedule space를 구성하면 이는 1차원 리스트로 표현된다. 하지만, 리스트로 표현되는 데이터의 경우에는 탐색할 수 있는 경로가 두 개 밖에 없기 때문에 locality가 떨어진다. 그래서 1차원 리스트로 표현된 것을 고차원 공간에서 표현될 수 있게끔 변환을 시켜준다. 고차원 공간의 local region은 더 많은 point를 갖고 있기 때문에 locality를 잘 활용할 수 있다.
백엔드에서는 프론트에서 생성한 schedule space를 탐색한다. 이때 heuristic method와 machine learning method 두 가지가 사용되는데, 전자는 탐색의 시작점을 고르기 위한 방법이고 후자는 space를 효율적으로 탐색하기 위한 방법이다. 먼저 heuristic method부터 알아보자. 우선, space내의 점 p가 있으면 이것의 성능을 E_p라고 한다. 탐색하는 중에 FlexTensor는 space내에 높은 E값을 갖는 점을 찾고, 이 값을 E*라고 한다. FlexTensors는 p를 exp(-gamma*(E*-E_p)/E*)의 확률로 그 다음 스텝의 시작점으로 선택한다. E_p가 E*에 가까워질 수록 지수가 0에 가까워지고, 이에 따라 확률이 1에 가까워지니 시작점으로 선택될 확률이 높아진다고 생각할 수 있다. 더불어, 한번데 둘 이상의 시작점을 선택할 수도 있다. 이와 같은 방법을 'simulated annealing'이라고 한다.
이제 시작점을 골랐으니 space를 탐색해야할텐데, 가능한 경로가 너무 많다. 이를 효율적으로 탐색하기 위해 강화학습 모델 중 하나인 DQN을 에이전트로 사용해 탐색을 했다. 강화학습 모델은 기본적으로 state, action, reward가 있어야 학습이 가능하다. 본 논문에서는 p를 state, direction d를 action, 그리고 p에서 출발해 direction d를 따라 가면 얻어지는 점을 e로 나타냈다. 여기서 우리가 찾고 싶은 것은 최적의 point에 갈 수 있는 좋은 direction d의 series를 얻는 것이다. 그래서 reward function을 r(p,d)=E_e-E_p로 설정하고 sum(t<T)r(s_t,a_t)=E(s_T)로 잡았다. 그러면, E_(s_T)를 최대화하는 것이 곧, 최적의 point를 찾는 것과 동치가 되기 때문에 에이전트는 학습을 통해 최적의 schedule을 찾을 수 있을 것이다.
heuristic method를 이야기할 때 점 p에 대한 성능을 E_p로 나타낸다고 하였다. 이 성능을 알기 위해서는 두 가지 방법이 있다. 첫 번째는 타겟 디바이스에 실행시켜 실제 성능을 측정하는 방법인데, 이 접근 방법은 실제로 시간이 너무 오래 걸린다. 두 번째는 성능을 추정하기 위해 analytical model을 사용하는 방법이다. 이 방법은 빠르지만 하드웨어에 따라 다양한 옵션들이 있어 model이 다양하게 변할 수 있는 단점이 있다. 본 논문에서는 analytical model을 사용해 탐색 시간을 최대한 줄이려고 했다. 실제로 쓰인 모델의 detail은 논문에 적혀있지 않으나, 대략적으로 workload, processing element의 개수, data write/read time, computation time과 같은 것을 팩터로 사용하였다.
이는 실제로 찾은 schedule에 대해서 이야기하는 부분
Experiment
