//Mesh Laplacian
//Author: Weiyue Wang
//Reference: https://github.com/charlesq34/pointnet-autoencoder/blob/master/tf_ops/nn_distance/tf_nndistance_g.cu
//           https://github.com/PointCloudLibrary/pcl/blob/master/tools/mesh_sampling.cpp

#if GOOGLE_CUDA
#define EIGEN_USE_GPU
#include <stdio.h>
#include <assert.h>


__device__ void getPoint(const float *vertices, int v_id, float *p){
    p[0] = vertices[3* v_id];
    p[1] = vertices[3* v_id+1];
    p[2] = vertices[3* v_id+2];
}

__device__ void getTriangle(const int *triangles, int t_id, int &v1, int &v2, int &v3){
    v1 = triangles[3 * t_id];
    v2 = triangles[3 * t_id + 1];
    v3 = triangles[3 * t_id + 2];
}

__host__ void getTriangle_cpu(const int *triangles, int t_id, int &v1, int &v2, int &v3){
    v1 = triangles[3 * t_id];
    v2 = triangles[3 * t_id + 1];
    v3 = triangles[3 * t_id + 2];
}

__host__ bool findnb_cpu(const int *nb, int v_id){
    bool flag = false;
    for (int i=0; i<20; i++)
        if (nb[i] == v_id){
            flag = true;
            break;
        }
    return flag;
}
__device__  bool findnb(const int *nb, int v_id){
    bool flag = false;
    for (int i=0; i<20; i++)
        if (nb[i] == v_id){
            flag = true;
            break;
        }
    return flag;
}

__device__ void setLaplacian(float *laplacian, int *count, int v_id1, int v_id2, float * p1, float * p2){

    atomicAdd(&laplacian[3 * v_id1], p2[0] - p1[0]);
    atomicAdd(&laplacian[3 * v_id1+1], p2[1] - p1[1]);
    atomicAdd(&laplacian[3 * v_id1+2], p2[2] - p1[2]);

    atomicAdd(&count[v_id1], 1);


    atomicAdd(&laplacian[3 * v_id2], p1[0] - p2[0]);
    atomicAdd(&laplacian[3 * v_id2+1], p1[1] - p2[1]);
    atomicAdd(&laplacian[3 * v_id2+2], p1[2] - p2[2]);

    atomicAdd(&count[v_id2], 1);
}


__global__ void InitMeshLaplacianKernel(const int b, const int maxnverts, float* laplacian, int* count){
    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        for (int v_id=threadIdx.x+blockIdx.y*blockDim.x; v_id<maxnverts; v_id+=blockDim.x*gridDim.y){
            for (int i_c = 0; i_c < 3; i_c++){
                laplacian[i*maxnverts*3+v_id*3+i_c] = 0;
                count[i*maxnverts+v_id] = 0;
            }
        }
    }
}

__global__ void AvgMeshLaplacianKernel(const int b, const int maxnverts, float* laplacian, int* count){
    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        for (int v_id=threadIdx.x+blockIdx.y*blockDim.x; v_id<maxnverts; v_id+=blockDim.x*gridDim.y){
            for (int i_c = 0; i_c < 3; i_c++)
                if (count[i*maxnverts+v_id]!=0)
                    laplacian[i*maxnverts*3+v_id*3+i_c] /= count[i*maxnverts+v_id];
        }
    }
}
__global__ void MeshLaplacianKernel(const int b, const int * nverts, const int maxnverts, const float * vertices, const int * ntriangles, const int maxntriangles, const int * triangles,
                                    float * laplacian, int* nb, int* count){

    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        int n_triangles = ntriangles[i];
        for (int triangle_id=threadIdx.x+blockIdx.y*blockDim.x; triangle_id < n_triangles; triangle_id+=blockDim.x*gridDim.y){

            int v1, v2, v3;
            float p1[3], p2[3], p3[3];
            getTriangle(&triangles[i*maxntriangles*3], triangle_id, v1, v2, v3);

            getPoint(&vertices[i*maxnverts*3], v1, p1);
            getPoint(&vertices[i*maxnverts*3], v2, p2);
            getPoint(&vertices[i*maxnverts*3], v3, p3);

            if (!findnb(&nb[(i*maxnverts+v1)*20], v2)){
                nb[(i*maxnverts+v1)*20+count[i*maxnverts+v1]] = v2;
                nb[(i*maxnverts+v2)*20+count[i*maxnverts+v2]] = v1;
                setLaplacian(&laplacian[i*maxnverts*3], &count[i*maxnverts], v1, v2, p1, p2);
            }
            if (!findnb(&nb[(i*maxnverts+v1)*20], v3)){
                nb[(i*maxnverts+v1)*20+count[i*maxnverts+v1]] = v3;
                nb[(i*maxnverts+v3)*20+count[i*maxnverts+v3]] = v1;
                setLaplacian(&laplacian[i*maxnverts*3], &count[i*maxnverts], v1, v3, p1, p3);
            }
            if (!findnb(&nb[(i*maxnverts+v2)*20], v3)){
                nb[(i*maxnverts+v2)*20+count[i*maxnverts+v2]] = v3;
                nb[(i*maxnverts+v3)*20+count[i*maxnverts+v3]] = v2;
                setLaplacian(&laplacian[i*maxnverts*3], &count[i*maxnverts], v2, v3, p2, p3);
            }

        }
        __syncthreads();
    }
}


void MeshLaplacianKernelLauncher( \
    /*inputs*/  const int b, const int * n_verts, const int maxn_verts, const float * vertices, const int * n_triangles, const int maxn_triangles, const int * triangles, \
    /*outputs*/ float * laplacian, int * count, int * nb){

    InitMeshLaplacianKernel<<<dim3(2,8,1),512>>>(b, maxn_verts, laplacian, count);
    cudaMemset(nb, -1, b*maxn_verts*20*sizeof(int));
    MeshLaplacianKernel<<<64,512>>>(b, n_verts, maxn_verts, vertices, n_triangles, maxn_triangles, triangles, laplacian, nb, count);
    AvgMeshLaplacianKernel<<<dim3(32,16,1),512>>>(b, maxn_verts, laplacian, count);


}

/****************** Gradient ******************/
__device__ void setLaplaciangrad (const float * grad_laplacian, float *grad_verts, int v_id1, int v_id2){

    atomicAdd(&grad_verts[3 * v_id1], - 1 * grad_laplacian[3 * v_id1]);
    atomicAdd(&grad_verts[3 * v_id1+1], - 1 * grad_laplacian[3 * v_id1+1]);
    atomicAdd(&grad_verts[3 * v_id1+2], - 1 * grad_laplacian[3 * v_id1+2]);

    atomicAdd(&grad_verts[3 * v_id2], grad_laplacian[3 * v_id1]);
    atomicAdd(&grad_verts[3 * v_id2+1], grad_laplacian[3 * v_id1+1]);
    atomicAdd(&grad_verts[3 * v_id2+2], grad_laplacian[3 * v_id1+2]);
}


__global__ void MeshLaplacianGradKernel(const int b, const int maxnverts, const int *nverts, const int * nb, const float * grad_laplacian, float* grad_verts){
    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        int n_verts = nverts[i];
        for (int v1=threadIdx.x+blockIdx.y*blockDim.x; v1 < n_verts; v1+=blockDim.x*gridDim.y){
            for (int v2=0; v2<20; v2++){
                if (nb[(i*maxnverts+v1)*20 + v2] == -1)
                    break;
                else{
                    setLaplaciangrad(&grad_laplacian[i*maxnverts*3], &grad_verts[i*maxnverts*3], v1, nb[(i*maxnverts+v1)*20 + v2]);
                }
            }
        }
    }
}

__global__ void AvgGradKernel(const int b, const int maxnverts, float* grad_verts, const int* cumulativeCounts){
    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        for (int v_id=threadIdx.x+blockIdx.y*blockDim.x; v_id<maxnverts; v_id+=blockDim.x*gridDim.y){
            for (int i_c = 0; i_c < 3; i_c++){
                if (cumulativeCounts[i*maxnverts+v_id]!=0){
                    grad_verts[i*maxnverts*3+v_id*3+i_c] /= cumulativeCounts[i*maxnverts+v_id];
                }
            }
        }
    }
}

void MeshLaplacianGradKernelLauncher(const int b, const int maxnverts, const int * nverts, const float * grad_laplacian, const int * count, const int * nb,
                                    float* grad_verts){
    cudaMemset(grad_verts, 0, b*maxnverts*3*sizeof(float));
    MeshLaplacianGradKernel<<<dim3(32,16,1),512>>>(b, maxnverts, nverts, nb, grad_laplacian, grad_verts);
    AvgGradKernel<<<dim3(32,16,1),512>>>(b, maxnverts, grad_verts, count);
}

#endif


	.file	"tmpxft_0029542b_00000000-6_tf_meshlaplacian_g.cudafe1.cpp"
	.text
#APP

#NO_APP
	.type	_ZL26__cudaUnregisterBinaryUtilv, @function
_ZL26__cudaUnregisterBinaryUtilv:
.LFB2010:
	.cfi_startproc
	endbr64
	movq	_ZL20__cudaFatCubinHandle(%rip), %rdi
	jmp	__cudaUnregisterFatBinary@PLT
	.cfi_endproc
.LFE2010:
	.size	_ZL26__cudaUnregisterBinaryUtilv, .-_ZL26__cudaUnregisterBinaryUtilv
	.section	.text.startup,"ax",@progbits
	.type	_ZL24__sti____cudaRegisterAllv, @function
_ZL24__sti____cudaRegisterAllv:
.LFB2033:
	.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
.LFE2033:
	.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:

//Mesh Laplacian
//Author: Weiyue Wang
//Reference: https://github.com/charlesq34/pointnet-autoencoder/blob/master/tf_ops/nn_distance/tf_nndistance_g.cu
//           https://github.com/PointCloudLibrary/pcl/blob/master/tools/mesh_sampling.cpp

#if GOOGLE_CUDA
#define EIGEN_USE_GPU

#include <hip/hip_runtime.h>
#include <stdio.h>
#include <assert.h>


__device__ void getPoint(const float *vertices, int v_id, float *p){
    p[0] = vertices[3* v_id];
    p[1] = vertices[3* v_id+1];
    p[2] = vertices[3* v_id+2];
}

__device__ void getTriangle(const int *triangles, int t_id, int &v1, int &v2, int &v3){
    v1 = triangles[3 * t_id];
    v2 = triangles[3 * t_id + 1];
    v3 = triangles[3 * t_id + 2];
}

__host__ void getTriangle_cpu(const int *triangles, int t_id, int &v1, int &v2, int &v3){
    v1 = triangles[3 * t_id];
    v2 = triangles[3 * t_id + 1];
    v3 = triangles[3 * t_id + 2];
}

__host__ bool findnb_cpu(const int *nb, int v_id){
    bool flag = false;
    for (int i=0; i<20; i++)
        if (nb[i] == v_id){
            flag = true;
            break;
        }
    return flag;
}
__device__  bool findnb(const int *nb, int v_id){
    bool flag = false;
    for (int i=0; i<20; i++)
        if (nb[i] == v_id){
            flag = true;
            break;
        }
    return flag;
}

__device__ void setLaplacian(float *laplacian, int *count, int v_id1, int v_id2, float * p1, float * p2){

    atomicAdd(&laplacian[3 * v_id1], p2[0] - p1[0]);
    atomicAdd(&laplacian[3 * v_id1+1], p2[1] - p1[1]);
    atomicAdd(&laplacian[3 * v_id1+2], p2[2] - p1[2]);

    atomicAdd(&count[v_id1], 1);


    atomicAdd(&laplacian[3 * v_id2], p1[0] - p2[0]);
    atomicAdd(&laplacian[3 * v_id2+1], p1[1] - p2[1]);
    atomicAdd(&laplacian[3 * v_id2+2], p1[2] - p2[2]);

    atomicAdd(&count[v_id2], 1);
}


__global__ void InitMeshLaplacianKernel(const int b, const int maxnverts, float* laplacian, int* count){
    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        for (int v_id=threadIdx.x+blockIdx.y*blockDim.x; v_id<maxnverts; v_id+=blockDim.x*gridDim.y){
            for (int i_c = 0; i_c < 3; i_c++){
                laplacian[i*maxnverts*3+v_id*3+i_c] = 0;
                count[i*maxnverts+v_id] = 0;
            }
        }
    }
}

__global__ void AvgMeshLaplacianKernel(const int b, const int maxnverts, float* laplacian, int* count){
    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        for (int v_id=threadIdx.x+blockIdx.y*blockDim.x; v_id<maxnverts; v_id+=blockDim.x*gridDim.y){
            for (int i_c = 0; i_c < 3; i_c++)
                if (count[i*maxnverts+v_id]!=0)
                    laplacian[i*maxnverts*3+v_id*3+i_c] /= count[i*maxnverts+v_id];
        }
    }
}
__global__ void MeshLaplacianKernel(const int b, const int * nverts, const int maxnverts, const float * vertices, const int * ntriangles, const int maxntriangles, const int * triangles,
                                    float * laplacian, int* nb, int* count){

    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        int n_triangles = ntriangles[i];
        for (int triangle_id=threadIdx.x+blockIdx.y*blockDim.x; triangle_id < n_triangles; triangle_id+=blockDim.x*gridDim.y){

            int v1, v2, v3;
            float p1[3], p2[3], p3[3];
            getTriangle(&triangles[i*maxntriangles*3], triangle_id, v1, v2, v3);

            getPoint(&vertices[i*maxnverts*3], v1, p1);
            getPoint(&vertices[i*maxnverts*3], v2, p2);
            getPoint(&vertices[i*maxnverts*3], v3, p3);

            if (!findnb(&nb[(i*maxnverts+v1)*20], v2)){
                nb[(i*maxnverts+v1)*20+count[i*maxnverts+v1]] = v2;
                nb[(i*maxnverts+v2)*20+count[i*maxnverts+v2]] = v1;
                setLaplacian(&laplacian[i*maxnverts*3], &count[i*maxnverts], v1, v2, p1, p2);
            }
            if (!findnb(&nb[(i*maxnverts+v1)*20], v3)){
                nb[(i*maxnverts+v1)*20+count[i*maxnverts+v1]] = v3;
                nb[(i*maxnverts+v3)*20+count[i*maxnverts+v3]] = v1;
                setLaplacian(&laplacian[i*maxnverts*3], &count[i*maxnverts], v1, v3, p1, p3);
            }
            if (!findnb(&nb[(i*maxnverts+v2)*20], v3)){
                nb[(i*maxnverts+v2)*20+count[i*maxnverts+v2]] = v3;
                nb[(i*maxnverts+v3)*20+count[i*maxnverts+v3]] = v2;
                setLaplacian(&laplacian[i*maxnverts*3], &count[i*maxnverts], v2, v3, p2, p3);
            }

        }
        __syncthreads();
    }
}


void MeshLaplacianKernelLauncher( \
    /*inputs*/  const int b, const int * n_verts, const int maxn_verts, const float * vertices, const int * n_triangles, const int maxn_triangles, const int * triangles, \
    /*outputs*/ float * laplacian, int * count, int * nb){

    InitMeshLaplacianKernel<<<dim3(2,8,1),512>>>(b, maxn_verts, laplacian, count);
    hipMemset(nb, -1, b*maxn_verts*20*sizeof(int));
    MeshLaplacianKernel<<<64,512>>>(b, n_verts, maxn_verts, vertices, n_triangles, maxn_triangles, triangles, laplacian, nb, count);
    AvgMeshLaplacianKernel<<<dim3(32,16,1),512>>>(b, maxn_verts, laplacian, count);


}

/****************** Gradient ******************/
__device__ void setLaplaciangrad (const float * grad_laplacian, float *grad_verts, int v_id1, int v_id2){

    atomicAdd(&grad_verts[3 * v_id1], - 1 * grad_laplacian[3 * v_id1]);
    atomicAdd(&grad_verts[3 * v_id1+1], - 1 * grad_laplacian[3 * v_id1+1]);
    atomicAdd(&grad_verts[3 * v_id1+2], - 1 * grad_laplacian[3 * v_id1+2]);

    atomicAdd(&grad_verts[3 * v_id2], grad_laplacian[3 * v_id1]);
    atomicAdd(&grad_verts[3 * v_id2+1], grad_laplacian[3 * v_id1+1]);
    atomicAdd(&grad_verts[3 * v_id2+2], grad_laplacian[3 * v_id1+2]);
}


__global__ void MeshLaplacianGradKernel(const int b, const int maxnverts, const int *nverts, const int * nb, const float * grad_laplacian, float* grad_verts){
    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        int n_verts = nverts[i];
        for (int v1=threadIdx.x+blockIdx.y*blockDim.x; v1 < n_verts; v1+=blockDim.x*gridDim.y){
            for (int v2=0; v2<20; v2++){
                if (nb[(i*maxnverts+v1)*20 + v2] == -1)
                    break;
                else{
                    setLaplaciangrad(&grad_laplacian[i*maxnverts*3], &grad_verts[i*maxnverts*3], v1, nb[(i*maxnverts+v1)*20 + v2]);
                }
            }
        }
    }
}

__global__ void AvgGradKernel(const int b, const int maxnverts, float* grad_verts, const int* cumulativeCounts){
    for (int i=blockIdx.x;i<b;i+=gridDim.x){
        for (int v_id=threadIdx.x+blockIdx.y*blockDim.x; v_id<maxnverts; v_id+=blockDim.x*gridDim.y){
            for (int i_c = 0; i_c < 3; i_c++){
                if (cumulativeCounts[i*maxnverts+v_id]!=0){
                    grad_verts[i*maxnverts*3+v_id*3+i_c] /= cumulativeCounts[i*maxnverts+v_id];
                }
            }
        }
    }
}

void MeshLaplacianGradKernelLauncher(const int b, const int maxnverts, const int * nverts, const float * grad_laplacian, const int * count, const int * nb,
                                    float* grad_verts){
    hipMemset(grad_verts, 0, b*maxnverts*3*sizeof(float));
    MeshLaplacianGradKernel<<<dim3(32,16,1),512>>>(b, maxnverts, nverts, nb, grad_laplacian, grad_verts);
    AvgGradKernel<<<dim3(32,16,1),512>>>(b, maxnverts, grad_verts, count);
}

#endif


	.text
	.file	"tf_meshlaplacian_g.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
