8f43927575
examples/ is based on KasperskyLab/knp at 55470ed6 with local changes to the mnist-learn altai/blifat models, cuda-backend CMake and model_loader. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
248 lines
11 KiB
Text
248 lines
11 KiB
Text
/**
|
|
* @file projections_impl.cu
|
|
* @brief Contains functions for projection calculation.
|
|
* @kaspersky_support A. Vartenkov
|
|
* @date 14.08.2025
|
|
* @license Apache 2.0
|
|
* @copyright © 2025 AO Kaspersky Lab
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include "projections_impl.cuh"
|
|
#include "projection.cuh"
|
|
#include <cuda_runtime.h>
|
|
#include <thrust/binary_search.h>
|
|
#include <thrust/sort.h>
|
|
|
|
|
|
/**
|
|
* @brief Namespace for CUDA backend.
|
|
*/
|
|
namespace knp::backends::gpu::cuda
|
|
{
|
|
template<>
|
|
void gpu_insert<ProjectionVariants>(const ProjectionVariants &cpu_source, ProjectionVariants *gpu_target)
|
|
{
|
|
::cuda::std::visit([gpu_target](const auto &val)
|
|
{
|
|
using ValueType = std::decay_t<decltype(val)>;
|
|
ValueType *buffer;
|
|
call_and_check(cudaMalloc(&buffer, sizeof(ValueType)));
|
|
gpu_insert(val, buffer);
|
|
device_lib::make_variant_kernel<<<1, 1>>>(gpu_target, buffer);
|
|
call_and_check(cudaFree(buffer));
|
|
}, cpu_source);
|
|
}
|
|
|
|
|
|
__global__ void get_projection_uids_kernel(const ProjectionVariants *projection,
|
|
cuda::UID *pre_uid,
|
|
cuda::UID *post_uid,
|
|
cuda::UID *self_uid)
|
|
{
|
|
::cuda::std::visit([pre_uid, post_uid, self_uid](const auto &proj)
|
|
{
|
|
*pre_uid = proj.presynaptic_uid_;
|
|
*post_uid = proj.postsynaptic_uid_;
|
|
*self_uid = proj.uid_;
|
|
}, *projection);
|
|
}
|
|
|
|
|
|
auto get_projection_uids(const ProjectionVariants *proj)
|
|
{
|
|
cuda::UID *pre_uid_gpu;
|
|
cuda::UID *post_uid_gpu;
|
|
cuda::UID *self_uid_gpu;
|
|
cuda::UID pre_uid, post_uid, self_uid;
|
|
call_and_check(cudaMalloc(&pre_uid_gpu, sizeof(cuda::UID)));
|
|
call_and_check(cudaMalloc(&post_uid_gpu, sizeof(cuda::UID)));
|
|
call_and_check(cudaMalloc(&self_uid_gpu, sizeof(cuda::UID)));
|
|
get_projection_uids_kernel<<<1, 1>>>(proj, pre_uid_gpu, post_uid_gpu, self_uid_gpu);
|
|
call_and_check(cudaMemcpy(&pre_uid, pre_uid_gpu, sizeof(cuda::UID), cudaMemcpyDeviceToHost));
|
|
call_and_check(cudaMemcpy(&post_uid, post_uid_gpu, sizeof(cuda::UID), cudaMemcpyDeviceToHost));
|
|
call_and_check(cudaMemcpy(&self_uid, self_uid_gpu, sizeof(cuda::UID), cudaMemcpyDeviceToHost));
|
|
call_and_check(cudaFree(pre_uid_gpu));
|
|
call_and_check(cudaFree(post_uid_gpu));
|
|
call_and_check(cudaFree(self_uid_gpu));
|
|
return std::make_tuple(pre_uid, post_uid, self_uid);
|
|
}
|
|
|
|
|
|
__global__ void calculate_synaptic_impact(
|
|
const device_lib::CUDAVectorView<CUDAProjection<knp::synapse_traits::DeltaSynapse>::Synapse> synapses,
|
|
const device_lib::LongIndex *synapse_indices, size_t size, StepIndex current_step,
|
|
device_lib::LongIndex *results, device_lib::LongIndex *send_steps)
|
|
{
|
|
const size_t i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= size) return;
|
|
const device_lib::LongIndex synapse_id = synapse_indices[i];
|
|
if (synapse_id >= synapses.size_) return;
|
|
results[i] = synapse_id;
|
|
auto delay = ::cuda::std::get<0>(synapses.data_[synapse_id]).delay_;
|
|
send_steps[i] = delay + current_step - 1;
|
|
}
|
|
|
|
|
|
/**
|
|
* @brief Calculate the series of impacts based on a single spike. Using index from projection we know the number of
|
|
* synapses per each neuron, therefore we can know the output synapse index.
|
|
* @param synapses all synapses, we use them for indexing.
|
|
* @param spike_ids spike indexes, also the indexes of activated neurons
|
|
* @param index synaptic index, see projection for more information
|
|
* @param start_offsets output starting points for the synaptic impacts originating at the neuron.
|
|
* @param current_step current step
|
|
* @param results synapse indexes, output
|
|
* @param send_steps sending steps, output
|
|
*/
|
|
__global__ void calculate_impacts_per_spike(
|
|
const device_lib::CUDAVectorView<CUDAProjection<knp::synapse_traits::DeltaSynapse>::Synapse> synapses,
|
|
device_lib::CUDAVectorView<SpikeIndex> spike_ids, device_lib::IndexView index,
|
|
device_lib::CUDAVectorView<device_lib::LongIndex> start_offsets,
|
|
StepIndex current_step, device_lib::LongIndex *results, device_lib::LongIndex *send_steps)
|
|
{
|
|
const size_t i = blockIdx.x * blockDim.x + threadIdx.x; // Spike index.
|
|
if (i >= spike_ids.size_ || i >= start_offsets.size_) return;
|
|
const SpikeIndex neuron_id = spike_ids.data_[i];
|
|
if (neuron_id + 1 >= index.offsets_size_) return;
|
|
const device_lib::LongIndex start = index.offsets_ptr_[neuron_id];
|
|
const device_lib::LongIndex size = index.offsets_ptr_[neuron_id + 1] - index.offsets_ptr_[neuron_id];
|
|
const device_lib::LongIndex output_start = start_offsets.data_[i];
|
|
const auto [num_blocks, num_threads] = device_lib::get_blocks_config(size);
|
|
// printf("Calc impacts: start %lu, size %lu, out_start %lu\n", start, size, output_start);
|
|
calculate_synaptic_impact<<<num_blocks, num_threads>>>(synapses, index.indices_ptr_ + start, size, current_step,
|
|
results + output_start, send_steps + output_start);
|
|
__syncthreads(); // TODO TEMP
|
|
}
|
|
|
|
|
|
__global__ void delta_indices_to_impacts_kernel(device_lib::LongIndex *indices_begin,
|
|
device_lib::LongIndex *indices_end,
|
|
cuda::device_lib::CUDAVectorView<CUDAProjection<knp::synapse_traits::DeltaSynapse>::Synapse> synapses,
|
|
SynapticImpact *impacts_out)
|
|
{
|
|
constexpr int data_index = core::SynapseElementAccess::synapse_data;
|
|
constexpr int source_id_index = core::SynapseElementAccess::source_neuron_id;
|
|
constexpr int target_id_index = core::SynapseElementAccess::target_neuron_id;
|
|
|
|
const size_t i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
device_lib::LongIndex *index = indices_begin + i;
|
|
if (index >= indices_end) return;
|
|
const device_lib::LongIndex synapse_id = *index;
|
|
if (synapse_id >= synapses.size_) return;
|
|
auto &synapse = ::cuda::std::get<data_index>(synapses.data_[synapse_id]);
|
|
SynapticImpact impact_out;
|
|
impact_out.connection_index_ = synapse_id;
|
|
impact_out.impact_value_ = synapse.weight_;
|
|
impact_out.synapse_type_ = synapse.output_type_;
|
|
impact_out.presynaptic_neuron_index_ = ::cuda::std::get<source_id_index>(synapses.data_[synapse_id]);
|
|
impact_out.postsynaptic_neuron_index_ = ::cuda::std::get<target_id_index>(synapses.data_[synapse_id]);
|
|
impacts_out[i] = impact_out;
|
|
}
|
|
|
|
|
|
template<>
|
|
void CUDAProjection<knp::synapse_traits::DeltaSynapse>::form_message(StepIndex current_step)
|
|
{
|
|
auto iter = thrust::upper_bound(thrust::device, sending_steps_.begin(), sending_steps_.end(), current_step);
|
|
if (iter == sending_steps_.begin())
|
|
{
|
|
// message_buf_.impacts_.clear();
|
|
return;
|
|
}
|
|
device_lib::LongIndex num_impacts = iter - sending_steps_.begin();
|
|
|
|
SynapticImpact *impacts;
|
|
auto [num_blocks, num_threads] = device_lib::get_blocks_config(num_impacts);
|
|
call_and_check(cudaMalloc(&impacts, sizeof(SynapticImpact) * num_impacts));
|
|
delta_indices_to_impacts_kernel<<<num_blocks, num_threads>>>(impact_indexes_.data(),
|
|
impact_indexes_.data() + num_impacts, synapses_.view(),
|
|
impacts);
|
|
MessageHeader header{uid_, current_step};
|
|
message_buf_.header_ = header;
|
|
message_buf_.presynaptic_population_uid_ = presynaptic_uid_;
|
|
message_buf_.postsynaptic_population_uid_ = postsynaptic_uid_;
|
|
message_buf_.impacts_ = device_lib::CUDAVector<SynapticImpact>{impacts, num_impacts};
|
|
message_buf_.is_forcing_ = true;
|
|
sending_steps_.erase(sending_steps_.begin(), sending_steps_.begin() + num_impacts);
|
|
impact_indexes_.erase(impact_indexes_.begin(), impact_indexes_.begin() + num_impacts);
|
|
}
|
|
|
|
|
|
/**
|
|
* @brief Calculate delta projection, host gpu-using variant.
|
|
*/
|
|
__host__ void calculate_projection(
|
|
CUDAProjection<knp::synapse_traits::DeltaSynapse> &projection,
|
|
const CUDAMessageBus &device_message_bus,
|
|
const std::vector<device_lib::LongIndex> &message_ids,
|
|
StepIndex step_n)
|
|
{
|
|
const auto &messages = device_message_bus.all_messages<SpikeMessage>();
|
|
for (size_t i = 0; i < message_ids.size(); ++i)
|
|
{
|
|
const device_lib::LongIndex msg_index = message_ids[i];
|
|
const size_t data_size = messages[msg_index].neuron_indexes_.size();
|
|
auto msg_data_pointer_cpu = messages[msg_index].neuron_indexes_.data();
|
|
SPDLOG_TRACE("Got message data: pointer {}, size {}", reinterpret_cast<const void*>(msg_data_pointer_cpu),
|
|
data_size);
|
|
|
|
if (data_size)
|
|
{
|
|
device_lib::LongIndex impacts_count = count_values_by_indexes(projection.index_,
|
|
device_lib::CUDAVectorView<SpikeIndex>{msg_data_pointer_cpu, data_size});
|
|
|
|
device_lib::LongIndex *impacts_buffer;
|
|
device_lib::LongIndex *delay_buffer;
|
|
call_and_check(cudaMalloc(&impacts_buffer, sizeof(device_lib::LongIndex) * impacts_count));
|
|
call_and_check(cudaMalloc(&delay_buffer, sizeof(device_lib::LongIndex) * impacts_count));
|
|
|
|
// 2. For each active synapse calculate its impact.
|
|
auto [num_blocks, num_threads] = device_lib::get_blocks_config(data_size);
|
|
auto output_start_indices = device_lib::calculate_neuron_scan(projection.index_,
|
|
device_lib::CUDAVectorView<SpikeIndex>{msg_data_pointer_cpu, data_size});
|
|
|
|
calculate_impacts_per_spike<<<num_blocks, num_threads>>>(projection.synapses_.view(),
|
|
device_lib::CUDAVectorView<SpikeIndex>{msg_data_pointer_cpu, data_size}, projection.index_.view(),
|
|
output_start_indices.view(), step_n, impacts_buffer, delay_buffer);
|
|
|
|
cudaDeviceSynchronize();
|
|
// 3. Sort impacts by time
|
|
thrust::sort_by_key(thrust::device, delay_buffer, delay_buffer + impacts_count, impacts_buffer);
|
|
projection.add_impacts(device_lib::CUDAVector<device_lib::LongIndex>{impacts_buffer, impacts_count},
|
|
device_lib::CUDAVector<device_lib::LongIndex>{delay_buffer, impacts_count});
|
|
}
|
|
// Make messages
|
|
projection.form_message(step_n);
|
|
}
|
|
}
|
|
|
|
|
|
__host__ void calculate_projection(
|
|
CUDAProjection<knp::synapse_traits::AdditiveSTDPDeltaSynapse> &projection,
|
|
const std::vector<device_lib::LongIndex> &message_ids,
|
|
StepIndex step_n)
|
|
{
|
|
//SPDLOG_TRACE("Calculate AdditiveSTDPDelta synapse projection {}.", std::string(projection.get_uid()));
|
|
}
|
|
|
|
|
|
__host__ void calculate_projection(
|
|
CUDAProjection<knp::synapse_traits::SynapticResourceSTDPDeltaSynapse> &projection,
|
|
const std::vector<device_lib::LongIndex> &message_ids,
|
|
StepIndex step_n)
|
|
{
|
|
}
|
|
|
|
} // namespace knp::backends::gpu::cuda
|