<commit_before>#include <stdint.h>
#include <vector>
#include "gninasrc/lib/gpu_util.h"

#include "caffe/layers/molgrid_data_layer.hpp"


//gridding is implemented in gridmaker
#include "gninasrc/lib/gridmaker.cu"
#define THREADS_PER_BLOCK 1024

namespace caffe {

template <typename Dtype>
void MolGridDataLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
		const vector<Blob<Dtype>*>& top)
{
	forward(bottom, top, true);
}

template <typename Dtype>
void MolGridDataLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
		const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom)
{
	backward(top, bottom, true);
}

template <typename Dtype>
void MolGridDataLayer<Dtype>::setAtomGradientsGPU(GridMaker& gmaker, Dtype
        *diff, unsigned batch_size)  {

  //launch a kernel for each batch
  for (int item_id = 0; item_id < batch_size; ++item_id) {
    int offset = item_id*example_size;
    //malloc and copy batch data
    float4* atoms;
    short* whichGrid;
    float3* gradient;

    mol_transform& transform = batch_transform[item_id];
    int natoms = transform.mol.atoms.size();
    cudaMalloc(&atoms, sizeof(float4)*natoms);
    cudaMemcpy(atoms, &transform.mol.atoms[0],
            sizeof(float4)*transform.mol.atoms.size(), cudaMemcpyHostToDevice);
    cudaMalloc(&whichGrid, sizeof(short)*transform.mol.whichGrid.size());
    cudaMemcpy(whichGrid, &transform.mol.whichGrid[0],
            sizeof(short)*transform.mol.whichGrid.size(),
            cudaMemcpyHostToDevice);
    cudaMalloc(&gradient, sizeof(float3)*transform.mol.gradient.size());
    cudaMemset(gradient, 0, sizeof(float3)*transform.mol.gradient.size());
    gmaker.setCenter(transform.center[0], transform.center[1], transform.center[2]);

    quaternion& cpu_q = transform.Q;
    qt gpu_q(cpu_q.R_component_2(), cpu_q.R_component_3(),
            cpu_q.R_component_4(), cpu_q.R_component_1());
    vec& molcenter = transform.mol.center;
    //diff is batch x channel x X x Y x Z
	unsigned nfull_blocks = natoms / THREADS_PER_BLOCK;
	unsigned nthreads_remain = natoms % THREADS_PER_BLOCK;
    setAtomGradientGPU<<<nfull_blocks, THREADS_PER_BLOCK>>>(gmaker, atoms, whichGrid, gradient, 
            make_float3(molcenter[0], molcenter[1], molcenter[2]), gpu_q, 
            make_float3(transform.center[0], transform.center[1],
            transform.center[2]), diff, offset);
    setAtomGradientGPU<<<1, nthreads_remain>>>(gmaker, atoms, whichGrid, gradient, 
            make_float3(molcenter[0], molcenter[1], molcenter[2]), gpu_q, 
            make_float3(transform.center[0], transform.center[1],
            transform.center[2]), diff, offset);

    cudaStreamSynchronize(cudaStreamPerThread);
    cudaMemcpy(&transform.mol.gradient[0], gradient,
            sizeof(float3)*transform.mol.gradient.size(),
            cudaMemcpyDeviceToHost);
    cudaFree(atoms);
    cudaFree(whichGrid);<commit_msg>Not sure if cudaStreamSynchronize plays nice with caffe<commit_after>#include <stdint.h>
#include <vector>
#include "gninasrc/lib/gpu_util.h"

#include "caffe/layers/molgrid_data_layer.hpp"


//gridding is implemented in gridmaker
#include "gninasrc/lib/gridmaker.cu"
#define THREADS_PER_BLOCK 1024

namespace caffe {

template <typename Dtype>
void MolGridDataLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
		const vector<Blob<Dtype>*>& top)
{
	forward(bottom, top, true);
}

template <typename Dtype>
void MolGridDataLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
		const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom)
{
	backward(top, bottom, true);
}

template <typename Dtype>
void MolGridDataLayer<Dtype>::setAtomGradientsGPU(GridMaker& gmaker, Dtype
        *diff, unsigned batch_size)  {

  //launch a kernel for each batch
  for (int item_id = 0; item_id < batch_size; ++item_id) {
    int offset = item_id*example_size;
    //malloc and copy batch data
    float4* atoms;
    short* whichGrid;
    float3* gradient;

    mol_transform& transform = batch_transform[item_id];
    int natoms = transform.mol.atoms.size();
    cudaMalloc(&atoms, sizeof(float4)*natoms);
    cudaMemcpy(atoms, &transform.mol.atoms[0],
            sizeof(float4)*transform.mol.atoms.size(), cudaMemcpyHostToDevice);
    cudaMalloc(&whichGrid, sizeof(short)*transform.mol.whichGrid.size());
    cudaMemcpy(whichGrid, &transform.mol.whichGrid[0],
            sizeof(short)*transform.mol.whichGrid.size(),
            cudaMemcpyHostToDevice);
    cudaMalloc(&gradient, sizeof(float3)*transform.mol.gradient.size());
    cudaMemset(gradient, 0, sizeof(float3)*transform.mol.gradient.size());
    gmaker.setCenter(transform.center[0], transform.center[1], transform.center[2]);

    quaternion& cpu_q = transform.Q;
    qt gpu_q(cpu_q.R_component_2(), cpu_q.R_component_3(),
            cpu_q.R_component_4(), cpu_q.R_component_1());
    vec& molcenter = transform.mol.center;
    //diff is batch x channel x X x Y x Z
	unsigned nfull_blocks = natoms / THREADS_PER_BLOCK;
	unsigned nthreads_remain = natoms % THREADS_PER_BLOCK;
    setAtomGradientGPU<<<nfull_blocks, THREADS_PER_BLOCK>>>(gmaker, atoms, whichGrid, gradient, 
            make_float3(molcenter[0], molcenter[1], molcenter[2]), gpu_q, 
            make_float3(transform.center[0], transform.center[1],
            transform.center[2]), diff, offset);
    setAtomGradientGPU<<<1, nthreads_remain>>>(gmaker, atoms, whichGrid, gradient, 
            make_float3(molcenter[0], molcenter[1], molcenter[2]), gpu_q, 
            make_float3(transform.center[0], transform.center[1],
            transform.center[2]), diff, offset);

    //could probably be a StreamSync instead
    cudaDeviceSynchronize();
    cudaMemcpy(&transform.mol.gradient[0], gradient,
            sizeof(float3)*transform.mol.gradient.size(),
            cudaMemcpyDeviceToHost);
    cudaFree(atoms);
    cudaFree(whichGrid);
{'fraction_non_alphanumeric': 0.130420054200542, 'fraction_numerical': 0.013550135501355014, 'mean_word_length': 4.4373848987108655, 'pattern_counts': {'":': 0, '<': 51, '<?xml version=': 0, '>': 49, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5991122', 'n_tokens_mistral': 2044, 'n_tokens_neox': 1956, 'n_words': 366}