Print XLA Kernel Traces | Relevant env vars

[Rohanjames1997]

Hi all,
Does anyone know how I can get the compiled XLA graph dump of an ML model?

I'm trying to print the resulting XLA graph of a Flax BERT model using TF_CPP_MAX_VLOG_LEVEL=1. I'm following steps present in Intel's guide to debug AI model performance, and it seems from other online sources that TF_CPP_MAX_VLOG_LEVEL was enough to print the model graph.
However, it seems like this environment variable does not have any effect in Jax / OpenXLA. Maybe there are alternative variables(?)

My goal is to compare the resulting graphs on different platforms

[yashk2810]

You can use the JAX AOT API: print(jax.jit(f).lower(*args, **kwargs).compile().as_text())

[Rohanjames1997]

Thanks @yashk2810 !
With AOT Compilation on BERT, I now get a lowered graph with nodes like this:

Node

%fused_computation.167 (param_0.722: f32[12,3072]) -> f32[12,3072] {
  %param_0.722 = f32[12,3072]{1,0} parameter(0)
  %bitcast.679 = f32[1,12,3072]{2,1,0} bitcast(f32[12,3072]{1,0} %param_0.722), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/dense/dot_general[dimension_numbers=(((2,), (0,)), ((), ())) precision=None preferred_element_type=None]" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/flax/linen/linear.py" source_line=276}
  %constant.1621 = f32[3072]{0} constant({...}), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/dense/add" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/flax/linen/linear.py" source_line=283}
  %broadcast.1545 = f32[1,12,3072]{2,1,0} broadcast(f32[3072]{0} %constant.1621), dimensions={2}, metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/dense/add" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/flax/linen/linear.py" source_line=283}
  %add.604 = f32[1,12,3072]{2,1,0} add(f32[1,12,3072]{2,1,0} %bitcast.679, f32[1,12,3072]{2,1,0} %broadcast.1545), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/dense/add" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/flax/linen/linear.py" source_line=283}
  %constant.1624 = f32[] constant(0.707106769)
  %broadcast.1548 = f32[1,12,3072]{2,1,0} broadcast(f32[] %constant.1624), dimensions={}
  %multiply.435 = f32[1,12,3072]{2,1,0} multiply(f32[1,12,3072]{2,1,0} %add.604, f32[1,12,3072]{2,1,0} %broadcast.1548), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/div" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/bert/modeling_flax_bert.py" source_line=468}
  %erf.12 = f32[1,12,3072]{2,1,0} erf(f32[1,12,3072]{2,1,0} %multiply.435), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/erf" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/bert/modeling_flax_bert.py" source_line=468}
  %constant.1623 = f32[] constant(1)
  %broadcast.1547 = f32[1,12,3072]{2,1,0} broadcast(f32[] %constant.1623), dimensions={}
  %add.603 = f32[1,12,3072]{2,1,0} add(f32[1,12,3072]{2,1,0} %erf.12, f32[1,12,3072]{2,1,0} %broadcast.1547), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/add" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/bert/modeling_flax_bert.py" source_line=468}
  %multiply.434 = f32[1,12,3072]{2,1,0} multiply(f32[1,12,3072]{2,1,0} %add.604, f32[1,12,3072]{2,1,0} %add.603), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/mul" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/bert/modeling_flax_bert.py" source_line=468}
  %constant.1622 = f32[] constant(0.5)
  %broadcast.1546 = f32[1,12,3072]{2,1,0} broadcast(f32[] %constant.1622), dimensions={}
  %multiply.433 = f32[1,12,3072]{2,1,0} multiply(f32[1,12,3072]{2,1,0} %multiply.434, f32[1,12,3072]{2,1,0} %broadcast.1546), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/div" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/bert/modeling_flax_bert.py" source_line=468}
  ROOT %bitcast.678 = f32[12,3072]{1,0} bitcast(f32[1,12,3072]{2,1,0} %multiply.433), metadata={op_name="jit(bertInfer)/jit(main)/FlaxBertForSequenceClassificationModule/bert/encoder/layer/1/intermediate/div" source_file="/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/bert/modeling_flax_bert.py" source_line=468}
}

I'm interested in knowing even lower level details of the graph - which XLA kernels are being invoked at runtime, for multiply, add, etc.

Is there a way to log the runtime execution flow of the low level kernel invocations?

[Rohanjames1997]

TF_CPP_MAX_VLOG_LEVEL=2 TF_CPP_MIN_LOG_LEVEL=0 python script.py