#include <cassert>
#include <cstdio>
#include "sgemm_kernel.cuh"

__global__ void fill(float* data, const int size, const float value)
{
    int idx = blockDim.x * blockIdx.x + threadIdx.x;
    while (idx < size)
    {
        data[idx] = value;
        idx += blockDim.x * gridDim.x;
    }
}

__global__ void sgemm(const float* a, const float* b, float* result, const int size, const int stride)
{
    const int thread_idx = threadIdx.y * blockDim.x + threadIdx.x;
    const int num_thread = blockDim.x * blockDim.y;
    const int large_ty = threadIdx.y * small_block_size;
    const int large_tx = threadIdx.x * small_block_size;

    // thread 単位で small_balock_size^2 だけ要素を持っている時に確保できる block_size_y
    const int block_size_y = num_thread * small_block_size * small_block_size / block_size_x;
    assert(block_size_y <= block_size_x);

    constexpr int block_k_size = 16;

    for (int i = blockIdx.y * block_size_y; i < size; i += gridDim.y * block_size_y)
    {
        for (int j = blockIdx.x * block_size_x; j < size; j += gridDim.x * block_size_x)
        {
            const bool has_result = (i + large_ty < size && j + large_tx < size);

            // 単一スレッドの結果保存用
            float local_result[small_block_size][small_block_size];
            for (int ii = 0; ii < small_block_size; ii++)
            {
                for (int jj = 0; jj < small_block_size; jj++)
                {
                    local_result[ii][jj] = 0;
                }
            }

            const int height = min(block_size_y, size - i);
            const int width = min(block_size_x, size - j);

            for (int k = 0; k < size; k += block_k_size)
            {
                __shared__ float temp_a[block_k_size][block_size_x + 1], temp_b[block_k_size][block_size_x + 1];
                {
                    const int l = thread_idx % block_size_x;
                    if (l < width)
                    {
                        for (int kk = thread_idx / block_size_x; kk < block_k_size; kk += num_thread / block_size_x)
                        {
                            temp_b[kk][l] = b[(k + kk) * stride + (j + l)];
                        }
                    }
                }
                {
                    const int kk = thread_idx % block_k_size;
                    for (int l = thread_idx / block_k_size; l < height; l += num_thread / block_k_size)
                    {
                        temp_a[kk][l] = a[(i + l) * stride + k + kk];
                    }
                }
                __syncthreads();

                if (has_result)
                {
                    for (int kk = 0; kk < block_k_size; kk++)
                    {
                        for (int ii = 0; ii < small_block_size; ii++)
                        {
                            for (int jj = 0; jj < small_block_size; jj++)
                            {
                                local_result[ii][jj] += temp_a[kk][large_ty + ii] * temp_b[kk][large_tx + jj];
                            }
                        }
                    }
                }
                __syncthreads();
            }
            if (has_result)
            {
                for (int ii = 0; ii < small_block_size; ii++)
                {
                    for (int jj = 0; jj < small_block_size; jj++)
                    {
                        result[(i + large_ty + ii) * stride + (j + large_tx + jj)] = local_result[ii][jj];
                    }
                }
            }
            __syncthreads();
        }
    }
}
	.file	"tmpxft_002dd87f_00000000-6_sgemm_kernel.cudafe1.cpp"
	.text
#APP

#NO_APP
	.type	_ZL26__cudaUnregisterBinaryUtilv, @function
_ZL26__cudaUnregisterBinaryUtilv:
.LFB2029:
	.cfi_startproc
	endbr64
	movq	_ZL20__cudaFatCubinHandle(%rip), %rdi
	jmp	__cudaUnregisterFatBinary@PLT
	.cfi_endproc
.LFE2029:
	.size	_ZL26__cudaUnregisterBinaryUtilv, .-_ZL26__cudaUnregisterBinaryUtilv
	.globl	_Z25__device_stub__Z4fillPfifPfif
	.type	_Z25__device_stub__Z4fillPfifPfif, @function
_Z25__device_stub__Z4fillPfifPfif:
.LFB2051:
	.cfi_startproc
	endbr64
	subq	$120, %rsp
	.cfi_def_cfa_offset 128
	movq	%rdi, 8(%rsp)
	leaq	24(%rsp), %rcx
	leaq	16(%rsp), %rdx
	movl	%esi, 4(%rsp)
	leaq	32(%rsp), %rdi
	leaq	44(%rsp), %rsi
	movss	%xmm0, (%rsp)
	movq	%fs:40, %rax
	movq	%rax, 104(%rsp)
	xorl	%eax, %eax
	leaq	8(%rsp), %rax
	movl	$1, 40(%rsp)
	movq	%rax, 80(%rsp)
	leaq	4(%rsp), %rax
	movq	%rax, 88(%rsp)
	movq	%rsp, %rax
	movq	%rax, 96(%rsp)
	movabsq	$4294967297, %rax
	movq	%rax, 32(%rsp)
	movq	%rax, 44(%rsp)
	movl	$1, 52(%rsp)
	call	__cudaPopCallConfiguration@PLT
	testl	%eax, %eax
	jne	.L2
	pushq	24(%rsp)
	.cfi_def_cfa_offset 136
	leaq	_Z4fillPfif(%rip), %rdi
	pushq	24(%rsp)
	.cfi_def_cfa_offset 144
	movq	60(%rsp), %rcx
	movl	68(%rsp), %r8d
	movq	48(%rsp), %rsi
	movl	56(%rsp), %edx
	leaq	96(%rsp), %r9
	call	cudaLaunchKernel@PLT
	popq	%rax
	.cfi_def_cfa_offset 136
	popq	%rdx
	.cfi_def_cfa_offset 128
.L2:
	movq	104(%rsp), %rax
	subq	%fs:40, %rax
	je	.L4
	call	__stack_chk_fail@PLT
.L4:
	addq	$120, %rsp
	.cfi_def_cfa_offset 8
	ret
	.cfi_endproc
.LFE2051:
	.size	_Z25__device_stub__Z4fillPfifPfif, .-_Z25__device_stub__Z4fillPfifPfif
	.globl	_Z4fillPfif
	.type	_Z4fillPfif, @function
_Z4fillPfif:
.LFB2052:
	.cfi_startproc
	endbr64
	jmp	_Z25__device_stub__Z4fillPfifPfif
	.cfi_endproc
.LFE2052:
	.size	_Z4fillPfif, .-_Z4fillPfif
	.globl	_Z32__device_stub__Z5sgemmPKfS0_PfiiPKfS0_Pfii
	.type	_Z32__device_stub__Z5sgemmPKfS0_PfiiPKfS0_Pfii, @function
_Z32__device_stub__Z5sgemmPKfS0_PfiiPKfS0_Pfii:
.LFB2053:
	.cfi_startproc
	endbr64
	subq	$152, %rsp
	.cfi_def_cfa_offset 160
	movq	%rdi, 24(%rsp)
	leaq	48(%rsp), %rdi
	movq	%rsi, 16(%rsp)
	leaq	60(%rsp), %rsi
	movq	%rdx, 8(%rsp)
	leaq	32(%rsp), %rdx
	movl	%ecx, 4(%rsp)
	leaq	40(%rsp), %rcx
	movl	%r8d, (%rsp)
	movq	%fs:40, %rax
	movq	%rax, 136(%rsp)
	xorl	%eax, %eax
	leaq	24(%rsp), %rax
	movl	$1, 56(%rsp)
	movq	%rax, 96(%rsp)
	leaq	16(%rsp), %rax
	movq	%rax, 104(%rsp)
	leaq	8(%rsp), %rax
	movq	%rax, 112(%rsp)
	leaq	4(%rsp), %rax
	movq	%rax, 120(%rsp)
	movq	%rsp, %rax
	movq	%rax, 128(%rsp)
	movabsq	$4294967297, %rax
	movq	%rax, 48(%rsp)
	movq	%rax, 60(%rsp)
	movl	$1, 68(%rsp)
	call	__cudaPopCallConfiguration@PLT
	testl	%eax, %eax
	jne	.L8
	pushq	40(%rsp)
	.cfi_def_cfa_offset 168
	leaq	_Z5sgemmPKfS0_Pfii(%rip), %rdi
	pushq	40(%rsp)
	.cfi_def_cfa_offset 176
	movq	76(%rsp), %rcx
	movl	84(%rsp), %r8d
	movq	64(%rsp), %rsi
	movl	72(%rsp), %edx
	leaq	112(%rsp), %r9
	call	cudaLaunchKernel@PLT
	popq	%rax
	.cfi_def_cfa_offset 168
	popq	%rdx
	.cfi_def_cfa_offset 160
.L8:
	movq	136(%rsp), %rax
	subq	%fs:40, %rax
	je	.L10
	call	__stack_chk_fail@PLT
.L10:
	addq	$152, %rsp
	.cfi_def_cfa_offset 8
	ret
	.cfi_endproc
.LFE2053:
	.size	_Z32__device_stub__Z5sgemmPKfS0_PfiiPKfS0_Pfii, .-_Z32__device_stub__Z5sgemmPKfS0_PfiiPKfS0_Pfii
	.globl	_Z5sgemmPKfS0_Pfii
	.type	_Z5sgemmPKfS0_Pfii, @function
_Z5sgemmPKfS0_Pfii:
.LFB2054:
	.cfi_startproc
	endbr64
	jmp	_Z32__device_stub__Z5sgemmPKfS0_PfiiPKfS0_Pfii
	.cfi_endproc
.LFE2054:
	.size	_Z5sgemmPKfS0_Pfii, .-_Z5sgemmPKfS0_Pfii
	.section	.rodata.str1.1,"aMS",@progbits,1
.LC0:
	.string	"_Z5sgemmPKfS0_Pfii"
.LC1:
	.string	"_Z4fillPfif"
	.section	.text.startup,"ax",@progbits
	.type	_ZL24__sti____cudaRegisterAllv, @function
_ZL24__sti____cudaRegisterAllv:
.LFB2056:
	.cfi_startproc
	endbr64
	pushq	%rbx
	.cfi_def_cfa_offset 16
	.cfi_offset 3, -16
	leaq	_ZL15__fatDeviceText(%rip), %rdi
	call	__cudaRegisterFatBinary@PLT
	pushq	$0
	.cfi_def_cfa_offset 24
	xorl	%r9d, %r9d
	orl	$-1, %r8d
	pushq	$0
	.cfi_def_cfa_offset 32
	leaq	.LC0(%rip), %rdx
	movq	%rax, %rdi
	movq	%rax, %rbx
	pushq	$0
	.cfi_def_cfa_offset 40
	movq	%rdx, %rcx
	leaq	_Z5sgemmPKfS0_Pfii(%rip), %rsi
	pushq	$0
	.cfi_def_cfa_offset 48
	movq	%rax, _ZL20__cudaFatCubinHandle(%rip)
	call	__cudaRegisterFunction@PLT
	addq	$32, %rsp
	.cfi_def_cfa_offset 16
	movq	%rbx, %rdi
	xorl	%r9d, %r9d
	pushq	$0
	.cfi_def_cfa_offset 24
	leaq	.LC1(%rip), %rdx
	orl	$-1, %r8d
	leaq	_Z4fillPfif(%rip), %rsi
	pushq	$0
	.cfi_def_cfa_offset 32
	movq	%rdx, %rcx
	pushq	$0
	.cfi_def_cfa_offset 40
	pushq	$0
	.cfi_def_cfa_offset 48
	call	__cudaRegisterFunction@PLT
	movq	_ZL20__cudaFatCubinHandle(%rip), %rdi
	addq	$32, %rsp
	.cfi_def_cfa_offset 16
	call	__cudaRegisterFatBinaryEnd@PLT
	leaq	_ZL26__cudaUnregisterBinaryUtilv(%rip), %rdi
	popq	%rbx
	.cfi_def_cfa_offset 8
	jmp	atexit@PLT
	.cfi_endproc
.LFE2056:
	.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
	.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:

#pragma once
#include <hip/hip_runtime.h>


constexpr int small_block_size = 4;

constexpr int block_size_x = 128;

__global__ void fill(float* data, int size, float value);

__global__ void sgemm(const float* a, const float* b, float* result, int size, int stride);
	.text
	.file	"sgemm_kernel.hip"
	.type	__hip_cuid_,@object             # @__hip_cuid_
	.bss
	.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
