그림 (참고1)

TensorFlow runtime 내의 원기둥은 tensorflow graph → XLA graph 가 가능하도록 준비된 연산들이다. 그림에서 add 는 1:1 대응되는 커널이고, softmax 는 1:3 대응되는 커널이라고 생각할 수 있다.

파이프라인을 정리하면 다음과 같다.

Python 등 API 로 TensorFlow Graph 생성 → 그래프의 일부 또는 전체를 clustering (clustering 된 부분을 하나의 노드라고 여김) → 런타임에 클러스터링된 노드를 마주치는 경우 클러스터를 JIT 컴파일하여 실행하고 캐싱

모든 TensorFlow ops 에 대해서 컴파일이 불가능하기 때문에 XLA 가 고민해야 하는 양이 많아진다. 그래프의 크기가 크고 복잡해질수록 이 문제는 더욱 복잡해질 것 같다 (참고4). 게다가 클러스터를 잘못 잡으면 자칫 사이클 그래프를 만들어 데드락이 걸릴수도 있다 (참고3).

XLA 가 운영하지 못하는 연산에는 어떤 것이 있을까? 내가 TPU 구현체를 만드는 과정에서 어려웠던 점이 바로 이것과 관련이 있지는 않을까?

왜 그래프의 일부가 아닌 전체 그래프를 XLA 컴파일러로 컴파일하면 안되는걸까?

최적화할 코드가 복잡하면 왜 JIT 의 퍼포먼스가 잘 나오지 않는 것일까?

참고

1. 22:20, 모든 TensorFlow 연산을 XLA 컴파일 가능하도록 준비해놓지는 않았다. 너무 동적으로 움직이기 때문에 바꾸기 까다롭기 때문이다.
2. 23:25, 일반적인 TensorFlow Graph 가 있다고 쳐 보자. Graph 에서 일부 노드들을 골라 묶어버린다 (cluster). 이 클러스터만 XLA 에 의해 컴파일될 수 있게 할 것이다. 다른 노드들이 선택되지 않은 이유는, 뭐 많이 있겠지만, XLA 로 컴파일되기 어렵기 때문이라고 치자. TensorFlow 는 그냥 평소처럼 하나씩 실행시켜 나가는데, 런타임에 XLA 로 컴파일되도록 묶인 부 그래프를 하나로 보고, 부 그래프를 마주쳤을 때 JIT 컴파일하고 머신 코드를 캐싱한다. 캐싱함으로써 부 그래프가 하나의 모델에서 반복적으로 등장할 때 계속 사용할 수 있게 된다.
3. 26:00, XLA JIT 컴파일 시 생길 수 있는 문제로 데드락이 있다. 클러스터링을 잘못하게 되면, 그래프가 사이클을 만들어서 컴파일자체를 할 수 없게 될지도 모른다.
4. ‣