mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
1433160a36
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/66742 Modified loops in files under fbsource/fbcode/caffe2/ from the format `for(TYPE var=x0;var<x_max;x++)` to the format `for(const auto var: irange(xmax))` This was achieved by running r-barnes's loop upgrader script (D28874212) with some modification to exclude all files under /torch/jit and a number of reversions or unused variable suppression warnings added by hand. Test Plan: Sandcastle Reviewed By: malfet Differential Revision: D31705366 fbshipit-source-id: be58222426c192406a7f93c21582c3f6f2082401
55 lines
1.5 KiB
C++
55 lines
1.5 KiB
C++
#pragma once
|
|
|
|
#include <c10/util/irange.h>
|
|
#include "caffe2/core/context.h"
|
|
#include "caffe2/core/operator.h"
|
|
#include "caffe2/utils/math.h"
|
|
|
|
#include <vector>
|
|
|
|
namespace caffe2 {
|
|
template <typename T, class Context>
|
|
class KeySplitOp : public Operator<Context> {
|
|
public:
|
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
|
|
|
template <class... Args>
|
|
explicit KeySplitOp(Args&&... args)
|
|
: Operator<Context>(std::forward<Args>(args)...),
|
|
categorical_limit_(
|
|
this->template GetSingleArgument<int>("categorical_limit", 0)) {
|
|
CAFFE_ENFORCE_GT(categorical_limit_, 0);
|
|
}
|
|
|
|
bool RunOnDevice() override {
|
|
auto& keys = Input(0);
|
|
const auto N = keys.numel();
|
|
const T *const keys_data = keys.template data<T>();
|
|
std::vector<int> counts(categorical_limit_);
|
|
std::vector<int*> eids(categorical_limit_);
|
|
for (const auto k : c10::irange(categorical_limit_)) {
|
|
counts[k] = 0;
|
|
}
|
|
for (const auto i : c10::irange(N)) {
|
|
const auto k = keys_data[i];
|
|
CAFFE_ENFORCE_GT(categorical_limit_, k);
|
|
CAFFE_ENFORCE_GE(k, 0);
|
|
counts[k]++;
|
|
}
|
|
for (const auto k : c10::irange(categorical_limit_)) {
|
|
auto *const eid = Output(k, {counts[k]}, at::dtype<int>());
|
|
eids[k] = eid->template mutable_data<int>();
|
|
counts[k] = 0;
|
|
}
|
|
for (const auto i : c10::irange(N)) {
|
|
const auto k = keys_data[i];
|
|
eids[k][counts[k]++] = i;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
private:
|
|
int categorical_limit_;
|
|
};
|
|
} // namespace caffe2
|