Adam Paszke ded9272a5b [Mosaic GPU] Implement a simple profilng tool using CUDA events ...

The other JAX profiling tools are a little heavyweight when we only care about
timing a single kernel programatically.

Also adapt wgmma.py to match failures triggered by upstream MLIR changes.

PiperOrigin-RevId: 628096973

2024-04-25 09:18:39 -07:00

..

_triton_ext.pyi

Added ir.Value-based versions of load and store in triton.compat

2024-02-13 06:13:31 -08:00

BUILD.bazel

[Mosaic GPU] Implement a simple profilng tool using CUDA events

2024-04-25 09:18:39 -07:00

jaxlib_mlir_capi_shims.cc

Initial commit for Mosaic GPU

2024-04-18 04:04:10 -07:00

jaxlib_mlir_capi_shims.h

Initial commit for Mosaic GPU

2024-04-18 04:04:10 -07:00

mosaic_gpu_ext.cc

[Mosaic GPU] Implement a simple profilng tool using CUDA events

2024-04-25 09:18:39 -07:00

register_jax_dialects.cc

Switch Mosaic GPU to a custom pass pipeline and improve the lowering of GPU launch

2024-04-24 03:27:45 -07:00

tpu_ext.cc

[Mosaic] Remove Python implementation of apply_vector_layout and infer_memref_layout.

2024-01-10 13:00:21 -08:00

triton_ext.cc

Added ir.Value-based versions of load and store in triton.compat

2024-02-13 06:13:31 -08:00