chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,6 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from .lightconv_layer import LightconvLayer # noqa
|
||||
@@ -0,0 +1,289 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
|
||||
def gen_forward():
|
||||
|
||||
kernels = [3, 5, 7, 15, 31, 63, 127, 255]
|
||||
seqs = [32 * x for x in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]]
|
||||
|
||||
head = """
|
||||
/**
|
||||
* Copyright (c) Facebook, Inc. and its affiliates.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include "lightconv_cuda.cuh"
|
||||
|
||||
std::vector<at::Tensor> lightconv_cuda_forward(at::Tensor input, at::Tensor filters, int padding_l) {
|
||||
|
||||
at::DeviceGuard g(input.device());
|
||||
const auto minibatch = input.size(0);
|
||||
const auto numFeatures = input.size(1);
|
||||
const auto sequenceLength = input.size(2);
|
||||
|
||||
const auto numHeads = filters.size(0);
|
||||
const auto filterSize = filters.size(1);
|
||||
|
||||
const auto numFiltersInBlock = numFeatures / numHeads;
|
||||
|
||||
const dim3 blocks(minibatch, numFeatures);
|
||||
|
||||
auto output = at::zeros_like(input);
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
"""
|
||||
|
||||
sequence_if = """
|
||||
if (sequenceLength <= {seq}) {{
|
||||
switch(filterSize) {{
|
||||
"""
|
||||
|
||||
case_k = """
|
||||
case {k}:
|
||||
"""
|
||||
|
||||
main_block = """
|
||||
if (padding_l == {pad}) {{
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "lightconv_forward", ([&] {{
|
||||
lightconv_forward_kernel<{k}, {b_size}, {pad}, scalar_t>
|
||||
<<<blocks, {b_size}, 0, stream>>>(
|
||||
input.data<scalar_t>(),
|
||||
filters.data<scalar_t>(),
|
||||
minibatch,
|
||||
sequenceLength,
|
||||
numFeatures,
|
||||
numFiltersInBlock,
|
||||
output.data<scalar_t>());
|
||||
}}));
|
||||
}} else
|
||||
"""
|
||||
|
||||
bad_padding = """
|
||||
{
|
||||
std::cout << "WARNING: Unsupported padding size - skipping forward pass" << std::endl;
|
||||
}
|
||||
break;
|
||||
"""
|
||||
|
||||
bad_filter = """
|
||||
default:
|
||||
std::cout << "WARNING: Unsupported filter length passed - skipping forward pass" << std::endl;
|
||||
}
|
||||
"""
|
||||
|
||||
con_else = """
|
||||
} else
|
||||
"""
|
||||
|
||||
final_else = """
|
||||
{
|
||||
switch(filterSize) {
|
||||
"""
|
||||
|
||||
final_return = """
|
||||
}
|
||||
|
||||
return {output};
|
||||
}
|
||||
"""
|
||||
|
||||
with open("lightconv_cuda_forward.cu", "w") as forward:
|
||||
forward.write(head)
|
||||
for seq in seqs:
|
||||
forward.write(sequence_if.format(seq=seq))
|
||||
for k in kernels:
|
||||
forward.write(case_k.format(k=k))
|
||||
for pad in [k // 2, k - 1]:
|
||||
forward.write(main_block.format(k=k, b_size=seq, pad=pad))
|
||||
forward.write(bad_padding)
|
||||
forward.write(bad_filter)
|
||||
forward.write(con_else)
|
||||
|
||||
forward.write(final_else)
|
||||
for k in kernels:
|
||||
forward.write(case_k.format(k=k))
|
||||
for pad in [k // 2, k - 1]:
|
||||
forward.write(main_block.format(k=k, b_size=seq, pad=pad))
|
||||
forward.write(bad_padding)
|
||||
forward.write(bad_filter)
|
||||
forward.write(final_return)
|
||||
|
||||
|
||||
def gen_backward():
|
||||
|
||||
head = """
|
||||
/**
|
||||
* Copyright (c) Facebook, Inc. and its affiliates.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include "lightconv_cuda.cuh"
|
||||
|
||||
std::vector<at::Tensor> lightconv_cuda_backward(
|
||||
at::Tensor gradOutput,
|
||||
int padding_l,
|
||||
at::Tensor input,
|
||||
at::Tensor filters) {
|
||||
|
||||
// gradWrtInput
|
||||
const int minibatch = input.size(0);
|
||||
const int numFeatures = input.size(1);
|
||||
const int sequenceLength = input.size(2);
|
||||
|
||||
const int numHeads = filters.size(0);
|
||||
const int filterSize = filters.size(1);
|
||||
|
||||
const dim3 gradBlocks(minibatch, numFeatures);
|
||||
const dim3 weightGradFirstpassShortBlocks(minibatch, numHeads);
|
||||
const dim3 weightGradSecondpassBlocks(numHeads, filterSize);
|
||||
|
||||
const int numFiltersInBlock = numFeatures / numHeads;
|
||||
|
||||
auto gradInput = at::zeros_like(input);
|
||||
auto gradFilters = at::zeros_like(filters);
|
||||
|
||||
at::DeviceGuard g(input.device());
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
switch(filterSize) {
|
||||
"""
|
||||
|
||||
sequence_if = """
|
||||
if (sequenceLength <= {seq}) {{
|
||||
"""
|
||||
|
||||
case_k = """
|
||||
case {k}:
|
||||
"""
|
||||
|
||||
main_block = """
|
||||
if (padding_l == {p}) {{
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "lightconv_backward", ([&] {{
|
||||
lightconv_grad_wrt_input_kernel<{k}, {b_size}, {p}, scalar_t>
|
||||
<<<gradBlocks, {b_size}, 0, stream>>>(
|
||||
gradOutput.data<scalar_t>(),
|
||||
filters.data<scalar_t>(),
|
||||
minibatch,
|
||||
sequenceLength,
|
||||
numFeatures,
|
||||
numFiltersInBlock,
|
||||
gradInput.data<scalar_t>());
|
||||
|
||||
"""
|
||||
|
||||
weight_grad_short = """
|
||||
at::Tensor tempSumGradFilters = at::zeros({{minibatch, numHeads, filterSize}}, input.options().dtype(at::kFloat));
|
||||
lightconv_grad_wrt_weights_firstpass_short_kernel<{k}, {b_size}, {p}, scalar_t>
|
||||
<<<weightGradFirstpassShortBlocks, {b_size}, 0, stream>>>(
|
||||
input.data<scalar_t>(),
|
||||
gradOutput.data<scalar_t>(),
|
||||
minibatch,
|
||||
sequenceLength,
|
||||
numFeatures,
|
||||
numFiltersInBlock,
|
||||
numHeads,
|
||||
tempSumGradFilters.data<float>()
|
||||
);
|
||||
|
||||
lightconv_grad_wrt_weights_secondpass_short_kernel<{k}, {b_size}, scalar_t>
|
||||
<<<weightGradSecondpassBlocks, {b_size}, 0, stream>>>(
|
||||
tempSumGradFilters.data<float>(),
|
||||
minibatch,
|
||||
numFiltersInBlock,
|
||||
gradFilters.data<scalar_t>()
|
||||
);
|
||||
}}));
|
||||
}} else
|
||||
"""
|
||||
|
||||
weight_grad = """
|
||||
at::Tensor tempSumGradFilters = at::zeros({{minibatch, numFeatures, filterSize}}, input.options().dtype(at::kFloat));
|
||||
lightconv_grad_wrt_weights_firstpass_kernel<{k}, {b_size}, {p}, scalar_t>
|
||||
<<<gradBlocks, {b_size}, 0, stream>>>(
|
||||
input.data<scalar_t>(),
|
||||
gradOutput.data<scalar_t>(),
|
||||
minibatch,
|
||||
sequenceLength,
|
||||
numFeatures,
|
||||
numFiltersInBlock,
|
||||
tempSumGradFilters.data<float>()
|
||||
);
|
||||
|
||||
lightconv_grad_wrt_weights_secondpass_kernel<{k}, {b_size}, scalar_t>
|
||||
<<<weightGradSecondpassBlocks, {b_size}, 0, stream>>>(
|
||||
tempSumGradFilters.data<float>(),
|
||||
minibatch,
|
||||
numFiltersInBlock,
|
||||
gradFilters.data<scalar_t>()
|
||||
);
|
||||
}}));
|
||||
}} else
|
||||
"""
|
||||
|
||||
bad_padding = """
|
||||
{
|
||||
std::cout << "WARNING: Unsupported padding size - skipping backward pass" << std::endl;
|
||||
}
|
||||
"""
|
||||
|
||||
breakout = """
|
||||
break;
|
||||
"""
|
||||
|
||||
bad_filter = """
|
||||
default:
|
||||
std::cout << "WARNING: Unsupported filter length passed - skipping backward pass" << std::endl;
|
||||
"""
|
||||
|
||||
con_else = """
|
||||
} else
|
||||
"""
|
||||
|
||||
final_else = """
|
||||
{
|
||||
switch(filterSize) {
|
||||
"""
|
||||
|
||||
last_return = """
|
||||
}
|
||||
return {gradInput, gradFilters};
|
||||
}
|
||||
"""
|
||||
|
||||
kernels = [3, 5, 7, 15, 31, 63, 127, 255]
|
||||
seqs = [32 * x for x in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]]
|
||||
thresh = [32, 32, 64, 128, 256, -1, -1, -1]
|
||||
max_mem = [-1, -1, -1, -1, -1, 192, 96, 64]
|
||||
|
||||
with open("lightconv_cuda_backward.cu", "w") as backward:
|
||||
backward.write(head)
|
||||
for (k, t, mem) in zip(kernels, thresh, max_mem):
|
||||
backward.write(case_k.format(k=k))
|
||||
for seq in seqs:
|
||||
if (t == -1 or seq <= t) and (mem == -1 or seq < mem):
|
||||
backward.write(sequence_if.format(seq=seq))
|
||||
for p in [k // 2, k - 1]:
|
||||
backward.write(main_block.format(k=k, b_size=seq, p=p))
|
||||
backward.write(weight_grad_short.format(k=k, b_size=seq, p=p))
|
||||
backward.write(bad_padding)
|
||||
else:
|
||||
for p in [k // 2, k - 1]:
|
||||
backward.write(main_block.format(k=k, b_size=32, p=p))
|
||||
backward.write(weight_grad.format(k=k, b_size=32, p=p))
|
||||
backward.write(bad_padding)
|
||||
backward.write(breakout)
|
||||
break
|
||||
backward.write(con_else)
|
||||
backward.write(bad_filter)
|
||||
backward.write(last_return)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
gen_forward()
|
||||
gen_backward()
|
||||
@@ -0,0 +1,54 @@
|
||||
/**
|
||||
* Copyright (c) Facebook, Inc. and its affiliates.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
std::vector<at::Tensor> lightconv_cuda_forward(
|
||||
at::Tensor input,
|
||||
at::Tensor filters,
|
||||
int padding_l);
|
||||
|
||||
std::vector<at::Tensor> lightconv_cuda_backward(
|
||||
at::Tensor gradOutput,
|
||||
int padding_l,
|
||||
at::Tensor input,
|
||||
at::Tensor filters);
|
||||
|
||||
|
||||
#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
|
||||
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
|
||||
|
||||
std::vector<at::Tensor> lightconv_forward(
|
||||
at::Tensor input,
|
||||
at::Tensor filters,
|
||||
int padding_l) {
|
||||
|
||||
CHECK_INPUT(input);
|
||||
CHECK_INPUT(filters);
|
||||
|
||||
return lightconv_cuda_forward(input, filters, padding_l);
|
||||
}
|
||||
|
||||
std::vector<at::Tensor> lightconv_backward(
|
||||
at::Tensor gradOutput,
|
||||
int padding_l,
|
||||
at::Tensor input,
|
||||
at::Tensor filters) {
|
||||
|
||||
CHECK_INPUT(gradOutput);
|
||||
CHECK_INPUT(input);
|
||||
CHECK_INPUT(filters);
|
||||
|
||||
return lightconv_cuda_backward(gradOutput, padding_l, input, filters);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("forward", &lightconv_forward, "lighconv forward (CUDA)");
|
||||
m.def("backward", &lightconv_backward, "lighconv backward (CUDA)");
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
/**
|
||||
* Copyright (c) Facebook, Inc. and its affiliates.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
#include <iostream>
|
||||
#include <stdexcept>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include <stdlib.h>
|
||||
#include <assert.h>
|
||||
|
||||
#define SHFL_MASK 0xffffffff
|
||||
|
||||
template<int FS, int SB, int padding_l, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_forward_kernel(const scalar_t* input,
|
||||
const scalar_t* filters,
|
||||
int minibatch, int sequenceLength,
|
||||
int numFeatures, int numFiltersInBlock,
|
||||
scalar_t* output);
|
||||
|
||||
template<int FS, int SB, int padding_l, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_input_kernel(
|
||||
const scalar_t* input,
|
||||
const scalar_t* filters,
|
||||
int minibatch,
|
||||
int sequenceLength,
|
||||
int numFeatures,
|
||||
int numFiltersInBlock,
|
||||
scalar_t* output);
|
||||
|
||||
template<int FS, int SB, int padding_l, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_weights_firstpass_short_kernel(
|
||||
const scalar_t* input,
|
||||
const scalar_t* gradInput,
|
||||
int minibatch,
|
||||
int sequenceLength,
|
||||
int numFeatures,
|
||||
int numFiltersInBlock,
|
||||
int numHeads,
|
||||
float* output);
|
||||
|
||||
template<int FS, int SB, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_weights_secondpass_short_kernel(
|
||||
const float* input,
|
||||
const int minibatch,
|
||||
const int numFiltersInBlock,
|
||||
scalar_t* output);
|
||||
|
||||
template<int FS, int SB, int padding_l, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_weights_firstpass_kernel(
|
||||
const scalar_t* input,
|
||||
const scalar_t* gradInput,
|
||||
int minibatch,
|
||||
int sequenceLength,
|
||||
int numFeatures,
|
||||
int numFiltersInBlock,
|
||||
float* output);
|
||||
|
||||
template<int FS, int SB, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_weights_secondpass_kernel(
|
||||
const float* input,
|
||||
const int minibatch,
|
||||
const int numFiltersInBlock,
|
||||
scalar_t* output);
|
||||
|
||||
@@ -0,0 +1,375 @@
|
||||
/**
|
||||
* Copyright (c) Facebook, Inc. and its affiliates.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include "lightconv_cuda.cuh"
|
||||
#include "lightconv_cuda_forward.cu"
|
||||
#include "lightconv_cuda_backward.cu"
|
||||
#include "../cuda_utils.cu"
|
||||
|
||||
template<int FS, int SB, int padding_l, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_forward_kernel(const scalar_t* input,
|
||||
const scalar_t* filters,
|
||||
int minibatch, int sequenceLength,
|
||||
int numFeatures, int numFiltersInBlock,
|
||||
scalar_t* output) {
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
const int batchIdx = blockIdx.x;
|
||||
const int featureIdx = blockIdx.y;
|
||||
const int filterIdx = featureIdx / numFiltersInBlock;
|
||||
|
||||
const int IOOffset = numFeatures * sequenceLength * batchIdx + featureIdx * sequenceLength;
|
||||
const scalar_t* inputFeature = &input[IOOffset];
|
||||
scalar_t* outputFeature = &output[IOOffset];
|
||||
const scalar_t* inputFilter = &filters[filterIdx * FS];
|
||||
|
||||
assert(blockDim.x == SB);
|
||||
|
||||
scalar_t filter[FS];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FS; ++i) {
|
||||
filter[i] = inputFilter[i];
|
||||
}
|
||||
|
||||
__shared__ scalar_t temp[SB + FS];
|
||||
zeroSharedMem<FS, SB, padding_l>(temp);
|
||||
|
||||
const int numIterations = divUp<int, int>(sequenceLength, SB);
|
||||
|
||||
for (int i = 0; i < numIterations; ++i) {
|
||||
// Read input into shared memory
|
||||
const int inputOffset = i * SB;
|
||||
|
||||
load_input_to_shared<FS, SB, padding_l>(inputFeature, inputOffset, sequenceLength,
|
||||
i, numIterations, (numIterations == 1), temp);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
scalar_t out = 0;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < FS; ++j) {
|
||||
out += filter[j] * temp[tid + j];
|
||||
}
|
||||
|
||||
// Write output
|
||||
const int outputOffset = inputOffset;
|
||||
if ((outputOffset + tid) < sequenceLength) {
|
||||
outputFeature[outputOffset + tid] = out;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template<int FS, int SB, int padding_l, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_input_kernel(
|
||||
const scalar_t* input,
|
||||
const scalar_t* filters,
|
||||
int minibatch,
|
||||
int sequenceLength,
|
||||
int numFeatures,
|
||||
int numFiltersInBlock,
|
||||
scalar_t* output) {
|
||||
|
||||
// input grad kernel is similar to forward kernel
|
||||
const int tid = threadIdx.x;
|
||||
const int batchIdx = blockIdx.x;
|
||||
const int featureIdx = blockIdx.y;
|
||||
const int filterIdx = featureIdx / numFiltersInBlock;
|
||||
|
||||
const int IOOffset = numFeatures * sequenceLength * batchIdx + featureIdx * sequenceLength;
|
||||
const scalar_t* inputFeature = &input[IOOffset];
|
||||
scalar_t* outputFeature = &output[IOOffset];
|
||||
const scalar_t* inputFilter = &filters[filterIdx * FS];
|
||||
|
||||
assert(blockDim.x == SB);
|
||||
|
||||
scalar_t filter[FS];
|
||||
|
||||
// The only change is loading the filter in reverse
|
||||
#pragma unroll
|
||||
for (int i = 0; i < FS; ++i) {
|
||||
filter[i] = inputFilter[FS - i - 1];
|
||||
}
|
||||
|
||||
__shared__ scalar_t temp[SB + FS];
|
||||
const int padding = FS - padding_l - 1;
|
||||
zeroSharedMem<FS, SB, padding>(temp);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
const int numIterations = divUp<int, int>(sequenceLength, SB);
|
||||
|
||||
for (int i = 0; i < numIterations; ++i) {
|
||||
// Read input into shared memory
|
||||
const int inputOffset = i * SB;
|
||||
|
||||
load_input_to_shared<FS, SB, padding>(inputFeature, inputOffset, sequenceLength,
|
||||
i, numIterations, false, temp);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
scalar_t out = 0;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < FS; ++j) {
|
||||
out += filter[j] * temp[tid + j];
|
||||
}
|
||||
|
||||
// Write output
|
||||
const int outputOffset = inputOffset;
|
||||
if ((outputOffset + tid) < sequenceLength) {
|
||||
outputFeature[outputOffset + tid] = out;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
// This is by far the most expensive kernel in terms of time taken.
|
||||
// Can be 16x slower than the forward or grad_wrt_input when filter size is 31
|
||||
template<int FS, int SB, int padding_l, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_weights_firstpass_short_kernel(
|
||||
const scalar_t* input,
|
||||
const scalar_t* gradInput,
|
||||
int minibatch,
|
||||
int sequenceLength,
|
||||
int numFeatures,
|
||||
int numFiltersInBlock,
|
||||
int numHeads,
|
||||
float* output) {
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
const int batchIdx = blockIdx.x;
|
||||
const int filterIdx = blockIdx.y;
|
||||
|
||||
const int numIterations = divUp<int, int>(sequenceLength, SB);
|
||||
|
||||
float* tempOutputGradWeight = &output[filterIdx * FS * minibatch];
|
||||
|
||||
assert(blockDim.x == SB);
|
||||
|
||||
__shared__ scalar_t tempInput[SB + FS];
|
||||
__shared__ scalar_t tempGradInput[SB + FS];
|
||||
|
||||
// local weight accumulation
|
||||
float accumWeights[FS];
|
||||
|
||||
// Initialize memory
|
||||
for (int i = 0; i < FS; ++i) {
|
||||
accumWeights[i] = float(0.0);
|
||||
}
|
||||
|
||||
|
||||
// loop over each sequence within filterblock
|
||||
for (int idxInFilterBlock = 0; idxInFilterBlock < numFiltersInBlock; ++idxInFilterBlock) {
|
||||
|
||||
const int featureOffset = batchIdx * numFeatures * sequenceLength + (filterIdx * numFiltersInBlock + idxInFilterBlock) * sequenceLength;
|
||||
const scalar_t* inputFeature = &input[featureOffset];
|
||||
const scalar_t* gradInputFeature = &gradInput[featureOffset];
|
||||
|
||||
zeroSharedMem<FS, SB, padding_l>(tempInput);
|
||||
zeroSharedMem<FS, SB, (FS/2)>(tempGradInput);
|
||||
__syncthreads();
|
||||
|
||||
for (int i = 0; i < numIterations; ++i) {
|
||||
|
||||
const int inputOffset = i * SB;
|
||||
|
||||
load_input_to_shared<FS, SB, padding_l>(inputFeature, inputOffset, sequenceLength,
|
||||
i, numIterations, false, tempInput);
|
||||
load_input_to_shared<FS, SB, (FS/2)>(gradInputFeature, inputOffset, sequenceLength,
|
||||
i, numIterations, false, tempGradInput);
|
||||
|
||||
__syncthreads();
|
||||
|
||||
const int gradIndex = (FS/2) + tid;
|
||||
scalar_t tempGrad = tempGradInput[gradIndex];
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < FS; j++) {
|
||||
const int inputIndex = tid + j;
|
||||
accumWeights[j] += tempInput[inputIndex] * tempGrad;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// Row-major sum
|
||||
for (int filterWeightIdx = 0; filterWeightIdx < FS; ++filterWeightIdx) {
|
||||
|
||||
float temp;
|
||||
if (tid < sequenceLength) {
|
||||
temp = accumWeights[filterWeightIdx];
|
||||
} else {
|
||||
temp = float(0.0);
|
||||
}
|
||||
|
||||
const int outputOffset = filterWeightIdx * minibatch + batchIdx;
|
||||
|
||||
temp = blockReduce(temp);
|
||||
|
||||
if (tid == 0) {
|
||||
tempOutputGradWeight[outputOffset] = temp;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<int FS, int SB, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_weights_secondpass_short_kernel(
|
||||
const float* input,
|
||||
const int minibatch,
|
||||
const int numFiltersInBlock,
|
||||
scalar_t* output) {
|
||||
|
||||
assert(blockDim.x == SB);
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
|
||||
const int filterIdx = blockIdx.x;
|
||||
const int filterWeightIdx = blockIdx.y;
|
||||
|
||||
const int inputOffset = filterIdx * FS * minibatch +
|
||||
filterWeightIdx * minibatch;
|
||||
const float* tempInput = &input[inputOffset];
|
||||
|
||||
// read into shared memory for reduction
|
||||
int readIndex = tid;
|
||||
|
||||
float sum = 0.0;
|
||||
while (readIndex < minibatch) {
|
||||
sum += tempInput[readIndex];
|
||||
readIndex += SB;
|
||||
}
|
||||
|
||||
float temp = blockReduce(sum);
|
||||
|
||||
if (tid == 0) {
|
||||
output[blockIdx.x * FS + blockIdx.y] = temp;
|
||||
}
|
||||
}
|
||||
|
||||
// This is by far the most expensive kernel in terms of time taken.
|
||||
// Can be 16x slower than the forward or grad_wrt_input when filter size is 31
|
||||
template<int FS, int SB, int padding_l, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_weights_firstpass_kernel(
|
||||
const scalar_t* input,
|
||||
const scalar_t* gradInput,
|
||||
int minibatch,
|
||||
int sequenceLength,
|
||||
int numFeatures,
|
||||
int numFiltersInBlock,
|
||||
float* output) {
|
||||
|
||||
assert(blockDim.x == SB);
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
const int batchIdx = blockIdx.x;
|
||||
const int featureIdx = blockIdx.y;
|
||||
const int filterIdx = featureIdx / numFiltersInBlock;
|
||||
const int idxInFilterBlock = featureIdx % numFiltersInBlock;
|
||||
|
||||
const int numIterations = divUp<int, int>(sequenceLength, SB);
|
||||
|
||||
float temp;
|
||||
|
||||
__shared__ scalar_t tempInput[SB + FS];
|
||||
__shared__ scalar_t tempGradInput[SB + FS];
|
||||
zeroSharedMem<FS, SB, padding_l>(tempInput);
|
||||
zeroSharedMem<FS, SB, (FS/2)>(tempGradInput);
|
||||
__syncthreads();
|
||||
|
||||
float accumWeights[FS];
|
||||
|
||||
for (int i = 0; i < FS; ++i) {
|
||||
accumWeights[i] = float(0.0);
|
||||
}
|
||||
|
||||
const int IOOffset = batchIdx * numFeatures * sequenceLength + featureIdx * sequenceLength;
|
||||
const scalar_t* inputFeature = &input[IOOffset];
|
||||
const scalar_t* gradInputFeature = &gradInput[IOOffset];
|
||||
float* tempOutputGradWeight = &output[filterIdx * FS * minibatch * numFiltersInBlock];
|
||||
|
||||
for (int i = 0; i < numIterations; ++i) {
|
||||
const int inputOffset = i * SB;
|
||||
|
||||
load_input_to_shared<FS, SB, padding_l>(inputFeature, inputOffset, sequenceLength,
|
||||
i, numIterations, false, tempInput);
|
||||
load_input_to_shared<FS, SB, (FS/2)>(gradInputFeature, inputOffset, sequenceLength,
|
||||
i, numIterations, false, tempGradInput);
|
||||
__syncthreads();
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < FS; ++j) {
|
||||
accumWeights[j] += tempInput[tid + j] * tempGradInput[tid + (FS/2)];
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Row-major sum
|
||||
for (int filterWeightIdx = 0; filterWeightIdx < FS; ++filterWeightIdx) {
|
||||
|
||||
// Write to shared memory before reduction
|
||||
if (tid < sequenceLength) {
|
||||
temp = accumWeights[filterWeightIdx];
|
||||
} else {
|
||||
temp = float(0.0);
|
||||
}
|
||||
|
||||
temp = blockReduce(temp);
|
||||
|
||||
const int outputOffset = filterWeightIdx * minibatch * numFiltersInBlock +
|
||||
batchIdx * numFiltersInBlock +
|
||||
idxInFilterBlock;
|
||||
|
||||
if (tid == 0) {
|
||||
tempOutputGradWeight[outputOffset] = temp;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<int FS, int SB, typename scalar_t>
|
||||
__global__
|
||||
void lightconv_grad_wrt_weights_secondpass_kernel(
|
||||
const float* input,
|
||||
const int minibatch,
|
||||
const int numFiltersInBlock,
|
||||
scalar_t* output) {
|
||||
|
||||
assert(blockDim.x == SB);
|
||||
const int tid = threadIdx.x;
|
||||
|
||||
// What is the id within a minibatch
|
||||
const int filterIdx = blockIdx.x;
|
||||
const int filterWeightIdx = blockIdx.y;
|
||||
|
||||
const int inputOffset = filterIdx * FS * minibatch * numFiltersInBlock +
|
||||
filterWeightIdx * minibatch * numFiltersInBlock;
|
||||
const float* tempInput = &input[inputOffset];
|
||||
|
||||
int readIndex = tid;
|
||||
|
||||
float sum = float(0.0);
|
||||
while (readIndex < (minibatch * numFiltersInBlock)) {
|
||||
sum += tempInput[readIndex];
|
||||
readIndex += SB;
|
||||
}
|
||||
|
||||
float temp = blockReduce(sum);
|
||||
|
||||
if (tid == 0) {
|
||||
output[blockIdx.x * FS + blockIdx.y] = temp;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import lightconv_cuda
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from fairseq import utils
|
||||
from fairseq.incremental_decoding_utils import with_incremental_state
|
||||
from fairseq.modules.fairseq_dropout import FairseqDropout
|
||||
from torch import nn
|
||||
from torch.autograd import Function
|
||||
|
||||
|
||||
class lightconvFunction(Function):
|
||||
@staticmethod
|
||||
def forward(ctx, x, weights, padding_l):
|
||||
ctx.padding_l = padding_l
|
||||
outputs = lightconv_cuda.forward(x, weights, padding_l)
|
||||
variables = [x, weights]
|
||||
ctx.save_for_backward(*variables)
|
||||
return outputs[0]
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
outputs = lightconv_cuda.backward(
|
||||
grad_output.contiguous(), ctx.padding_l, *ctx.saved_tensors
|
||||
)
|
||||
grad_input, grad_weights = outputs
|
||||
return grad_input, grad_weights, None
|
||||
|
||||
|
||||
@with_incremental_state
|
||||
class LightconvLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_size,
|
||||
kernel_size=1,
|
||||
padding_l=None,
|
||||
weight_softmax=False,
|
||||
num_heads=1,
|
||||
weight_dropout=0.0,
|
||||
bias=False,
|
||||
):
|
||||
super(LightconvLayer, self).__init__()
|
||||
self.input_size = input_size
|
||||
self.kernel_size = kernel_size
|
||||
self.padding_l = padding_l
|
||||
self.num_heads = num_heads
|
||||
self.weight_softmax = weight_softmax
|
||||
self.weight_dropout_module = FairseqDropout(
|
||||
weight_dropout, module_name=self.__class__.__name__
|
||||
)
|
||||
|
||||
self.weight = nn.Parameter(torch.Tensor(num_heads, kernel_size))
|
||||
if bias:
|
||||
self.bias = nn.Parameter(torch.Tensor(input_size))
|
||||
else:
|
||||
self.bias = None
|
||||
self.reset_parameters()
|
||||
|
||||
def upgrade_state_dict_named(self, state_dict, name):
|
||||
prefix = name + "." if name != "" else ""
|
||||
for k, v in state_dict.items():
|
||||
if k.endswith(prefix + "weight"):
|
||||
if v.dim() == 3 and v.size(1) == 1:
|
||||
state_dict[k] = v.squeeze(1)
|
||||
|
||||
def reset_parameters(self):
|
||||
nn.init.xavier_uniform_(self.weight)
|
||||
if self.bias is not None:
|
||||
nn.init.constant_(self.bias, 0.0)
|
||||
|
||||
def forward(self, x, incremental_state=None):
|
||||
|
||||
# during inference time, incremental BMM is faster
|
||||
if incremental_state is not None:
|
||||
T, B, C = x.size()
|
||||
K, H = self.kernel_size, self.num_heads
|
||||
R = C // H
|
||||
input_buffer = self._get_input_buffer(incremental_state)
|
||||
if input_buffer is None:
|
||||
input_buffer = x.new()
|
||||
x_unfold = torch.cat([input_buffer, x.unsqueeze(3)], dim=3)
|
||||
if self.kernel_size > 1:
|
||||
self._set_input_buffer(
|
||||
incremental_state, x_unfold[:, :, :, -self.kernel_size + 1 :]
|
||||
)
|
||||
x_unfold = x_unfold.view(T * B * H, R, -1)
|
||||
|
||||
weight = self.weight
|
||||
if self.weight_softmax:
|
||||
weight = F.softmax(weight.float(), dim=1).type_as(weight)
|
||||
|
||||
weight = weight[:, -x_unfold.size(2) :]
|
||||
|
||||
K = weight.size(1)
|
||||
|
||||
weight = (
|
||||
weight.view(1, H, K)
|
||||
.expand(T * B, H, K)
|
||||
.contiguous()
|
||||
.view(T * B * H, K, 1)
|
||||
)
|
||||
|
||||
weight = self.weight_dropout_module(weight)
|
||||
output = torch.bmm(x_unfold, weight) # T*B*H x R x 1
|
||||
output = output.view(T, B, C)
|
||||
return output
|
||||
|
||||
# during training time, use CUDA kernel
|
||||
else:
|
||||
x = x.permute(1, 2, 0).contiguous()
|
||||
weight = self.weight
|
||||
if self.weight_softmax:
|
||||
weight = F.softmax(self.weight, -1)
|
||||
if self.weight_dropout_module.p:
|
||||
weight = self.weight_dropout_module(weight)
|
||||
return lightconvFunction.apply(x, weight, self.padding_l).permute(2, 0, 1)
|
||||
|
||||
def reorder_incremental_state(self, incremental_state, new_order):
|
||||
input_buffer = self._get_input_buffer(incremental_state)
|
||||
if input_buffer is not None:
|
||||
input_buffer = input_buffer.index_select(1, new_order)
|
||||
self._set_input_buffer(incremental_state, input_buffer)
|
||||
|
||||
def _get_input_buffer(self, incremental_state):
|
||||
return utils.get_incremental_state(self, incremental_state, "input_buffer")
|
||||
|
||||
def _set_input_buffer(self, incremental_state, new_buffer):
|
||||
return utils.set_incremental_state(
|
||||
self, incremental_state, "input_buffer", new_buffer
|
||||
)
|
||||
|
||||
def half(self):
|
||||
return self._apply(lambda t: t.half() if t.is_floating_point() else t)
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
|
||||
|
||||
setup(
|
||||
name="lightconv_layer",
|
||||
ext_modules=[
|
||||
CUDAExtension(
|
||||
"lightconv_cuda",
|
||||
[
|
||||
"lightconv_cuda.cpp",
|
||||
"lightconv_cuda_kernel.cu",
|
||||
],
|
||||
),
|
||||
],
|
||||
cmdclass={"build_ext": BuildExtension},
|
||||
)
|
||||
Reference in New Issue
Block a user