XLA Architecture & Tutorials

Why did we build XLA?

XLA가 TensorFlow와 함께 작동하기 위한 몇 가지 목표

• Improve execution speed (실행 속도 향상) : 서브그래프를 컴파일하여 짧은 지속 시간을 가진 연산의 실행 시간을 줄여 TensorFlow 런타임의 오버헤드를 제거하고, 파이프라인 연산을 융합하여 메모리 오버헤드를 줄이고, 알려진 텐서 형태에 특화되어 더 공격적인 상수 전파를 허용합니다.
• Improve memory usage (메모리 사용량 개선) : 메모리 사용량을 분석하고 스케줄링하여 원칙적으로 많은 중간 저장 버퍼를 제거합니다.
• Reduce reliance on custom Ops (사용자 지정 Ops에 대한 의존성 줄이기) : 자동으로 융합된 저수준 Ops의 성능을 개선하여 수동으로 융합된 사용자 지정 Ops의 성능과 일치하게 하여 많은 사용자 지정 Ops가 필요하지 않도록 합니다.
• Reduce mobile footprint (모바일 풋프린트 축소) : 서브그래프를 미리 컴파일하고 개체/헤더 파일 쌍으로 출력하여 TensorFlow 런타임을 제거합니다. 이 결과로 모바일 추론의 풋프린트가 수십 개 정도 줄어들 수 있습니다.
• Improve portability (이식성 개선) : 새로운 하드웨어에 대해 비교적 쉽게 새로운 백엔드를 작성할 수 있도록 하여 TensorFlow 프로그램의 상당 부분이 해당 하드웨어에서 수정 없이 실행될 수 있게 합니다. 이는 새로운 하드웨어에 대해 개별적인 단일 뭉치 연산을 특수화하는 방법과 대조적으로, TensorFlow 프로그램을 그러한 연산을 사용하도록 다시 작성해야 하는 접근 방식과 달라집니다.

 

How does XLA work?

XLA의 입력 언어는 "HLO IR" 또는 간단히 HLO(고수준 연산)라고 불립니다. HLO의 의미론은 Operation Semantics 페이지에서 설명되어 있습니다. HLO는 주로 컴파일러 IR(중간 표현)로 생각하는 것이 가장 편리합니다.

 

XLA는 HLO로 정의된 그래프("계산")를 다양한 아키텍처의 기계 명령어로 컴파일합니다. XLA는 새로운 하드웨어 아키텍처를 대상으로 대체 백엔드를 쉽게 추가할 수 있도록 모듈화되어 있습니다. x64 및 ARM64용 CPU 백엔드와 NVIDIA GPU 백엔드는 TensorFlow 소스 트리에 포함되어 있습니다.

 

다음 다이어그램은 XLA의 컴파일 과정을 보여줍니다.

 

XLA에는 CSE, 타겟 독립적인 연산 융합(target-independent operation fusion) 및 계산을 위해 런타임 메모리를 할당하기 위한 버퍼 분석과 같은 여러 최적화 및 분석 패스들이 탑재되어 있습니다.

 

비타겟 종속적인 단계 이후에, XLA는 HLO 계산을 백엔드에 보냅니다. 백엔드는 추가적인 HLO 수준의 최적화를 수행하는데, 이번에는 대상 특정 정보와 요구사항을 고려합니다. 예를 들어, XLA GPU 백엔드는 특히 GPU 프로그래밍 모델에 유리한 연산 융합을 수행하고, 계산을 스트림으로 분할하는 방법을 결정할 수 있습니다. 이 단계에서 백엔드는 특정 연산이나 그 조합을 패턴 매치하여 최적화된 라이브러리 호출을 할 수도 있습니다.

 

다음 단계는 대상 특정 코드 생성입니다. XLA와 함께 제공되는 CPU 및 GPU 백엔드는 저수준 IR, 최적화 및 코드 생성을 위해 LLVM을 사용합니다. 이러한 백엔드는 효율적인 방식으로 XLA HLO 계산을 표현하는 데 필요한 LLVM IR을 생성하고, 그런 다음 LLVM을 호출하여 이 LLVM IR에서 네이티브 코드를 생성합니다.

 

현재 GPU 백엔드는 LLVM NVPTX 백엔드를 통해 NVIDIA GPU를 지원하며, CPU 백엔드는 여러 CPU ISA를 지원합니다.

 

Tutorials

Use XLA with tf.function

이 튜토리얼은 MNIST 데이터셋을 분류하기 위해 TensorFlow 모델을 훈련시키는데, 이때 훈련 함수는 XLA를 사용하여 컴파일됩니다.

 

먼저 TensorFlow를 로드하고 eager execution을 활성화합니다.

import tensorflow as tf

 

그런 다음 필요한 상수들을 정의하고 MNIST 데이터셋을 준비합니다.

# Size of each input image, 28 x 28 pixels
IMAGE_SIZE = 28 * 28
# Number of distinct number labels, [0..9]
NUM_CLASSES = 10
# Number of examples in each training batch (step)
TRAIN_BATCH_SIZE = 100
# Number of training steps to run
TRAIN_STEPS = 1000

# Loads MNIST dataset.
train, test = tf.keras.datasets.mnist.load_data()
train_ds = tf.data.Dataset.from_tensor_slices(train).batch(TRAIN_BATCH_SIZE).repeat()

# Casting from raw data to the required datatypes.
def cast(images, labels):
  images = tf.cast(
      tf.reshape(images, [-1, IMAGE_SIZE]), tf.float32)
  labels = tf.cast(labels, tf.int64)
  return (images, labels)

 

마지막으로, 모델과 옵티마이저를 정의합니다. 모델은 하나의 dense layer를 사용합니다.

layer = tf.keras.layers.Dense(NUM_CLASSES)
optimizer = tf.keras.optimizers.Adam()

 

Define the training function

훈련 함수에서는 위에서 정의한 레이어를 사용하여 예측된 레이블을 얻은 다음, 옵티마이저를 사용하여 손실의 그래디언트를 최소화합니다. XLA를 사용하여 계산을 컴파일하려면 jit_compile=True로 설정하여 tf.function 내에 배치합니다.

@tf.function(jit_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))

 

Train and test the model

훈련 함수를 정의한 후에, 모델을 정의합니다.

for images, labels in train_ds:
  if optimizer.iterations > TRAIN_STEPS:
    break
  train_mnist(images, labels)

 

마지막으로 정확도를 확인합니다.

images, labels = cast(test[0], test[1])
predicted_labels = layer(images)
correct_prediction = tf.equal(tf.argmax(predicted_labels, 1), labels)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
print("Prediction accuracy after training: %s" % accuracy)

Prediction accuracy after training: tf.Tensor(0.8773, shape=(), dtype=float32)

뒤에서 XLA 컴파일러는 전체 TF 함수를 HLO로 컴파일했으며, 이로 인해 융합 최적화(fusion optimization)가 가능해졌습니다. 내부 검사 기능을 사용하여 HLO 코드를 확인할 수 있습니다. ('stage'에 대해 다른 흥미로운 가능한 값으로는 optimized_hlo는 최적화 후 HLO이며 optimized_hlo_dot은 Graphviz 그래프입니다)

print(train_mnist.experimental_get_compiler_ir(images, labels)(stage='hlo'))