ROCm/jax

Pull Requests Commits

Arech8 committed 55 days ago

394ac94e

async_copy() stub

Arech8 committed 55 days ago

37459f91

make kernel launcher use proper WORKGROUP_SIZE number of threads

Arech8 committed 55 days ago

67fe951b

enable IR verification for debug builds

Arech8 committed 55 days ago

a3f406dd

bugfix amdgpu dialect handling

Arech8 committed 55 days ago

0192a988

Arech8 committed 55 days ago

ca9bf05f

Arech8 committed 55 days ago

a1bddea2

Improvements and cleanup

Arech8 committed 55 days ago

245d0839

Arech8 committed 55 days ago

c7bc5894

Fix a fence in the init section

Arech8 committed 55 days ago

9f8746a7

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

Arech8 committed 55 days ago

230b777e

fix warp_barrier()

Arech8 committed 55 days ago

3d15b761

improve comments

Arech8 committed 55 days ago

0fd36905

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

Arech8 committed 55 days ago

f2f0f7ef

Implement warpgroup_barrier() with s_waitcnt + s_barrier

Arech8 committed 55 days ago

bb96dc60

compile and launch now works

Arech8 committed 55 days ago

1f067fbb

stubify bar.sync too for ROCm

Arech8 committed 55 days ago

ff14795b

support for the runtime

Arech8 committed 55 days ago

3e312474

WIP emrgency saving - working on C++ runtime support

Arech8 committed 55 days ago

af41b375

format runtime.cc

Arech8 committed 55 days ago

af39c8ce

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

Arech8 committed 55 days ago

d66802bc

Lift amdgpu dialect into Python/JAX

Arech8 committed 55 days ago

6f530407

fix WARPGROUP_SIZE for lowering.py

Arech8 committed 55 days ago

c3da8177

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

Arech8 committed 55 days ago

d7ca8cac

DISABLE important assertion in fragmented_array. RESTORE IT!

Arech8 committed 55 days ago

481439f0

RocmModuleToBinaryPass WIP broken

Arech8 committed 55 days ago

c7e38206

lift ROCDL dialect to Python/JAX/jaxlib

Arech8 committed 55 days ago

ebd7ff3d

what? it works already?

Arech8 committed 55 days ago

77fbdfc1

WIP build + custom_call

Arech8 committed 55 days ago

ed2577aa

reformat + todo

Arech8 committed 55 days ago

be3cd930

Older