chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,195 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
__global__ static void splitCuda(const void* vx, const LongType* xShapeInfo, void* pVz,
|
||||
const LongType* zTadShapeInfo, const LongType axis) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
|
||||
// Shared memory for caching shape information and related variables
|
||||
extern __shared__ unsigned char shmem[];
|
||||
LongType* sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
// Shared variables
|
||||
__shared__ LongType shared_xLen;
|
||||
__shared__ LongType shared_totalThreads;
|
||||
__shared__ int shared_xRank;
|
||||
__shared__ LongType shared_zDim;
|
||||
|
||||
// Cached shape and stride pointers
|
||||
__shared__ const LongType* shared_xShape;
|
||||
__shared__ const LongType* shared_xStride;
|
||||
__shared__ const LongType* shared_zTadShape;
|
||||
__shared__ const LongType* shared_zTadStride;
|
||||
__shared__ int shared_zTadRank;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache shape and stride information for xShapeInfo
|
||||
shared_xRank = shape::rank(xShapeInfo);
|
||||
shared_xShape = shape::shapeOf(xShapeInfo);
|
||||
shared_xStride = shape::stride(xShapeInfo);
|
||||
|
||||
// Cache shape and stride information for zTadShapeInfo
|
||||
shared_zTadRank = shape::rank(zTadShapeInfo);
|
||||
shared_zTadShape = shape::shapeOf(zTadShapeInfo);
|
||||
shared_zTadStride = shape::stride(zTadShapeInfo);
|
||||
shared_zDim = shared_zTadShape[axis]; // Assuming zDim is constant across splits
|
||||
|
||||
// Cache length and total threads
|
||||
shared_xLen = shape::length(xShapeInfo);
|
||||
shared_totalThreads = gridDim.x * blockDim.x;
|
||||
}
|
||||
|
||||
// Ensure all threads have access to the cached values
|
||||
__syncthreads();
|
||||
|
||||
const LongType tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
// Allocate space in shared memory for coordinates
|
||||
LongType* coords = sharedMem + threadIdx.x * shared_xRank;
|
||||
|
||||
for (LongType i = tid; i < shared_xLen; i += shared_totalThreads) {
|
||||
// Convert linear index to multi-dimensional coordinates
|
||||
INDEX2COORDS(i, shared_xRank, shared_xShape, coords);
|
||||
|
||||
LongType xOffset;
|
||||
// Convert coordinates to linear index for x
|
||||
COORDS2INDEX(shared_xRank, shared_xStride, coords, xOffset);
|
||||
|
||||
// Determine the split index along the specified axis
|
||||
LongType splitIndex = coords[axis] / shared_zDim;
|
||||
|
||||
// Retrieve the pointer to the target output tensor
|
||||
T* z = reinterpret_cast<T*>(reinterpret_cast<void**>(pVz)[splitIndex]);
|
||||
|
||||
// Update the coordinate along the split axis
|
||||
coords[axis] %= shared_zDim;
|
||||
|
||||
LongType zOffset;
|
||||
// Convert updated coordinates to linear index for z
|
||||
COORDS2INDEX(shared_zTadRank, shared_zTadStride, coords, zOffset);
|
||||
|
||||
// Perform the split operation
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST static void splitCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const cudaStream_t* stream,
|
||||
const void* vx, const LongType* xShapeInfo, void* pVz,
|
||||
const LongType* zTadShapeInfo, const LongType axis) {
|
||||
splitCuda<T><<<blocksPerGrid, threadsPerBlock, 256, *stream>>>(vx, xShapeInfo, pVz, zTadShapeInfo, axis);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "splitCuda failed");
|
||||
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void splitCudaLauncher,
|
||||
(const int blocksPerGrid, const int threadsPerBlock, const cudaStream_t* stream, const void* vx,
|
||||
const sd::LongType* xShapeInfo, void* pVz, const sd::LongType* zTadShapeInfo, const sd::LongType axis),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void split(LaunchContext* context, NDArray& input, std::vector<NDArray*>& outArrs, const LongType axis) {
|
||||
const int numOfSubArrs = outArrs.size();
|
||||
const auto sizeofT = input.sizeOfT();
|
||||
|
||||
for (int i = 0; i < numOfSubArrs; ++i) outArrs[i]->syncToDevice();
|
||||
input.syncToDevice();
|
||||
|
||||
bool luckCase1 = false;
|
||||
|
||||
if (luckCase1) {
|
||||
for (LongType i = 0; i < numOfSubArrs; ++i) {
|
||||
luckCase1 &= outArrs[i]->ordering() == input.ordering();
|
||||
if (!luckCase1) break;
|
||||
}
|
||||
}
|
||||
|
||||
if (luckCase1) { // for example {1,10} + {2,10} + {3,10} = {6, 10} order c; or {10,1} + {10,2} + {10,3} = {10, 6}
|
||||
// order f
|
||||
|
||||
auto x = static_cast<const int8_t*>(input.specialBuffer());
|
||||
|
||||
for (LongType i = 0; i < numOfSubArrs; ++i) {
|
||||
const auto memAmountToCopy = outArrs[i]->lengthOf() * sizeofT;
|
||||
cudaMemcpyAsync(static_cast<int8_t*>(outArrs[i]->specialBuffer()), x, memAmountToCopy, cudaMemcpyDeviceToDevice,
|
||||
*context->getCudaStream());
|
||||
x = static_cast<const int8_t*>(x) + memAmountToCopy;
|
||||
}
|
||||
|
||||
if (cudaStreamSynchronize(*context->getCudaStream()) != 0)
|
||||
THROW_EXCEPTION("split cuda: luckCase1 failed!");
|
||||
|
||||
for (int i = 0; i < numOfSubArrs; ++i) outArrs[i]->tickWriteDevice();
|
||||
input.tickReadDevice();
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (input.lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
|
||||
|
||||
// prepare arrays of pointers on buffers and shapes
|
||||
std::vector<void*> hOutBuffers(numOfSubArrs);
|
||||
|
||||
for (int i = 0; i < numOfSubArrs; ++i) hOutBuffers[i] = outArrs[i]->specialBuffer();
|
||||
|
||||
PointersManager manager(context, "helpers::split");
|
||||
|
||||
void* dOutBuffers = manager.replicatePointer(hOutBuffers.data(), hOutBuffers.size() * sizeof(void*));
|
||||
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), splitCudaLauncher,
|
||||
(blocksPerGrid, threadsPerBlock, context->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), dOutBuffers, outArrs[0]->specialShapeInfo(), axis),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
manager.synchronize();
|
||||
// }
|
||||
|
||||
for (int i = 0; i < numOfSubArrs; ++i) outArrs[i]->tickWriteDevice();
|
||||
input.tickReadDevice();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
Reference in New Issue
Block a user