Merge #225

225: Forward port of #224 r=vchuravy a=vchuravy

(cherry picked from commit 7318a2c)

bors r+


Co-authored-by: Valentin Churavy <[email protected]>

Author: bors[bot]

src/compiler.jl

@@ -60,5 +60,7 @@ function generate_overdubs(mod, Ctx)
             end
             Base.steprangelen_hp(T, start, step, 0, len, 1)
         end
+
+        @inline Cassette.overdub(::$Ctx, ::typeof(Base.literal_pow), f::F, x, p) where F = Base.literal_pow(f, x, p)
     end
 end

test/compiler.jl

@@ -8,10 +8,19 @@ import KernelAbstractions.NDIteration: NDRange, StaticSize, DynamicSize
     @show A[I...]
 end
 
+@kernel function literal_pow(A)
+    A[1] = 2^11
+end
+
 function compiler_testsuite()
     kernel = index(CPU(), DynamicSize(), DynamicSize())
     iterspace = NDRange{1, StaticSize{(128,)}, StaticSize{(8,)}}();
     ctx = KernelAbstractions.mkcontext(kernel, 1, nothing, iterspace, Val(KernelAbstractions.NoDynamicCheck()))
 
     @test KernelAbstractions.Cassette.overdub(ctx, KernelAbstractions.__index_Global_NTuple, CartesianIndex(1)) == (1,)
+
+    CI, rt = @ka_code_typed literal_pow(CPU())(zeros(Int,1), ndrange=1)
+    # test that there is no invoke of overdub
+    @test !any(stmt->(stmt isa Expr) && stmt.head == :invoke, CI.code)
 end
+