Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add fused GEMM + SiLU kernel #180

Open
wants to merge 12 commits into
base: main
Choose a base branch
from
Open

Add fused GEMM + SiLU kernel #180

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
0327398
Reorder iter args to match ordering of init args and outputs (#161)
harsh-nod Sep 24, 2024
4e95351
[ExportedProgram] Add mutable attribute to buffer (#123)
chrsmcgrr Sep 24, 2024
65eb532
[TKW] Add xfail decorator for unaligned shape (#163)
raikonenfnu Sep 24, 2024
909411a
[TKW] Fix indexing of Reduction and GetResult to enable post-tile op.…
raikonenfnu Sep 24, 2024
d37c6a4
Get GEMMs working without minimize_global_loads (#167)
harsh-nod Sep 26, 2024
04a4ba5
Add first draft of introduction (#168)
harsh-nod Sep 26, 2024
7686157
[TKW] igemm shared mem tests (#171)
Hardcode84 Sep 26, 2024
0e16d54
[TKW] Implement support for multiple iter args on Reduction (#166)
raikonenfnu Sep 27, 2024
192a786
Handle complex element type in torch.vtensor conversion (#175)
sogartar Sep 27, 2024
92ad900
[TKW] Rework vector mask generation (#172)
Hardcode84 Sep 30, 2024
621cbe1
Enable import_symbolic_shape_expressions in the FxImporter. (#179)
stellaraccident Sep 30, 2024
7de2e0b
Add fused GEMM + SiLU kernel
harsh-nod Oct 1, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion .github/workflows/ci.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,6 @@ jobs:
if: "contains(matrix.os, 'mi300') && !cancelled()"
run: |
export WAVE_RUN_E2E_TESTS=1
export TEST_PARAMS_PATH=./tests/kernel/wave/test_param.json
pytest -n 4 ./tests/kernel/wave/

- name: Run LIT tests
Expand Down
4 changes: 2 additions & 2 deletions lit_tests/kernel/wave/barriers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,9 +207,9 @@ def test_gemm():

# Reduction subgraph:
# CHECK: %acc_0_0_0
# CHECK-NEXT: %acc_1_1_0
# CHECK-NEXT: %acc_1_0_0
# CHECK-NEXT: %acc_0_1_0
# CHECK-NEXT: %acc_1_0_0
# CHECK-NEXT: %acc_1_1_0
# CHECK-NEXT: %a
# CHECK-NEXT: %read_0_0_0
# CHECK-NEXT: %read_0_0_1
Expand Down
457 changes: 323 additions & 134 deletions lit_tests/kernel/wave/codegen.py
View file
Open in desktop

Large diffs are not rendered by default.

32 changes: 16 additions & 16 deletions lit_tests/kernel/wave/expansion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,31 +243,31 @@ def test_gemm():
# CHECK-NEXT: placeholder(_name=a
# CHECK-NEXT: placeholder(_name=b
# CHECK-NEXT: placeholder(_name=c
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N})
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + 16})
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N})
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + 16})
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16)})
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16) + 16})
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16)})
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16) + 16})
# CHECK-NEXT: reduction(axis=K, init_args=[register_0_0_0, register_0_1_0, register_1_0_0, register_1_1_0], subgraph_name=region_0, implicit_captures=[a, b])
# CHECK-NEXT: get_result(value=reduction, res_idx=3)
# CHECK-NEXT: get_result(value=reduction, res_idx=2)
# CHECK-NEXT: get_result(value=reduction, res_idx=1)
# CHECK-NEXT: get_result(value=reduction, res_idx=0)
# CHECK-NEXT: write(register_=getresult_0_0_0
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N}
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16)}
# CHECK-NEXT: write(register_=getresult_1_1_0
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + 16}
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16) + 16}
# CHECK-NEXT: write(register_=getresult_1_0_0
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N}
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16)}
# CHECK-NEXT: write(register_=getresult_0_1_0
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + 16}
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16) + 16}
# CHECK-NEXT: output

# Reduction subgraph:

# CHECK: %acc_0_0_0
# CHECK-NEXT: %acc_1_1_0
# CHECK-NEXT: %acc_1_0_0
# CHECK-NEXT: %acc_0_1_0
# CHECK-NEXT: %acc_1_0_0
# CHECK-NEXT: %acc_1_1_0

# CHECK-NEXT: %a
# CHECK-NEXT: %read_0_0_0
Expand Down Expand Up @@ -305,13 +305,13 @@ def test_gemm():
# CHECK-SAME: (%read_0_0_0, %read_0_1_0, %acc_0_1_0)
# CHECK-NEXT: %mma_0_1_1
# CHECK-SAME: (%read_0_0_1, %read_0_1_1, %mma_0_1_0)
# CHECK-NEXT: return [mma_0_0_1, mma_1_1_1, mma_1_0_1, mma_0_1_1]
# CHECK-NEXT: return [mma_0_0_1, mma_0_1_1, mma_1_0_1, mma_1_1_1]

# Custom format:
# CHECK-NEXT: placeholder(_name=acc_0_0_0
# CHECK-NEXT: placeholder(_name=acc_1_1_0
# CHECK-NEXT: placeholder(_name=acc_1_0_0
# CHECK-NEXT: placeholder(_name=acc_0_1_0
# CHECK-NEXT: placeholder(_name=acc_1_0_0
# CHECK-NEXT: placeholder(_name=acc_1_1_0
# CHECK-NEXT: placeholder(_name=a
# CHECK-NEXT: read(memory=a, elements_per_thread=4, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), K: ARGK*BLOCK_K + 4*floor((Mod($T0, 64))/16) : 4 : 1})
# CHECK-NEXT: read(memory=a, elements_per_thread=4, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), K: ARGK*BLOCK_K + 4*floor((Mod($T0, 64))/16) + 16 : 4 : 1})
Expand Down Expand Up @@ -346,7 +346,7 @@ def test_gemm():
# CHECK-NEXT: mma(lhs=read_0_0_1 (index = {M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), K: ARGK*BLOCK_K + 4*floor((Mod($T0, 64))/16) + 16 : 4 : 1})
# CHECK-SAME: rhs=read_0_1_1 (index = {N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16) + 16, K: ARGK*BLOCK_K + 4*floor((Mod($T0, 64))/16) + 16 : 4 : 1})
# CHECK-SAME: acc=mma_0_1_0 (index = {M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + 4*floor((Mod($T0, 64))/16) : 4 : 16, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16) + 16}))
# CHECK-NEXT: output(return_vals=([mma_0_0_1, mma_1_1_1, mma_1_0_1, mma_0_1_1],))
# CHECK-NEXT: output(return_vals=([mma_0_0_1, mma_0_1_1, mma_1_0_1, mma_1_1_1],))

# CHECK-NEXT: -----

Expand Down Expand Up @@ -389,11 +389,11 @@ def test_gemm_reduction_expansion_only():
# CHECK-NEXT: placeholder(_name=a
# CHECK-NEXT: placeholder(_name=b
# CHECK-NEXT: placeholder(_name=c
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N})
# CHECK-NEXT: register(shape=(M, N), dtype=f32, value=0.0, index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16)})
# CHECK-NEXT: reduction(axis=K, init_args=[register_0_0_0]
# CHECK-NEXT: get_result(value=reduction, res_idx=0)
# CHECK-NEXT: write(register_=getresult_0_0_0
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M, N: $T1*BLOCK_N/2 + $WG1*BLOCK_N})
# CHECK-SAME: index={M: $T0*BLOCK_M/128 + $WG0*BLOCK_M + Mod($T0, 16), N: $T1*BLOCK_N/2 + $WG1*BLOCK_N + Mod($T0, 16)})
# CHECK-NEXT: output(return_vals=(None,))

# Reduction subgraph:
Expand Down
32 changes: 16 additions & 16 deletions lit_tests/kernel/wave/index_sequence_analysis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,13 +182,13 @@ def test_gemm():
# CHECK-NEXT: placeholder(_name=b
# CHECK-NEXT: placeholder(_name=c
# CHECK-NEXT: register
# CHECK-SAME: index={M: $WG0*BLOCK_M, N: $WG1*BLOCK_N + BLOCK_N/2})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16), N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16)})
# CHECK-NEXT: register(
# CHECK-SAME: index={M: $WG0*BLOCK_M + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + 16})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16) + 16})
# CHECK-NEXT: register(
# CHECK-SAME: index={M: $WG0*BLOCK_M + 16, N: $WG1*BLOCK_N + BLOCK_N/2})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16)})
# CHECK-NEXT: register(
# CHECK-SAME: index={M: $WG0*BLOCK_M, N: $WG1*BLOCK_N + BLOCK_N/2 + 16})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16), N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16) + 16})
# CHECK-NEXT: allocate(
# CHECK-NEXT: allocate(
# CHECK-NEXT: reduction(
Expand All @@ -210,16 +210,16 @@ def test_gemm():
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 3, N: 64*$WG1 + Mod($T0, 16) + 32})
# CHECK-NEXT: extract_slice(register_=getresult_1_1_0, offset=[0], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_4, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16), N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 16, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-NEXT: extract_slice(register_=getresult_1_1_0, offset=[1], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_5, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 1, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 17, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-NEXT: extract_slice(register_=getresult_1_1_0, offset=[2], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_6, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 2, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 18, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-NEXT: extract_slice(register_=getresult_1_1_0, offset=[3], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_7, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 3, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 19, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-NEXT: extract_slice(register_=getresult_1_0_0, offset=[0], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_8, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 16, N: 64*$WG1 + Mod($T0, 16) + 32})
Expand All @@ -234,22 +234,22 @@ def test_gemm():
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 19, N: 64*$WG1 + Mod($T0, 16) + 32})
# CHECK-NEXT: extract_slice(register_=getresult_0_1_0, offset=[0], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_12, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 16, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16), N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-NEXT: extract_slice(register_=getresult_0_1_0, offset=[1], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_13, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 17, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 1, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-NEXT: extract_slice(register_=getresult_0_1_0, offset=[2], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_14, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 18, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 2, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-NEXT: extract_slice(register_=getresult_0_1_0, offset=[3], size=[1], stride=[1])
# CHECK-NEXT: write(register_=extract_slice_15, memory=c, elements_per_thread=1,
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 19, N: 64*$WG1 + Mod($T0, 16) + 48})
# CHECK-SAME: index={M: 64*$WG0 + 4*floor((Mod($T0, 64))/16) + 3, N: 64*$WG1 + Mod($T0, 16) + 48})

# Reduction subgraph:
# CHECK: %acc_0_0_0
# CHECK-NEXT: %acc_1_1_0
# CHECK-NEXT: %acc_1_0_0
# CHECK-NEXT: %acc_0_1_0
# CHECK-NEXT: %acc_1_0_0
# CHECK-NEXT: %acc_1_1_0
# CHECK-NEXT: %a
# CHECK-NEXT: %read_4
# CHECK-SAME: (%a, 8, None, None)
Expand Down Expand Up @@ -303,9 +303,9 @@ def test_gemm():

# Reduction subgraph (custom format):
# CHECK: placeholder(_name=acc_0_0_0
# CHECK-NEXT: placeholder(_name=acc_1_1_0
# CHECK-NEXT: placeholder(_name=acc_1_0_0
# CHECK-NEXT: placeholder(_name=acc_0_1_0
# CHECK-NEXT: placeholder(_name=acc_1_0_0
# CHECK-NEXT: placeholder(_name=acc_1_1_0
# CHECK-NEXT: placeholder(_name=a
# CHECK-NEXT: read(memory=a, elements_per_thread=8,
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod(16*$T1 + 32*$T2 + floor($T0/8), 64), K: ARGK*BLOCK_K + 8*(Mod($T0, 8)) : 8 : 1})
Expand Down
24 changes: 12 additions & 12 deletions lit_tests/kernel/wave/minimize_global_loads.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,13 +131,13 @@ def test_gemm():
# CHECK-NEXT: placeholder(_name=b
# CHECK-NEXT: placeholder(_name=c
# CHECK-NEXT: register
# CHECK-SAME: index={M: $WG0*BLOCK_M, N: $WG1*BLOCK_N + BLOCK_N/2})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16), N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16)})
# CHECK-NEXT: register(
# CHECK-SAME: index={M: $WG0*BLOCK_M + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + 16})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16) + 16})
# CHECK-NEXT: register(
# CHECK-SAME: index={M: $WG0*BLOCK_M + 16, N: $WG1*BLOCK_N + BLOCK_N/2})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16)})
# CHECK-NEXT: register(
# CHECK-SAME: index={M: $WG0*BLOCK_M, N: $WG1*BLOCK_N + BLOCK_N/2 + 16})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16), N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16) + 16})
# CHECK-NEXT: allocate(
# CHECK-NEXT: allocate(
# CHECK-NEXT: reduction(
Expand All @@ -146,19 +146,19 @@ def test_gemm():
# CHECK-NEXT: get_result(value=reduction, res_idx=1)
# CHECK-NEXT: get_result(value=reduction, res_idx=0)
# CHECK-NEXT: write(register_=getresult_0_0_0, memory=c
# CHECK-SAME: index={M: $WG0*BLOCK_M, N: $WG1*BLOCK_N + BLOCK_N/2})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16), N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16)})
# CHECK-NEXT: write(register_=getresult_1_1_0, memory=c
# CHECK-SAME: index={M: $WG0*BLOCK_M + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + 16})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16) + 16})
# CHECK-NEXT: write(register_=getresult_1_0_0, memory=c
# CHECK-SAME: index={M: $WG0*BLOCK_M + 16, N: $WG1*BLOCK_N + BLOCK_N/2})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16) + 16, N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16)})
# CHECK-NEXT: write(register_=getresult_0_1_0, memory=c
# CHECK-SAME: index={M: $WG0*BLOCK_M, N: $WG1*BLOCK_N + BLOCK_N/2 + 16})
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod($T0, 16), N: $WG1*BLOCK_N + BLOCK_N/2 + Mod($T0, 16) + 16})

# Reduction subgraph:
# CHECK: %acc_0_0_0
# CHECK-NEXT: %acc_1_1_0
# CHECK-NEXT: %acc_1_0_0
# CHECK-NEXT: %acc_0_1_0
# CHECK-NEXT: %acc_1_0_0
# CHECK-NEXT: %acc_1_1_0
# CHECK-NEXT: %a
# CHECK-NEXT: %read_4
# CHECK-SAME: (%a, 8, None, None)
Expand Down Expand Up @@ -215,9 +215,9 @@ def test_gemm():

# Reduction subgraph (custom format):
# CHECK: placeholder(_name=acc_0_0_0
# CHECK-NEXT: placeholder(_name=acc_1_1_0
# CHECK-NEXT: placeholder(_name=acc_1_0_0
# CHECK-NEXT: placeholder(_name=acc_0_1_0
# CHECK-NEXT: placeholder(_name=acc_1_0_0
# CHECK-NEXT: placeholder(_name=acc_1_1_0
# CHECK-NEXT: placeholder(_name=a
# CHECK-NEXT: read(memory=a, elements_per_thread=8,
# CHECK-SAME: index={M: $WG0*BLOCK_M + Mod(16*$T1 + 32*$T2 + floor($T0/8), 64), K: ARGK*BLOCK_K + 8*(Mod($T0, 8)) : 8 : 1})
Expand Down
26 changes: 21 additions & 5 deletions shark_turbine/aot/compiled_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@

from .support.ir_utils import (
ModuleBuilder,
ModuleBuilderOptions,
)


Expand Down Expand Up @@ -162,11 +163,13 @@ class CompiledModuleClassInfo:
__slots__ = [
"all_exports",
"ir_module_name",
"options",
]

def __init__(self, *, ir_module_name: str):
def __init__(self, *, ir_module_name: str, options: ModuleBuilderOptions):
self.ir_module_name = ir_module_name
self.all_exports: Dict[str, Exportable] = dict()
self.options = options

def add_export(self, key: str, value: Exportable):
if key in self.all_exports:
Expand Down Expand Up @@ -370,13 +373,23 @@ class CompiledModuleMeta(type):
# It is passed the dictionary of declared attributes and any keyword
# arguments from the class declaration:
# class Foo(Bar, kwarg="you probably just learned this is possible"):
def __new__(mcls, name: str, bases, dct, *, export_name: Optional[str] = None):
def __new__(
mcls,
name: str,
bases,
dct,
*,
export_name: Optional[str] = None,
options: Optional[ModuleBuilderOptions] = None,
):
if not _metaclass_setup_complete:
return type.__new__(mcls, name, bases, dct)

ir_module_name = _derive_ir_module_name(name, export_name)
logger.debug("Create new CompiledModule: %s", ir_module_name)
info = CompiledModuleClassInfo(ir_module_name=ir_module_name)
info = CompiledModuleClassInfo(
ir_module_name=ir_module_name, options=options or ModuleBuilderOptions()
)

# Process that attributes that were set as part of class definition.
# Any attributes that we decide are part of the compiled module
Expand Down Expand Up @@ -436,6 +449,7 @@ def create_from_dict(
dct: dict,
*,
export_name: Optional[str] = None,
options: Optional[ModuleBuilderOptions] = None,
) -> CompiledModuleMeta:
"""Creates a CompiledModule subclass with an explicit dictionary of members.

Expand All @@ -446,7 +460,9 @@ class Foo(CompiledModule, export_name="bar"):
def member(): ...
```
"""
return CompiledModuleMeta(name, (cls,), dct, export_name=export_name)
return CompiledModuleMeta(
name, (cls,), dct, export_name=export_name, options=options
)

@staticmethod
def get_class_info(cls: CompiledModuleMeta) -> CompiledModuleClassInfo:
Expand Down Expand Up @@ -596,7 +612,7 @@ def __new__(
module_op.attributes["sym_name"] = StringAttr.get(
class_info.ir_module_name, context=context
)
module_builder = ModuleBuilder(module_op)
module_builder = ModuleBuilder(module_op, options=class_info.options)
info = CompiledModuleInstanceInfo(class_info, module_builder=module_builder)
_all_compiled_module_instance_infos[self] = info

Expand Down
Loading
Loading