wehub/deeplearning4j--deeplearning4j
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:05 +08:00
commit 4f3b7da785
7394 changed files with 2005594 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libnd4j/include/array/cuda/CudaPointerDeallocator.cu
+47
View File
@@ -0,0 +1,47 @@
/*
* ******************************************************************************
* *
* *
* * This program and the accompanying materials are made available under the
* * terms of the Apache License, Version 2.0 which is available at
* * https://www.apache.org/licenses/LICENSE-2.0.
* *
* * See the NOTICE file distributed with this work for additional
* * information regarding copyright ownership.
* * 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.
* *
* * SPDX-License-Identifier: Apache-2.0
* *****************************************************************************
*/
//
// @author raver119@gmail.com
//
#include <array/CudaPointerDeallocator.h>
namespace sd {
void CudaPointerDeallocator::release(void *ptr) {
if (ptr == nullptr) return;
// Check if this is a valid device pointer before freeing
cudaPointerAttributes attributes;
cudaError_t result = cudaPointerGetAttributes(&attributes, ptr);
if (result == cudaSuccess) {
// Only free if it's a regular device pointer
// cudaMemoryTypeDevice is for regular allocations we can free
if (attributes.type == cudaMemoryTypeDevice) {
cudaFree(ptr);
}
// Don't free other types (like constant memory)
} else {
// Clear the error and don't try to free this pointer
cudaGetLastError(); // Clear the error state
}
}
} // namespace sd
libnd4j/include/array/cuda/DataBuffer.cu
+608
View File
@@ -0,0 +1,608 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* 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.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author raver119@gmail.com
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <array/DataTypeUtils.h>
#include <exceptions/allocation_exception.h>
#include <exceptions/cuda_exception.h>
#include <execution/AffinityManager.h>
#include <memory/MemoryCounter.h>
#include <system/op_boilerplate.h>
#include <system/type_boilerplate.h>
#include "../DataBuffer.h"
#include "helpers/DebugHelper.h"
#if defined(SD_GCC_FUNCTRACE)
#include <array/DataBufferLifecycleTracker.h>
#endif
namespace sd {
void DataBuffer::expand(const uint64_t size) {
if (size > _lenInBytes) {
// allocate new buffer
int8_t* newBuffer = nullptr;
int8_t* newSpecialBuffer = nullptr;
ALLOCATE_SPECIAL(newSpecialBuffer, _workspace, size, int8_t);
// copy data from existing buffer
if (_primaryBuffer != nullptr) {
// there's non-zero chance that primary buffer doesn't exist yet
ALLOCATE(newBuffer, _workspace, size, int8_t);
std::memcpy(newBuffer, _primaryBuffer, _lenInBytes);
if (_isOwnerPrimary) {
auto ipb = reinterpret_cast<int8_t*>(_primaryBuffer);
RELEASE(ipb, _workspace);
}
_primaryBuffer = newBuffer;
_isOwnerPrimary = true;
}
cudaMemcpy(newSpecialBuffer, _specialBuffer, _lenInBytes, cudaMemcpyDeviceToDevice);
if (_isOwnerSpecial) {
auto isb = reinterpret_cast<int8_t*>(_specialBuffer);
RELEASE_SPECIAL(isb, _workspace);
}
_specialBuffer = newSpecialBuffer;
_lenInBytes = size;
_isOwnerSpecial = true;
}
}
DataBuffer DataBuffer::dup() {
DataBuffer result;
result._dataType = _dataType;
result._lenInBytes = _lenInBytes;
result._primaryBuffer = _primaryBuffer;
result._specialBuffer = _specialBuffer;
result._isOwnerPrimary = _isOwnerPrimary;
result._isOwnerSpecial = _isOwnerSpecial;
result.allocateBuffers(true);
result.copyCounters(*this);
result.copyBufferFrom(*this);
return result;
}
template <typename T>
void* DataBuffer::primaryAtOffset(const LongType offset) {
if(_primaryBuffer == nullptr)
return nullptr;
T *type = reinterpret_cast<T*>(_primaryBuffer);
return reinterpret_cast<void *>(type + offset);
}
template <typename T>
void* DataBuffer::specialAtOffset(const LongType offset) {
if(_specialBuffer == nullptr)
return nullptr;
T *type = reinterpret_cast<T*>(_specialBuffer);
return reinterpret_cast<void *>(type + offset);
}
#define PRIMARYOFFSET(T) template SD_LIB_EXPORT void* DataBuffer::primaryAtOffset<GET_SECOND(T)>(sd::LongType offset);
ITERATE_LIST((SD_COMMON_TYPES),PRIMARYOFFSET)
#define SPECIALOFFSET(T) template SD_LIB_EXPORT void* DataBuffer::specialAtOffset<GET_SECOND(T)>(sd::LongType offset);
ITERATE_LIST((SD_COMMON_TYPES),SPECIALOFFSET)
template <typename T>
void _printHostBuffer(DataBuffer* buffer, long offset) {
sd::LongType len = buffer->getNumElements();
auto buff = buffer->template primaryAsT<T>();
sd::LongType limit = len;
if (limit == -1 || limit >= buffer->getNumElements()) {
limit = buffer->getNumElements();
}
const char* msg = nullptr;
if (msg != nullptr) {
printf("%s: ", msg);
} else {
printf("[");
}
sd::DataType dataType = buffer->getDataType();
auto baseOffset = offset;
if (dataType == sd::DataType::DOUBLE || dataType == sd::DataType::FLOAT32) {
for (sd::LongType e = baseOffset; e < limit; e++) {
if (e > offset) printf(", ");
if (dataType == sd::DataType::DOUBLE) {
printf("%.15f", buff[e]);
} else {
printf("%.15f", static_cast<float>(buff[e]));
}
}
} else if (dataType == sd::DataType::INT64 || dataType == sd::DataType::UINT64 ||
dataType == sd::DataType::INT32 || dataType == sd::DataType::UINT32) {
for (sd::LongType e = baseOffset; e < limit; e++) {
if (dataType == sd::DataType::INT64 || dataType == sd::DataType::UINT64) {
printf("%lld", static_cast<long long>(buff[e]));
} else {
printf("%d", static_cast<int>(buff[e]));
}
if (e < limit - 1) {
printf(", ");
}
}
} else if (dataType == sd::DataType::BOOL) {
for (sd::LongType e = baseOffset; e < limit; e++) {
if (static_cast<bool>(buff[e])) {
printf("true");
} else {
printf("false");
}
if (e < limit - 1) {
printf(", ");
}
}
} else if (dataType == sd::DataType::UTF8 || dataType == sd::DataType::UTF16 ||
dataType == sd::DataType::UTF32) {
for (sd::LongType e = baseOffset; e < limit; e++) {
printf("\"%s\"", reinterpret_cast<const char*>(&buff[e]));
if (e < limit - 1) {
printf(", ");
}
}
}
printf("]\n");
fflush(stdout);
}
void DataBuffer::printHostDevice(long offset) {
THROW_EXCEPTION("");
}
void DataBuffer::printSpecialAllocationTraces() {
//no op on purpose
}
void DataBuffer::showBufferLimited() {
}
template <typename T>
SD_KERNEL void printDeviceBufferKernel(void* buffer, sd::LongType offset, sd::LongType length) {
T* typedBuffer = reinterpret_cast<T*>(buffer);
if (threadIdx.x == 0 && blockIdx.x == 0) {
printf("[ ");
for (sd::LongType i = offset; i < offset + length; i++) {
// Cast to double for consistent formatting
printf("%g ", (double)typedBuffer[i]);
}
printf("]");
}
}
BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT SD_KERNEL void printDeviceBufferKernel,(void* buffer, sd::LongType offset, sd::LongType length),SD_COMMON_TYPES);
// Wrapper function to launch the kernel
template <typename T>
void launchPrintDeviceBufferKernel(void* buffer, sd::LongType offset, sd::LongType length) {
printDeviceBufferKernel<T><<<1, 1, 32*1024, *LaunchContext::defaultContext()->getCudaStream()>>>(
buffer, offset, length);
cudaStreamSynchronize(*LaunchContext::defaultContext()->getCudaStream());
sd::DebugHelper::checkErrorCode(LaunchContext::defaultContext()->getCudaStream(),
"printBufferDebug kernel failed");
}
BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void launchPrintDeviceBufferKernel,(void* buffer, sd::LongType offset, sd::LongType length),SD_COMMON_TYPES);
template <typename T>
void DataBuffer::printHostBufferContent(void* buffer, sd::LongType offset, sd::LongType length) {
T* typedBuffer = reinterpret_cast<T*>(buffer);
sd_printf("[ ", 0);
for (sd::LongType i = offset; i < offset + length; i++) {
// For numeric types, cast to double for consistent formatting
if (std::is_arithmetic<T>::value) {
sd_printf("%g ", (double)typedBuffer[i]);
} else {
// For non-numeric types, print as hex
sd_printf("0x%x ", *reinterpret_cast<int*>(&typedBuffer[i]));
}
}
sd_printf("]", 0);
}
BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void DataBuffer::printHostBufferContent,(void* buffer, sd::LongType offset, sd::LongType length),SD_COMMON_TYPES);
// DataBuffer implementation for .cu file
void DataBuffer::printBufferDebug(const char* msg, sd::LongType offset, sd::LongType limit) {
if (msg) sd_printf("%s:\n", msg);
// Print metadata
sd_printf("DataBuffer: DataType=%s, Length=%lld elements, DeviceId=%d\n",
DataTypeUtils::asString(_dataType).c_str(), (long long)getNumElements(), deviceId());
// Print host buffer content
if (_primaryBuffer != nullptr) {
sd_printf("Host buffer (@%p): ", _primaryBuffer);
sd::LongType len = getNumElements();
sd::LongType printLen = limit < 0 ? len : std::min(len - offset, limit);
// Print based on datatype
BUILD_SINGLE_SELECTOR(_dataType, printHostBufferContent,
(_primaryBuffer, offset, printLen), SD_COMMON_TYPES);
if (offset + printLen < len) sd_printf("... ", 0);
sd_printf("\n", 0);
} else {
sd_printf("Host buffer: nullptr\n", 0);
}
// Print device buffer using kernel
if (_specialBuffer != nullptr) {
sd_printf("Device buffer (@%p): ", _specialBuffer);
sd::LongType len = getNumElements();
sd::LongType printLen = limit < 0 ? len : std::min(len - offset, limit);
// Launch kernel through wrapper function
BUILD_SINGLE_SELECTOR(_dataType, launchPrintDeviceBufferKernel,
(_specialBuffer, offset, printLen), SD_COMMON_TYPES);
sd_printf("\n", 0);
} else {
sd_printf("Device buffer: nullptr\n", 0);
}
// Print sync state counters
sd_printf("Sync state: _counter=%lld, _writePrimary=%lld, _writeSpecial=%lld, _readPrimary=%lld, _readSpecial=%lld\n",
(long long)_counter.load(), (long long)_writePrimary.load(), (long long)_writeSpecial.load(),
(long long)_readPrimary.load(), (long long)_readSpecial.load());
sd_printf("isPrimaryActual=%d, isSpecialActual=%d\n", isPrimaryActual(), isSpecialActual());
}
void DataBuffer::showCounters(const char* msg1, const char* msg2) {
sd_debug("%s %s || primary %p special %p :: wP: %d wS: %d rP: %d rS: %d\n", msg1, msg2, _primaryBuffer,
_specialBuffer, (int)_writePrimary.load(), (int)_writeSpecial.load(), (int)_readPrimary.load(),
(int)_readSpecial.load());
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::allocateSpecial() {
if (_specialBuffer != nullptr) {
return;
}
if (_lenInBytes == 0) {
std::string errorMessage;
errorMessage += "DataBuffer::allocateSpecial: ";
errorMessage += "Special buffer is already allocated";
errorMessage += " or length is 0";
errorMessage += "Length is: ";
errorMessage += std::to_string(getLenInBytes());
errorMessage += "Special buffer is nullptr : ";
errorMessage += std::to_string(_specialBuffer == nullptr);
THROW_EXCEPTION(errorMessage.c_str());
}
#if defined(SD_GCC_FUNCTRACE)
if(Environment::getInstance().isFuncTracePrintAllocate()) {
allocationStackTraceSpecial = new StackTrace();
allocationStackTraceSpecial->load_here();
}
#endif
if (_specialBuffer == nullptr) {
auto deviceId = AffinityManager::currentDeviceId();
if (_workspace == nullptr) {
if (!memory::MemoryCounter::getInstance().validate(getLenInBytes())) {
std::string errorMessage;
errorMessage += "DataBuffer::allocateSpecial: ";
errorMessage += "Requested amount exceeds device limits";
errorMessage += "DeviceId: ";
errorMessage += std::to_string(deviceId);
errorMessage += "Device limit: ";
errorMessage += std::to_string(memory::MemoryCounter::getInstance().deviceLimit(deviceId));
errorMessage += "Requested amount: ";
errorMessage += std::to_string(getLenInBytes());
errorMessage += "Special buffer is nullptr : ";
errorMessage += std::to_string(_specialBuffer == nullptr);
THROW_EXCEPTION(errorMessage.c_str());
}
}
ALLOCATE_SPECIAL(_specialBuffer, _workspace, getLenInBytes(), int8_t);
_isOwnerSpecial = true;
#if defined(SD_GCC_FUNCTRACE)
// Record SPECIAL (device) buffer allocation
array::DataBufferLifecycleTracker::getInstance().recordAllocation(
_specialBuffer, getLenInBytes(), getDataType(),
array::BufferType::SPECIAL, this, _workspace != nullptr);
#endif
if (_workspace == nullptr) {
memory::MemoryCounter::getInstance().countIn(deviceId, getLenInBytes());
memory::MemoryCounter::getInstance().countIn(memory::MemoryType::DEVICE, getLenInBytes());
}
} else if(getLenInBytes() == 0) {
std::string errorMessage;
errorMessage += "DataBuffer::allocateSpecial: ";
errorMessage += "Special buffer is already allocated";
errorMessage += " or length is 0";
errorMessage += "Length is: ";
errorMessage += std::to_string(getLenInBytes());
errorMessage += "Special buffer is nullptr : ";
errorMessage += std::to_string(_specialBuffer == nullptr);
THROW_EXCEPTION(errorMessage.c_str());
}
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::syncToPrimary(const LaunchContext* context, const bool forceSync) {
if (isPrimaryActual() && !forceSync) {
return;
}
allocatePrimary();
auto res = cudaStreamSynchronize(*context->getCudaStream());
if (res != 0) {
std::string errorMessage;
errorMessage += "DataBuffer::syncToPrimary: cudaStreamSynchronize failed: ";
errorMessage += std::to_string(getLenInBytes());
errorMessage += cudaGetErrorString(res);
errorMessage += "Special buffer is nullptr : ";
THROW_EXCEPTION(errorMessage.c_str());
}
res = cudaMemcpy(_primaryBuffer, _specialBuffer, getLenInBytes(), cudaMemcpyDeviceToHost);
if (res != 0) {
std::string errorMessage;
errorMessage += "DataBuffer::syncToPrimary: cudaMemcpy failed: ";
errorMessage += std::to_string(getLenInBytes());
errorMessage += cudaGetErrorString(res);
errorMessage += "Special buffer is nullptr : ";
errorMessage += std::to_string(_specialBuffer == nullptr);
THROW_EXCEPTION(errorMessage.c_str());
}
readPrimary();
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::syncToSpecial(const bool forceSync) {
// in this case there's nothing to do here
if (_primaryBuffer == nullptr) return;
if (isSpecialActual() && !forceSync) {
return;
}
allocateSpecial();
auto res = cudaMemcpy(_specialBuffer, _primaryBuffer, getLenInBytes(), cudaMemcpyHostToDevice);
if (res != 0) {
std::string errorMessage;
errorMessage += "Failed to copy dataBuffer::syncToSpecial: ";
errorMessage += std::to_string(getLenInBytes());
errorMessage += cudaGetErrorString(res);
THROW_EXCEPTION(errorMessage.c_str());
}
readSpecial();
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::deleteSpecial() {
if (_isOwnerSpecial && _specialBuffer != nullptr && getLenInBytes() != 0) {
auto p = reinterpret_cast<int8_t*>(_specialBuffer);
#if defined(SD_GCC_FUNCTRACE)
// Record SPECIAL (device) buffer deallocation before releasing
array::DataBufferLifecycleTracker::getInstance().recordDeallocation(
_specialBuffer, array::BufferType::SPECIAL);
#endif
RELEASE_SPECIAL(p, _workspace);
_specialBuffer = nullptr;
_isOwnerSpecial = false;
// count out towards DataBuffer device, only if we're not in workspace
if (_workspace == nullptr) {
sd::memory::MemoryCounter::getInstance().countOut(_deviceId, getLenInBytes());
sd::memory::MemoryCounter::getInstance().countOut(sd::memory::MemoryType::DEVICE, getLenInBytes());
}
}
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::setCountersToZero() {
_counter.store(0L);
_writePrimary.store(0L);
_writeSpecial.store(0L);
_readPrimary.store(0L);
_readSpecial.store(0L);
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::copyCounters(const DataBuffer& other) {
_counter.store(other._counter);
_writePrimary.store(other._readSpecial);
_writeSpecial.store(other._readPrimary);
_readPrimary.store(other._writeSpecial);
_readSpecial.store(other._writePrimary);
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::copyBufferFrom(const DataBuffer& other, size_t sizeToCopyinBytes, const sd::LongType offsetThis,
const sd::LongType offsetOther) { // copies only to special buffer
if (other._primaryBuffer == nullptr && other._specialBuffer == nullptr) {
return;
}
if (sizeToCopyinBytes == 0) {
sizeToCopyinBytes = other.getLenInBytes();
}
if (sizeToCopyinBytes == 0) {
return;
}
if (other.isPrimaryActual()) {
auto res = cudaMemcpy(
static_cast<int8_t*>(_specialBuffer) + offsetThis * DataTypeUtils::sizeOfElement(_dataType),
static_cast<const int8_t*>(other._primaryBuffer) + offsetOther * DataTypeUtils::sizeOfElement(other._dataType),
sizeToCopyinBytes, cudaMemcpyHostToDevice);
if (res != 0)
throw cuda_exception::build("DataBuffer::copyBufferFrom: cudaMemcpy_cudaMemcpyHostToDevice failed!", res);
other.readPrimary();
} else {
auto res = cudaMemcpy(
static_cast<int8_t*>(_specialBuffer) + offsetThis * DataTypeUtils::sizeOfElement(_dataType),
static_cast<const int8_t*>(other._specialBuffer) + offsetOther * DataTypeUtils::sizeOfElement(other._dataType),
sizeToCopyinBytes, cudaMemcpyDeviceToDevice);
if (res != 0)
throw cuda_exception::build("DataBuffer::copyBufferFrom: cudaMemcpy_cudaMemcpyDeviceToDevice failed!", res);
other.readSpecial();
}
writeSpecial();
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::copyBufferFromHost(const void* hostBuffer, size_t sizeToCopyinBytes, const sd::LongType offsetThis,
const sd::LongType offsetHostBuffer) { // copies only to special buffer
if (hostBuffer == nullptr) return;
if (sizeToCopyinBytes == 0) sizeToCopyinBytes = getLenInBytes();
if (sizeToCopyinBytes == 0) return;
auto res =
cudaMemcpy(static_cast<int8_t*>(_specialBuffer) + offsetThis * DataTypeUtils::sizeOfElement(_dataType),
static_cast<const int8_t*>(hostBuffer) + offsetHostBuffer * DataTypeUtils::sizeOfElement(_dataType),
sizeToCopyinBytes, cudaMemcpyHostToDevice);
if (res != 0)
throw cuda_exception::build("DataBuffer::copyBufferFromHost: cudaMemcpy_cudaMemcpyHostToDevice failed!", res);
writeSpecial();
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::setSpecial(void* special, const bool isOwnerSpecial) {
deleteSpecial();
_specialBuffer = special;
_isOwnerSpecial = isOwnerSpecial;
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::allocateBuffers(const bool allocBoth) { // always allocate special buffer only (cuda case)
allocateSpecial();
if (allocBoth) allocatePrimary();
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::setToZeroBuffers(const bool both) {
if(getLenInBytes() < 1 || special() == nullptr)
return;
cudaMemsetAsync(special(), 0, getLenInBytes(), *LaunchContext::defaultContext()->getCudaStream());
auto res = cudaStreamSynchronize(*LaunchContext::defaultContext()->getCudaStream());
if (res != 0) throw cuda_exception::build("DataBuffer::setToZeroBuffers: streamSync failed!", res);
writeSpecial();
if (both) {
memset(primary(), 0, getLenInBytes());
readPrimary();
}
}
/////////////////////////
template <typename T>
void memcpyWithT(DataBuffer* dst, DataBuffer* src, sd::LongType startingOffset, sd::LongType dstOffset) {
if (src->getLenInBytes() > dst->getLenInBytes())
THROW_EXCEPTION("DataBuffer::memcpy: Source data buffer is larger than destination");
int res = 0;
if (src->isSpecialActual()) {
res = cudaMemcpyAsync(dst->specialAtOffset<T>(dstOffset), src->specialAtOffset<T>(startingOffset), src->getLenInBytes(), cudaMemcpyDeviceToDevice,
*LaunchContext::defaultContext()->getCudaStream());
} else if (src->isPrimaryActual()) {
res = cudaMemcpyAsync(dst->specialAtOffset<T>(dstOffset), src->specialAtOffset<T>(startingOffset), src->getLenInBytes(), cudaMemcpyHostToDevice,
*LaunchContext::defaultContext()->getCudaStream());
}
if (res != 0) throw cuda_exception::build("DataBuffer::memcpy: cudaMemcpyAsync failed!", res);
res = cudaStreamSynchronize(*LaunchContext::defaultContext()->getCudaStream());
if (res != 0) throw cuda_exception::build("DataBuffer::memcpy: streamSync failed!", res);
dst->writeSpecial();
}
void DataBuffer::memcpy(DataBuffer* dst, DataBuffer* src,
sd::LongType startingOffset, sd::LongType dstOffset) {
BUILD_SINGLE_TEMPLATE(memcpyWithT,(dst, src, startingOffset, dstOffset),
SD_COMMON_TYPES);
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::migrate() {
memory::Workspace* newWorkspace = nullptr;
void* newBuffer;
ALLOCATE_SPECIAL(newBuffer, newWorkspace, getLenInBytes(), int8_t);
auto res = cudaMemcpy(newBuffer, _specialBuffer, getLenInBytes(), cudaMemcpyDeviceToDevice);
if (res != 0) throw cuda_exception::build("DataBuffer::migrate: cudaMemcpyAsync failed!", res);
if (_isOwnerSpecial) {
// now we're releasing original buffer
RELEASE_SPECIAL(_specialBuffer, _workspace);
}
_isOwnerSpecial = true;
_specialBuffer = newBuffer;
}
////////////////////////////////////////////////////////////////////////
void DataBuffer::writePrimary() const { _writePrimary = ++_counter; }
void DataBuffer::writeSpecial() const { _writeSpecial = ++_counter; }
void DataBuffer::readPrimary() const { _readPrimary = ++_counter; }
void DataBuffer::readSpecial() const { _readSpecial = ++_counter; }
bool DataBuffer::isPrimaryActual() const {
return (_writePrimary.load() > _writeSpecial.load() || _readPrimary.load() > _writeSpecial.load());
}
bool DataBuffer::isSpecialActual() const {
return (_writeSpecial.load() > _writePrimary.load() || _readSpecial.load() > _writePrimary.load());
}
} // namespace sd
libnd4j/include/array/cuda/NDArray.cu
+720
View File
@@ -0,0 +1,720 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* 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.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
#ifndef NDARRAY_CPP
#define NDARRAY_CPP
#include <array/NDArray.h>
#include <array/NDArrayFactory.h>
#include <exceptions/cuda_exception.h>
#include <exceptions/datatype_exception.h>
#include <helpers/ArrayUtils.h>
#include <helpers/ConstantShapeHelper.h>
#include <helpers/MmulHelper.h>
#include <helpers/PointersManager.h>
#include <helpers/ShapeUtils.h>
#include <helpers/logger.h>
#include <helpers/threshold.h>
#include <indexing/IndicesList.h>
#include <indexing/NDIndex.h>
#include <legacy/NativeOpExecutioner.h>
#include <loops/broadcasting.h>
#include <loops/pairwise_transform.h>
#include <loops/random.h>
#include <loops/special_kernels.h>
#include <loops/transform_same.h>
#include <memory/MemoryRegistrator.h>
#include <memory/Workspace.h>
#include <ops/ops.h>
#include <ops/specials_cuda.h>
#include <array/NDArray.hXX>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <system/selective_rendering.h>
#include "execution/cuda/LaunchDims.h"
namespace sd {
void* NDArray::platformBuffer() { return specialBuffer(); }
void NDArray::syncToDevice() {
auto currentDeviceId = AffinityManager::currentDeviceId();
if (currentDeviceId != _deviceId) {
// first of all we update shapeInfo
const_cast<NDArray*>(this)->setShapeInfo(this->shapeInfo());
// now we actually migrate data buffer
_buffer->migrate();
}
_buffer->syncToSpecial();
}
void NDArray::syncToHost() { if(!isEmpty()) _buffer->syncToPrimary(getContext()); }
void NDArray::tickWriteHost() { if(!isEmpty()) _buffer->writePrimary(); }
void NDArray::tickWriteDevice() { if(!isEmpty()) _buffer->writeSpecial(); }
void NDArray::tickReadHost() { if(!isEmpty()) _buffer->readPrimary(); }
void NDArray::tickReadDevice() { if(!isEmpty()) _buffer->readSpecial(); }
void NDArray::tickBothActual() {
_buffer->writePrimary();
_buffer->readSpecial();
}
bool NDArray::isActualOnHostSide() { return _buffer->isPrimaryActual(); }
bool NDArray::isActualOnDeviceSide() { return _buffer->isSpecialActual(); }
void NDArray::makeBothBuffersActual() {
if (!isActualOnHostSide()) syncToHost();
if (!isActualOnDeviceSide()) syncToDevice();
}
///////////////////////////////////////////////////////////////////
template <typename T>
SD_KERNEL static void fillAsTriangularCuda(const void* vx, const LongType* xShapeInfo, void* vz,
const LongType* zShapeInfo, const T val, const int lower,
const int upper, char direction, bool includeEdges) {
const auto x = reinterpret_cast<const T*>(vx);
auto z = reinterpret_cast<T*>(vz);
__shared__ LongType zRank, xRank, areSameOffsets, *sharedMem; // xRank == zRank always, except when xRank = 1, in this case zRank = 2
__shared__ LongType zLen, totalThreads; // xLen == zLen, except when xRank = 1, in this case zLen = 2*xLen
__shared__ LongType *zShape;
__shared__ LongType *zStride;
__shared__ LongType *xShape;
__shared__ LongType *xStride;
if (threadIdx.x == 0) {
extern __shared__ unsigned char shmem[];
sharedMem = reinterpret_cast<LongType*>(shmem);
areSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
xRank = shape::rank(xShapeInfo);
zRank = shape::rank(zShapeInfo);
zLen = shape::length(zShapeInfo);
totalThreads = gridDim.x * blockDim.x;
zShape = shape::shapeOf(zShapeInfo);
zStride = shape::stride(zShapeInfo);
xShape = shape::shapeOf(xShapeInfo);
xStride = shape::stride(xShapeInfo);
}
__syncthreads();
auto coords = sharedMem + threadIdx.x * zRank;
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
bool dirU = direction == 'u';
bool dirL = direction == 'l';
for (LongType i = tid; i < zLen; i += totalThreads) {
INDEX2COORDS(i, zRank, zShape, coords);
LongType zOffset;
COORDS2INDEX(zRank, zStride, coords, zOffset);
auto row = coords[zRank - 2];
auto col = coords[zRank - 1];
auto lCompare = includeEdges ? row + lower <= col : row + lower < col;
auto uCompare = includeEdges ? row + upper >= col : row + upper > col;
if (dirU && lCompare || dirL && uCompare) {
z[zOffset] = val;
} else if (vx != vz) { // when x and z are different arrays
if (xRank != zRank) coords[0] = coords[1];
LongType xOffset;
COORDS2INDEX(xRank, xStride, coords, xOffset);
z[zOffset] = x[xOffset];
}
}
}
///////////////////////////////////////////////////////////////////
template <typename T>
void NDArray::fillAsTriangular(const float val, int lower, int upper, NDArray& target, const char direction,
const bool includeEdges) {
if (isS()) THROW_EXCEPTION("NDArray::fillAsTriangular: you can't use this method on String array!");
if (!isSameShape(target) &&
!(rankOf() == 1 && target.rankOf() == 2 && sizeAt(0) == target.sizeAt(0) && sizeAt(0) == target.sizeAt(1)))
throw std::string("NDArray::fillAsTriangular method: wrong shape of target array !");
const int threadsPerBlock = SD_MAX_NUM_THREADS / 4;
int len = target.isScalar() ? 1 : target.lengthOf();
const int blocksPerGrid = (len + threadsPerBlock - 1) / threadsPerBlock;
const int sharedMem = threadsPerBlock * sizeof(int) * target.rankOf() + 128;
dim3 launchDims = getFillTriLaunchDims(target.lengthOf(), target.rankOf());
PointersManager manager(getContext(), "NDArray::fillAsTriangular");
prepareSpecialUse({&target}, {this});
fillAsTriangularCuda<T><<<launchDims.y, launchDims.x, launchDims.z, *getContext()->getCudaStream()>>>(
platformBuffer(), specialShapeInfo(), target.platformBuffer(), target.specialShapeInfo(), static_cast<T>(val),
lower, upper, direction, includeEdges);
registerSpecialUse({&target}, {this});
sd::DebugHelper::checkGlobalErrorCode("fillTriangular failed");
manager.synchronize();
}
BUILD_SINGLE_TEMPLATE( SD_LIB_EXPORT void NDArray::fillAsTriangular,
(const float val, int lower, int upper, NDArray& target, const char direction,
const bool includeEdges),
SD_COMMON_TYPES);
////////////////////////////////////////////////////////////////////////
template <typename T>
SD_KERNEL static void identityMatrixCuda(void* vx, const LongType* xShapeInfo, const T val) {
auto x = reinterpret_cast<T*>(vx);
// Shared memory variables
__shared__ LongType rank;
__shared__ LongType len;
__shared__ LongType totalThreads;
__shared__ const LongType* shapePtr;
__shared__ const LongType* stridePtr;
__shared__ LongType* sharedMem;
// Initialize shared variables in thread 0
if (threadIdx.x == 0) {
extern __shared__ unsigned char shmem[];
sharedMem = reinterpret_cast<LongType*>(shmem);
// Cache rank and length
rank = shape::rank(xShapeInfo);
len = shape::length(xShapeInfo);
// Cache pointers to shape and stride arrays
shapePtr = shape::shapeOf(xShapeInfo);
stridePtr = shape::stride(xShapeInfo);
// Calculate total number of threads
totalThreads = gridDim.x * blockDim.x;
}
__syncthreads();
// Each thread has its own coordinates array in shared memory
auto coords = sharedMem + threadIdx.x * rank;
// Calculate global thread ID
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
// Iterate over assigned elements
for (LongType i = tid; i < len; i += totalThreads) {
// Convert linear index to multi-dimensional coordinates using cached shape
INDEX2COORDS(i, rank, shapePtr, coords);
// Compute linear offset from coordinates using cached stride
LongType offset;
COORDS2INDEX(rank, stridePtr, coords, offset);
// Check if the current position is on the diagonal (row == col)
if (coords[rank - 2] == coords[rank - 1]) { // Assuming 0-based indexing
x[offset] = val;
}
else {
x[offset] = static_cast<T>(0);
}
}
}
///////////////////////////////////////////////////////////////////
template <typename T>
static void identityMatrixCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
const cudaStream_t* stream, void* vx, const LongType* xShapeInfo,
const float val) {
identityMatrixCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, static_cast<T>(val));
sd::DebugHelper::checkGlobalErrorCode("identityMatrix failed");
}
BUILD_SINGLE_TEMPLATE( void identityMatrixCudaLauncher,
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
const cudaStream_t* stream, void* vx, const sd::LongType* xShapeInfo, const float val),
SD_COMMON_TYPES);
////////////////////////////////////////////////////////////////////////
void NDArray::setIdentity() {
if (isS()) THROW_EXCEPTION("NDArray::setIdentity: you can't use this method on String array!");
int len = isScalar() ? 1 : lengthOf();
dim3 launchDims = getIdentityLaunchDims(len, rankOf());
PointersManager manager(getContext(), "NDArray::setIdentity");
syncToDevice();
BUILD_SINGLE_SELECTOR(dataType(), identityMatrixCudaLauncher,
(launchDims.y, launchDims.x,launchDims.z, getContext()->getCudaStream(), platformBuffer(),
specialShapeInfo(), 1.f),
SD_COMMON_TYPES);
tickWriteDevice();
manager.synchronize();
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void NDArray::swapUnsafe(NDArray& other) {
auto xType = this->dataType();
if (xType != other.dataType())
THROW_EXCEPTION("NDArray::swapUnsage method: both arrays must have the same data type");
if (specialBuffer() == nullptr || other.specialBuffer() == nullptr)
THROW_EXCEPTION("NDArray::swapUnsafe method: input array should not be empty!");
if (lengthOf() != other.lengthOf())
THROW_EXCEPTION("NDArray::swapUnsafe method: input arrays should have the same length!");
PointersManager manager(getContext(), "NDArray::swapUnsafe");
prepareSpecialUse({&other, this}, {&other, this});
BUILD_SINGLE_SELECTOR(xType, templatedSwapUnsafe,
(specialBuffer(), specialShapeInfo(), other.specialBuffer(), other.specialShapeInfo(),
getContext()->getCudaStream()),
SD_COMMON_TYPES);
registerSpecialUse({&other, this}, {&other, this});
manager.synchronize();
}
////////////////////////////////////////////////////////////////////////
void NDArray::synchronize(const char* msg) {
auto res = cudaStreamSynchronize(*(getContext()->getCudaStream()));
if (res != 0) {
std::string message = msg + std::string(": synchronization failed !");
THROW_EXCEPTION(message.c_str());
}
}
// NDArray implementation for .cu file
void NDArray::printBufferDebug(const char* msg, sd::LongType offset, sd::LongType limit) {
if (msg) sd_printf("%s:\n", msg);
if(limit < 0) limit = lengthOf();
// Print array info
sd_printf("NDArray: Shape=[", 0);
for (int i = 0; i < rankOf(); i++) {
sd_printf("%lld", (long long)sizeAt(i));
if (i < rankOf() - 1) sd_printf(",", 0);
}
sd_printf("], DataType=%s, Order=%c\n",
DataTypeUtils::asString(dataType()).c_str(), ordering());
#if defined(SD_GCC_FUNCTRACE)
printf("========================================================\n");
Printer p;
StackTrace st;
st.load_here();
p.print(st);
printf("========================================================\n");
fflush(stdout);
#endif
// Print buffer state
if (_buffer != nullptr) {
_buffer->printBufferDebug("Buffer contents", offset, limit);
} else {
sd_printf("Buffer is nullptr\n", 0);
}
}
////////////////////////////////////////////////////////////////////////
void NDArray::prepareSpecialUse(const std::vector<NDArray*>& writeList,
const std::vector<NDArray*>& readList, bool synchronizeWritables) {
for (const auto& a : readList)
if (a != nullptr) a->syncToDevice();
for (const auto& a : writeList) {
if (a != nullptr) {
a->getDataBuffer()->allocateSpecial();
if (synchronizeWritables) a->syncToDevice();
}
}
}
////////////////////////////////////////////////////////////////////////
void NDArray::registerSpecialUse(const std::vector<NDArray*>& writeList,
const std::vector<NDArray*>& readList) {
for (const auto& p : readList)
if (p != nullptr) p->tickReadDevice();
for (const auto& p : writeList)
if (p != nullptr) p->tickWriteDevice();
}
////////////////////////////////////////////////////////////////////////
void NDArray::preparePrimaryUse(const std::vector<NDArray*>& writeList,
const std::vector<NDArray*>& readList, bool synchronizeWritables) {
for (const auto& a : readList)
if (a != nullptr) a->syncToHost();
for (const auto& a : writeList) {
if (a != nullptr) {
a->getDataBuffer()->allocatePrimary();
if (synchronizeWritables) a->syncToHost();
}
}
}
////////////////////////////////////////////////////////////////////////
void NDArray::registerPrimaryUse(const std::vector<NDArray*>& writeList,
const std::vector<NDArray*>& readList) {
for (const auto& p : readList)
if (p != nullptr) p->tickReadHost();
for (const auto& p : writeList)
if (p != nullptr) p->tickWriteHost();
}
//////////////////////////////////////////////////////////////////////////
void NDArray::syncShape() {
cudaMemcpy(const_cast<LongType*>(specialShapeInfo()), shapeInfo(), shape::shapeInfoByteLength(shapeInfo()),
cudaMemcpyHostToDevice);
}
//////////////////////////////////////////////////////////////////////////
//////////////////////////////////////////////////////////////////////////
// change an array by repeating it the number of times given by reps.
NDArray NDArray::tile(const std::vector<LongType>& reps) {
int dim = reps.size();
LongType product = 1;
for (const auto& item : reps) product *= item;
if (product < 1) THROW_EXCEPTION("NDArray::tile method: one of the elements in reps array is zero !");
int rankOld = rankOf();
int diff = rankOld - dim;
if (product == 1) { // in this case 2 possibilities are present: just reshape or nothing to do
NDArray result(*this);
if (diff < 0) { // reshape to higher dimension
std::vector<LongType> shapeNew = reps; // need to have unities at first "diff" positions of new shape
memcpy(&shapeNew[-diff], result.shapeInfo() + 1,
rankOld * sizeof(LongType)); // put old shape numbers at rest of positions
result.reshapei(ordering(), shapeNew);
}
return result; // nothing to do, if diff >= 0 -> identity tile
}
// evaluate shapeInfo for resulting array
auto newShapeInfo = ShapeUtils::evalTileShapeInfo(*this, reps, getContext()->getWorkspace());
// create new buffer, in any case the memory amount new buffer points to is bigger then those for old _buffer
DataBuffer * newBuff = new DataBuffer(shape::length(newShapeInfo) * sizeOfT(),
dataType(), getContext()->getWorkspace(), true);
// assign new shape and new buffer to resulting array
NDArray result(newBuff,const_cast<sd::LongType *>(newShapeInfo) , getContext());
// fill newBuff, loop through all elements of newBuff
// looping through buffer() goes automatically by means of getSubArrayIndex applying
const auto resultLen = result.lengthOf();
auto xType = this->dataType();
auto stream = getContext()->getCudaStream();
prepareSpecialUse({&result}, {this});
BUILD_SINGLE_SELECTOR(xType, tileKernelH,
(this->specialBuffer(), this->specialShapeInfo(), result.specialBuffer(),
result.specialShapeInfo(), resultLen, stream),
SD_COMMON_TYPES);
registerSpecialUse({&result}, {this});
return result;
}
//////////////////////////////////////////////////////////////////////////
// change an array by repeating it the number of times given by reps.
void NDArray::tile(const std::vector<LongType>& reps, NDArray& target) {
auto repProd = shape::prodLong(reps.data(), reps.size());
if (repProd < 1) THROW_EXCEPTION("NDArray::tile: reps can't contain 0s");
// evaluate true tile shapeInfo for comparison with target shapeInfo
auto newShapeInfo = ShapeUtils::evalTileShapeInfo(*this, reps, getContext()->getWorkspace());
if (!shape::equalsSoft(newShapeInfo, target.shapeInfo())) {
THROW_EXCEPTION("NDArray::tile method - shapeInfo of target array is not suitable for tile operation !");
}
// fill newBuff, loop through all elements of newBuff
// looping through buffer() goes automatically by means of getSubArrayIndex applying
const int ews = target.ews();
const int targetLen = target.lengthOf();
auto stream = getContext()->getCudaStream();
prepareSpecialUse({&target}, {this});
BUILD_SINGLE_SELECTOR_TWICE(
target.dataType(), tileKernelHH,
(specialBuffer(), specialShapeInfo(), target.specialBuffer(), target.specialShapeInfo(), targetLen, stream),
SD_COMMON_TYPES);
registerSpecialUse({&target}, {this});
}
//////////////////////////////////////////////////////////////////////////
void NDArray::tile(NDArray& target) {
if (rankOf() > target.rankOf())
THROW_EXCEPTION(
"NDArray::tile method - rank of target array must be bigger or equal to the rank of this array !");
if (!ShapeUtils::areShapesBroadcastable(*this, target))
THROW_EXCEPTION("NDArray::tile method - shapeInfo of target array is not suitable for tile operation !");
// fill newBuff, loop through all elements of newBuff
// looping through getBuffer() goes automatically by means of getSubArrayIndex applying
const auto ews = target.ews();
const auto targetLen = target.lengthOf();
auto stream = getContext()->getCudaStream();
prepareSpecialUse({&target}, {this});
BUILD_SINGLE_SELECTOR_TWICE(
target.dataType(), tileKernelHH,
(specialBuffer(), specialShapeInfo(), target.specialBuffer(), target.specialShapeInfo(), targetLen, stream),
SD_COMMON_TYPES);
registerSpecialUse({&target}, {this});
}
////////////////////////////////////////////////////////////////////////
template <typename X, typename Z>
SD_KERNEL static void repeatCuda(const void* vx, const LongType* xShapeInfo, void* vz,
const LongType* zShapeInfo, const LongType* repeats, const LongType repSize,
const int axis) {
const X* x = reinterpret_cast<const X*>(vx);
Z* z = reinterpret_cast<Z*>(vz);
__shared__ LongType rank, *sharedMem;
__shared__ LongType zLen, totalThreads; // xLen = zLen
if (threadIdx.x == 0) {
extern __shared__ unsigned char shmem[];
sharedMem = reinterpret_cast<LongType*>(shmem);
rank = shape::rank(zShapeInfo); // xRank = zRank
zLen = shape::length(zShapeInfo); // xLen <= zLen
totalThreads = gridDim.x * blockDim.x;
}
__syncthreads();
auto coords = sharedMem + threadIdx.x * rank;
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
for (LongType i = tid; i < zLen; i += totalThreads) {
INDEX2COORDS(i, rank, shape::shapeOf(zShapeInfo), coords);
LongType zOffset;
COORDS2INDEX(rank, shape::stride(zShapeInfo), coords, zOffset);
if (repSize > 1) {
for (LongType j = 0; j < repSize; ++j) {
coords[axis] -= repeats[j];
if (coords[axis] < 0) {
coords[axis] = j;
break;
}
}
} else
coords[axis] /= repeats[0];
LongType xOffset;
COORDS2INDEX(rank, shape::stride(xShapeInfo), coords, xOffset);
z[zOffset] = x[xOffset];
}
}
//////////////////////////////////////////////////////////////////////////
template <typename X, typename Z>
static void repeatCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo, void* vz,
const LongType* zShapeInfo, const LongType* repeats, const LongType repSize, const LongType axis) {
repeatCuda<X, Z>
<<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, repeats, repSize, axis);
DebugHelper::checkGlobalErrorCode("NDArray repeat cuda failed(...) failed");
}
BUILD_DOUBLE_TEMPLATE( void repeatCudaLauncher,
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
const cudaStream_t* stream, const void* vx, const sd::LongType* xShapeInfo, void* vz,
const sd::LongType* zShapeInfo, const sd::LongType* repeats, const sd::LongType repSize, const sd::LongType axis),
SD_COMMON_TYPES, SD_COMMON_TYPES);
//////////////////////////////////////////////////////////////////////////
// create new array by repeating it the number of times given by repeats
NDArray NDArray::repeat(const int axis, const std::vector<LongType>& repeats) {
auto nonConst = const_cast<NDArray *>(this);
std::vector<sd::LongType> shape = ShapeUtils::evalRepeatShape(axis, repeats, *nonConst);
NDArray output('c',shape, dataType(), getContext());
dim3 launchDims = getRepeatLaunchDims(output.lengthOf(), output.rankOf());
PointersManager manager(getContext(), "NDArray::repeat(const int axis, const std::vector<int>& repeats)");
const LongType* reps = reinterpret_cast<LongType*>(manager.replicatePointer(repeats.data(), repeats.size() * sizeof(LongType)));
prepareSpecialUse({&output}, {this});
BUILD_SINGLE_SELECTOR_TWICE(
dataType(), repeatCudaLauncher,
(launchDims.y, launchDims.x, launchDims.z, getContext()->getCudaStream(), specialBuffer(), specialShapeInfo(),
output.specialBuffer(), output.specialShapeInfo(), reps, repeats.size(), axis),
SD_COMMON_TYPES);
prepareSpecialUse({&output}, {this});
manager.synchronize();
return output;
}
//////////////////////////////////////////////////////////////////////////
// fill array by repeating it the number of times given by repeats
void NDArray::repeat(const int axis, const std::vector<LongType>& repeats, NDArray& target) {
auto nonConst = const_cast<NDArray *>(this);
std::vector<sd::LongType> shape = ShapeUtils::evalRepeatShape(axis, repeats, *nonConst);
if (!target.isSameShape(shape))
THROW_EXCEPTION(
"NDArray::repeat(const int axis, const std::vector<int>& repeats, NDArray& target) method: wrong shape of "
"target array!");
dim3 launchDims = getRepeatLaunchDims(target.lengthOf(), target.rankOf());
PointersManager manager(getContext(), "NDArray::repeat(const int axis, const std::vector<int>& repeats)");
const LongType* reps = reinterpret_cast<LongType*>(manager.replicatePointer(repeats.data(), repeats.size() * sizeof(LongType)));
auto targetDataType = target.dataType();
auto selfDType = dataType();
prepareSpecialUse({&target}, {this});
BUILD_DOUBLE_SELECTOR(
dataType(), target.dataType(), repeatCudaLauncher,
(launchDims.y, launchDims.x, launchDims.z, getContext()->getCudaStream(), specialBuffer(), specialShapeInfo(),
target.specialBuffer(), target.specialShapeInfo(), reps, repeats.size(), axis),
SD_COMMON_TYPES, SD_COMMON_TYPES);
prepareSpecialUse({&target}, {this});
manager.synchronize();
}
////////////////////////////////////////////////////////////////////////
void* NDArray::specialBuffer() {
if (_buffer == nullptr) {
THROW_EXCEPTION("NDArray::specialBuffer(): _buffer is nullptr - array not properly initialized");
}
void* specialBuf = _buffer->special();
if (specialBuf == nullptr) {
syncToDevice();
tickReadHost();
specialBuf = _buffer->special();
if (specialBuf == nullptr) {
THROW_EXCEPTION("NDArray::specialBuffer(): _buffer->special() returned nullptr even after syncToDevice - buffer not allocated");
}
}
// FIXME: this should be fixed once CUDA backend added
return static_cast<int8_t*>(specialBuf) + (offset() * sizeOfT());
}
//////////////////////////////////////////////////////////////////////////
template <typename T>
void NDArray::printCurrentBuffer(const bool host, const char* msg, const int precision) {
if (!isScalar() && _length == 0) {
printf("NDArray::printActualBuffer: array length is zero !\n");
return;
}
if(isScalar()) {
if(host) {
if (msg) printf("%s", msg);
if (buffer() == nullptr ) {
printf("NDArray::printActualBuffer: host buffer is nullptr !\n");
return;
}
const T* buff = bufferAsT<T>();
if (msg) printf("%s", msg);
printf("%.*f\n", precision, (double)buff[getOffset(0)]);
return;
} else {
if (msg) printf("%s", msg);
if (specialBuffer() == nullptr) {
printf("NDArray::printSpecialBuffer: special buffer is nullptr !\n");
return;
}
const auto sizeOfBuffer = sizeOfT();
void* pHost = operator new(sizeOfBuffer);
cudaMemcpyAsync(pHost, specialBuffer(), sizeOfBuffer, cudaMemcpyDeviceToHost, *getContext()->getCudaStream());
cudaDeviceSynchronize();
cudaError_t cudaResult = cudaStreamSynchronize(*getContext()->getCudaStream());
auto cast = reinterpret_cast<T*>(pHost);
if (cudaResult != 0) THROW_EXCEPTION("NDArray::printSpecialBuffer: cudaStreamSynchronize failed!");
printf("%.*f\n", precision, (double)cast[0]);
return;
}
}
if (msg) printf("%s", msg);
if (host) {
if (buffer() == nullptr || _length == 0) {
printf("NDArray::printActualBuffer: host buffer is nullptr !\n");
return;
}
const T* buff = bufferAsT<T>();
for (LongType i = 0; i < _length; i++) printf("%.*f, ", precision, (double)buff[getOffset(i)]);
printf("\n");
} else {
if (specialBuffer() == nullptr) {
printf("NDArray::printSpecialBuffer: special buffer is nullptr !\n");
return;
}
const auto sizeOfBuffer = sizeOfT() * (getOffset(_length - 1) + 1);
void* pHost = operator new(sizeOfBuffer);
cudaMemcpyAsync(pHost, specialBuffer(), sizeOfBuffer, cudaMemcpyDeviceToHost, *getContext()->getCudaStream());
cudaError_t cudaResult = cudaStreamSynchronize(*getContext()->getCudaStream());
if (cudaResult != 0) THROW_EXCEPTION("NDArray::printSpecialBuffer: cudaStreamSynchronize failed!");
for (LongType i = 0; i < _length; i++)
printf("%.*f, ", precision, (double)reinterpret_cast<T*>(pHost)[getOffset(i)]);
printf("\n");
operator delete(pHost);
}
}
#define PRINT_BUFFER(T) template void NDArray::printCurrentBuffer<GET_SECOND(T)>(const bool host, const char* msg, const int precision);
ITERATE_LIST((SD_COMMON_TYPES),PRINT_BUFFER)
} // end namespace sd
#endif
libnd4j/include/array/cuda/NDArrayLambda.cu
+624
View File
@@ -0,0 +1,624 @@
/*
* ******************************************************************************
* *
* *
* * This program and the accompanying materials are made available under the
* * terms of the Apache License, Version 2.0 which is available at
* * https://www.apache.org/licenses/LICENSE-2.0.
* *
* * See the NOTICE file distributed with this work for additional
* * information regarding copyright ownership.
* * 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.
* *
* * SPDX-License-Identifier: Apache-2.0
* *****************************************************************************
*/
//
// @author raver119@gmail.com
//
#include <array/DataType.h>
#include <array/DataTypeUtils.h>
#include <array/NDArray.h>
#include <exceptions/cuda_exception.h>
#include <execution/ThreadPool.h>
#include <helpers/DebugHelper.h>
#include <loops/legacy_ops.h>
#include <system/Environment.h>
#include <system/op_boilerplate.h>
#include <types/types.h>
#include "helpers/ShapeUtils.h"
namespace sd {
// ----------- Unary Lambda Operations ----------------
template <typename T>
SD_KERNEL void applyLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
void* vz, const sd::LongType* zShapeInfo,
void* vextraParams) {
// Cast input and output pointers
auto x = reinterpret_cast<const T*>(vx);
auto z = reinterpret_cast<T*>(vz);
auto extraParams = reinterpret_cast<void*>(vextraParams);
// Cache shape information for x buffer
__shared__ sd::LongType length;
__shared__ sd::LongType xRank;
__shared__ const sd::LongType* xShapePtr;
__shared__ const sd::LongType* xStridePtr;
// Cache shape information for z buffer
__shared__ sd::LongType zRank;
__shared__ const sd::LongType* zShapePtr;
__shared__ const sd::LongType* zStridePtr;
if (threadIdx.x == 0) {
length = shape::length(xShapeInfo);
// Cache x shape information
xRank = shape::rank(xShapeInfo);
xShapePtr = shape::shapeOf(xShapeInfo);
xStridePtr = shape::stride(xShapeInfo);
// Cache z shape information
zRank = shape::rank(zShapeInfo);
zShapePtr = shape::shapeOf(zShapeInfo);
zStridePtr = shape::stride(zShapeInfo);
}
__syncthreads();
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
int totalThreads = gridDim.x * blockDim.x;
for (sd::LongType i = tid; i < length; i += totalThreads) {
sd::LongType xCoords[SD_MAX_RANK];
sd::LongType zCoords[SD_MAX_RANK];
sd::LongType xOffset;
sd::LongType zOffset;
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
// Apply the function using extraParams (this will be handled in the wrapper function)
// For now, using a placeholder
z[zOffset] = x[xOffset]; // This will be replaced with the actual lambda function call
}
}
// ----------- Indexed Lambda Operations ----------------
template <typename T>
SD_KERNEL void applyIndexedLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
void* vz, const sd::LongType* zShapeInfo,
void* vextraParams) {
// Cast input and output pointers
auto x = reinterpret_cast<const T*>(vx);
auto z = reinterpret_cast<T*>(vz);
auto extraParams = reinterpret_cast<void*>(vextraParams);
// Cache shape information for x buffer
__shared__ sd::LongType length;
__shared__ sd::LongType xRank;
__shared__ const sd::LongType* xShapePtr;
__shared__ const sd::LongType* xStridePtr;
// Cache shape information for z buffer
__shared__ sd::LongType zRank;
__shared__ const sd::LongType* zShapePtr;
__shared__ const sd::LongType* zStridePtr;
if (threadIdx.x == 0) {
length = shape::length(xShapeInfo);
// Cache x shape information
xRank = shape::rank(xShapeInfo);
xShapePtr = shape::shapeOf(xShapeInfo);
xStridePtr = shape::stride(xShapeInfo);
// Cache z shape information
zRank = shape::rank(zShapeInfo);
zShapePtr = shape::shapeOf(zShapeInfo);
zStridePtr = shape::stride(zShapeInfo);
}
__syncthreads();
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
int totalThreads = gridDim.x * blockDim.x;
for (sd::LongType i = tid; i < length; i += totalThreads) {
sd::LongType xCoords[SD_MAX_RANK];
sd::LongType zCoords[SD_MAX_RANK];
sd::LongType xOffset;
sd::LongType zOffset;
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
// Apply the indexed function - placeholder for actual lambda call
z[zOffset] = x[xOffset]; // This will be replaced with the actual indexed lambda function call
}
}
// ----------- Pairwise Lambda Operations ----------------
template <typename T>
SD_KERNEL void applyPairwiseLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
const void* vy, const sd::LongType* yShapeInfo,
void* vz, const sd::LongType* zShapeInfo,
void* vextraParams, bool isScalar) {
// Cast input and output pointers
auto x = reinterpret_cast<const T*>(vx);
auto y = reinterpret_cast<const T*>(vy);
auto z = reinterpret_cast<T*>(vz);
auto extraParams = reinterpret_cast<void*>(vextraParams);
// Cache shape information
__shared__ sd::LongType length;
__shared__ sd::LongType xRank;
__shared__ sd::LongType yRank;
__shared__ sd::LongType zRank;
__shared__ const sd::LongType* xShapePtr;
__shared__ const sd::LongType* yShapePtr;
__shared__ const sd::LongType* zShapePtr;
__shared__ const sd::LongType* xStridePtr;
__shared__ const sd::LongType* yStridePtr;
__shared__ const sd::LongType* zStridePtr;
__shared__ T scalarValue;
__shared__ sd::LongType yOffset0;
if (threadIdx.x == 0) {
length = shape::length(xShapeInfo);
// Cache shape information
xRank = shape::rank(xShapeInfo);
yRank = shape::rank(yShapeInfo);
zRank = shape::rank(zShapeInfo);
xShapePtr = shape::shapeOf(xShapeInfo);
yShapePtr = shape::shapeOf(yShapeInfo);
zShapePtr = shape::shapeOf(zShapeInfo);
xStridePtr = shape::stride(xShapeInfo);
yStridePtr = shape::stride(yShapeInfo);
zStridePtr = shape::stride(zShapeInfo);
if (isScalar) {
sd::LongType yCoords[SD_MAX_RANK];
for (int i = 0; i < yRank; i++) {
yCoords[i] = 0;
}
COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset0);
scalarValue = y[yOffset0];
}
}
__syncthreads();
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
int totalThreads = gridDim.x * blockDim.x;
for (sd::LongType i = tid; i < length; i += totalThreads) {
sd::LongType xCoords[SD_MAX_RANK];
sd::LongType yCoords[SD_MAX_RANK];
sd::LongType zCoords[SD_MAX_RANK];
sd::LongType xOffset;
sd::LongType yOffset;
sd::LongType zOffset;
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
if (isScalar) {
// Apply the pairwise function with scalar - placeholder
z[zOffset] = x[xOffset]; // Will be replaced with actual function call using scalarValue
} else {
INDEX2COORDS(i, yRank, yShapePtr, yCoords);
COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
// Apply the pairwise function - placeholder
z[zOffset] = x[xOffset]; // Will be replaced with actual function call using y[yOffset]
}
}
}
// ----------- Indexed Pairwise Lambda Operations ----------------
template <typename T>
SD_KERNEL void applyIndexedPairwiseLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
const void* vy, const sd::LongType* yShapeInfo,
void* vz, const sd::LongType* zShapeInfo,
void* vextraParams) {
// Cast input and output pointers
auto x = reinterpret_cast<const T*>(vx);
auto y = reinterpret_cast<const T*>(vy);
auto z = reinterpret_cast<T*>(vz);
auto extraParams = reinterpret_cast<void*>(vextraParams);
// Cache shape information
__shared__ sd::LongType length;
__shared__ sd::LongType xRank;
__shared__ sd::LongType yRank;
__shared__ sd::LongType zRank;
__shared__ const sd::LongType* xShapePtr;
__shared__ const sd::LongType* yShapePtr;
__shared__ const sd::LongType* zShapePtr;
__shared__ const sd::LongType* xStridePtr;
__shared__ const sd::LongType* yStridePtr;
__shared__ const sd::LongType* zStridePtr;
if (threadIdx.x == 0) {
length = shape::length(xShapeInfo);
// Cache shape information
xRank = shape::rank(xShapeInfo);
yRank = shape::rank(yShapeInfo);
zRank = shape::rank(zShapeInfo);
xShapePtr = shape::shapeOf(xShapeInfo);
yShapePtr = shape::shapeOf(yShapeInfo);
zShapePtr = shape::shapeOf(zShapeInfo);
xStridePtr = shape::stride(xShapeInfo);
yStridePtr = shape::stride(yShapeInfo);
zStridePtr = shape::stride(zShapeInfo);
}
__syncthreads();
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
int totalThreads = gridDim.x * blockDim.x;
for (sd::LongType i = tid; i < length; i += totalThreads) {
sd::LongType xCoords[SD_MAX_RANK];
sd::LongType yCoords[SD_MAX_RANK];
sd::LongType zCoords[SD_MAX_RANK];
sd::LongType xOffset;
sd::LongType yOffset;
sd::LongType zOffset;
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
INDEX2COORDS(i, yRank, yShapePtr, yCoords);
COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
// Apply the indexed pairwise function - placeholder
z[zOffset] = x[xOffset]; // Will be replaced with actual function call
}
}
// ----------- Triplewise Lambda Operations ----------------
template <typename T>
SD_KERNEL void applyTriplewiseLambdaKernel(const void* vx, const sd::LongType* xShapeInfo,
const void* vy, const sd::LongType* yShapeInfo,
const void* vt, const sd::LongType* tShapeInfo,
void* vz, const sd::LongType* zShapeInfo,
void* vextraParams) {
// Cast input and output pointers
auto x = reinterpret_cast<const T*>(vx);
auto y = reinterpret_cast<const T*>(vy);
auto t = reinterpret_cast<const T*>(vt);
auto z = reinterpret_cast<T*>(vz);
auto extraParams = reinterpret_cast<void*>(vextraParams);
// Cache shape information
__shared__ sd::LongType length;
__shared__ sd::LongType xRank;
__shared__ sd::LongType yRank;
__shared__ sd::LongType tRank;
__shared__ sd::LongType zRank;
__shared__ const sd::LongType* xShapePtr;
__shared__ const sd::LongType* yShapePtr;
__shared__ const sd::LongType* tShapePtr;
__shared__ const sd::LongType* zShapePtr;
__shared__ const sd::LongType* xStridePtr;
__shared__ const sd::LongType* yStridePtr;
__shared__ const sd::LongType* tStridePtr;
__shared__ const sd::LongType* zStridePtr;
if (threadIdx.x == 0) {
length = shape::length(xShapeInfo);
// Cache shape information
xRank = shape::rank(xShapeInfo);
yRank = shape::rank(yShapeInfo);
tRank = shape::rank(tShapeInfo);
zRank = shape::rank(zShapeInfo);
xShapePtr = shape::shapeOf(xShapeInfo);
yShapePtr = shape::shapeOf(yShapeInfo);
tShapePtr = shape::shapeOf(tShapeInfo);
zShapePtr = shape::shapeOf(zShapeInfo);
xStridePtr = shape::stride(xShapeInfo);
yStridePtr = shape::stride(yShapeInfo);
tStridePtr = shape::stride(tShapeInfo);
zStridePtr = shape::stride(zShapeInfo);
}
__syncthreads();
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
int totalThreads = gridDim.x * blockDim.x;
for (sd::LongType i = tid; i < length; i += totalThreads) {
sd::LongType xCoords[SD_MAX_RANK];
sd::LongType yCoords[SD_MAX_RANK];
sd::LongType tCoords[SD_MAX_RANK];
sd::LongType zCoords[SD_MAX_RANK];
sd::LongType xOffset;
sd::LongType yOffset;
sd::LongType tOffset;
sd::LongType zOffset;
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
INDEX2COORDS(i, yRank, yShapePtr, yCoords);
COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
INDEX2COORDS(i, tRank, tShapePtr, tCoords);
COORDS2INDEX(tRank, tStridePtr, tCoords, tOffset);
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
// Apply the triplewise function - placeholder
z[zOffset] = x[xOffset]; // Will be replaced with actual function call
}
}
// ---------------------- Wrapper functions -----------------------
// Helper class for CUDA Lambda operations
template <typename T>
class NDArrayLambdaCuda {
public:
static int constexpr LAMBDA_THREADS = 256;
static int constexpr LAMBDA_BLOCKS = 512;
// Unary lambda wrapper
static void executeLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
void* z, const sd::LongType* zShapeInfo, void* extraParams) {
if(stream == nullptr) {
THROW_EXCEPTION("executeLambda: Stream must not be nullptr!");
}
dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
applyLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
x, xShapeInfo, z, zShapeInfo, extraParams);
sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executeLambda failed");
}
// Indexed lambda wrapper
static void executeIndexedLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
void* z, const sd::LongType* zShapeInfo, void* extraParams) {
if(stream == nullptr) {
THROW_EXCEPTION("executeIndexedLambda: Stream must not be nullptr!");
}
dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
applyIndexedLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
x, xShapeInfo, z, zShapeInfo, extraParams);
sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executeIndexedLambda failed");
}
// Pairwise lambda wrapper
static void executePairwiseLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
const void* y, const sd::LongType* yShapeInfo,
void* z, const sd::LongType* zShapeInfo, void* extraParams, bool isScalar) {
dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
if(stream == nullptr) {
THROW_EXCEPTION("executePairwiseLambda: Stream must not be nullptr!");
}
applyPairwiseLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, isScalar);
sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executePairwiseLambda failed");
}
// Indexed pairwise lambda wrapper
static void executeIndexedPairwiseLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
const void* y, const sd::LongType* yShapeInfo,
void* z, const sd::LongType* zShapeInfo, void* extraParams) {
dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
applyIndexedPairwiseLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executeIndexedPairwiseLambda failed");
}
// Triplewise lambda wrapper
static void executeTriplewiseLambda(cudaStream_t* stream, const void* x, const sd::LongType* xShapeInfo,
const void* y, const sd::LongType* yShapeInfo,
const void* t, const sd::LongType* tShapeInfo,
void* z, const sd::LongType* zShapeInfo, void* extraParams) {
if(stream == nullptr) {
THROW_EXCEPTION("executeTriplewiseLambda: Stream must not be nullptr!");
}
dim3 launchDims(LAMBDA_BLOCKS, LAMBDA_THREADS, 1024);
applyTriplewiseLambdaKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
x, xShapeInfo, y, yShapeInfo, t, tShapeInfo, z, zShapeInfo, extraParams);
sd::DebugHelper::checkErrorCode(stream, "NDArrayLambdaCuda::executeTriplewiseLambda failed");
}
};
// Implementation of the NDArray Lambda methods for CUDA
template <typename T>
SD_LIB_EXPORT void NDArray::applyLambda(std::function<T(T)>& func, NDArray* target) {
// Validate types
if (dataType() != DataTypeUtils::fromT<T>())
THROW_EXCEPTION(
"NDArray::applyLambdaCuda<T> method: wrong template parameter T, its type should be the same as type of this "
"array!");
if (dataType() != target->dataType())
THROW_EXCEPTION("NDArray::applyLambdaCuda<T> method: types of this and target array should match!");
// Get device pointers and stream
auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
auto x = this->specialBuffer();
auto z = target->specialBuffer();
auto xShapeInfo = this->specialShapeInfo();
auto zShapeInfo = target->specialShapeInfo();
// Create and set up extraParams
void* extraParams = nullptr; // This would hold the function pointer for the lambda
// Execute the CUDA kernel
NDArrayLambdaCuda<T>::executeLambda(stream, x, xShapeInfo, z, zShapeInfo, extraParams);
}
template <typename T>
SD_LIB_EXPORT void NDArray::applyIndexedLambda(std::function<T(sd::LongType, T)>& func, NDArray* target) {
// Validate types
if (dataType() != DataTypeUtils::fromT<T>())
THROW_EXCEPTION(
"NDArray::applyIndexedLambdaCuda<T> method: wrong template parameter T, its type should be the same as type of "
"this array!");
if (dataType() != target->dataType())
THROW_EXCEPTION("NDArray::applyIndexedLambdaCuda<T> method: types of this and target array should match!");
// Get device pointers and stream
auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
auto x = this->specialBuffer();
auto z = target->specialBuffer();
auto xShapeInfo = this->specialShapeInfo();
auto zShapeInfo = target->specialShapeInfo();
// Create and set up extraParams
void* extraParams = nullptr; // This would hold the function pointer for the lambda
// Execute the CUDA kernel
NDArrayLambdaCuda<T>::executeIndexedLambda(stream, x, xShapeInfo, z, zShapeInfo, extraParams);
}
template <typename T>
SD_LIB_EXPORT void NDArray::applyPairwiseLambda(NDArray* other, std::function<T(T, T)>& func,
NDArray* target) {
// Validate types
if (dataType() != DataTypeUtils::fromT<T>())
THROW_EXCEPTION(
"NDArray::applyPairwiseLambdaCuda<T> method: wrong template parameter T, its type should be the same as type of "
"this array!");
if (dataType() != other->dataType() || dataType() != target->dataType())
THROW_EXCEPTION(
"NDArray::applyPairwiseLambdaCuda<T> method: all three arrays (this, other, target) must have the same type!");
// Check for scalar or same length
bool isScalar = other->isScalar();
if (this->lengthOf() != other->lengthOf() && !this->isScalar() && !isScalar) {
THROW_EXCEPTION("applyPairwiseLambdaCuda requires both operands to have the same shape or one to be a scalar");
}
// Get device pointers and stream
auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
auto x = this->specialBuffer();
auto y = other->specialBuffer();
auto z = target->specialBuffer();
auto xShapeInfo = this->specialShapeInfo();
auto yShapeInfo = other->specialShapeInfo();
auto zShapeInfo = target->specialShapeInfo();
// Create and set up extraParams
void* extraParams = nullptr; // This would hold the function pointer for the lambda
// Execute the CUDA kernel
NDArrayLambdaCuda<T>::executePairwiseLambda(stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams, isScalar);
}
template <typename T>
SD_LIB_EXPORT void NDArray::applyIndexedPairwiseLambda(NDArray* other, std::function<T(sd::LongType, T, T)>& func,
NDArray* target) {
// Validate types
if (dataType() != DataTypeUtils::fromT<T>())
THROW_EXCEPTION(
"NDArray::applyIndexedPairwiseLambdaCuda<T> method: wrong template parameter T, its type should be the same as "
"type of this array!");
if (dataType() != target->dataType())
THROW_EXCEPTION(
"NDArray::applyIndexedPairwiseLambdaCuda<T> method: types of this and target array should match!");
if (this->lengthOf() != other->lengthOf()) {
THROW_EXCEPTION("applyIndexedPairwiseLambdaCuda requires both operands to have the same shape");
}
// Get device pointers and stream
auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
auto x = this->specialBuffer();
auto y = other->specialBuffer();
auto z = target->specialBuffer();
auto xShapeInfo = this->specialShapeInfo();
auto yShapeInfo = other->specialShapeInfo();
auto zShapeInfo = target->specialShapeInfo();
// Create and set up extraParams
void* extraParams = nullptr; // This would hold the function pointer for the lambda
// Execute the CUDA kernel
NDArrayLambdaCuda<T>::executeIndexedPairwiseLambda(stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams);
}
template <typename T>
SD_LIB_EXPORT void NDArray::applyTriplewiseLambda(NDArray* second, NDArray* third,
std::function<T(T, T, T)>& func, NDArray* target) {
// Validate types
if (dataType() != DataTypeUtils::fromT<T>())
THROW_EXCEPTION(
"NDArray::applyTriplewiseLambdaCuda<T> method: wrong template parameter T, its type should be the same as type of "
"this array!");
if (dataType() != second->dataType() || dataType() != third->dataType() || dataType() != target->dataType())
THROW_EXCEPTION(
"NDArray::applyTriplewiseLambdaCuda<T> method: all four arrays (this, second, third, target) should have the "
"same type!");
if (this->lengthOf() != second->lengthOf() || this->lengthOf() != third->lengthOf() || !this->isSameShape(second) ||
!this->isSameShape(third)) {
std::string errorMessage;
errorMessage += "applyTriplewiseLambdaCuda requires all operands to have the same shape\n";
errorMessage += "this shape: " + ShapeUtils::shapeAsString(this->shapeInfo()) + "\n";
errorMessage += "second shape: " + ShapeUtils::shapeAsString(second->shapeInfo()) + "\n";
errorMessage += "third shape: " + ShapeUtils::shapeAsString(third->shapeInfo()) + "\n";
errorMessage += "target shape: " + ShapeUtils::shapeAsString(target->shapeInfo()) + "\n";
THROW_EXCEPTION(errorMessage.c_str());
}
// Get device pointers and stream
auto stream = LaunchContext::defaultContext()->getCudaStream(); // Get the CUDA stream
auto x = this->specialBuffer();
auto y = second->specialBuffer();
auto t = third->specialBuffer();
auto z = target->specialBuffer();
auto xShapeInfo = this->specialShapeInfo();
auto yShapeInfo = second->specialShapeInfo();
auto tShapeInfo = third->specialShapeInfo();
auto zShapeInfo = target->specialShapeInfo();
// Create and set up extraParams
void* extraParams = nullptr; // This would hold the function pointer for the lambda
// Execute the CUDA kernel
NDArrayLambdaCuda<T>::executeTriplewiseLambda(stream, x, xShapeInfo, y, yShapeInfo, t, tShapeInfo,
z, zShapeInfo, extraParams);
}
#define INSTANTIATE_LAMBDA_METHODS(T) template SD_LIB_EXPORT void NDArray::applyLambda( std::function<GET_SECOND(T)(GET_SECOND(T))>& func, NDArray* target);
ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS);
#define INSTANTIATE_LAMBDA_METHODS_INDEXED(T) template SD_LIB_EXPORT void NDArray::applyIndexedLambda( std::function<GET_SECOND(T)(sd::LongType, GET_SECOND(T))>& func, NDArray* target);
ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS_INDEXED);
#define INSTANTIATE_LAMBDA_METHODS_PAIRWISE(T) template SD_LIB_EXPORT void NDArray::applyPairwiseLambda(NDArray* other, std::function<GET_SECOND(T)(GET_SECOND(T), GET_SECOND(T))>& func, NDArray* target);
ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS_PAIRWISE);
#define INSTANTIATE_LAMBDA_METHODS_INDEX_PAIR(T) template SD_LIB_EXPORT void NDArray::applyIndexedPairwiseLambda(NDArray* other, std::function<GET_SECOND(T)(sd::LongType, GET_SECOND(T), GET_SECOND(T))>& func, NDArray* target);
ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS_INDEX_PAIR);
#define INSTANTIATE_LAMBDA_METHODS_TRIPLE(T) template void NDArray::applyTriplewiseLambda(NDArray* second, NDArray* third, std::function<GET_SECOND(T)(GET_SECOND(T), GET_SECOND(T), GET_SECOND(T))>& func, NDArray* target);
ITERATE_LIST((SD_COMMON_TYPES),INSTANTIATE_LAMBDA_METHODS_TRIPLE);
} // namespace sd