ROCm/jax

Pull Requests Commits

[ROCm] Add ROCm platform support for psend/precv collective operations

AratiGanesh committed 9 hours ago

d548f5c8

[Mosaic GPU] Fix the return type of `_lift_fast_packed_instr` to match its argument types.

bchetioui committed 10 hours ago

ed12030c

[Mosaic GPU][NFC] Replace usages of `_gpu_ops_gen` with `gpu` dialect whenever possible.

allanrenucci committed 10 hours ago

35a14b8b

[Mosaic GPU] Fix issue in `vector_dim` reduction with `vec_len > 2`.

bchetioui committed 11 hours ago

678adf59

[Pallas:MGPU][NFC] Update outdated docstring for `scratch_view`.

allanrenucci committed 13 hours ago

f25011f0

[Pallas/interpreter] Add a prototype for a GPU kernel interpreter.

Google-ML-Automation committed 14 hours ago

6eb71e90

[Mosaic GPU] Support more reduction kinds and layouts in `MultiDimReductionOp` lowering.

allanrenucci committed 14 hours ago

e5faf1d2

Add pallas lowering for jnp.tile using tpu.repeat.

levskaya committed 16 hours ago

f0ca449a

Update XLA dependency to use revision http://github.com/openxla/xla/commit/9258aaa31683b9f5e5a1cc21c6c2a3c769df4bfb

Google-ML-Automation committed 17 hours ago

9a243d76

[Mosaic:TPU] Clean up tpu.memref_slice verifier

tlongeri committed 18 hours ago

4691dfee

Automated Code Change

Google-ML-Automation committed 20 hours ago

bd07769f

levskaya committed 20 hours ago

c396a367

[Pallas] Allowlist semaphore/prng effects under remat and custom

sharadmv committed 23 hours ago

c5e42402

Merge pull request #33974 from Prakharprasun:fix-padtype-to-pads

Google-ML-Automation committed 1 day ago

a10149b6

[sc] Removed `infer_kernel_arguments` from infer-memref-layout

superbobry committed 1 day ago

c984def2

[Pallas/TPU] Don't lower eqns that have all dropvar outputs (aka DCE at

sharadmv committed 1 day ago

9e8fae66

[sc] Generalized infer-memref-layout to support SC tiling

superbobry committed 1 day ago

06d75423

Plumb prim params for call discharge rule (to handle named_computation_p

sharadmv committed 1 day ago

96787dd8

[jax.collect_profile] Allow arbitrary options to be passed to XProf

Matt-Hurd committed 1 day ago

a4036e4a

lax: ensure padtype_to_pads returns Python ints

Prakharprasun committed 1 day ago

6ce4314f

[Pallas MGPU] Clip the size of contiguous TMA transfers to the size of the ref. This avoids OOB accesses and is more in-line with the tensor map TMAs.

Rifur13 committed 1 day ago

27a568a3

[Mosaic GPU] Add support for f8 types for WGMMA with lhs in registers under WG semantics.

allanrenucci committed 1 day ago

7cee4db6

Add `thread_guard` to the public API.

emilyfertig committed 1 day ago

1242ba32

[Pallas:MGPU] Add Pallas lowering for `pl.debug_check` under WG semantics.

allanrenucci committed 1 day ago

a070482d

[Pallas:MGPU] Lower `lax.sign` consistently for LANE and WG semantics.

allanrenucci committed 1 day ago

e4fa0259

[mosaic] infer-memref-layout now accepts target shape as a span

superbobry committed 1 day ago

dcb011a1

[XLA:MGPU] Port Replicated wrapper to C++.

golechwierowicz committed 1 day ago

658e6505

[Mosaic GPU] Enable redux.sync.f32 on Blackwell

apaszke committed 1 day ago

098e953a

[Mosaic GPU] Test approximate math functions properly + eta reduce args

apaszke committed 1 day ago

0ad11175

Merge pull request #34120 from oulgen:mgpu-ops

Google-ML-Automation committed 1 day ago

1ef117f3

Older