/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/operators/lengths_reducer_ops.h (23315B)
#pragma once
#include "caffe2/core/context.h"
#include "caffe2/core/operator.h"
#include "caffe2/perfkernels/embedding_lookup.h"
#ifdef USE_FBGEMM
#include "fbgemm/Fbgemm.h"
#endif
#include
#include
namespace caffe2 {
// A templated class that implements SparseLengths[Sum,WeightedSum,Mean].
template <
typename T, // output type
class InputTypes, // supported input types, such as TensorTypes
bool USE_WEIGHT = false, // Whether it is SparseLengthsWeightedSum
bool USE_MEAN = false, // Whether this is SparseLengthsMean
bool USE_POSITIONAL_WEIGHT = false
// USE_WEIGHT = true and USE_POSITIONAL_WEIGHT = true
// -> SparseLengthsPositionalWeightedSum
>
class CPUSparseLengthsReductionOp : public Operator {
public:
USE_OPERATOR_FUNCTIONS(CPUContext);
template
explicit CPUSparseLengthsReductionOp(Args&&... args)
: Operator(std::forward(args)...) {
static_assert(
!(USE_WEIGHT & USE_MEAN), "Cannot both specify weight and mean.");
}
~CPUSparseLengthsReductionOp() {}
// Currently, we support float and at::Half inputs for input data type, and
// int32_t and int64_t for the index type.
bool RunOnDevice() override {
return DispatchHelper::call(this, Input(DATA));
}
template
bool DoRunWithType() {
return DispatchHelper, InputType>::call(
this, Input(INDICES));
}
template
bool DoRunWithType2() {
auto& dataInput = Input(DATA);
auto& indicesInput = Input(INDICES);
auto& lengthsInput = Input(LENGTHS);
const int64_t M = lengthsInput.size(0);
const int64_t indices_size = indicesInput.numel();
auto shape = dataInput.sizes().vec();
shape[0] = M;
auto* output = Output(0, shape, at::dtype());
T* out_data = output->template mutable_data();
if (indices_size == 0) {
if (M > 0) {
memset(out_data, 0, output->numel() * sizeof(T));
}
return true;
}
CAFFE_ENFORCE_EQ(1, indicesInput.dim(), "INDICES must be a vector");
CAFFE_ENFORCE_EQ(1, lengthsInput.dim(), "LENGTHS must be a vector");
const int64_t N = dataInput.size(0);
const int D = dataInput.size_from_dim(1);
const InputType* in_data = dataInput.template data();
const IndexType* indices = indicesInput.template data();
const int* lengths = lengthsInput.template data();
const T* in_weight = nullptr;
if (USE_WEIGHT) {
// static if
auto& weightInput = Input(WEIGHT);
CAFFE_ENFORCE_EQ(1, weightInput.dim(), "WEIGHT must be a vector");
if (!USE_POSITIONAL_WEIGHT) {
CAFFE_ENFORCE_EQ(
weightInput.numel(),
indices_size,
"Weight should have the same length as indices.");
}
in_weight = weightInput.template data();
}
#ifdef USE_FBGEMM
// If this is the first call or block size has changed (should never
// happen actually), generate a kernel.
if (D != last_block_size) {
last_block_size = D;
if (std::is_same::value) {
if (std::is_same::value) {
kernel_fp32_i32_ =
fbgemm::GenerateEmbeddingSpMDM(
D,
USE_WEIGHT,
USE_MEAN,
/*prefetch distance*/ 16,
USE_POSITIONAL_WEIGHT,
/*use_offsets*/ false);
} else {
CAFFE_ENFORCE((std::is_same::value));
kernel_fp32_i64_ =
fbgemm::GenerateEmbeddingSpMDM(
D,
USE_WEIGHT,
USE_MEAN,
/*prefetch distance*/ 16,
USE_POSITIONAL_WEIGHT,
/*use_offsets*/ false);
}
} else {
CAFFE_ENFORCE((std::is_same::value));
if (std::is_same::value) {
kernel_fp16_i32_ =
fbgemm::GenerateEmbeddingSpMDM(
D,
USE_WEIGHT,
USE_MEAN,
/*prefetch distance*/ 16,
USE_POSITIONAL_WEIGHT,
/*use_offsets*/ false);
} else {
CAFFE_ENFORCE((std::is_same::value));
kernel_fp16_i64_ =
fbgemm::GenerateEmbeddingSpMDM(
D,
USE_WEIGHT,
USE_MEAN,
/*prefetch distance*/ 16,
USE_POSITIONAL_WEIGHT,
/*use_offsets*/ false);
}
}
}
bool success;
if (std::is_same::value) {
if (std::is_same::value) {
success = kernel_fp32_i32_(
M,
indices_size,
N,
reinterpret_cast(in_data),
indicesInput.template data(),
lengths,
in_weight,
out_data);
} else {
success = kernel_fp32_i64_(
M,
indices_size,
N,
reinterpret_cast(in_data),
indicesInput.template data(),
lengths,
in_weight,
out_data);
}
} else {
if (std::is_same::value) {
success = kernel_fp16_i32_(
M,
indices_size,
N,
reinterpret_cast(in_data),
indicesInput.template data(),
lengths,
in_weight,
out_data);
} else {
success = kernel_fp16_i64_(
M,
indices_size,
N,
reinterpret_cast(in_data),
indicesInput.template data(),
lengths,
in_weight,
out_data);
}
}
if (success) {
return true;
}
int64_t current = 0;
for (int m = 0; m < M; ++m) {
for (int i = 0; i < lengths[m]; ++i) {
CAFFE_ENFORCE_LT(
current,
indices_size,
"Your input seems to be incorrect: the sum of lengths values "
"should be the size of the indices tensor, but it appears not.");
IndexType idx = indices[current];
CAFFE_ENFORCE(
0 <= idx && idx < N,
"Index ",
current,
" is out of bounds: ",
idx,
", range 0 to ",
N,
", actual batch length is ",
M);
++current;
}
}
CAFFE_ENFORCE_EQ(
current,
indices_size,
"Your input seems to be incorrect: the sum of lengths values should be "
"the size of the indices tensor, but it appears not.");
return false;
#endif
// delegate work to perfkernel that branches based on architecture
EmbeddingLookup(
D,
M,
indices_size,
N,
in_data,
indices,
lengths,
in_weight,
nullptr, // scale_bias field is only used in SparseLengths8BitsRowwiseOp
USE_MEAN,
out_data);
return true;
}
enum {
DATA = 0, // Data input.
WEIGHT = 1, // Weight input used in SparseLengthsWeightedSum
INDICES = 1 + USE_WEIGHT, // 1 in SparseLengths[Sum,Mean] and
// 2 in SparseLengthsWeightedSum
LENGTHS = 2 + USE_WEIGHT, // 2 in SparseLengths[Sum, Mean],
// 3 in SparseLengthsWeightedSum
};
#ifdef USE_FBGEMM
private:
std::int64_t last_block_size{-1};
fbgemm::EmbeddingSpMDMKernelSignature::Type
kernel_fp32_i32_;
fbgemm::EmbeddingSpMDMKernelSignature::Type
kernel_fp32_i64_;
fbgemm::EmbeddingSpMDMKernelSignature::Type
kernel_fp16_i32_;
fbgemm::EmbeddingSpMDMKernelSignature::Type
kernel_fp16_i64_;
#endif
};
template
class TTSparseLengthsSumOp final : public Operator {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
template
explicit TTSparseLengthsSumOp(Args&&... args)
: Operator(std::forward(args)...),
factor_i(this->template GetRepeatedArgument(
"factor_i",
vector{1, 1, 1})),
factor_j(this->template GetRepeatedArgument(
"factor_j",
vector{1, 1, 1})),
ranks(this->template GetRepeatedArgument(
"ranks",
vector{1, 1, 1, 1})),
emb_size(this->template GetSingleArgument("emb_size", 64)) {
// cumprod of i, used for index slice
l_cumprod.push_back(1);
for (size_t i = 1; i < factor_i.size(); ++i) {
l_cumprod.push_back(l_cumprod[i - 1] * factor_i[i - 1]);
}
}
~TTSparseLengthsSumOp() {}
void Ind2Sub(int64_t* out_factor_index, const int64_t* indices, int len) {
// TODO: vectorization
auto N = factor_i.size();
for (int j = 0; j < len; j++) {
auto idx = indices[j];
for (int i = N; i > 0; i--) {
out_factor_index[j * N + i - 1] = idx / l_cumprod[i - 1];
idx = idx % l_cumprod[i - 1];
}
}
}
bool GetSlice(
std::vector>& tgt_slice,
const T* core,
const vector& ind_slice,
int bs,
int idx) {
// implement the functinality index_select(core, 1, ind_slice)
auto num_of_elements = ranks[idx] * factor_j[idx] * ranks[idx + 1];
for (int i = 0; i < bs; i++) {
memcpy(
tgt_slice[i].data(),
core + ind_slice[i] * num_of_elements,
num_of_elements * sizeof(T));
}
return true;
}
// ind: it stores the index to each tensor core
// bs: the number of indices
// GatherAllRows uses two steps to calculate the lengthsum functionality: 1) it uses tensor train
// to calculate the embedding for each index. 2) it sums the embedding for each bag.
// In Step 1), it batches all the indices together. Specifically, for every index, it uses the pre-computed
// ind of each tensor core to extract the corresponding slice of the core. Then it does gemm operation
// sequentially on the slices to produce the embedding result for each index.
// In Step 2), it takes the embedding computed in step 1) and apply the sum operation for each bag.
bool GatherAllRows(
int64_t* ind,
int bs,
int x_len,
vector cores,
int segments,
const int* lengths,
T* out_data) {
// compute the largest memory consumption of intermediate result
// TODO: dynamic allocation size: cur_rows*factor_j[i]*ranks[i+1]
// and also explore the contiguous memory storage for res and int_res
int max_rank = *max_element(ranks.begin(), ranks.end());
std::vector> res(bs, std::vector(emb_size * max_rank, 0));
std::vector> int_res(
bs, std::vector(emb_size * max_rank, 0));
// Store the matrix A
vector Y_ptr(bs);
// Store the intermediate result in each layer
vector Z_ptr(bs);
for (int b = 0; b < bs; b++) {
Y_ptr[b] = res[b].data();
Z_ptr[b] = int_res[b].data();
}
vector ind_slice(bs);
int rows = 0;
for (int i = 0; i < x_len; i++) {
// slice cur
for (int j = 0; j < bs; j++) {
ind_slice[j] = ind[x_len * j + i];
}
if (i == 0) {
GetSlice(res, cores[i], ind_slice, bs, i);
rows = factor_j[0];
} else {
std::vector> slice(
bs, std::vector(ranks[i] * factor_j[i] * ranks[i + 1], 0));
vector X_ptr(bs);
for (int b = 0; b < bs; b++) {
X_ptr[b] = slice[b].data();
}
GetSlice(slice, cores[i], ind_slice, bs, i);
math::GemmBatched(
CblasNoTrans,
CblasNoTrans,
bs,
rows,
factor_j[i] * ranks[i + 1],
ranks[i],
1.0f,
const_cast(Y_ptr.data()),
X_ptr.data(),
0.0f,
Z_ptr.data(),
&context_);
for (int b = 0; b < bs; b++) {
std::memcpy(Y_ptr[b], Z_ptr[b], (emb_size * max_rank) * sizeof(T));
}
rows *= factor_j[i];
}
// save the intermediate output for backward path
// shape for the core
auto shape = vector({bs, rows, ranks[i + 1]});
if (i < 2) {
auto* core_data = Output(i + 1, shape, at::dtype());
T* out_core = core_data->template mutable_data();
for (int b = 0; b < bs; b++) {
std::memcpy(
out_core + b * rows * ranks[i + 1],
Y_ptr[b],
rows * ranks[i + 1] * sizeof(T));
}
}
}
// reduction and store back to output
vector cum_lengths(segments);
for (int seg = 0; seg < segments; seg++) {
cum_lengths[seg] =
seg == 0 ? lengths[0] : lengths[seg] + cum_lengths[seg - 1];
}
int length_idx = 0;
vector tmp_sum(emb_size, 0.0f);
for (int i = 0; i <= bs; i++) {
while ((length_idx < segments) && (i == cum_lengths[length_idx])) {
// store the tmp_sum into output
memcpy(
&out_data[length_idx * emb_size],
tmp_sum.data(),
emb_size * sizeof(T));
length_idx++;
fill(tmp_sum.begin(), tmp_sum.end(), 0.0f);
}
if (i == bs) {
break;
}
transform(
res[i].begin(),
res[i].begin() + emb_size,
tmp_sum.begin(),
tmp_sum.begin(),
std::plus());
}
return true;
}
bool RunOnDevice() override {
const auto& dataInput0 = Input(0);
const auto& dataInput1 = Input(1);
const auto& dataInput2 = Input(2);
const auto& indicesInput = Input(3);
const auto& lengthsInput = Input(4);
CAFFE_ENFORCE_EQ(1, indicesInput.dim(), "INDICES must be a vector");
CAFFE_ENFORCE_EQ(1, lengthsInput.dim(), "LENGTHS must be a vector");
int N = factor_i.size();
const int64_t M = lengthsInput.size(0);
auto shape = vector({M, emb_size});
auto* output = Output(0, shape, at::dtype());
T* out_data = output->template mutable_data();
const T* core0 = dataInput0.template data();
const T* core1 = dataInput1.template data();
const T* core2 = dataInput2.template data();
const int* lengths = lengthsInput.template data();
vector cores = {core0, core1, core2};
const int64_t* indices = indicesInput.template data();
// Store the factor index for backward path
auto index_shape = vector({indicesInput.size(), N});
auto* index_data = Output(3, index_shape, at::dtype());
int64_t* out_factor_index = index_data->template mutable_data();
// Store the factorized index for each core
Ind2Sub(out_factor_index, indices, indicesInput.size());
return GatherAllRows(
out_factor_index, indicesInput.size(), N, cores, M, lengths, out_data);
}
protected:
vector factor_i;
vector factor_j;
vector ranks;
vector l_cumprod;
int emb_size;
};
template
class TTSparseLengthsSumGradientOp final : public Operator {
public:
USE_OPERATOR_CONTEXT_FUNCTIONS;
template
explicit TTSparseLengthsSumGradientOp(Args&&... args)
: Operator(std::forward(args)...) {}
bool RunOnDevice() override;
~TTSparseLengthsSumGradientOp() {}
};
// implement the graident op for TTLengthSumGradient op
template
bool TTSparseLengthsSumGradientOp::RunOnDevice() {
const auto& core0 = Input(0);
const auto& core1 = Input(1);
const auto& core2 = Input(2);
const auto& lengths = Input(3);
const auto& core0_out = Input(4);
const auto& core1_out = Input(5);
const auto& index_out = Input(6);
const auto& dY = Input(7);
const int* lengths_data = lengths.template data();
const T* dY_data = dY.template data();
// restore the arguments from shape
const int64_t bs = index_out.size(0);
const int64_t emb_size = dY.size(1);
const int64_t num_segments = lengths.size(0);
auto core0_shape = core0.sizes().vec();
auto core1_shape = core1.sizes().vec();
auto core2_shape = core2.sizes().vec();
auto core0_out_shape = core0_out.sizes().vec();
auto core1_out_shape = core1_out.sizes().vec();
auto* dCore0 = Output(0, core0_shape, at::dtype());
auto* dCore1 = Output(1, core1_shape, at::dtype());
auto* dCore2 = Output(2, core2_shape, at::dtype());
T* dCore0_data = dCore0->template mutable_data();
T* dCore1_data = dCore1->template mutable_data();
T* dCore2_data = dCore2->template mutable_data();
memset(
dCore0_data,
0.0f,
sizeof(T) *
accumulate(
core0_shape.begin(), core0_shape.end(), 1, std::multiplies()));
memset(
dCore1_data,
0.0f,
sizeof(T) *
accumulate(
core1_shape.begin(), core1_shape.end(), 1, std::multiplies()));
memset(
dCore2_data,
0.0f,
sizeof(T) *
accumulate(
core2_shape.begin(), core2_shape.end(), 1, std::multiplies()));
int64_t* index_out_data = index_out.template mutable_data();
vector> index_slice(bs, vector(3, 0));
for (int64_t b = 0; b < bs; b++) {
memcpy(index_slice[b].data(), index_out_data + b * 3, 3 * sizeof(int64_t));
}
vector A_ptr(bs);
vector B_ptr(bs);
vector C_ptr(bs);
// size of each batch
int64_t num_of_elements = 0;
// construct the ranks
// expand the gradient into all indices
vector> core2_out_grad(bs, vector(emb_size, 0));
int64_t data_index = 0;
for (int64_t range_index = 0; range_index < num_segments; ++range_index) {
for (int64_t start = data_index;
data_index < start + lengths_data[range_index];
++data_index) {
memcpy(
core2_out_grad[data_index].data(),
dY_data + range_index * emb_size,
emb_size * sizeof(T));
}
}
// =======================================================
// Calculate dCore2_data:
// 1) Transpose core1_out and multiply iwth core2_out_grad
// 2) add to dCore2_data
vector> dCore2_data_slice_grad(
bs, vector(core2_shape[1] * core2_shape[2] * core2_shape[3], 0));
const T* core1_out_data = core1_out.template data();
// const T* core1_out_p[bs];
for (int64_t b = 0; b < bs; b++) {
A_ptr[b] = core1_out_data + b * core1_out.size(1) * core1_out.size(2);
B_ptr[b] = core2_out_grad[b].data();
C_ptr[b] = dCore2_data_slice_grad[b].data();
}
math::GemmBatched(
CblasTrans,
CblasNoTrans,
bs,
core2.size(1), // M
core2.size(2) * core2.size(3), // N
core1_out.size(1), // K
1.0f,
const_cast(A_ptr.data()),
const_cast(B_ptr.data()),
0.0f,
C_ptr.data(),
&context_);
// update the corresponding slice
num_of_elements = core2_shape[1] * core2_shape[2] * core2_shape[3];
T* core2_data = core2.template mutable_data();
vector> core2_slice(
bs, vector(core2_shape[1] * core2_shape[2] * core2_shape[3], 0));
for (int64_t b = 0; b < bs; b++) {
for (int i = 0; i < num_of_elements; i++) {
dCore2_data[index_slice[b][2] * num_of_elements + i] += C_ptr[b][i];
}
memcpy(
core2_slice[b].data(),
core2_data + index_slice[b][2] * num_of_elements,
sizeof(T) * num_of_elements);
}
// Calculate core1_out_grad
vector> core1_out_grad(
bs, vector(core1_out_shape[1] * core1_out_shape[2], 0));
for (int64_t b = 0; b < bs; b++) {
A_ptr[b] = core2_out_grad[b].data();
B_ptr[b] = core2_slice[b].data();
C_ptr[b] = core1_out_grad[b].data();
}
math::GemmBatched(
CblasNoTrans,
CblasTrans,
bs,
core1_out.size(1), // M
core2_shape[1], // N
core2_shape[2] * core2_shape[3], // K
1.0f,
const_cast(A_ptr.data()),
const_cast(B_ptr.data()),
0.0f,
C_ptr.data(),
&context_);
// =======================================================
// Calcuate dCore1_data:
// 1) Transpose core1_out_grad and multiply with core0_out
// 2) Transpose the result and then add to dCore1_data
vector> dCore1_data_slice_grad(
bs, vector(core1_shape[1] * core1_shape[2] * core1_shape[3], 0));
const T* core0_out_data = core0_out.template data();
for (int64_t b = 0; b < bs; b++) {
A_ptr[b] = core0_out_data + b * core0_out.size(1) * core0_out.size(2);
B_ptr[b] = core1_out_grad[b].data();
C_ptr[b] = dCore1_data_slice_grad[b].data();
}
math::GemmBatched(
CblasTrans,
CblasNoTrans,
bs,
core1.size(1), // M
core1.size(2) * core1.size(3), // N
core0_out.size(1), // K
1.0f,
const_cast(A_ptr.data()),
const_cast(B_ptr.data()),
0.0f,
C_ptr.data(),
&context_);
// update the corresponding slice
num_of_elements = core1_shape[1] * core1_shape[2] * core1_shape[3];
T* core1_data = core1.template mutable_data();
vector> core1_slice(
bs, vector(core1_shape[1] * core1_shape[2] * core1_shape[3], 0));
for (int64_t b = 0; b < bs; b++) {
for (int i = 0; i < num_of_elements; i++) {
dCore1_data[index_slice[b][1] * num_of_elements + i] += C_ptr[b][i];
}
memcpy(
core1_slice[b].data(),
core1_data + index_slice[b][1] * num_of_elements,
sizeof(T) * num_of_elements);
}
// Calcuate core0_out_grad
vector> core0_out_grad(
bs, vector(core0_out_shape[1] * core0_out_shape[2], 0));
for (int64_t b = 0; b < bs; b++) {
A_ptr[b] = core1_out_grad[b].data();
B_ptr[b] = core1_slice[b].data();
C_ptr[b] = core0_out_grad[b].data();
}
math::GemmBatched(
CblasNoTrans,
CblasTrans,
bs,
core0_out.size(1), // M
core1_shape[1], // N
core1_shape[2] * core1_shape[3], // K
1.0f,
const_cast(A_ptr.data()),
const_cast(B_ptr.data()),
0.0f,
C_ptr.data(),
&context_);
num_of_elements = core0_shape[1] * core0_shape[2] * core0_shape[3];
for (int64_t b = 0; b < bs; b++) {
for (int i = 0; i < num_of_elements; i++) {
dCore0_data[index_slice[b][0] * num_of_elements + i] += C_ptr[b][i];
}
}
return true;
}
} // namespace caffe2