Sergei Lebedev
f3b91b2042
Export PointerType and register_dialect from jaxlib.triton.dialect
...
The `... as ...` form tells the type checker that the name is exported.
See #7570 .
PiperOrigin-RevId: 671318047
2024-09-05 04:15:32 -07:00
Sergei Lebedev
778933dfda
Removed inspect.signature() call from jaxlib.triton.dialect.ScanOp
...
PiperOrigin-RevId: 614772594
2024-03-11 13:30:41 -07:00
Goran Flegar
53364b438c
Integrate Triton up to [bfb8e413]( bfb8e413b0)
...
PiperOrigin-RevId: 614740360
2024-03-11 11:43:46 -07:00
Sergei Lebedev
5e2e609a9b
_triton_ext no longer links in MLIR C APIs
...
I re-used the same trick we do for the TPU dialect. Specifically, _triton_ext no longer depends on :triton_dialect_capi. Instead
* we include Triton dialect C bindings into :jaxlib_mlir_capi_objects
* and _triton_ext depends on :jaxlib_mlir_capi_objects and a header-only cc_library providing Triton dialect C bindings
This is a fork of #19680 with a few internal-only fixes.
PiperOrigin-RevId: 604929377
2024-02-07 03:39:29 -08:00
Jake VanderPlas
4d5c557264
jaxlib/triton: ensure python 3.9 compatibility
2024-02-06 11:39:51 -08:00
Sergei Lebedev
0596c804cb
Inlined all builder.create_* calls
...
The lowering code now directly calls into the auto-generated MLIR bindings
for the respective dialects.
PiperOrigin-RevId: 604260837
2024-02-05 03:00:33 -08:00
Sergei Lebedev
5b7a0d9c91
Pallas now uses MLIR Python builders to lower to Triton IR
...
This allows us to drop a dependency on the Triton Python package in the future,
and delegate ->ptx compilation to XLA.
PiperOrigin-RevId: 597640756
2024-01-11 13:33:26 -08:00
Sergei Lebedev
f219482212
The Triton MLIR bindings now include auto-generated wrappers for enums
...
PiperOrigin-RevId: 596873541
2024-01-09 03:00:47 -08:00
Sergei Lebedev
e6c890171b
Generate Python bindings for the Triton MLIR dialect
...
The bindings are not yet included in the jaxlib wheel. I will do that in a
follow up PR.
PiperOrigin-RevId: 595174466
2024-01-02 11:55:05 -08:00
