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.

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

Test more WMMA configurations #171

Merged
merged 3 commits into from
Jan 2, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
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
19 changes: 17 additions & 2 deletions configs/configs.jl
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,7 @@ macro get_wmma_config()
mul!,
Epilogue.Default(),
verify_default,
Kernel.matmul_pipelined,
kernel,
wmma_baseline)
end end)
end
Expand Down Expand Up @@ -520,7 +520,22 @@ function get_configs()
[2, 2, 1],
[1, 1, 2],
[2, 2, 2]], [[2048, 2048, 2048]]),
zero_c in [false]
zero_c in [false],
kernel in [Kernel.matmul_pipelined]

push!(rv, @get_wmma_config)
end

# WMMA GEMM parameters
for (M, N, K) in [(256, 256, 256)],
(AB_type, CD_type) in [(Float16, Float32)],
transpose_a in [false, true],
transpose_b in [false, true],
(BLOCK_M, BLOCK_N, BLOCK_K) in filter(x -> prod(x[1:2]) <= 128*128, collect(Iterators.product([64, 128, 256], [64, 128, 256], [16, 32, 64]))[:]),
(WARPS_M, WARPS_N) in filter(x -> prod(x) >= 4, collect(Iterators.product([1, 2, 4], [1, 2, 4]))[:]),
zero_c in [false, true],
(OP_M, OP_N, OP_K) in [(16, 16, 16)],
kernel in [Kernel.matmul_singlestage, Kernel.matmul_pipelined]

push!(rv, @get_wmma_config)
end
Expand Down
33 changes: 33 additions & 0 deletions src/config.jl
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,39 @@
is_b_col_major
end

function Base.show(io::IO, config::Config)
println(io, "matmul_shape: $(config.matmul_shape)")
println(io, "block_shape: $(config.block_shape)")
println(io, "warps_per_block: $(config.warps_per_block)")

println(io, "mem_a_warp: $(config.mem_a_warp)")
println(io, "mem_a_thread: $(config.mem_a_thread)")

println(io, "mem_b_warp: $(config.mem_b_warp)")
println(io, "mem_b_thread: $(config.mem_b_thread)")

println(io, "mem_cd_warp: $(config.mem_cd_warp)")
println(io, "mem_cd_thread: $(config.mem_cd_thread)")

println(io, "compute_warp: $(config.compute_warp)")
println(io, "compute_op_shape: $(config.compute_op_shape)")

println(io, "global_a_layout: $(config.global_a_layout)")
println(io, "global_b_layout: $(config.global_b_layout)")
println(io, "global_c_layout: $(config.global_c_layout)")
println(io, "global_d_layout: $(config.global_d_layout)")

println(io, "shared_a_layout: $(config.shared_a_layout)")
println(io, "shared_b_layout: $(config.shared_b_layout)")
println(io, "shared_c_layout: $(config.shared_c_layout)")
println(io, "shared_d_layout: $(config.shared_d_layout)")

println(io, "operator: $(config.operator)")

println(io, "is_a_col_major: $(config.is_a_col_major)")
println(io, "is_b_col_major: $(config.is_b_col_major)")
end

struct ConfigError <: Exception
message::String
end
Expand Down
15 changes: 12 additions & 3 deletions test/matmul.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,17 @@ include("../configs/configs.jl")

@testset "Matrix multiplication" begin
@testcase "$( cf.name )" for cf in get_configs()
c_h, a, b, c, d = generate_inputs(cf)
run_gemm(cf, a, b, c, d)
@test verify(cf, c_h, d)
try
c_h, a, b, c, d = generate_inputs(cf)
run_gemm(cf, a, b, c, d)
@test verify(cf, c_h, d)
catch err
# Count tests with config errors as "broken".
if isa(err, GemmKernels.ConfigError)
@test true skip=true
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@maleadt For now, I've marked unsupported configurations as "Broken". The term is not really appropriate, as broken tests should be tests that ought to pass, but do not currently. We could also just mark these as "pass", but I quite like that they are reported separately so we can easily see how many configurations are skipped due to ConfigErrors. WDYT?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ideally we hard-code whether a configuration is unsupported, i.e., which configuration we know is supposed to throw a ConfigError. That way the broken would work as intended, resulting in test failures when it anything starts failing/passing without updating the tests. However, I guess that would be hard to do, given how we now generate configurations using multiple loops. Maybe we could maintain a separate list of known-broken configurations? I guess it would also depend on the device though (because of the shmem limitation), so maybe that's not feasible either...

else
rethrow()
end
end
end
end