ROCm/jax

Pull Requests Commits

Arech8 committed 3 days ago

6c1bf731

Arech8 committed 3 days ago

bc3ddb7f

Improvements and cleanup

Arech8 committed 3 days ago

f8a99810

Arech8 committed 4 days ago

b11e343b

Fix a fence in the init section

Arech8 committed 4 days ago

5b4f9397

Make empty _commit_group_lowering() and await_async_copy() impls (don't need more)

Arech8 committed 4 days ago

3ca06cd8

fix warp_barrier()

Arech8 committed 4 days ago

935d04d4

improve comments

Arech8 committed 4 days ago

e547b697

implement commit_shared() as pass b/c of a barrier directly after it

Arech8 committed 5 days ago

85ff6681

Implement warpgroup_barrier() with s_waitcnt + s_barrier

Arech8 committed 5 days ago

ceffd7ca

compile and launch now works

Arech8 committed 6 days ago

077109ed

stubify bar.sync too for ROCm

Arech8 committed 6 days ago

0b179f52

support for the runtime

Arech8 committed 6 days ago

89155363

WIP emrgency saving - working on C++ runtime support

Arech8 committed 6 days ago

acc395bb

format runtime.cc

Arech8 committed 7 days ago

2340e532

Add stubs on the rest of nvvm mbarrier & cp.asynch calls

Arech8 committed 7 days ago

d2fef606

Lift amdgpu dialect into Python/JAX

Arech8 committed 10 days ago

5c81b229

fix WARPGROUP_SIZE for lowering.py

Arech8 committed 12 days ago

3ecdcc2b

rocdl alternatives for nvvm: elect.sync, shfl.sync

Arech8 committed 12 days ago

a9770af7

DISABLE important assertion in fragmented_array. RESTORE IT!

Arech8 committed 12 days ago

369d4086

RocmModuleToBinaryPass WIP broken

Arech8 committed 12 days ago

5da96ca1

lift ROCDL dialect to Python/JAX/jaxlib

Arech8 committed 12 days ago

7631868e

what? it works already?

Arech8 committed 12 days ago

5e0aca40

WIP build + custom_call

Arech8 committed 12 days ago

02879111

reformat + todo

Arech8 committed 12 days ago

a0ea27de

moved all to the standard custom_call.cc instead of customized file

Arech8 committed 12 days ago

da7c25eb

reformat custom_call.cc

Arech8 committed 12 days ago

35ca7d48

move rocm/custom_call back to main ./gpu dir and drop ./rocm

Arech8 committed 12 days ago

4ba2ac49

first e2e work scaffolding

Arech8 committed 12 days ago

67ccae6b

reformat my custom_call

Arech8 committed 12 days ago

afc85c62

Older