<div id="readability-page-1" class="page"><div id="content">
		<main id="main" role="main">

		
<article id="post-3022">
	
	

	<div>
		최근 업무 관련해서 XLA를 알게 되었는데, 재밌는 프로젝트인데 반해 관련된 자료를 찾는게 쉽지 않길래 한번 소개를 해보면 좋을 것 같다는 생각이 들었습니다.
<img loading="lazy" src="https://b.mytears.org/wp-content/uploads/2019/08/xla-logo.png" alt="" width="455" height="456" sizes="(max-width: 455px) 100vw, 455px">
우선 <a href="https://www.tensorflow.org/xla">XLA</a> (Accelerated Linear Algebra)는 <a href="https://www.tensorflow.org/">Tensorflow</a>의 서브 프로젝트로 그래프 연산의 최적화 / 바이너리 사이즈의 최소화 등을 목적으로 하는 컴파일러입니다.
결과를 먼저 말씀드리자면, XLA를 활용하는 것을 통해 그래프 연산 과정에 필요한 임시 버퍼도 덜 사용할 수 있고, 그래프 연산 속도도 단축시킬 수 있습니다. (<a href="https://developers.googleblog.com/2017/03/xla-tensorflow-compiled.html">구글 측 자료</a>에 의하면 Nvidia GPU를 사용할 때 50%까지도 성능 향상이 있었다고 합니다.)
예를 들어 아래와 같은 그래프가 있다고 했을때…
<img src="https://b.mytears.org/wp-content/uploads/2019/08/xlaGraphExample.png" alt="">
각 연산마다 룹을 돌리는 것이 아니라, 아래와 같은 형식으로 연산을 fuse해서 처리하게 될 경우, 임시 데이터 저장도 줄어들고, 룹도 하나의 룹으로 묶어서 처리가 가능하기 때문에 연산 속도나 임시 버퍼 사용량(및 read/write bandwidth), 캐쉬 효율 측면에서 여러가지 이점을 얻을 수 있습니다.
<img src="https://b.mytears.org/wp-content/uploads/2019/08/xlaGraphExampleFused.png" alt="">
XLA에 있는 그래프 최적화기가 하는 일이 위와 같은 것들이며, XLA에서 하는 일을 간단히 도식화 하면 아래와 같습니다.
<img loading="lazy" src="https://b.mytears.org/wp-content/uploads/2019/08/xlaCompilerFlow.png" alt="" width="507" height="57" sizes="(max-width: 507px) 100vw, 507px">
XLA도 일종의 컴파일러이기 때문에 일반적인 컴파일러 들이 동작하는 것처럼 FrontEnd (Parser) / MiddleEnd (High-level Optimizer) / BackEnd (Target Code Generator)의 구조를 가지고 있으며, 주어진 타겟 플랫폼에 맞춰 그래프를 최적화 한 후 실행가능한 형태의 무언가를 생성하는 역할을 수행합니다.
BackEnd에서는 LLVM을 사용하고 있으며, LLVM과 마찬가지로 JiT compilation과 AoT compilation 모두를 지원하는게 특징입니다. 두가지 중 어떤 방식을 사용하느냐에 따라 실제 사용 시나리오가 달라지게 되며,
AoT compilation을 사용하는 경우 아래와 같이 컴파일의 결과로 object file 및 이를 사용하기 위한 helper class에 대한 header file이 생성됩니다.
<img loading="lazy" width="702" height="218" src="https://b.mytears.org/wp-content/uploads/2019/08/xlaAotFlow.png" alt="" sizes="(max-width: 702px) 100vw, 702px">
컴파일된 그래프는 아래와 같은 형태의 코드와 링크를 시켜 바이너리로 만든 후 실행할 수 있습니다. (참고: <a href="https://www.tensorflow.org/xla/tfcompile">Using AOT compilation</a>)


		<div id="crayon-60af7430b6d28676298697" data-settings=" touchscreen minimize scroll-mouseover">
				<table>
					<tbody><tr>
				<td data-settings="show">
					
				</td>
						<td><div>&nbsp;&nbsp;BlahGraph graph;&nbsp;&nbsp;// Set up args and run the computation.&nbsp;&nbsp;const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};&nbsp;&nbsp;std::copy(args + 0, args + 6, graph.arg0_data());&nbsp;&nbsp;std::copy(args + 6, args + 12, graph.arg1_data());&nbsp;&nbsp;graph.Run();&nbsp;&nbsp;std::cout &amp;lt;&amp;lt; graph.result0(0, 0) &amp;lt;&amp;lt; std::endl;</div></td>
					</tr>
				</tbody></table>
			</div>


실행 바이너리를 만들기 위해 XLA 라이브러리나 Tensorflow 라이브러리 등을 포함할 필요가 없기 때문에 최종 타겟 바이너리 사이즈를 굉장히 작게 만들 수 있으나 기본 CPU backend의 경우 convolution, matmul 등의 연산에 대해 XLA runtime function을 호출하는 코드를 생성하도록 되어 있기 때문에 일부 XLA runtime function들과 링크되어야 할 수도 있습니다.
JiT compilation을 사용하는 경우는 compile의 결과로 XlaExecutable 데이터가 생성되며, StreamExecutor를 통해 input parameter를 feeding하고 그래프를 실행시킬 수 있는 형태입니다.
<img loading="lazy" width="702" height="341" src="https://b.mytears.org/wp-content/uploads/2019/08/xlaJitFlow.png" alt="" sizes="(max-width: 702px) 100vw, 702px">
JitCompilation을 사용하는 경우 런타임에 그래프를 입력으로 받아 실행할 수 있다는 장점이 있으나 애플리케이션이 XLA 런타임을 포함해야하기 때문에 eigen, protobuf, LLVM 등등 의존성들이 줄줄이 엮인다는 단점이 있습니다. (타겟 플랫폼에 해당 라이브러리들을 줄줄이 포팅해야할 수도 있습니다.)
애플리케이션을 배포하기 전 단계에 타겟 플랫폼이나 네트워크 형태에 대해 미리 알고 있는 경우에는 AoT compilation을 활용하는 방안을 고려해볼 수 있으며, 애플리케이션은 네트워크를 모르는 상태로 배포되어야 하는 환경이라면 JiT compilation을 사용해야합니다.
Tensorflow 외에도 다양한 머신러닝 프레임웍의 백엔드로 XLA를 활용할 수 있으며, <a href="https://pytorch.org/">PyTorch</a>를 위해선 <a href="https://github.com/pytorch/xla">PyTorch XLA</a>, <a href="https://github.com/onnx">ONNX (Open Neural Network Exchange )</a>를 위해선 <a href="https://github.com/onnx/onnx-xla">onnx xla</a> 프로젝트가 개발되고 있습니다..
<img loading="lazy" src="https://b.mytears.org/wp-content/uploads/2019/08/xlaOverview.png" alt="" width="545" height="587" sizes="(max-width: 545px) 100vw, 545px">
Service / Client 구조나 기타 여러가지 재밌는 부분들이 많이 남아있지만 소개글이니 우선은 이 정도로 시작해볼까 합니다.
관심 있으신 분들 있으면 이어서 다른 내용들도 정리를 해보도록 하겠습니다. 😀
추가: 얼마전 <a href="https://github.com/ConstantPark/DL_Compiler">ML Compiler 스터디</a>에서 관련 내용을 발표한 자료를 하나 진행했으니 관심있으신 분들은 참고하셔도 좋을 것 같군요.


<ul><li>PPT: <a href="https://github.com/.../XLA101.pdf">https://github.com/…/XLA101.pdf</a></li><li>Video: <a href="https://youtu.be/_3ykXQH5h2o" data-type="URL" data-id="https://youtu.be/_3ykXQH5h2o">https://youtu.be/_3ykXQH5h2o</a></li></ul>




<h3>Similar Posts:</h3></div>

	
	

</article>



	
		</main>
	</div></div>

내 맘대로 보는 세상

Tee Jung

최근 업무 관련해서 XLA를 알게 되었는데, 재밌는 프로젝트인데 반해 관련된 자료를 찾는게 쉽지 않길래 한번 소개를 해보면 좋을 것 같다는 생각이 들었습니다.
우선 XLA (Accelerated Linear Algebra)는 Tensorflow의 서브 프로젝트로 그래프 연산의 최적화 / 바이너리 사이즈의 최소화 등을 목적으로 하는 컴파일러입니다.
결과를 먼저 말씀드리자면, XLA를 활용하는 것을 통해 그래프 연산 과정에 필요한 임시 버퍼도 덜 사용할 수 있고, 그래프 연산 속도도 단축시킬 수 있습니다. (구글 측 자료에 의하면 Nvidia GPU를 사용할 때 50%까지도 성능 향상이 있었다고 합니다.)
예를 들어 아래와 같은 그래프가 있다고 했을때…
각 연산마다 룹을 돌리는 것이 아니라, 아래와 같은 형식으로 연산을 fuse해서 처리하게 될 경우, 임시 데이터 저장도 줄어들고, 룹도 하나의 룹으로 묶어서 처리가 가능하기 때문에 연산 속도나 임시 버퍼 사용량(및 read/write bandwidth), 캐쉬 효율 측면에서 여러가지 이점을 얻을 수 있습니다.
XLA에 있는 그래프 최적화기가 하는 일이 위와 같은 것들이며, XLA에서 하는 일을 간단히 도식화 하면 아래와 같습니다.
XLA도 일종의 컴파일러이기 때문에 일반적인 컴파일러 들이 동작하는 것처럼 FrontEnd (Parser) / MiddleEnd (High-level Optimizer) / BackEnd (Target Code Generator)의 구조를 가지고 있으며, 주어진 타겟 플랫폼에 맞춰 그래프를 최적화 한 후 실행가능한 형태의 무언가를 생성하는 역할을 수행합니다.
BackEnd에서는 LLVM을 사용하고 있으며, LLVM과 마찬가지로 JiT compilation과 AoT compilation 모두를 지원하는게 특징입니다. 두가지 중 어떤 방식을 사용하느냐에 따라 실제 사용 시나리오가 달라지게 되며,
AoT compilation을 사용하는 경우 아래와 같이 컴파일의 결과로 object file 및 이를 사용하기 위한 helper class에 대한 header file이 생성됩니다.
컴파일된 그래프는 아래와 같은 형태의 코드와 링크를 시켜 바이너리로 만든 후 실행할 수 있습니다. (참고: Using AOT compilation)
BlahGraph graph;
  // Set up args and run the computation.
  const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  std::copy(args + 0, args + 6, graph.arg0_data());
  std::copy(args + 6, args + 12, graph.arg1_data());
  graph.Run();
  std::cout &lt;&lt; graph.result0(0, 0) &lt;&lt; std::endl;
실행 바이너리를 만들기 위해 XLA 라이브러리나 Tensorflow 라이브러리 등을 포함할 필요가 없기 때문에 최종 타겟 바이너리 사이즈를 굉장히 작게 만들 수 있으나 기본 CPU backend의 경우 convolution, matmul 등의 연산에 대해 XLA runtime function을 호출하는 코드를 생성하도록 되어 있기 때문에 일부 XLA runtime function들과 링크되어야 할 수도 있습니다.
JiT compilation을 사용하는 경우는 compile의 결과로 XlaExecutable 데이터가 생성되며, StreamExecutor를 통해 input parameter를 feeding하고 그래프를 실행시킬 수 있는 형태입니다.
JitCompilation을 사용하는 경우 런타임에 그래프를 입력으로 받아 실행할 수 있다는 장점이 있으나 애플리케이션이 XLA 런타임을 포함해야하기 때문에 eigen, protobuf, LLVM 등등 의존성들이 줄줄이 엮인다는 단점이 있습니다. (타겟 플랫폼에 해당 라이브러리들을 줄줄이 포팅해야할 수도 있습니다.)
애플리케이션을 배포하기 전 단계에 타겟 플랫폼이나 네트워크 형태에 대해 미리 알고 있는 경우에는 AoT compilation을 활용하는 방안을 고려해볼 수 있으며, 애플리케이션은 네트워크를 모르는 상태로 배포되어야 하는 환경이라면 JiT compilation을 사용해야합니다.
Tensorflow 외에도 다양한 머신러닝 프레임웍의 백엔드로 XLA를 활용할 수 있으며, PyTorch를 위해선 PyTorch XLA, ONNX (Open Neural Network Exchange )를 위해선 onnx xla 프로젝트가 개발되고 있습니다..
Service / Client 구조나 기타 여러가지 재밌는 부분들이 많이 남아있지만 소개글이니 우선은 이 정도로 시작해볼까 합니다.
관심 있으신 분들 있으면 이어서 다른 내용들도 정리를 해보도록 하겠습니다. 
추가: 얼마전 ML Compiler 스터디에서 관련 내용을 발표한 자료를 하나 진행했으니 관심있으신 분들은 참고하셔도 좋을 것 같군요.
PPT: https://github.com/…/XLA101.pdf
Video: https://youtu.be/_3ykXQH5h2o

XLA를 소개합니다.

댓글