support emitting a generated function

[vchuravy]

Motivated by #665

Proposed syntax

@kernel generated=true function f(::Val{N}) where N
           KernelAbstractions.Extras.@unroll $N for i in 1:10
           end
       end

Sadly this doesn't quite work yet, since I need to handle the $N correctly

Currently:

    function cpu_f(__ctx__, ::Val{N}; ) where N
        if $(Expr(:generated))
            $(Expr(:copyast, :($(QuoteNode(:(let
      begin
          var"##N#249" = length((KernelAbstractions.__workitems_iterspace)(__ctx__))
          begin
              #= /home/vchuravy/src/KernelAbstractions/src/macros.jl:317 =#
              for var"##I#248" = (KernelAbstractions.__workitems_iterspace)(__ctx__)
                  #= /home/vchuravy/src/KernelAbstractions/src/macros.jl:318 =#
                  (KernelAbstractions.__validindex)(__ctx__, var"##I#248") || continue
                  #= /home/vchuravy/src/KernelAbstractions/src/macros.jl:319 =#
                  #= /home/vchuravy/src/KernelAbstractions/src/macros.jl:320 =#
                  #= REPL[19]:2 =# KernelAbstractions.Extra.@unroll $(Expr(:$, :N)) for i = 1:10
                          #= REPL[19]:2 =#
                          #= REPL[19]:3 =#
                      end
                  #= /home/vchuravy/src/KernelAbstractions/src/macros.jl:321 =#
              end
          end
      end
      return nothing
  end))))))
        else
            $(Expr(:meta, :generated_only))
        end
    end

Whereas

julia> @macroexpand @generated function(::Val{N}) where N
           quote
               KernelAbstractions.Extras.@unroll $N for i in 1:10
               end
           end
       end
:(function (::Val{N},) where N
      #= REPL[18]:1 =#
      if $(Expr(:generated))
          #= REPL[18]:1 =#
          #= REPL[18]:2 =#
          Core._expr(:block, $(QuoteNode(:(#= REPL[18]:3 =#))), Core._expr(:macrocall, $(Expr(:copyast, :($(QuoteNode(:(KernelAbstractions.Extras.var"@unroll")))))), $(QuoteNode(:(#= REPL[18]:3 =#))), N, $(Expr(:copyast, :($(QuoteNode(:(for i = 1:10
      #= REPL[18]:3 =#
      #= REPL[18]:4 =#
  end))))))))
      else
          $(Expr(:meta, :generated_only))
          return
      end
  end)

Note that the QuoteNode got broken into smaller pieces with the interpolated variable being left alone and the rest being passed to _expr and QuoteNode

[github-actions]

Benchmark Results

	main	9195d94...	main / 9195d94...
saxpy/default/Float32/1024	0.0421 ± 0.03 ms	0.65 ± 0.0089 μs	64.8 ± 46
saxpy/default/Float32/1048576	0.441 ± 0.019 ms	0.192 ± 0.022 ms	2.3 ± 0.29
saxpy/default/Float32/16384	0.0543 ± 0.024 ms	2.81 ± 0.3 μs	19.4 ± 8.9
saxpy/default/Float32/2048	0.0443 ± 0.025 ms	0.779 ± 0.022 μs	56.9 ± 32
saxpy/default/Float32/256	0.0407 ± 0.027 ms	0.595 ± 0.0062 μs	68.5 ± 46
saxpy/default/Float32/262144	0.152 ± 0.025 ms	0.0447 ± 0.0029 ms	3.41 ± 0.6
saxpy/default/Float32/32768	0.0591 ± 0.025 ms	5.4 ± 0.63 μs	10.9 ± 4.7
saxpy/default/Float32/4096	0.048 ± 0.023 ms	1.16 ± 0.14 μs	41.6 ± 20
saxpy/default/Float32/512	0.0542 ± 0.029 ms	0.628 ± 0.0073 μs	86.3 ± 47
saxpy/default/Float32/64	0.0404 ± 0.025 ms	0.583 ± 0.0072 μs	69.3 ± 43
saxpy/default/Float32/65536	0.0718 ± 0.025 ms	12.2 ± 1.1 μs	5.91 ± 2.1
saxpy/default/Float64/1024	0.0449 ± 0.029 ms	0.753 ± 0.035 μs	59.6 ± 38
saxpy/default/Float64/1048576	0.523 ± 0.059 ms	0.478 ± 0.044 ms	1.09 ± 0.16
saxpy/default/Float64/16384	0.0559 ± 0.025 ms	5.39 ± 0.6 μs	10.4 ± 4.8
saxpy/default/Float64/2048	0.0447 ± 0.025 ms	1.15 ± 0.14 μs	38.9 ± 23
saxpy/default/Float64/256	0.0418 ± 0.029 ms	0.578 ± 0.0084 μs	72.3 ± 50
saxpy/default/Float64/262144	0.166 ± 0.027 ms	0.0916 ± 0.0084 ms	1.81 ± 0.34
saxpy/default/Float64/32768	0.064 ± 0.025 ms	12.2 ± 1.1 μs	5.24 ± 2.1
saxpy/default/Float64/4096	0.0483 ± 0.023 ms	1.75 ± 0.26 μs	27.6 ± 14
saxpy/default/Float64/512	0.049 ± 0.029 ms	0.628 ± 0.012 μs	78.1 ± 47
saxpy/default/Float64/64	0.0402 ± 0.023 ms	0.557 ± 0.0067 μs	72 ± 42
saxpy/default/Float64/65536	0.0833 ± 0.025 ms	24 ± 1.9 μs	3.47 ± 1.1
saxpy/static workgroup=(1024,)/Float32/1024	0.0412 ± 0.03 ms	2.17 ± 0.028 μs	19 ± 14
saxpy/static workgroup=(1024,)/Float32/1048576	0.443 ± 0.019 ms	0.201 ± 0.021 ms	2.2 ± 0.24
saxpy/static workgroup=(1024,)/Float32/16384	0.051 ± 0.023 ms	4.54 ± 0.51 μs	11.2 ± 5.2
saxpy/static workgroup=(1024,)/Float32/2048	0.0431 ± 0.026 ms	2.33 ± 0.079 μs	18.5 ± 11
saxpy/static workgroup=(1024,)/Float32/256	0.0412 ± 0.023 ms	2.65 ± 0.02 μs	15.5 ± 8.8
saxpy/static workgroup=(1024,)/Float32/262144	0.152 ± 0.026 ms	0.0475 ± 0.004 ms	3.2 ± 0.62
saxpy/static workgroup=(1024,)/Float32/32768	0.0559 ± 0.024 ms	7.72 ± 0.97 μs	7.25 ± 3.3
saxpy/static workgroup=(1024,)/Float32/4096	0.0465 ± 0.023 ms	2.64 ± 0.13 μs	17.7 ± 8.8
saxpy/static workgroup=(1024,)/Float32/512	0.0455 ± 0.028 ms	2.81 ± 0.027 μs	16.2 ± 9.9
saxpy/static workgroup=(1024,)/Float32/64	0.0407 ± 0.022 ms	2.56 ± 0.018 μs	15.9 ± 8.8
saxpy/static workgroup=(1024,)/Float32/65536	0.0692 ± 0.025 ms	14.5 ± 1.4 μs	4.77 ± 1.8
saxpy/static workgroup=(1024,)/Float64/1024	0.0434 ± 0.029 ms	2.23 ± 0.075 μs	19.5 ± 13
saxpy/static workgroup=(1024,)/Float64/1048576	0.479 ± 0.051 ms	0.493 ± 0.057 ms	0.972 ± 0.15
saxpy/static workgroup=(1024,)/Float64/16384	0.0526 ± 0.025 ms	7.41 ± 0.68 μs	7.09 ± 3.4
saxpy/static workgroup=(1024,)/Float64/2048	0.0426 ± 0.024 ms	2.54 ± 0.12 μs	16.8 ± 9.6
saxpy/static workgroup=(1024,)/Float64/256	0.0415 ± 0.025 ms	2.55 ± 0.023 μs	16.3 ± 9.8
saxpy/static workgroup=(1024,)/Float64/262144	0.163 ± 0.028 ms	0.0933 ± 0.0088 ms	1.75 ± 0.34
saxpy/static workgroup=(1024,)/Float64/32768	0.0602 ± 0.025 ms	14.7 ± 1.3 μs	4.09 ± 1.8
saxpy/static workgroup=(1024,)/Float64/4096	0.0454 ± 0.023 ms	3.13 ± 0.23 μs	14.5 ± 7.5
saxpy/static workgroup=(1024,)/Float64/512	0.0437 ± 0.028 ms	2.71 ± 0.032 μs	16.1 ± 10
saxpy/static workgroup=(1024,)/Float64/64	0.0403 ± 0.02 ms	2.45 ± 0.12 μs	16.5 ± 8.4
saxpy/static workgroup=(1024,)/Float64/65536	0.0801 ± 0.027 ms	26.5 ± 2.6 μs	3.02 ± 1
time_to_load	0.942 ± 0.0027 s	0.28 ± 0.0021 s	3.37 ± 0.028

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

[codecov]

Codecov Report

❌ Patch coverage is 46.15385% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 71.36%. Comparing base (1f84b17) to head (9195d94).

Files with missing lines	Patch %	Lines
src/macros.jl	42.85%	4 Missing ⚠️
src/KernelAbstractions.jl	50.00%	3 Missing ⚠️

Additional details and impacted files

@@               Coverage Diff               @@
##           release-0.9     #667      +/-   ##
===============================================
- Coverage        71.85%   71.36%   -0.49%     
===============================================
  Files               14       14              
  Lines              906      915       +9     
===============================================
+ Hits               651      653       +2     
- Misses             255      262       +7     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[vchuravy]

julia> @kernel generated=false function g(::Val{N}) where N
                  KernelAbstractions.Extras.@unroll $N for i in 1:10
                  end
              end
ERROR: LoadError: Syntax error: `@unroll N expr` needs a constant integer N
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] var"@unroll"(__source__::LineNumberNode, __module__::Module, N::Any, expr::Any)
   @ KernelAbstractions.Extras.LoopInfo ~/src/KernelAbstractions/src/extras/loopinfo.jl:60
in expression starting at REPL[5]:2

julia> @kernel generated=true function f(::Val{N}) where N
                  KernelAbstractions.Extras.@unroll $N for i in 1:10
                  end
              end

 @ka_code_llvm optimize=false raw=true dump_module = true f(CPU())(Val(17), ndrange=1)

; ...
L33:                                              ; preds = %L29
; │ @ none within `macro expansion` @ /home/vchuravy/src/KernelAbstractions/src/macros.jl:321
; │┌ @ multidimensional.jl:417 within `iterate`
    br label %L22, !dbg !87, !llvm.loop !100
; ...
!100 = distinct !{!100, !101}
!101 = !{!"llvm.loop.unroll.count", i64 17}

Seems to work!

[vchuravy]

Homework for myself:

Write the kernel below with @nif

f(x::Integer) = 1
f(x::AbstractFloat) = 2

@kernel function _kern!(A::AbstractArray, B::AbstractArray)
    idx = @index(Global, Cartesian) 
    tpl = B[idx]
    k = tpl[1] > 0 ? 2 : 1
    A[idx] = ntuple(Val(length(tpl))) do j
        if j == k
            f(tpl[j])
        end
    end[k]
    nothing
end

wtf(N=64) = let
    A = cu(zeros(N,N))
    B = cu(tuple.(randn(N,N), ones(Int,N,N)))
    _kern!(get_backend(A))(A, B; ndrange=size(A))
end