Seeing some weird erroring when branching (being fixed in #53):

        qk = if !EVEN_K[] && j >= mask_start
            offs_n = ((j-Int32(1)) * TILE_N[]) .+ offs_n_tile
            mask = ct.full((TILE_N[], TILE_M[]), true, Bool)
            mask = mask .& (offs_n .<= k_seqlen)
            mask = ct.where(mask, ct.zeros((TILE_N[], TILE_M[],), Float32), ct.full((TILE_N[], TILE_M[],), -Inf32, Float32))
            qk .+ mask
        else
            qk
        end

        if !EVEN_K[] && j >= mask_start
            offs_n = ((j-Int32(1)) * TILE_N[]) .+ offs_n_tile
            mask = ct.full((TILE_N[], TILE_M[]), true, Bool)
            mask = mask .& (offs_n .<= k_seqlen)
            mask = ct.where(mask, ct.zeros((TILE_N[], TILE_M[],), Float32), ct.full((TILE_N[], TILE_M[],), -Inf32, Float32))
            qk = qk .+ mask
        end

        qk = if !EVEN_K[] && j >= mask_start
            offs_n = ((j-Int32(1)) * TILE_N[]) .+ offs_n_tile
            mask = ct.full((TILE_N[], TILE_M[]), true, Bool)
            if !EVEN_K[]
                mask .& (offs_n .<= k_seqlen)
            end
            mask = ct.where(mask, ct.zeros((TILE_N[], TILE_M[],), Float32), ct.full((TILE_N[], TILE_M[],), -Inf32, Float32))
            qk .+ mask
        else
            qk
        end

        if !EVEN_K[]
            offs_n = ((j-Int32(1)) * TILE_N[]) .+ offs_n_tile
            mask = ct.full((TILE_N[], TILE_M[]), true, Bool)
            if !EVEN_K[]
                mask = mask .& (offs_n .<= k_seqlen)
            end
            mask = ct.where(mask, ct.zeros((TILE_N[], TILE_M[],), Float32), ct.full((TILE_N[], TILE_M[],), -Inf32, Float32))
            qk = qk .+ mask
        end

Does the if block need to depend on compile time constants?

I'd need this to make the padding and causal mask properly.

In the second and third block, I get "ERROR: SSAValue %___ not found in context"

That's an IRStructurizer error. Can you provide an MWE?

That's an IRStructurizer error. Can you provide an MWE?

I was able to narrow it down and believe it is covered by #53. See the added tests for MWE.

With #51 and #53, I now have forward and backward passes working locally!

Currently needing to wrap outside Float32 constants in Float32 within the kernel because MulF otherwise sees it as nothing:

qk_scale = Float32(qk_scale) * Float32(INV_LOG_2)

Another concern is whether I should convert to Int32 or not, essentially every time I do index arithmetic.

Currently needing to wrap outside Float32 constants in Float32 within the kernel because MulF otherwise sees it as nothing:

qk_scale = Float32(qk_scale) * Float32(INV_LOG_2)

Can you elaborate?

Another concern is whether I should convert to Int32 or not, essentially every time I do index arithmetic.

Yeah that's a common Julia pain point. It's why we have One(), and in CUDA.jl you can do e.g. 1i32.

Can you elaborate?

I define const INV_LOG_2 = Float32(1 / log(2)). If I use it without wrapping in Float32 within the kernel I get:

ERROR: LoadError: MethodError: no method matching encode_MulFOp!(::cuTile.CodeBuilder, ::cuTile.TypeId, ::cuTile.Value, ::Nothing)
The function `encode_MulFOp!` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  encode_MulFOp!(::cuTile.CodeBuilder, ::cuTile.TypeId, ::cuTile.Value, ::cuTile.Value; rounding_mode, flush_to_zero)
   @ cuTile ~/.julia/dev/cuTile/src/bytecode/encodings.jl:720

in CUDA.jl you can do e.g. 1i32.

Oh, neat. I didn't know. I considered maybe a @32 macro (macro var"32" ... end) to make all 64-bit integers wrapped in their 32-bit counterparts. Found that it wouldn't work within curly brackets like for type parameters since e.g. Array{T,Int32(1)} won't count as a vector, but the macro doesn't need to descend into :curly expressions. Problem is still that some functions actually only have methods for Int so it can't be applied to the entire function.

Problem is still that some functions actually only have methods for Int so it can't be applied to the entire function.

In general, Julia's array indexing requires Int. In CUDA.jl we've added some additional methods to override part of the getindex chain to support Int32, but it's tricky...

Constants should work without the type conversion now.

See #77 (comment)

See #40 (comment). Required for full parity with the python example which has a function that uses autotuned launch.

Superceded by #170