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	d4ffa3b...	main / d4ffa3b...
saxpy/default/Float32/1024	0.0728 ± 0.029 ms	0.627 ± 0.0079 μs	116 ± 46
saxpy/default/Float32/1048576	0.454 ± 0.022 ms	0.274 ± 0.02 ms	1.65 ± 0.14
saxpy/default/Float32/16384	0.0702 ± 0.03 ms	2.84 ± 0.31 μs	24.7 ± 11
saxpy/default/Float32/2048	0.0737 ± 0.03 ms	0.76 ± 0.059 μs	96.9 ± 40
saxpy/default/Float32/256	0.0721 ± 0.028 ms	0.565 ± 0.0053 μs	128 ± 49
saxpy/default/Float32/262144	0.167 ± 0.029 ms	0.0676 ± 0.0046 ms	2.47 ± 0.47
saxpy/default/Float32/32768	0.0713 ± 0.03 ms	5.49 ± 0.59 μs	13 ± 5.7
saxpy/default/Float32/4096	0.0755 ± 0.031 ms	1.13 ± 0.14 μs	66.6 ± 28
saxpy/default/Float32/512	0.0711 ± 0.029 ms	0.598 ± 0.0056 μs	119 ± 49
saxpy/default/Float32/64	0.0721 ± 0.029 ms	0.555 ± 0.0052 μs	130 ± 51
saxpy/default/Float32/65536	0.0938 ± 0.03 ms	13.8 ± 1.1 μs	6.77 ± 2.2
saxpy/default/Float64/1024	0.0689 ± 0.03 ms	0.746 ± 0.046 μs	92.3 ± 41
saxpy/default/Float64/1048576	0.615 ± 0.099 ms	0.595 ± 0.052 ms	1.03 ± 0.19
saxpy/default/Float64/16384	0.0773 ± 0.03 ms	5.45 ± 0.5 μs	14.2 ± 5.6
saxpy/default/Float64/2048	0.0697 ± 0.029 ms	1.17 ± 0.15 μs	59.7 ± 26
saxpy/default/Float64/256	0.0637 ± 0.031 ms	0.58 ± 0.0071 μs	110 ± 53
saxpy/default/Float64/262144	0.196 ± 0.034 ms	0.135 ± 0.013 ms	1.45 ± 0.29
saxpy/default/Float64/32768	0.0844 ± 0.029 ms	13.8 ± 1.1 μs	6.1 ± 2.1
saxpy/default/Float64/4096	0.071 ± 0.03 ms	1.73 ± 0.2 μs	41 ± 18
saxpy/default/Float64/512	0.0688 ± 0.03 ms	0.634 ± 0.0093 μs	108 ± 47
saxpy/default/Float64/64	0.0653 ± 0.031 ms	0.554 ± 0.0069 μs	118 ± 57
saxpy/default/Float64/65536	0.104 ± 0.03 ms	0.0334 ± 0.0024 ms	3.11 ± 0.92
saxpy/static workgroup=(1024,)/Float32/1024	0.0706 ± 0.029 ms	2.06 ± 0.026 μs	34.3 ± 14
saxpy/static workgroup=(1024,)/Float32/1048576	0.457 ± 0.023 ms	0.281 ± 0.022 ms	1.63 ± 0.15
saxpy/static workgroup=(1024,)/Float32/16384	0.0679 ± 0.029 ms	4.33 ± 0.34 μs	15.7 ± 6.8
saxpy/static workgroup=(1024,)/Float32/2048	0.0659 ± 0.029 ms	2.22 ± 0.083 μs	29.7 ± 13
saxpy/static workgroup=(1024,)/Float32/256	0.0696 ± 0.029 ms	2.55 ± 0.022 μs	27.3 ± 11
saxpy/static workgroup=(1024,)/Float32/262144	0.165 ± 0.028 ms	0.0693 ± 0.013 ms	2.39 ± 0.6
saxpy/static workgroup=(1024,)/Float32/32768	0.0759 ± 0.03 ms	7.67 ± 0.64 μs	9.89 ± 4.1
saxpy/static workgroup=(1024,)/Float32/4096	0.0725 ± 0.031 ms	2.56 ± 0.13 μs	28.3 ± 12
saxpy/static workgroup=(1024,)/Float32/512	0.0704 ± 0.029 ms	2.7 ± 0.027 μs	26.1 ± 11
saxpy/static workgroup=(1024,)/Float32/64	0.0683 ± 0.03 ms	2.45 ± 0.021 μs	27.9 ± 12
saxpy/static workgroup=(1024,)/Float32/65536	0.0897 ± 0.031 ms	17 ± 2.3 μs	5.29 ± 2
saxpy/static workgroup=(1024,)/Float64/1024	0.0681 ± 0.029 ms	2.23 ± 0.075 μs	30.6 ± 13
saxpy/static workgroup=(1024,)/Float64/1048576	0.576 ± 0.077 ms	0.649 ± 0.11 ms	0.887 ± 0.19
saxpy/static workgroup=(1024,)/Float64/16384	0.0754 ± 0.029 ms	7.64 ± 0.66 μs	9.87 ± 3.9
saxpy/static workgroup=(1024,)/Float64/2048	0.0705 ± 0.03 ms	2.56 ± 0.13 μs	27.6 ± 12
saxpy/static workgroup=(1024,)/Float64/256	0.0656 ± 0.029 ms	2.55 ± 0.026 μs	25.7 ± 11
saxpy/static workgroup=(1024,)/Float64/262144	0.197 ± 0.031 ms	0.144 ± 0.025 ms	1.37 ± 0.32
saxpy/static workgroup=(1024,)/Float64/32768	0.0818 ± 0.028 ms	17.1 ± 1.8 μs	4.8 ± 1.7
saxpy/static workgroup=(1024,)/Float64/4096	0.071 ± 0.03 ms	3.11 ± 0.21 μs	22.8 ± 9.8
saxpy/static workgroup=(1024,)/Float64/512	0.0695 ± 0.029 ms	2.71 ± 0.031 μs	25.6 ± 11
saxpy/static workgroup=(1024,)/Float64/64	0.0682 ± 0.03 ms	2.43 ± 0.021 μs	28.1 ± 12
saxpy/static workgroup=(1024,)/Float64/65536	0.102 ± 0.031 ms	0.0362 ± 0.0054 ms	2.81 ± 0.95
time_to_load	1.01 ± 0.036 s	0.298 ± 0.0018 s	3.4 ± 0.12

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!