XLA:GPU Emitters

There are three ways how to generate code for HLO in XLA:GPU.

Replacing HLO with custom calls to external libraries, e.g. NVidia cuBLAS, cuDNN.
Tiling HLO to block-level and then using OpenAI Triton.
Using XLA Emitters to progressively lower HLO to LLVM IR.

This document is focused on XLA:GPU Emitters.

Hero-based codegen

There are 7 emitter types in XLA:GPU. Each emitter type corresponds to a "hero" of the fusion, i.e. the most important op in the fused computation that shapes the code generation for the whole fusion.

For example, the tranpose emitter will be selected if there is a HloTransposeInstruction within the fusion that requires using shared memory to improve the memory reading and writing patterns. The reduction emitter generates reductions using shuffles and shared memory. The loop emitter is the default emitter. If a fusion does not have a hero for which we have a special emitter, then the loop emitter will be used.

High-level overview

The code consists of the following big building blocks:

Computation partitioner - splitting an HLO fusion computation into functions
Emitters - converting partitioned HLO fusion to MLIR (xla_gpu, tensor, arith, math, scf dialects)
Compilation pipeline - optimizes and lowers IR to LLVM

Partitioning

See computation_partitioner.h.

Non-elementwise HLO instructions cannot always be emitted together. Consider the following HLO graph:

     param
       |
      log
      |  \
      |  transpose
      |  /
      add

If we emit this in a single function, the log will be accessed at two different indices for each element of the add. The old emitters solve this problem by generating the log twice. For this particular graph, this is not a problem, but when there are multiple splits, the code size grows exponentially.

Here, we solve this problem by partitioning the graph into pieces that can be safely emitted as one function. The criteria are:

Instructions that have only one user are safe to emit together with their user.
Instructions that have multiple users are safe to emit together with their users if they are accessed through the same indices by all users.

In the example above, the add and tranpose access different indices of the log, so it is not safe to emit it together with them.

The graph is therefore partitioned into three functions (each containing just one instruction).

The same is applicable to the following example with slice and pad of add.

Elemental emission

See elemental_hlo_to_mlir.h.

Elemental emission creates loops and math/arith ops for HloInstructions. For the most part, this is straightforward, but there are some interesting things going on here.

Indexing transformations

Some instructions (transpose, broadcast, reshape, slice, reverse and a few more) are purely transformations on indices: to produce an element of the result, we need to produce some other element of the input. For this, we can reuse XLA's indexing_analysis, which has functions to produce the output to input mapping for an instruction.

For example, for a transpose from [20,40] to [40,20], it will produce the following indexing map (one symbolic expression per input dimension; d0 and d1 are the output dimensions):

  (d0, d1) -> d1
  (d0, d1) -> d0

So for these pure index transformation instructions, we can simply get the map, apply it to the output indices, and produce the input at the resulting index.

Similarly, the pad op uses indexing maps and constraints for most of the implementation. pad is also an indexing transformation with some added checks to see if we return an element of the input or the padding value.

Tuples

We do not support internal tuples. We also do not support nested tuple outputs. All XLA graphs that use these features can be converted to graphs that do not.

Gather

We only support canonical gathers as produced by gather_simplifier.

Subgraph functions

For a subgraph of a computation with parameters %p0 to %p_n, and subgraph roots with r dimensions and element types (e0 to e_m), we use the following MLIR function signature:

(%p0: tensor<...>, %p1: tensor<...>, ..., %pn: tensor<...>,
 %i0: index, %i1: index, ..., %i_r-1: index) -> (e0, ..., e_m)

That is, we have one tensor input per computation parameter, one index input per dimension of the output, and one result per output.

To emit a function, we simply use the elemental emitter above, and recursively emit its operands until we reach the edge of the subgraph. Then, we:emit a tensor.extract for parameters or emit a func.call for other subgraphs

Entry function

Each emitter type differs in how it generates the entry function, i.e. the function for the hero. The entry function is different from the functions above, since it has no indices as inputs (just the thread and block IDs) and actually needs to write the output somewhere. For the loop emitter, this is fairly straightforward, but the transpose and reduction emitters have non-trivial write logic.

The signature of the entry computation is:

(%p0: tensor<...>, ..., %pn: tensor<...>,
 %r0: tensor<...>, ..., %rn: tensor<...>) -> (tensor<...>, ..., tensor<...>)

Where like before, the %pns are the parameters of the computation, and the %rns are the results of the computation. The entry computation takes the results as tensors, tensor.inserts updates into them, and then returns them. No other uses of the output tensors are allowed.

Compilation pipeline

Loop emitter

See loop.h.

Let's study the most important passes of the MLIR compilation pipeline using the HLO for the GELU function.

This HLO computation only has elementwise ops, constants and broadcasts. It will be emitted using the loop emitter.

MLIR Conversion

After conversion to MLIR we get an xla_gpu.loop that depends on %thread_id_x and %block_id_x and defines the loop that traverses all elements of the output linearly to guarantee coalesced writes.

On every iteration of this loop we call

   %pure_call = xla_gpu.pure_call @gelu(%input, %dim0, %dim1, %dim2)
      : (tensor<6x512x4096xbf16>, index, index, index) -> bf16

to compute elements of the root operation. Note, that we have only one outlined function for @gelu, because the partitioner did not detect a tensor that has 2 or more various access patterns.

#map = #xla_gpu.indexing_map<"(th_x, bl_x)[vector_index] -> ("
 "bl_x floordiv 4096, (bl_x floordiv 8) mod 512, (bl_x mod 8) * 512 + th_x * 4 + vector_index),"
 "domain: th_x in [0, 127], bl_x in [0, 24575], vector_index in [0, 3]">

func.func @main(%input: tensor<6x512x4096xbf16> , %output: tensor<6x512x4096xbf16>)
   -> tensor<6x512x4096xbf16> {
 %thread_id_x = gpu.thread_id  x {xla.range = [0 : index, 127 : index]}
 %block_id_x = gpu.block_id  x {xla.range = [0 : index, 24575 : index]}

 %xla_loop = xla_gpu.loop (%thread_id_x, %block_id_x)[%vector_index] -> (%dim0, %dim1, %dim2)
     in #map iter_args(%iter = %output) -> (tensor<6x512x4096xbf16>) {
   %pure_call = xla_gpu.pure_call @gelu(%input, %dim0, %dim1, %dim2)
      : (tensor<6x512x4096xbf16>, index, index, index) -> bf16
   %inserted = tensor.insert %pure_call into %iter[%dim0, %dim1, %dim2] : tensor<6x512x4096xbf16>
   xla_gpu.yield %inserted : tensor<6x512x4096xbf16>
 }
 return %xla_loop : tensor<6x512x4096xbf16>
}

func.func private @gelu(%arg0: tensor<6x512x4096xbf16>, %i: index, %j: index, %k: index) -> bf16 {
  %cst = arith.constant 5.000000e-01 : bf16
  %cst_0 = arith.constant 1.000000e+00 : bf16
  %cst_1 = arith.constant 7.968750e-01 : bf16
  %cst_2 = arith.constant 4.467770e-02 : bf16
  %extracted = tensor.extract %arg0[%i, %j, %k] : tensor<6x512x4096xbf16>
  %0 = arith.mulf %extracted, %extracted : bf16
  %1 = arith.mulf %0, %extracted : bf16
  %2 = arith.mulf %1, %cst_2 : bf16
  %3 = arith.addf %extracted, %2 : bf16
  %4 = arith.mulf %3, %cst_1 : bf16
  %5 = math.tanh %4 : bf16
  %6 = arith.addf %5, %cst_0 : bf16
  %7 = arith.mulf %6, %cst : bf16
  %8 = arith.mulf %extracted, %7 : bf16
  return %8 :

XLA:GPU Emitters Stay organized with collections Save and categorize content based on your preferences.