Skip to content

Update performant_matmul.jl #590

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

Closed
wants to merge 1 commit into from
Closed

Update performant_matmul.jl #590

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
53 changes: 23 additions & 30 deletions examples/performant_matmul.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,36 +21,31 @@ const TILE_DIM = 32

# private variable for tile output
outval = @private eltype(output) 1
@inbounds outval[1] = -zero(eltype(output))
@inbounds outval[1] = zero(eltype(output))

@uniform N = size(output, 1)
# number of tiles depends on inner dimension
@uniform NUM_TILES = div(R + TILE_DIM - 1, TILE_DIM)
@uniform NUM_TILES = cld(R, TILE_DIM)

# Can't use @index(Global), because we use a smaller ndrange
I = (gi - 1) * TILE_DIM + i
J = (gj - 1) * TILE_DIM + j
# loop over all tiles needed for this calculation
for t in 0:(NUM_TILES - 1)
# Can't use @index(Global), because we use a smaller ndrange
I = (gi - 1) * TILE_DIM + i
J = (gj - 1) * TILE_DIM + j

# load inputs into tiles, with bounds checking for non-square matrices
if I <= N && t * TILE_DIM + j <= R
@inbounds tile1[i, j] = input1[I, t * TILE_DIM + j]
else
@inbounds tile1[i, j] = 0.0
end

if t * TILE_DIM + i <= R && J <= M
@inbounds tile2[i, j] = input2[t * TILE_DIM + i, J]
else
@inbounds tile2[i, j] = 0.0
end

# wait for all tiles to be loaded
@synchronize

# get global values again
I = (gi - 1) * TILE_DIM + i
J = (gj - 1) * TILE_DIM + j
@synchronize(true)

# calculate value of spot in output, use temporary value to allow for vectorization
out = zero(eltype(output))
Expand All @@ -59,29 +54,27 @@ const TILE_DIM = 32
end
outval[1] += out

@synchronize
@synchronize(true)
end

# get global indices again
I = (gi - 1) * TILE_DIM + i
J = (gj - 1) * TILE_DIM + j

# save if inbounds
if I <= N && J <= M
@inbounds output[I, J] = outval[1]
end
end

N = 1024
R = 512
M = 2048
A = rand!(allocate(backend, Float32, N, R))
B = rand!(allocate(backend, Float32, R, M))
C = KernelAbstractions.zeros(backend, Float32, N, M)

kern = coalesced_matmul_kernel!(backend, (TILE_DIM, TILE_DIM))

kern(C, A, B, N, R, M, ndrange = size(C))
KernelAbstractions.synchronize(backend)

@test isapprox(A * B, C)
@testset "dims for $N, $R, $M" for (N,R,M) in [rand(500:1000,3) for _ in 1:10]
A = rand!(allocate(backend, Float32, N, R))
B = rand!(allocate(backend, Float32, R, M))
C = KernelAbstractions.zeros(backend, Float32, N, M)

kern = coalesced_matmul_kernel!(backend, (TILE_DIM, TILE_DIM))

group_size_x = cld(N, TILE_DIM)
group_size_y = cld(M, TILE_DIM)

kern(C, A, B, N, R, M, ndrange = (group_size_x * TILE_DIM, group_size_y * TILE_DIM))
KernelAbstractions.synchronize(backend)

@test isapprox(A * B, C)
end