Add translate variant for offset

Author: thomasfaingnaert

src/device/matmul_kernels/kernel.jl

@@ -45,7 +45,7 @@ function matmul_impl(a, b, c, d,
 
     @unroll for i = 1 : NUM_FRAGMENTS_M
         @unroll for j = 1 : NUM_FRAGMENTS_N
-            tile = translate(warp_tile, (M = (i-1)*COMPUTE_OP_SHAPE.M, N = (j-1)*COMPUTE_OP_SHAPE.N))
+            tile = translate_offset(warp_tile, (M = (i-1)*COMPUTE_OP_SHAPE.M, N = (j-1)*COMPUTE_OP_SHAPE.N))
             @inbounds c_frags[i, j] = transf_sh2rf_c(Operator.load_c(OPERATOR, SHARED_C_LAYOUT, shmem_c, tile), tile)
         end
     end
@@ -84,15 +84,15 @@ function matmul_impl(a, b, c, d,
             a_frags = MArray{Tuple{NUM_FRAGMENTS_M}, Operator.fragtype_a(OPERATOR, SHARED_A_LAYOUT)}(undef)
 
             @unroll for i = 1 : NUM_FRAGMENTS_M
-                a_tile = translate(warp_tile.MK, (M = (i-1)*COMPUTE_OP_SHAPE.M, K = 0))
+                a_tile = translate_offset(warp_tile.MK, (M = (i-1)*COMPUTE_OP_SHAPE.M, K = 0))
                 @inbounds a_frags[i] = transf_sh2rf_a(Operator.load_a(OPERATOR, SHARED_A_LAYOUT, shmem_a, a_tile), a_tile)
             end
 
             # (3.3.2) Load a COMPUTE_WARP.K x COMPUTE_WARP.N tile of B from shared memory into registers
             b_frags = MArray{Tuple{NUM_FRAGMENTS_N}, Operator.fragtype_b(OPERATOR, SHARED_B_LAYOUT)}(undef)
 
             @unroll for j = 1 : NUM_FRAGMENTS_N
-                b_tile = translate(warp_tile.KN, (K = 0, N = (j-1)*COMPUTE_OP_SHAPE.N))
+                b_tile = translate_offset(warp_tile.KN, (K = 0, N = (j-1)*COMPUTE_OP_SHAPE.N))
                 @inbounds b_frags[j] = transf_sh2rf_b(Operator.load_b(OPERATOR, SHARED_B_LAYOUT, shmem_b, b_tile), b_tile)
             end
 
@@ -114,7 +114,7 @@ function matmul_impl(a, b, c, d,
 
     @unroll for i = 1 : NUM_FRAGMENTS_M
         @unroll for j = 1 : NUM_FRAGMENTS_N
-            tile = translate(warp_tile, (M = (i-1)*COMPUTE_OP_SHAPE.M, N = (j-1)*COMPUTE_OP_SHAPE.N))
+            tile = translate_offset(warp_tile, (M = (i-1)*COMPUTE_OP_SHAPE.M, N = (j-1)*COMPUTE_OP_SHAPE.N))
             Operator.store_d(OPERATOR, SHARED_D_LAYOUT, shmem_d, transf_rf2sh_d(c_frags[i, j], tile), tile)
         end
     end

src/device/matmul_kernels/operator.jl

@@ -25,31 +25,43 @@ struct WMMAOp{M, N, K} end
 @inline fragtype_b(::Type{WMMAOp{16, 16, 16}}, ::Type{Layout.AlignedColMajor{Float16}}) = WMMA.Fragment{16, 16, 16, 16, Float16, WMMA.ColMajor, WMMA.MatrixB}
 @inline fragtype_accum(::Type{WMMAOp{16, 16, 16}}, ::Type{Layout.AlignedColMajor{Float32}}) = WMMA.Fragment{16, 16, 16, 8, Float32, WMMA.Unspecified, WMMA.Accumulator}
 
-function load_a(::Type{WMMAOp{M, N, K}}, ::Type{Layout.AlignedColMajor{Float16}}, workspace, tile::Tile) where {M, N, K}
+@inline function load_a(::Type{WMMAOp{M, N, K}}, ::Type{Layout.AlignedColMajor{Float16}}, workspace, tile::Tile) where {M, N, K}
     conf = WMMA.Config{M, N, K, Float32}
-    linear_index = linearise(tile.index, size(workspace))
-    ptr = pointer(workspace, linear_index)
+
+    linear_base = linearise(tile.base, size(workspace))
+    linear_offset = linearise(tile.offset, size(workspace))
+
+    ptr = pointer(workspace, linear_base) + (linear_offset - 1) * sizeof(Float16)
     return WMMA.load_a(ptr, size(workspace, 1), WMMA.ColMajor, conf)
 end
 
-function load_b(::Type{WMMAOp{M, N, K}}, ::Type{Layout.AlignedColMajor{Float16}}, workspace, tile::Tile) where {M, N, K}
+@inline function load_b(::Type{WMMAOp{M, N, K}}, ::Type{Layout.AlignedColMajor{Float16}}, workspace, tile::Tile) where {M, N, K}
     conf = WMMA.Config{M, N, K, Float32}
-    linear_index = linearise(tile.index, size(workspace))
-    ptr = pointer(workspace, linear_index)
+
+    linear_base = linearise(tile.base, size(workspace))
+    linear_offset = linearise(tile.offset, size(workspace))
+
+    ptr = pointer(workspace, linear_base) + (linear_offset - 1) * sizeof(Float16)
     return WMMA.load_b(ptr, size(workspace, 1), WMMA.ColMajor, conf)
 end
 
-function load_c(::Type{WMMAOp{M, N, K}}, ::Type{Layout.AlignedColMajor{Float32}}, workspace, tile::Tile) where {M, N, K}
+@inline function load_c(::Type{WMMAOp{M, N, K}}, ::Type{Layout.AlignedColMajor{Float32}}, workspace, tile::Tile) where {M, N, K}
     conf = WMMA.Config{M, N, K, Float32}
-    linear_index = linearise(tile.index, size(workspace))
-    ptr = pointer(workspace, linear_index)
+
+    linear_base = linearise(tile.base, size(workspace))
+    linear_offset = linearise(tile.offset, size(workspace))
+
+    ptr = pointer(workspace, linear_base) + (linear_offset - 1) * sizeof(Float32)
     return WMMA.load_c(ptr, size(workspace, 1), WMMA.ColMajor, conf)
 end
 
-function store_d(::Type{WMMAOp{M, N, K}}, ::Type{Layout.AlignedColMajor{Float32}}, workspace, frag, tile::Tile) where {M, N, K}
+@inline function store_d(::Type{WMMAOp{M, N, K}}, ::Type{Layout.AlignedColMajor{Float32}}, workspace, frag, tile::Tile) where {M, N, K}
     conf = WMMA.Config{M, N, K, Float32}
-    linear_index = linearise(tile.index, size(workspace))
-    ptr = pointer(workspace, linear_index)
+
+    linear_base = linearise(tile.base, size(workspace))
+    linear_offset = linearise(tile.offset, size(workspace))
+
+    ptr = pointer(workspace, linear_base) + (linear_offset - 1) * sizeof(Float32)
     WMMA.store_d(ptr, frag, size(workspace, 1), WMMA.ColMajor, conf)
 end
 

src/device/tiling.jl

@@ -132,6 +132,15 @@ end
 
 @inline translate(tile::Tile{size, names, T}, offset::Tuple) where {names, T, size} = translate(tile, NamedTuple{names}(offset))
 
+export translate_offset
+
+@inline function translate_offset(tile::Tile{size, names, T}, offset::NamedTuple{names, T}) where {names, T, size}
+    new_offset = map(+, tile.offset, offset)
+    return Tile{size, names, T}(tile.base, new_offset)
+end
+
+@inline translate_offset(tile::Tile{size, names, T}, offset::Tuple) where {names, T, size} = translate_offset(tile, NamedTuple{names}(offset))
+
 # -------------
 # TileIterators
 # -------------