Enzyme/AMDGPU function without sub-tapes

@wsmoses @michel2323

While trying to look into the compilation times of Enzyme, I was on my laptop where I can only use AMDGPU.

Besides running into https://github.com/EnzymeAD/Enzyme.jl/issues/2478 and I also encountered

ERROR: LoadError: AssertionError: ROCArray only supports bits types
Stacktrace:
  [1] ROCArray{nothing, 1, AMDGPU.Runtime.Mem.HIPBuffer}(::UndefInitializer, dims::Tuple{Int64})
    @ AMDGPU ~/.julia/packages/AMDGPU/IDGfT/src/array.jl:7
  [2] (ROCArray{nothing, 1})(::UndefInitializer, dims::Tuple{Int64})
    @ AMDGPU ~/.julia/packages/AMDGPU/IDGfT/src/array.jl:70
  [3] (ROCArray{nothing})(::UndefInitializer, dims::UInt32)
    @ AMDGPU ~/.julia/packages/AMDGPU/IDGfT/src/array.jl:88
  [4] #augmented_primal#25
    @ ~/.julia/packages/AMDGPU/IDGfT/ext/AMDGPUEnzymeCoreExt/AMDGPUEnzymeCoreExt.jl:203
  [5] augmented_primal
    @ ~/.julia/packages/AMDGPU/IDGfT/ext/AMDGPUEnzymeCoreExt/AMDGPUEnzymeCoreExt.jl:181 [inlined]

MWE:

using AMDGPU
using Enzyme

function square_kernel!(x)
    i =  workitemIdx().x
    x[i] *= x[i]
    sync_workgroup()
    return nothing
end

# basic squaring on GPU
function square!(x)
    @roc gridsize = length(x) groupsize = length(x) square_kernel!(x)
    return nothing
end

A = AMDGPU.ones(64)
dA = AMDGPU.ones(64)

    @time square!(A)

A .= (1:1:64)
dA .= 1

    @time Enzyme.autodiff(Reverse, square!, Duplicated(A, dA))

@show dA
@assert all(dA .≈ (2:2:128))

Start with the MWE and the augmented_primal stack frames at AMDGPUEnzymeCoreExt.jl:181 and :203. Reproduce the ROCArray assertion on AMDGPU, then trace how Enzyme.autodiff reaches augmented_primal. Done means the MWE completes without the assertion and the final dA values satisfy the shown assertion.