XLA GPU MLIR Codegen

[Md Faijul Amin]

Hi,

I have watched  XLA GPU MLIR Codegen presented by  Alexander Belyaev in the 2024 OpenXLA Fall Dev Lab. There was an example with mlir pass dumping. 

run_hlo_module \
--platform=CUDA --xla_disable_all_hlo_passes \
--reference_platform=""--v=5 /tmp/gelu.hlo 

However, --v=5 does not work. Also I am unable to see mlir pass dumps. Could anyone help getting mlir pass dumps?

Thanks,

Amin

[Sayce Falk]

Good questions, Amin - + @Alexander Belyaev to help answer!

--
You received this message because you are subscribed to the Google Groups "OpenXLA Discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to openxla-discu...@openxla.org.
To view this discussion visit https://groups.google.com/a/openxla.org/d/msgid/openxla-discuss/1fe9f9a4-7248-4110-9595-13d886f815f0n%40openxla.org.
For more options, visit https://groups.google.com/a/openxla.org/d/optout.

[Alexander Belyaev]

Hi Amin,

does the command not work at all? Or it runs the module, but you see no logs? What version are you using? The MLIR emitters were enabled by default only in the end of October. It could be that the compilation goes via the legacy emitters, that were removed in October/November. There used to be a flag --xla_gpu_mlir_emitter_level=4 that was necessary to enable the new emitters.

Alex

[Md Faijul Amin]

Hi Alex,

The command works if remove  --v=5, but does not show any log. I am using xla main branch. I can see some logs having TF_CPP_VMODULE=emitter_base=5, though.

~Amin

[Alexander Belyaev]

Nice! So, do you see the IR dumps after every pass now?

[Md Faijul Amin]

Yes, I can see MLIR passes up to https://github.com/openxla/xla/blob/main/xla/backends/gpu/codegen/emitters/emitter_base.cc#L311
But I would like to see also https://github.com/openxla/xla/blob/main/xla/backends/gpu/codegen/emitters/emitter_base.cc#L319

[Alexander Belyaev]

But that's just a translation to LLVM IR, there are no MLIR passes left.

[Md Faijul Amin]

I see, thanks for the explanation. Is there anyway to see the steps or logs for that translation to LLVM IR?

[Alexander Belyaev]

You can dump the llvm Module after the translation. But you are probably interested not in this, but in the passes that lower LLVM -> PTX. you can find those in, for example, xla/service/gpu/nvptx_compiler.cc

you can also try to dump different modules/stages by adding XLA_FLAGS=--xla_dump_to=/tmp/myfolder  as described in https://openxla.org/xla/tools

[Md Faijul Amin]

Hi Alex,
Is it possible set a breakpoint at cudaMalloc/cudaMemcpy for loop fusion while running with run_hlo_module? Trying to understand where host to memory transfer code is located while running with run_hlo_module.

Thanks,

Amin

[Alexander Belyaev]

 Hi Amin,

I think this logic is in xla/stream_executor/gpu.

--

Alexander Belyaev

[Md Faijul Amin]

Thank you very much! Alex for the pointers to gpu memory allocations. I found the CUDA allocations here https://github.com/openxla/xla/blob/main/xla/stream_executor/cuda/cuda_executor.cc
This has been very useful.

[Alexander Belyaev]

We can also chat over VC, if it is more convenient for you.