#include "memory.h"

#include <assert.h>
#include <stdlib.h>
#include <stdio.h>

#define CACHED


#ifdef CACHED
#define CACHELEN 128

typedef struct _tup
{
    size_t bytes;
    void* ptr;
    bool free;
} _tup;

_tup cache[CACHELEN];
bool initialized = false;
#endif // CACHED


void
cuda_malloc_clear(void** ptr, size_t bytes)
{
    cudaError_t err;
    
#ifdef CACHED
    if (!initialized)
    {
        for (int i = 0; i < CACHELEN; i++)
        {
            cache[i].bytes = 0;
            cache[i].ptr = 0;
            cache[i].free = true;
        }
        initialized = true;
    }
    
    for (int i = 0; i < CACHELEN; i++)
    {
        if (cache[i].free && cache[i].bytes == bytes)
        {
            // We don't have to remalloc, we already have a valid free ptr
            cache[i].free = false;
            *ptr = cache[i].ptr;
            err = cudaMemset(*ptr, 0, bytes);
            assert(err == cudaSuccess);
            return;
        }
    }
#endif // CACHED

    // Malloc to device, check for errors
    err = cudaMalloc(ptr, bytes);
    assert(err == cudaSuccess);

    // Set val to 0, check for errors
    err = cudaMemset(*ptr, 0, bytes);
    assert(err == cudaSuccess);

#ifdef CACHED
    for (int i = 0; i < CACHELEN; i++)
    {
        if (cache[i].free && cache[i].ptr == 0)
        {
            cache[i].free = false;
            cache[i].ptr = *ptr;
            cache[i].bytes = bytes;
            return;
        }
    }
#endif // CACHED
}

void
cuda_malloc_free(void* ptr)
{
#ifdef CACHED
    for (int i = 0; i < CACHELEN; i++)
    {
        if (cache[i].ptr == ptr)
        {
            cache[i].free = true;
            return;
        }
    }
#endif // CACHED

    cudaFree(ptr);
}

	.file	"tmpxft_0030e6ea_00000000-6_memory.cudafe1.cpp"
	.text
#APP

#NO_APP
	.type	_ZL26__cudaUnregisterBinaryUtilv, @function
_ZL26__cudaUnregisterBinaryUtilv:
.LFB2031:
	.cfi_startproc
	endbr64
	movq	_ZL20__cudaFatCubinHandle(%rip), %rdi
	jmp	__cudaUnregisterFatBinary@PLT
	.cfi_endproc
.LFE2031:
	.size	_ZL26__cudaUnregisterBinaryUtilv, .-_ZL26__cudaUnregisterBinaryUtilv
	.globl	_Z17cuda_malloc_clearPPvm
	.type	_Z17cuda_malloc_clearPPvm, @function
_Z17cuda_malloc_clearPPvm:
.LFB2027:
	.cfi_startproc
	endbr64
	pushq	%r13
	.cfi_def_cfa_offset 16
	.cfi_offset 13, -16
	leaq	cache(%rip), %r13
	pushq	%r12
	.cfi_def_cfa_offset 24
	.cfi_offset 12, -24
	pushq	%rbp
	.cfi_def_cfa_offset 32
	.cfi_offset 6, -32
	movq	%rdi, %rbp
	pushq	%rbx
	.cfi_def_cfa_offset 40
	.cfi_offset 3, -40
	movq	%rsi, %rbx
	pushq	%r8
	.cfi_def_cfa_offset 48
	cmpb	$0, initialized(%rip)
	jne	.L3
	movq	%r13, %rax
	leaq	3072(%r13), %rdx
.L4:
	xorl	%esi, %esi
	movb	$1, 16(%rax)
	addq	$24, %rax
	movq	%rsi, -24(%rax)
	movq	%rsi, -16(%rax)
	cmpq	%rdx, %rax
	jne	.L4
	movb	$1, initialized(%rip)
.L3:
	leaq	cache(%rip), %r12
	xorl	%eax, %eax
	movq	%r12, %rdx
.L6:
	cmpb	$0, 16(%rdx)
	je	.L5
	cmpq	%rbx, (%rdx)
	jne	.L5
	cltq
	movq	%rbx, %rdx
	xorl	%esi, %esi
	imulq	$24, %rax, %rax
	addq	%r13, %rax
	movq	8(%rax), %rdi
	movb	$0, 16(%rax)
	movq	%rdi, 0(%rbp)
	popq	%rcx
	.cfi_remember_state
	.cfi_def_cfa_offset 40
	popq	%rbx
	.cfi_def_cfa_offset 32
	popq	%rbp
	.cfi_def_cfa_offset 24
	popq	%r12
	.cfi_def_cfa_offset 16
	popq	%r13
	.cfi_def_cfa_offset 8
	jmp	cudaMemset@PLT
.L5:
	.cfi_restore_state
	incl	%eax
	addq	$24, %rdx
	cmpl	$128, %eax
	jne	.L6
	movq	%rbx, %rsi
	movq	%rbp, %rdi
	call	cudaMalloc@PLT
	movq	0(%rbp), %rdi
	movq	%rbx, %rdx
	xorl	%esi, %esi
	call	cudaMemset@PLT
	xorl	%eax, %eax
.L9:
	cmpb	$0, 16(%r12)
	je	.L7
	cmpq	$0, 8(%r12)
	jne	.L7
	cltq
	movq	0(%rbp), %rdx
	imulq	$24, %rax, %rax
	addq	%r13, %rax
	movb	$0, 16(%rax)
	movq	%rdx, 8(%rax)
	movq	%rbx, (%rax)
	jmp	.L2
.L7:
	incl	%eax
	addq	$24, %r12
	cmpl	$128, %eax
	jne	.L9
.L2:
	popq	%rax
	.cfi_def_cfa_offset 40
	popq	%rbx
	.cfi_def_cfa_offset 32
	popq	%rbp
	.cfi_def_cfa_offset 24
	popq	%r12
	.cfi_def_cfa_offset 16
	popq	%r13
	.cfi_def_cfa_offset 8
	ret
	.cfi_endproc
.LFE2027:
	.size	_Z17cuda_malloc_clearPPvm, .-_Z17cuda_malloc_clearPPvm
	.globl	_Z16cuda_malloc_freePv
	.type	_Z16cuda_malloc_freePv, @function
_Z16cuda_malloc_freePv:
.LFB2028:
	.cfi_startproc
	endbr64
	leaq	cache(%rip), %rdx
	xorl	%eax, %eax
	movq	%rdx, %rcx
.L18:
	cmpq	%rdi, 8(%rdx)
	jne	.L16
	cltq
	imulq	$24, %rax, %rax
	movb	$1, 16(%rcx,%rax)
	ret
.L16:
	incl	%eax
	addq	$24, %rdx
	cmpl	$128, %eax
	jne	.L18
	jmp	cudaFree@PLT
	.cfi_endproc
.LFE2028:
	.size	_Z16cuda_malloc_freePv, .-_Z16cuda_malloc_freePv
	.section	.text.startup,"ax",@progbits
	.type	_ZL24__sti____cudaRegisterAllv, @function
_ZL24__sti____cudaRegisterAllv:
.LFB2054:
	.cfi_startproc
	endbr64
	pushq	%rax
	.cfi_def_cfa_offset 16
	leaq	_ZL15__fatDeviceText(%rip), %rdi
	call	__cudaRegisterFatBinary@PLT
	movq	%rax, %rdi
	movq	%rax, _ZL20__cudaFatCubinHandle(%rip)
	call	__cudaRegisterFatBinaryEnd@PLT
	leaq	_ZL26__cudaUnregisterBinaryUtilv(%rip), %rdi
	popq	%rdx
	.cfi_def_cfa_offset 8
	jmp	atexit@PLT
	.cfi_endproc
.LFE2054:
	.size	_ZL24__sti____cudaRegisterAllv, .-_ZL24__sti____cudaRegisterAllv
	.section	.init_array,"aw"
	.align 8
	.quad	_ZL24__sti____cudaRegisterAllv
	.section	.nvFatBinSegment,"aw"
	.align 8
	.type	_ZL15__fatDeviceText, @object
	.size	_ZL15__fatDeviceText, 24
_ZL15__fatDeviceText:
	.long	1180844977
	.long	1
	.quad	fatbinData
	.quad	0
	.local	_ZL20__cudaFatCubinHandle
	.comm	_ZL20__cudaFatCubinHandle,8,8
	.globl	initialized
	.bss
	.type	initialized, @object
	.size	initialized, 1
initialized:
	.zero	1
	.globl	cache
	.align 32
	.type	cache, @object
	.size	cache, 3072
cache:
	.zero	3072
	.ident	"GCC: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0"
	.section	.note.GNU-stack,"",@progbits
	.section	.note.gnu.property,"a"
	.align 8
	.long	1f - 0f
	.long	4f - 1f
	.long	5
0:
	.string	"GNU"
1:
	.align 8
	.long	0xc0000002
	.long	3f - 2f
2:
	.long	0x3
3:
	.align 8
4:


#include <hip/hip_runtime.h>
#include "memory.h"

#include <assert.h>
#include <stdlib.h>
#include <stdio.h>

#define CACHED


#ifdef CACHED
#define CACHELEN 128

typedef struct _tup
{
    size_t bytes;
    void* ptr;
    bool free;
} _tup;

_tup cache[CACHELEN];
bool initialized = false;
#endif // CACHED


void
cuda_malloc_clear(void** ptr, size_t bytes)
{
    hipError_t err;
    
#ifdef CACHED
    if (!initialized)
    {
        for (int i = 0; i < CACHELEN; i++)
        {
            cache[i].bytes = 0;
            cache[i].ptr = 0;
            cache[i].free = true;
        }
        initialized = true;
    }
    
    for (int i = 0; i < CACHELEN; i++)
    {
        if (cache[i].free && cache[i].bytes == bytes)
        {
            // We don't have to remalloc, we already have a valid free ptr
            cache[i].free = false;
            *ptr = cache[i].ptr;
            err = hipMemset(*ptr, 0, bytes);
            assert(err == hipSuccess);
            return;
        }
    }
#endif // CACHED

    // Malloc to device, check for errors
    err = hipMalloc(ptr, bytes);
    assert(err == hipSuccess);

    // Set val to 0, check for errors
    err = hipMemset(*ptr, 0, bytes);
    assert(err == hipSuccess);

#ifdef CACHED
    for (int i = 0; i < CACHELEN; i++)
    {
        if (cache[i].free && cache[i].ptr == 0)
        {
            cache[i].free = false;
            cache[i].ptr = *ptr;
            cache[i].bytes = bytes;
            return;
        }
    }
#endif // CACHED
}

void
cuda_malloc_free(void* ptr)
{
#ifdef CACHED
    for (int i = 0; i < CACHELEN; i++)
    {
        if (cache[i].ptr == ptr)
        {
            cache[i].free = true;
            return;
        }
    }
#endif // CACHED

    hipFree(ptr);
}

	.text
	.file	"memory.hip"
	.globl	_Z17cuda_malloc_clearPPvm       # -- Begin function _Z17cuda_malloc_clearPPvm
	.type	_Z17cuda_malloc_clearPPvm,@function
_Z17cuda_malloc_clearPPvm:              # @_Z17cuda_malloc_clearPPvm
	.cfi_startproc
# %bb.0:
	pushq	%r14
	.cfi_def_cfa_offset 16
	pushq	%rbx
	.cfi_def_cfa_offset 24
	pushq	%rax
	.cfi_def_cfa_offset 32
	.cfi_offset %rbx, -24
	.cfi_offset %r14, -16
	movq	%rsi, %rbx
	movq	%rdi, %r14
	cmpb	$0, initialized(%rip)
	jne	.LBB0_4
# %bb.1:                                # %.preheader.preheader
	movq	$-3072, %rax                    # imm = 0xF400
	xorps	%xmm0, %xmm0
.LBB0_2:                                # %.preheader
                                        # =>This Inner Loop Header: Depth=1
	movups	%xmm0, cache+3072(%rax)
	movb	$1, cache+3088(%rax)
	addq	$24, %rax
	jne	.LBB0_2
# %bb.3:
	movb	$1, initialized(%rip)
.LBB0_4:                                # %.preheader76
	movq	$-3072, %rax                    # imm = 0xF400
.LBB0_5:                                # =>This Inner Loop Header: Depth=1
	cmpb	$0, cache+3088(%rax)
	je	.LBB0_7
# %bb.6:                                #   in Loop: Header=BB0_5 Depth=1
	cmpq	%rbx, cache+3072(%rax)
	je	.LBB0_14
.LBB0_7:                                #   in Loop: Header=BB0_5 Depth=1
	addq	$24, %rax
	jne	.LBB0_5
# %bb.8:                                # %.critedge
	movq	%r14, %rdi
	movq	%rbx, %rsi
	callq	hipMalloc
	movq	(%r14), %rdi
	xorl	%esi, %esi
	movq	%rbx, %rdx
	callq	hipMemset
	movq	$-3072, %rax                    # imm = 0xF400
.LBB0_9:                                # =>This Inner Loop Header: Depth=1
	cmpb	$0, cache+3088(%rax)
	je	.LBB0_12
# %bb.10:                               #   in Loop: Header=BB0_9 Depth=1
	cmpq	$0, cache+3080(%rax)
	je	.LBB0_11
.LBB0_12:                               #   in Loop: Header=BB0_9 Depth=1
	addq	$24, %rax
	jne	.LBB0_9
	jmp	.LBB0_13
.LBB0_14:
	movb	$0, cache+3088(%rax)
	movq	cache+3080(%rax), %rdi
	movq	%rdi, (%r14)
	xorl	%esi, %esi
	movq	%rbx, %rdx
	addq	$8, %rsp
	.cfi_def_cfa_offset 24
	popq	%rbx
	.cfi_def_cfa_offset 16
	popq	%r14
	.cfi_def_cfa_offset 8
	jmp	hipMemset                       # TAILCALL
.LBB0_11:
	.cfi_def_cfa_offset 32
	movb	$0, cache+3088(%rax)
	movq	(%r14), %rcx
	movq	%rcx, cache+3080(%rax)
	movq	%rbx, cache+3072(%rax)
.LBB0_13:                               # %.loopexit
	addq	$8, %rsp
	.cfi_def_cfa_offset 24
	popq	%rbx
	.cfi_def_cfa_offset 16
	popq	%r14
	.cfi_def_cfa_offset 8
	retq
.Lfunc_end0:
	.size	_Z17cuda_malloc_clearPPvm, .Lfunc_end0-_Z17cuda_malloc_clearPPvm
	.cfi_endproc
                                        # -- End function
	.globl	_Z16cuda_malloc_freePv          # -- Begin function _Z16cuda_malloc_freePv
	.type	_Z16cuda_malloc_freePv,@function
_Z16cuda_malloc_freePv:                 # @_Z16cuda_malloc_freePv
	.cfi_startproc
# %bb.0:
	movq	$-3072, %rax                    # imm = 0xF400
	jmp	.LBB1_2
.LBB1_1:                                #   in Loop: Header=BB1_2 Depth=1
	addq	$24, %rax
	je	hipFree                         # TAILCALL
.LBB1_2:                                # =>This Inner Loop Header: Depth=1
	cmpq	%rdi, cache+3080(%rax)
	jne	.LBB1_1
# %bb.3:
	movb	$1, cache+3088(%rax)
	retq
.Lfunc_end1:
	.size	_Z16cuda_malloc_freePv, .Lfunc_end1-_Z16cuda_malloc_freePv
	.cfi_endproc
                                        # -- End function
	.type	cache,@object                   # @cache
	.bss
	.globl	cache
	.p2align	4, 0x0
cache:
	.zero	3072
	.size	cache, 3072

	.type	initialized,@object             # @initialized
	.globl	initialized
initialized:
	.byte	0                               # 0x0
	.size	initialized, 1

	.type	__hip_cuid_,@object             # @__hip_cuid_
	.globl	__hip_cuid_
__hip_cuid_:
	.byte	0                               # 0x0
	.size	__hip_cuid_, 1

	.ident	"AMD clang version 18.0.0git (https://github.com/RadeonOpenCompute/llvm-project roc-6.3.2 25012 e5bf7e55c91490b07c49d8960fa7983d864936c4)"
	.section	".note.GNU-stack","",@progbits
	.addrsig
	.addrsig_sym __hip_cuid_

	.text
	.p2alignl 7, 3214868480
	.fill 96, 4, 3214868480
	.type	__hip_cuid_,@object             ; @__hip_cuid_
	.section	.bss,"aw",@nobits
	.globl	__hip_cuid_
__hip_cuid_:
	.byte	0                               ; 0x0
	.size	__hip_cuid_, 1

	.ident	"AMD clang version 18.0.0git (https://github.com/RadeonOpenCompute/llvm-project roc-6.3.2 25012 e5bf7e55c91490b07c49d8960fa7983d864936c4)"
	.section	".note.GNU-stack","",@progbits
	.addrsig
	.addrsig_sym __hip_cuid_
	.amdgcn_target "amdgcn-amd-amdhsa--gfx1100"
	.amdgpu_metadata
---
amdhsa.kernels:  []
amdhsa.target:   amdgcn-amd-amdhsa--gfx1100
amdhsa.version:
  - 1
  - 2
...

	.end_amdgpu_metadata
