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	7f8f67a...	main / 7f8f67a...
const/@Const/Float32/262144	0.341 ± 0.029 ms	0.245 ± 0.0098 ms	1.39 ± 0.13
const/@Const/Float32/65536	0.124 ± 0.025 ms	0.0617 ± 0.00016 ms	2.02 ± 0.41
const/@Const/Float64/262144	0.426 ± 0.041 ms	0.502 ± 0.011 ms	0.848 ± 0.083
const/@Const/Float64/65536	0.164 ± 0.013 ms	0.128 ± 0.00025 ms	1.28 ± 0.099
const/unmarked/Float32/262144	1.14 ± 0.037 ms	3 ± 0.005 ms	0.38 ± 0.012
const/unmarked/Float32/65536	0.322 ± 0.031 ms	0.734 ± 0.0071 ms	0.438 ± 0.043
const/unmarked/Float64/262144	1.34 ± 0.037 ms	2.99 ± 0.0031 ms	0.448 ± 0.012
const/unmarked/Float64/65536	0.395 ± 0.017 ms	0.734 ± 0.0079 ms	0.538 ± 0.024
launch/3D static workgroup, dynamic ndrange	0.0686 ± 0.032 ms	0.165 ± 0.0038 μs	417 ± 1.9e+02
launch/3D static workgroup, static ndrange	0.0674 ± 0.033 ms	0.0898 ± 0.0028 μs	751 ± 3.7e+02
launch/dynamic workgroup, dynamic ndrange	0.0726 ± 0.03 ms	0.0938 ± 0.0017 μs	773 ± 3.2e+02
launch/dynamic workgroup, dynamic ndrange, workgroupsize given	0.0667 ± 0.033 ms	0.0921 ± 0.0023 μs	725 ± 3.6e+02
launch/static workgroup, dynamic ndrange	0.0699 ± 0.032 ms	0.0576 ± 0.00091 μs	1.21e+03 ± 5.5e+02
launch/static workgroup, static ndrange	0.0688 ± 0.032 ms	0.0532 ± 0.0009 μs	1.29e+03 ± 6e+02
partition/dynamic workgroup, dynamic ndrange	0.0582 ± 0.0011 μs	0.0557 ± 0.0014 μs	1.05 ± 0.033
partition/static workgroup, dynamic ndrange	0.0653 ± 0.012 μs	0.0683 ± 0.00074 μs	0.955 ± 0.18
partition/static workgroup, static ndrange	1.55 ± 0.001 ns	16.8 ± 6.2 ns	0.0923 ± 0.034
saxpy/default/Float16/1024	0.0761 ± 0.0044 ms	0.634 ± 0.0063 μs	120 ± 7
saxpy/default/Float16/1048576	0.773 ± 0.028 ms	0.164 ± 0.012 ms	4.71 ± 0.38
saxpy/default/Float16/16384	0.0635 ± 0.033 ms	3.06 ± 0.083 μs	20.7 ± 11
saxpy/default/Float16/2048	0.0784 ± 0.02 ms	0.78 ± 0.0047 μs	100 ± 26
saxpy/default/Float16/256	0.0738 ± 0.014 ms	0.5 ± 0.0067 μs	148 ± 28
saxpy/default/Float16/262144	0.232 ± 0.029 ms	0.0417 ± 0.0021 ms	5.57 ± 0.75
saxpy/default/Float16/32768	0.0733 ± 0.031 ms	5.63 ± 0.063 μs	13 ± 5.5
saxpy/default/Float16/4096	0.0792 ± 0.03 ms	1.11 ± 0.016 μs	71.3 ± 27
saxpy/default/Float16/512	0.0754 ± 0.0036 ms	0.562 ± 0.006 μs	134 ± 6.6
saxpy/default/Float16/64	0.0752 ± 0.0069 ms	0.475 ± 0.051 μs	158 ± 22
saxpy/default/Float16/65536	0.094 ± 0.027 ms	10.9 ± 0.16 μs	8.59 ± 2.5
saxpy/default/Float32/1024	0.0749 ± 0.0033 ms	0.538 ± 0.007 μs	139 ± 6.5
saxpy/default/Float32/1048576	0.482 ± 0.024 ms	0.239 ± 0.012 ms	2.02 ± 0.14
saxpy/default/Float32/16384	0.0556 ± 0.03 ms	3.01 ± 0.48 μs	18.4 ± 10
saxpy/default/Float32/2048	0.075 ± 0.027 ms	0.614 ± 0.008 μs	122 ± 44
saxpy/default/Float32/256	0.0745 ± 0.011 ms	0.477 ± 0.0043 μs	156 ± 24
saxpy/default/Float32/262144	0.164 ± 0.028 ms	0.0608 ± 0.004 ms	2.7 ± 0.5
saxpy/default/Float32/32768	0.0626 ± 0.029 ms	6.45 ± 0.56 μs	9.7 ± 4.6
saxpy/default/Float32/4096	0.0694 ± 0.032 ms	1.02 ± 0.048 μs	68.3 ± 32
saxpy/default/Float32/512	0.0746 ± 0.0055 ms	0.512 ± 0.0066 μs	146 ± 11
saxpy/default/Float32/64	0.0738 ± 0.015 ms	0.468 ± 0.05 μs	158 ± 35
saxpy/default/Float32/65536	0.0757 ± 0.028 ms	12.9 ± 1.4 μs	5.85 ± 2.2
saxpy/default/Float64/1024	0.0751 ± 0.0058 ms	0.627 ± 0.014 μs	120 ± 9.6
saxpy/default/Float64/1048576	0.531 ± 0.071 ms	0.504 ± 0.047 ms	1.05 ± 0.17
saxpy/default/Float64/16384	0.0583 ± 0.029 ms	6.46 ± 0.14 μs	9.02 ± 4.4
saxpy/default/Float64/2048	0.0754 ± 0.03 ms	1.01 ± 0.03 μs	74.4 ± 30
saxpy/default/Float64/256	0.0751 ± 0.0051 ms	0.492 ± 0.042 μs	153 ± 17
saxpy/default/Float64/262144	0.179 ± 0.033 ms	0.123 ± 0.0048 ms	1.46 ± 0.28
saxpy/default/Float64/32768	0.0668 ± 0.028 ms	14.2 ± 1.5 μs	4.7 ± 2.1
saxpy/default/Float64/4096	0.0646 ± 0.031 ms	1.81 ± 0.12 μs	35.8 ± 17
saxpy/default/Float64/512	0.0752 ± 0.0053 ms	0.546 ± 0.0098 μs	138 ± 10
saxpy/default/Float64/64	0.0755 ± 0.0064 ms	0.472 ± 0.05 μs	160 ± 22
saxpy/default/Float64/65536	0.0845 ± 0.028 ms	0.0321 ± 0.0038 ms	2.64 ± 0.92
saxpy/static workgroup=(1024,)/Float16/1024	0.0737 ± 0.0035 ms	2.13 ± 0.018 μs	34.6 ± 1.7
saxpy/static workgroup=(1024,)/Float16/1048576	0.77 ± 0.03 ms	0.162 ± 0.0096 ms	4.75 ± 0.34
saxpy/static workgroup=(1024,)/Float16/16384	0.0592 ± 0.034 ms	4.48 ± 0.15 μs	13.2 ± 7.5
saxpy/static workgroup=(1024,)/Float16/2048	0.0761 ± 0.031 ms	2.29 ± 0.021 μs	33.2 ± 13
saxpy/static workgroup=(1024,)/Float16/256	0.0722 ± 0.027 ms	2.71 ± 0.019 μs	26.7 ± 10
saxpy/static workgroup=(1024,)/Float16/262144	0.227 ± 0.028 ms	0.0422 ± 0.0018 ms	5.38 ± 0.69
saxpy/static workgroup=(1024,)/Float16/32768	0.0689 ± 0.028 ms	6.99 ± 0.27 μs	9.86 ± 4
saxpy/static workgroup=(1024,)/Float16/4096	0.0777 ± 0.032 ms	2.63 ± 0.027 μs	29.6 ± 12
saxpy/static workgroup=(1024,)/Float16/512	0.0729 ± 0.015 ms	2.97 ± 0.02 μs	24.5 ± 5.2
saxpy/static workgroup=(1024,)/Float16/64	0.0678 ± 0.031 ms	2.51 ± 0.014 μs	27 ± 12
saxpy/static workgroup=(1024,)/Float16/65536	0.0897 ± 0.019 ms	13 ± 0.49 μs	6.91 ± 1.5
saxpy/static workgroup=(1024,)/Float32/1024	0.0721 ± 0.002 ms	2.08 ± 0.023 μs	34.7 ± 1
saxpy/static workgroup=(1024,)/Float32/1048576	0.481 ± 0.025 ms	0.246 ± 0.014 ms	1.95 ± 0.15
saxpy/static workgroup=(1024,)/Float32/16384	0.052 ± 0.027 ms	4.53 ± 0.6 μs	11.5 ± 6.1
saxpy/static workgroup=(1024,)/Float32/2048	0.0729 ± 0.028 ms	2.26 ± 0.025 μs	32.3 ± 12
saxpy/static workgroup=(1024,)/Float32/256	0.0722 ± 0.027 ms	2.57 ± 0.018 μs	28.1 ± 10
saxpy/static workgroup=(1024,)/Float32/262144	0.161 ± 0.028 ms	0.0633 ± 0.0039 ms	2.55 ± 0.47
saxpy/static workgroup=(1024,)/Float32/32768	0.0585 ± 0.024 ms	7.91 ± 0.91 μs	7.4 ± 3.1
saxpy/static workgroup=(1024,)/Float32/4096	0.0733 ± 0.033 ms	2.55 ± 0.08 μs	28.8 ± 13
saxpy/static workgroup=(1024,)/Float32/512	0.0723 ± 0.0067 ms	2.71 ± 0.022 μs	26.7 ± 2.5
saxpy/static workgroup=(1024,)/Float32/64	0.0695 ± 0.029 ms	2.47 ± 0.016 μs	28.1 ± 12
saxpy/static workgroup=(1024,)/Float32/65536	0.072 ± 0.024 ms	16 ± 1.4 μs	4.51 ± 1.5
saxpy/static workgroup=(1024,)/Float64/1024	0.0716 ± 0.0054 ms	2.26 ± 0.027 μs	31.7 ± 2.4
saxpy/static workgroup=(1024,)/Float64/1048576	0.538 ± 0.057 ms	0.518 ± 0.036 ms	1.04 ± 0.13
saxpy/static workgroup=(1024,)/Float64/16384	0.0545 ± 0.025 ms	7.87 ± 0.87 μs	6.92 ± 3.3
saxpy/static workgroup=(1024,)/Float64/2048	0.0709 ± 0.03 ms	2.55 ± 0.085 μs	27.8 ± 12
saxpy/static workgroup=(1024,)/Float64/256	0.0714 ± 0.024 ms	2.6 ± 0.02 μs	27.5 ± 9.3
saxpy/static workgroup=(1024,)/Float64/262144	0.175 ± 0.033 ms	0.126 ± 0.0043 ms	1.39 ± 0.27
saxpy/static workgroup=(1024,)/Float64/32768	0.0628 ± 0.026 ms	16 ± 1.5 μs	3.92 ± 1.7
saxpy/static workgroup=(1024,)/Float64/4096	0.0556 ± 0.031 ms	3.27 ± 0.25 μs	17 ± 9.6
saxpy/static workgroup=(1024,)/Float64/512	0.0721 ± 0.0054 ms	2.77 ± 0.027 μs	26 ± 2
saxpy/static workgroup=(1024,)/Float64/64	0.0708 ± 0.028 ms	2.49 ± 0.017 μs	28.4 ± 11
saxpy/static workgroup=(1024,)/Float64/65536	0.0804 ± 0.025 ms	0.0336 ± 0.0027 ms	2.39 ± 0.78
time_to_load	0.747 ± 0.0053 s	0.248 ± 0.0012 s	3.01 ± 0.026

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

[fjwillemsen]

@vchuravy can this be merged? We would like to publish the package soon 😄

[vchuravy]

@fjwillemsen nobody expressed to me after my proposal that they actually wanted this, and I don't like to merge functionality without at least a clear user.

[fjwillemsen]

I understand, I'll get @evelyne-ringoot in the loop and get back to you on this. Thank you!

[evelyne-ringoot]

Hi, the original motivation for this was the performance difference in all kernels where input values are being used for loop unrolls, versus where this values is defined as const, so the scope is not limited to specifying the loop unroll number, but general availability of loop unrolling constants during early compile times, in particular for more complex functions. However, I can no longer reproduce a substantial performance difference in all kernels but one, where the performance difference swings both ways depending on hyperparameters (better using ::Val vs better using const). Considering that, the master branch works fine for our application!