Lowering of Pallas Kernel and DMAs

[yuanfz98]

I am trying to understand how pallas kernel is integrated into a graph in mlir. I don't understand the stablehlo.custom_call op generated by hlo_to_mhlo pass. As we can see in the graph it accepts two tensors which is from HBM I suppose, but actually a pallas kernel accepts Vmem pointers. When compiling a graph with many stablehlo ops (whether from pallas or not), where is the dma from hbm to vmem, as pallas defined DMAs "outside" of kernel ?

import jax.numpy as jnp
import jax
from jax.experimental import pallas as pl


def kernel(x_ref, y_ref, o_ref):
    o_ref[...] = x_ref[...] + y_ref[...]


x = jnp.empty((8, 128), dtype=jnp.float32)
pl.pallas_call(
    kernel,
    out_shape=jax.ShapeDtypeStruct((8, 128), jnp.float32),
    debug=True,
    name="my_call",
)(x, x)

Gives:

The kernel jaxpr for pallas_call my_call for kernel function kernel at /workspace/codebase/test_pallas.py:6:
{ lambda ; a:MemRef<None>{float32[8,128]} b:MemRef<None>{float32[8,128]} c:MemRef<None>{float32[8,128]}. let
    d:f32[8,128] <- a[:,:]
    e:f32[8,128] <- b[:,:]
    f:f32[8,128] = add d e
    c[:,:] <- f
  in () }

The Mosaic module for pallas_call my_call for kernel function kernel at /workspace/codebase/test_pallas.py:6:
module @my_call {
  func.func @main(%arg0: memref<8x128xf32, #tpu.memory_space<vmem>>, %arg1: memref<8x128xf32, #tpu.memory_space<vmem>>, %arg2: memref<8x128xf32, #tpu.memory_space<vmem>>) attributes {dimension_semantics = [], scalar_prefetch = 0 : i64, scratch_operands = 0 : i64} {
    %c0 = arith.constant 0 : index
    %c0_0 = arith.constant 0 : index
    %0 = vector.load %arg0[%c0, %c0_0] : memref<8x128xf32, #tpu.memory_space<vmem>>, vector<8x128xf32>
    %c0_1 = arith.constant 0 : index
    %c0_2 = arith.constant 0 : index
    %1 = vector.load %arg1[%c0_1, %c0_2] : memref<8x128xf32, #tpu.memory_space<vmem>>, vector<8x128xf32>
    %2 = arith.addf %0, %1 : vector<8x128xf32>
    %c0_3 = arith.constant 0 : index
    %c0_4 = arith.constant 0 : index
    %3 = vector.load %arg2[%c0_3, %c0_4] : memref<8x128xf32, #tpu.memory_space<vmem>>, vector<8x128xf32>
    tpu.vector_store %arg2[%c0_3, %c0_4], %2 {strides = array<i32>} : memref<8x128xf32, #tpu.memory_space<vmem>>, vector<8x128xf32>, 
    return
  }
}

PjRtCApiClient::Compile: module.dump()
module @jit_wrapped attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<8x128xf32>, %arg1: tensor<8x128xf32>) -> (tensor<8x128xf32> {jax.result_info = ""}) {
    %0 = stablehlo.custom_call @tpu_custom_call(%arg0, %arg1) {backend_config = "{\22custom_call_config\22: {\22body\22: \22TUzvUgFNTElSMjEuMC4wZ2l0AAEhCQEDBQcBAwkDDQsNDxETFQN3VQ8BUQcTDw8PCw8LGwsPCzMLCwsLhQsLCwsPCw8LHw8LHxMLDwsLGwtTCxMFA2EHA1kBDwcjGw8HDxsC1gMfAwMrLR0vCR1FCRUxNwUXEQcBBRkDBRMVCxcFGxEHDQUdAwsbHR8hJQ0nDQspBR8BAQUhDQ1hZmZpbmVfbWFwPChkMCwgZDEpIC0+IChkMCwgZDEpPgAFIwUlBScFKREBAQUrHTM1BS0tDwkPAQEBHTk7BS8tDwkXAQEBAwM/UwUxHUMJBTMFNQMFSUtNTwU3IwsJIQEAAAABAAAAAgAAAAAAAAAFOSMLAQEjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPgAjYXJpdGguZmFzdG1hdGg8bm9uZT4AAxdRBSECBAkjJwUhAgQJAQIECwECAgUHAwMDAQTZBQERAREHAwEFBxEBGQcDGzMHAwEDAQMBAwMFAwMBAwMFAwMBBQYFAwUHAQcJAwMFAwMBAwMFAwMBBQYFAwUHAw0PCQdBPQMFBQsRAwMHAwMBAwMHAwMBBQYHAwUHBRUXCwUHRwkTBRUXDQABBgMBBQEAvgU7ESkNCxMTDwsNCyMhHSkRLUcTGSMXFRkfDw0JHRFidWlsdGluAHN0YWJsZV9tb3NhaWMAdHB1AGFyaXRoAG1vZHVsZQBhcml0aC5jb25zdGFudAB2ZWN0b3IubG9hZABmdW5jLmZ1bmMAYXJpdGguYWRkZgB0cHUudmVjdG9yX3N0b3JlAGZ1bmMucmV0dXJuAHN5bV9uYW1lAC93b3Jrc3BhY2UvY29kZWJhc2UvdGVzdF9wYWxsYXMucHkAc3RhYmxlX21vc2FpYy52ZXJzaW9uAG15X2NhbGwAZGltZW5zaW9uX3NlbWFudGljcwBmdW5jdGlvbl90eXBlAHNjYWxhcl9wcmVmZXRjaABzY3JhdGNoX29wZXJhbmRzAG1haW4AdmFsdWUAL2dldABrZXJuZWwAPG1vZHVsZT4AZmFzdG1hdGgAL2FkZAAvc3dhcABvcGVyYW5kU2VnbWVudFNpemVzAHN0cmlkZXMA\22, \22serialization_format\22: 1, \22needs_layout_passes\22: true}}", kernel_name = "my_call", operand_layouts = [dense<[1, 0]> : tensor<2xindex>, dense<[1, 0]> : tensor<2xindex>], result_layouts = [dense<[1, 0]> : tensor<2xindex>]} : (tensor<8x128xf32>, tensor<8x128xf32>) -> tensor<8x128xf32>
    return %0 : tensor<8x128xf32>
  }
}

[yuanfz98]

What’s actually happening here?
Thus far we’ve described how to think about Pallas kernels but what we’ve actually accomplished is we’re writing a function that’s executed very close to the compute units since values are loaded into the innermost (fastest) portion of the memory hierarchy.
On GPU, x_ref corresponds to a value in high-bandwidth memory (HBM) and when we do x_ref[...] we are copying the value from HBM into static RAM (SRAM) (this is a costly operation generally speaking!). We then use GPU vector compute to execute the addition, then copy the resulting value in SRAM back to HBM.
On TPU, we do something slightly different. Before the kernel is ever executed, we fetch the value from HBM into SRAM. x_ref therefore corresponds to a value in SRAM and when we do x_ref[...] we are copying the value from SRAM into a register. We then use TPU vector compute to execute the addition, then copy the resulting value back into SRAM. After the kernel is executed, the SRAM value is copied back into HBM.

Well, if DMA is defined outside of kernels, how to describe it in MLIR graph ? I really need your help, thanks.

[sharadmv]

The DMAs you are referring to are emitted as part of the lowering of the kernel, so it won't be visible in either the Mosaic or the HLO. Specifically, when the HLO is compiled into TPU assembly in the very final stage of compilation, it additionally inlines these DMAs.

[yuanfz98]

1. Could you please refer the code for DMA inlining ?
2. As I am dealing with PJRT Plugin, for the exported graph it only abstracts tensors as HBM, if DMAs are lowered late in the compilation, how to handle the coherence of graph ? Is there any way to compile Pallas kernel in advance then insert the binary into the graph ?