All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
BroadcastSum.cu
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the CC-by-NC license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#include <algorithm>
#include "../../FaissAssert.h"
#include "../utils/DeviceUtils.h"
#include "../utils/MathOperators.cuh"
#include "../utils/Tensor.cuh"
#include "../utils/StaticUtils.h"
namespace faiss { namespace gpu {
template <typename T, int kRowsPerBlock, int kRowUnroll, int kColLoad>
__global__ void sumAlongColumns(Tensor<T, 1, true> input,
Tensor<T, 2, true> output) {
static_assert(kRowsPerBlock % kRowUnroll == 0, "must fit rows");
// blockIdx.x: which chunk of rows we are responsible for updating
// blockIdx.y: which chunk of columns we are responsible for
// updating
int rowStart = blockIdx.x * kRowsPerBlock;
int rowEnd = rowStart + kRowsPerBlock;
int colStart = blockIdx.y * blockDim.x * kColLoad;
// FIXME: if we have exact multiples, don't need this
bool endRow = (blockIdx.x == gridDim.x - 1);
bool endCol = (blockIdx.y == gridDim.y - 1);
if (endRow) {
if (output.getSize(0) % kRowsPerBlock == 0) {
endRow = false;
}
}
if (endCol) {
for (int col = colStart + threadIdx.x;
col < input.getSize(0); col += blockDim.x) {
T val = input[col];
if (endRow) {
for (int row = rowStart; row < output.getSize(0); ++row) {
T out = output[row][col].ldg();
out = Math<T>::add(out, val);
output[row][col] = out;
}
} else {
T rows[kRowUnroll];
for (int row = rowStart; row < rowEnd; row += kRowUnroll) {
#pragma unroll
for (int i = 0; i < kRowUnroll; ++i) {
rows[i] = output[row + i][col].ldg();
}
#pragma unroll
for (int i = 0; i < kRowUnroll; ++i) {
rows[i] = Math<T>::add(rows[i], val);
}
#pragma unroll
for (int i = 0; i < kRowUnroll; ++i) {
output[row + i][col] = rows[i];
}
}
}
}
} else {
int col = colStart + threadIdx.x;
T val[kColLoad];
#pragma unroll
for (int i = 0; i < kColLoad; ++i) {
val[i] = input[col + i * blockDim.x];
}
if (endRow) {
for (int row = rowStart; row < output.getSize(0); ++row) {
#pragma unroll
for (int i = 0; i < kColLoad; ++i) {
T out = output[row][col + i * blockDim.x].ldg();
out = Math<T>::add(out, val[i]);
output[row][col + i * blockDim.x] = out;
}
}
} else {
T rows[kRowUnroll * kColLoad];
for (int row = rowStart; row < rowEnd; row += kRowUnroll) {
#pragma unroll
for (int i = 0; i < kRowUnroll; ++i) {
#pragma unroll
for (int j = 0; j < kColLoad; ++j) {
rows[i * kColLoad + j] =
output[row + i][col + j * blockDim.x].ldg();
}
}
#pragma unroll
for (int i = 0; i < kRowUnroll; ++i) {
#pragma unroll
for (int j = 0; j < kColLoad; ++j) {
rows[i * kColLoad + j] =
Math<T>::add(rows[i * kColLoad + j], val[j]);
}
}
#pragma unroll
for (int i = 0; i < kRowUnroll; ++i) {
#pragma unroll
for (int j = 0; j < kColLoad; ++j) {
output[row + i][col + j * blockDim.x] =
rows[i * kColLoad + j];
}
}
}
}
}
}
template <typename T, int kRowsPerBlock, int kRowUnroll, int kColLoad>
__global__ void assignAlongColumns(Tensor<T, 1, true> input,
Tensor<T, 2, true> output) {
static_assert(kRowsPerBlock % kRowUnroll == 0, "must fit rows");
// blockIdx.x: which chunk of rows we are responsible for updating
// blockIdx.y: which chunk of columns we are responsible for
// updating
int rowStart = blockIdx.x * kRowsPerBlock;
int rowEnd = rowStart + kRowsPerBlock;
int colStart = blockIdx.y * blockDim.x * kColLoad;
// FIXME: if we have exact multiples, don't need this
bool endRow = (blockIdx.x == gridDim.x - 1);
bool endCol = (blockIdx.y == gridDim.y - 1);
if (endRow) {
if (output.getSize(0) % kRowsPerBlock == 0) {
endRow = false;
}
}
if (endCol) {
for (int col = colStart + threadIdx.x;
col < input.getSize(0); col += blockDim.x) {
T val = input[col];
if (endRow) {
for (int row = rowStart; row < output.getSize(0); ++row) {
output[row][col] = val;
}
} else {
for (int row = rowStart; row < rowEnd; row += kRowUnroll) {
#pragma unroll
for (int i = 0; i < kRowUnroll; ++i) {
output[row + i][col] = val;
}
}
}
}
} else {
int col = colStart + threadIdx.x;
T val[kColLoad];
#pragma unroll
for (int i = 0; i < kColLoad; ++i) {
val[i] = input[col + i * blockDim.x];
}
if (endRow) {
for (int row = rowStart; row < output.getSize(0); ++row) {
#pragma unroll
for (int i = 0; i < kColLoad; ++i) {
output[row][col + i * blockDim.x] = val[i];
}
}
} else {
for (int row = rowStart; row < rowEnd; row += kRowUnroll) {
#pragma unroll
for (int i = 0; i < kRowUnroll; ++i) {
#pragma unroll
for (int j = 0; j < kColLoad; ++j) {
output[row + i][col + j * blockDim.x] = val[j];
}
}
}
}
}
}
template <typename T, typename TVec>
__global__ void sumAlongRows(Tensor<T, 1, true> input,
Tensor<TVec, 2, true> output) {
__shared__ T sval;
int row = blockIdx.x;
if (threadIdx.x == 0) {
sval = input[row];
}
__syncthreads();
T val = sval;
// FIXME: speed up
for (int i = threadIdx.x; i < output.getSize(1); i += blockDim.x) {
TVec out = output[row][i];
out = Math<TVec>::add(out, val);
output[row][i] = out;
}
}
template <typename T, typename TVec>
void runSumAlongColumns(Tensor<T, 1, true>& input,
Tensor<T, 2, true>& output,
cudaStream_t stream) {
FAISS_ASSERT(input.getSize(0) == output.getSize(1));
int threadsPerBlock = 256;
constexpr int kRowUnroll = 4;
constexpr int kRowsPerBlock = kRowUnroll * 4;
constexpr int kColLoad = 4;
auto block = dim3(threadsPerBlock);
if (input.template canCastResize<TVec>() &&
output.template canCastResize<TVec>()) {
auto inputV = input.template castResize<TVec>();
auto outputV = output.template castResize<TVec>();
auto grid =
dim3(utils::divUp(outputV.getSize(0), kRowsPerBlock),
utils::divUp(outputV.getSize(1), threadsPerBlock * kColLoad));
sumAlongColumns<TVec, kRowsPerBlock, kRowUnroll, kColLoad>
<<<grid, block, 0, stream>>>(inputV, outputV);
} else {
auto grid =
dim3(utils::divUp(output.getSize(0), kRowsPerBlock),
utils::divUp(output.getSize(1), threadsPerBlock * kColLoad));
sumAlongColumns<T, kRowsPerBlock, kRowUnroll, kColLoad>
<<<grid, block, 0, stream>>>(input, output);
}
CUDA_TEST_ERROR();
}
void runSumAlongColumns(Tensor<float, 1, true>& input,
Tensor<float, 2, true>& output,
cudaStream_t stream) {
runSumAlongColumns<float, float4>(input, output, stream);
}
#ifdef FAISS_USE_FLOAT16
void runSumAlongColumns(Tensor<half, 1, true>& input,
Tensor<half, 2, true>& output,
cudaStream_t stream) {
runSumAlongColumns<half, half2>(input, output, stream);
}
#endif
template <typename T, typename TVec>
void runAssignAlongColumns(Tensor<T, 1, true>& input,
Tensor<T, 2, true>& output,
cudaStream_t stream) {
FAISS_ASSERT(input.getSize(0) == output.getSize(1));
int threadsPerBlock = 256;
constexpr int kRowUnroll = 4;
constexpr int kRowsPerBlock = kRowUnroll * 4;
constexpr int kColLoad = 4;
auto block = dim3(threadsPerBlock);
if (input.template canCastResize<TVec>() &&
output.template canCastResize<TVec>()) {
auto inputV = input.template castResize<TVec>();
auto outputV = output.template castResize<TVec>();
auto grid =
dim3(utils::divUp(outputV.getSize(0), kRowsPerBlock),
utils::divUp(outputV.getSize(1), threadsPerBlock * kColLoad));
assignAlongColumns<TVec, kRowsPerBlock, kRowUnroll, kColLoad>
<<<grid, block, 0, stream>>>(inputV, outputV);
} else {
auto grid =
dim3(utils::divUp(output.getSize(0), kRowsPerBlock),
utils::divUp(output.getSize(1), threadsPerBlock * kColLoad));
assignAlongColumns<T, kRowsPerBlock, kRowUnroll, kColLoad>
<<<grid, block, 0, stream>>>(input, output);
}
CUDA_TEST_ERROR();
}
void runAssignAlongColumns(Tensor<float, 1, true>& input,
Tensor<float, 2, true>& output,
cudaStream_t stream) {
runAssignAlongColumns<float, float4>(input, output, stream);
}
#ifdef FAISS_USE_FLOAT16
void runAssignAlongColumns(Tensor<half, 1, true>& input,
Tensor<half, 2, true>& output,
cudaStream_t stream) {
runAssignAlongColumns<half, half2>(input, output, stream);
}
#endif
template <typename T, typename TVec>
void runSumAlongRows(Tensor<T, 1, true>& input,
Tensor<T, 2, true>& output,
cudaStream_t stream) {
FAISS_ASSERT(input.getSize(0) == output.getSize(0));
if (output.template canCastResize<TVec>()) {
auto outputV = output.template castResize<TVec>();
int threadsPerBlock =
std::min(outputV.getSize(1), getMaxThreadsCurrentDevice());
auto grid = dim3(outputV.getSize(0));
auto block = dim3(threadsPerBlock);
sumAlongRows<T, TVec><<<grid, block, 0, stream>>>(input, outputV);
} else {
int threadsPerBlock =
std::min(output.getSize(1), getMaxThreadsCurrentDevice());
auto grid = dim3(output.getSize(0));
auto block = dim3(threadsPerBlock);
sumAlongRows<T, T><<<grid, block, 0, stream>>>(input, output);
}
CUDA_TEST_ERROR();
}
void runSumAlongRows(Tensor<float, 1, true>& input,
Tensor<float, 2, true>& output,
cudaStream_t stream) {
runSumAlongRows<float, float4>(input, output, stream);
}
#ifdef FAISS_USE_FLOAT16
void runSumAlongRows(Tensor<half, 1, true>& input,
Tensor<half, 2, true>& output,
cudaStream_t stream) {
runSumAlongRows<half, half2>(input, output, stream);
}
#endif
} } // namespace