Merge pull request #120 from JuliaGPU/tb/zero_layout

Add zero layout to optimize alpha/beta=zero.

Author: maleadt

src/blas.jl

@@ -4,22 +4,6 @@ using CUDA
 using GemmKernels
 using LinearAlgebra
 
-# Global layouts
-global_layout(::Type{<:CuArray{T}}, ::Val{false}) where {T} = Layout.AlignedColMajor{T}
-global_layout(::Type{<:CuArray{T}}, ::Val{true}) where {T} = Layout.AlignedRowMajor{T}
-global_layout(::Type{<:Diagonal{Float16, <:CuArray{Float16}}}, transpose) = Layout.Diagonal{Float16}
-
-# Shared layouts for A / B
-shared_layout_ab(typ::Type{<:CuArray{Float16}}, transpose) = Layout.Padded{global_layout(typ, transpose), 8}
-shared_layout_ab(::Type{<:Diagonal{Float16, <:CuArray{Float16, N}}}, transpose) where {N, P} = shared_layout_ab(CuArray{Float16, N}, transpose)
-
-# Shared layouts for C / D
-shared_layout_cd(typ::Type{<:CuArray{T}}, transpose) where {T} = global_layout(typ, transpose)
-
-# Convert matrix to type compatible with kernel
-convert_matrix(mat) = mat
-convert_matrix(mat::Diagonal{T, A}) where {T, A} = mat.diag
-
 # Select the best kernel
 kernel(layout_a, layout_b) = Kernel.matmul_singlestage
 kernel(::Type{Layout.AlignedColMajor{T}}, ::Type{Layout.AlignedColMajor{T}}) where {T} = Kernel.matmul_pipelined
@@ -28,7 +12,8 @@ kernel(::Type{Layout.AlignedRowMajor{T}}, ::Type{Layout.AlignedColMajor{T}}) whe
 kernel(::Type{Layout.AlignedRowMajor{T}}, ::Type{Layout.AlignedRowMajor{T}}) where {T} = Kernel.matmul_pipelined
 
 # Based on https://github.com/JuliaGPU/CUDA.jl/blob/bd5a2a8800e91eb6a7df89eb5dd4bb8fc503541d/lib/cublas/wrappers.jl#L743-L769
-function gemmEx!(transA::Char, transB::Char, alpha::Number, A, B, beta::Number, C)
+function gemmEx!(transA::Char, transB::Char, alpha::Number, A::CuMatrix, B::CuMatrix,
+                 beta::Number, C::CuMatrix)
     m = size(A, transA == 'N' ? 1 : 2)
     k = size(A, transA == 'N' ? 2 : 1)
     n = size(B, transB == 'N' ? 2 : 1)
@@ -40,31 +25,48 @@ function gemmEx!(transA::Char, transB::Char, alpha::Number, A, B, beta::Number,
     transpose_a = (transA == 'T')
     transpose_b = (transB == 'T')
 
-    a_layout = global_layout(typeof(A), Val(transpose_a))
-    b_layout = global_layout(typeof(B), Val(transpose_b))
+    a_layout_base = transpose_a ? Layout.AlignedRowMajor : Layout.AlignedColMajor
+    b_layout_base = transpose_b ? Layout.AlignedRowMajor : Layout.AlignedColMajor
+
+    # determine global memory layouts
+    ## if alpha is zero, we don't need to load A or B
+    if iszero(alpha)
+        global_a_layout = Layout.Zero{eltype(A)}
+        global_b_layout = Layout.Zero{eltype(B)}
+    else
+        global_a_layout = a_layout_base{eltype(A)}
+        global_b_layout = b_layout_base{eltype(B)}
+    end
+    ## if beta is zero, we don't need to load C
+    global_c_layout = if iszero(beta)
+        Layout.Zero{eltype(C)}
+    else
+        Layout.AlignedColMajor{eltype(C)}
+    end
+    global_d_layout = Layout.AlignedColMajor{eltype(C)}
 
-    conf = GemmKernels.get_config(
+    # determine shared memory layouts
+    ## padded to avoid bank conflicts
+    shared_a_layout = Layout.Padded{a_layout_base{eltype(A)}, 8}
+    shared_b_layout = Layout.Padded{b_layout_base{eltype(B)}, 8}
+    ## outputs are never transposed, and padding them doesn't seem worth it
+    shared_c_layout = shared_d_layout = Layout.AlignedColMajor{eltype(C)}
+
+    conf = GemmKernels.get_config(;
             gemm_shape = (M = m, N = n, K = k),
             operator = Operator.WMMAOp{16, 16, 16, eltype(C)},
 
-            global_a_layout = a_layout,
-            global_b_layout = b_layout,
-            global_c_layout = global_layout(typeof(C), Val(false)),
-            global_d_layout = global_layout(typeof(C), Val(false)),
-
-            shared_a_layout = shared_layout_ab(typeof(A), Val(transpose_a)),
-            shared_b_layout = shared_layout_ab(typeof(B), Val(transpose_b)),
-            shared_c_layout = shared_layout_cd(typeof(C), Val(false)),
-            shared_d_layout = shared_layout_cd(typeof(C), Val(false)),
+            global_a_layout, global_b_layout, global_c_layout, global_d_layout,
+            shared_a_layout, shared_b_layout, shared_c_layout, shared_d_layout,
 
             is_a_col_major = !transpose_a,
             is_b_col_major = !transpose_b
                                 )
 
-    GemmKernels.matmul(convert_matrix(A), convert_matrix(B), convert_matrix(C), convert_matrix(C), conf;
+    GemmKernels.matmul(A, B, C, C, conf;
                        transform_shared_to_regs_a = Transform.Elementwise(x -> x * alpha),
                        transform_shared_to_regs_c = Transform.Elementwise(x -> x * beta),
-                       kernel = kernel(a_layout, b_layout)
+                       kernel = kernel(global_a_layout, global_b_layout)
                       )
 end
 

src/layout.jl

@@ -57,6 +57,19 @@ abstract type LayoutBase{T} end
 @inline eltype(::Type{<:LayoutBase{T}}) where {T} = T
 @inline physical_size(::Type{<:LayoutBase{T}}, logical_size::NamedTuple) where {T} = Tuple(logical_size)
 
+# ----
+# Zero
+# ----
+
+abstract type Zero{T} <: LayoutBase{T} end
+
+@inline function load(::Type{<:Zero{T}}, workspace, tile::Tile{size}) where {T, size}
+    N = 16 ÷ sizeof(T)
+    return ntuple(i -> VecElement{T}(zero(T)), Val(N))
+end
+
+@inline store!(::Type{<:Zero{T}}, workspace, value, tile::Tile) where {T} = return
+
 # --------------
 # Padded layouts
 # --------------

test/blas.jl

@@ -36,35 +36,4 @@ CUDA.CUBLAS.cublasSetMathMode(CUBLAS.handle(), CUBLAS.CUBLAS_TENSOR_OP_MATH)
             @test all(isapprox.(Array(c_gemmkernels), Array(c_cublas); rtol=sqrt(eps(A_type))));
         end
     end
-
-    @testset "WMMA GEMM (A = diagonal, B = $( !transpose_b ? 'N' : 'T' ))" for transpose_b = [false, true]
-        @testset "(M = $M, N = $N, K = $K)" for M in [128, 256],
-            N in [128, 256],
-            K in [M]
-
-            transpose_a = false
-
-            alpha = rand(Float32)
-            beta = rand(Float32)
-
-            a_h = rand(Float16, M);
-            b_h = rand(Float16, (K, N)) / sqrt(Float16(K))
-            c_h = rand(Float32, (M, N))
-
-            # Transpose input if necessary
-            a_h = transpose_a ? transpose(a_h) : a_h
-            b_h = transpose_b ? transpose(b_h) : b_h
-
-            a   = Diagonal(CuArray(a_h))
-            b   = CuArray(b_h)
-
-            c_gemmkernels = CuArray(c_h)
-            GemmKernels.BLAS.gemmEx!(!transpose_a ? 'N' : 'T', !transpose_b ? 'N' : 'T', alpha, a, b, beta, c_gemmkernels)
-
-            c_cublas = CuArray(c_h)
-            CUDA.CUBLAS.gemmEx!(!transpose_a ? 'N' : 'T', !transpose_b ? 'N' : 'T', alpha, CuArray(Array(Diagonal(a_h))), b, beta, c_cublas)
-
-            @test all(isapprox.(Array(c_gemmkernels), Array(c_cublas); rtol=sqrt(eps(Float16))));
-        end
-    end
 end