How does constant propagation impact latency with a Val-based dispatch?

jishnub · May 22, 2024, 5:49am

The example here is schematic, and isn’t really an MWE.

Suppose I have a set of functions like

f() = f('A')
Base.@constprop :aggressive f(a, args...) = g(Val(a == 'A'), args...)
Base.@constprop :aggressive g(::Val{true}, args...) = 1
Base.@constprop :aggressive g(::Val{false}, args...) = 2

In this case, code_typed for f() detects that g(::Val{false}) is unreachable, and inlines g(::Val{true}), to produce

julia> @code_typed f()
CodeInfo(
1 ─     return 1
) => Int64

However, the first call still seems to propagate constants into g(::Val{false}). In my use case, propagating constants into g(::Val{false}) is necessary for type-inference within the method, but this adds quite a bit to the latency of f(). I would love it if the method was entirely ignored in this call, as it is not being compiled subsequently.

Simplifying the method definition of g(::Val{false}, args...) appears to reduce TTFX for f(), which I had not expected, as this method is not being used subsequently in the call.

The actual use case that I’m looking at is in matmul, where the latency in generic_matmul!

github.com

JuliaLang/julia/blob/263928f9ad450a28d601c3f185040987e0bc45d5/stdlib/LinearAlgebra/src/matmul.jl#L372-L418


      
          Base.@constprop :aggressive function generic_matmatmul!(C::StridedMatrix{T}, tA, tB, A::StridedVecOrMat{T}, B::StridedVecOrMat{T},
                                              α::Number, β::Number) where {T<:BlasFloat}
              mA, nA = lapack_size(tA, A)
              mB, nB = lapack_size(tB, B)
              if any(iszero, size(A)) || any(iszero, size(B)) || iszero(α)
                  if size(C) != (mA, nB)
                      throw(DimensionMismatch(lazy"C has dimensions $(size(C)), should have ($mA,$nB)"))
                  end
                  return _rmul_or_fill!(C, β)
              end
              if size(C) == size(A) == size(B) == (2,2)
                  return matmul2x2!(C, tA, tB, A, B, α, β)
              end
              if size(C) == size(A) == size(B) == (3,3)
                  return matmul3x3!(C, tA, tB, A, B, α, β)
              end
              # We convert the chars to uppercase to potentially unwrap a WrapperChar,
              # and extract the char corresponding to the wrapper type
              tA_uc, tB_uc = uppercase(tA), uppercase(tB)
              # the map in all ensures constprop by acting on tA and tB individually, instead of looping over them.

This file has been truncated. show original

may be reduced considerably if the last call

github.com

JuliaLang/julia/blob/263928f9ad450a28d601c3f185040987e0bc45d5/stdlib/LinearAlgebra/src/matmul.jl#L417


      
                  if tA_uc == 'S' && tB_uc == 'N'
                      return BLAS.symm!('L', tA == 'S' ? 'U' : 'L', alpha, A, B, beta, C)
                  elseif tA_uc == 'N' && tB_uc == 'S'
                      return BLAS.symm!('R', tB == 'S' ? 'U' : 'L', alpha, B, A, beta, C)
                  elseif tA_uc == 'H' && tB_uc == 'N'
                      return BLAS.hemm!('L', tA == 'H' ? 'U' : 'L', alpha, A, B, beta, C)
                  elseif tA_uc == 'N' && tB_uc == 'H'
                      return BLAS.hemm!('R', tB == 'H' ? 'U' : 'L', alpha, B, A, beta, C)
                  end
              end
              return _generic_matmatmul!(C, wrap(A, tA), wrap(B, tB), MulAddMul(α, β))
          end
          # legacy method
          Base.@constprop :aggressive generic_matmatmul!(C::StridedMatrix{T}, tA, tB, A::StridedVecOrMat{T}, B::StridedVecOrMat{T},
                  _add::MulAddMul = MulAddMul()) where {T<:BlasFloat} =
              generic_matmatmul!(C, tA, tB, A, B, _add.alpha, _add.beta)
          
          # Complex matrix times (transposed) real matrix. Reinterpret the first matrix to real for efficiency.
          Base.@constprop :aggressive function generic_matmatmul!(C::StridedVecOrMat{Complex{T}}, tA, tB, A::StridedVecOrMat{Complex{T}}, B::StridedVecOrMat{T},
                              α::Number, β::Number) where {T<:BlasReal}
              # We convert the chars to uppercase to potentially unwrap a WrapperChar,

on line 417 were absent. This line is not reachable in many common cases, but still contributes to the latency of matrix multiplication. The constant-propagation into wrap appears to lead to this latency, and removing the aggressive constprop annotations reduces TTFX. Ideally, I would only want this constprop to happen if the line is reachable.

The connection with a Val-based dispatch is that I was trying to hide this line behind a Val(all(map(in(('N', 'T', 'C')), (tA_uc, tB_uc))), which may be evaluated at compile-time, but doesn’t seem to impact latency.