mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
4404762d7d
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/16751 This was made more complicated by the fact that ivalue::IntList is a thing. So I had to fix all of the sites where we referring to IValue post facto. The following codemods were run, in this order: ``` codemod -m -d . --extensions cc,cpp,cu,cuh,h,hpp,py,cwrap,yaml,in IntList IntArrayRef codemod -m -d . --extensions cc,cpp,cu,cuh,h,hpp,py,cwrap,yaml,in IntArrayRef::create IntList::create codemod -m -d . --extensions cc,cpp,cu,cuh,h,hpp,py,cwrap,yaml,in ivalue::IntArrayRef ivalue::IntList codemod -m -d . --extensions cc,cpp,cu,cuh,h,hpp,py,cwrap,yaml,in Tag::IntArrayRef Tag::IntList codemod -m -d . --extensions cc,cpp,cu,cuh,h,hpp,py,cwrap,yaml,in isIntArrayRef isIntList codemod -m -d . --extensions cc,cpp,cu,cuh,h,hpp,py,cwrap,yaml,in toIntArrayRef toIntList codemod -m -d . --extensions cc,cpp,cu,cuh,h,hpp,py,cwrap,yaml,in 'Shared<IntArrayRef>' 'Shared<IntList>' codemod -m -d . --extensions cc,cpp,cu,cuh,h,hpp,py,cwrap,yaml,in 'intrusive_ptr<IntArrayRef>' 'intrusive_ptr<IntList>' ``` Some manual fixups were done afterwards; they can be reviewed separately at https://github.com/pytorch/pytorch/pull/16752 Reviewed By: dzhulgakov Differential Revision: D13954363 fbshipit-source-id: b5c40aacba042402155a2f5a229fa6db7992ac64
66 lines
1.9 KiB
C++
66 lines
1.9 KiB
C++
#ifndef CAFFE2_OPERATORS_CONV_OP_CACHE_H_
|
|
#define CAFFE2_OPERATORS_CONV_OP_CACHE_H_
|
|
|
|
#include <functional>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
#include "caffe2/core/logging.h"
|
|
#include "caffe2/core/tensor.h"
|
|
|
|
namespace caffe2 {
|
|
template <typename TAlgorithm>
|
|
class AlgorithmsCache {
|
|
public:
|
|
// Caches the best algorithm for a given
|
|
// combination of tensor dimensions & compute data type.
|
|
//
|
|
TAlgorithm getAlgorithm(
|
|
at::IntArrayRef tensorDimensions1,
|
|
at::IntArrayRef tensorDimensions2,
|
|
int algorithmFlags, // Differentiate between algorithms with different
|
|
// parameters in a generic way
|
|
std::function<TAlgorithm()> generatingFunc);
|
|
|
|
private:
|
|
std::unordered_map<int64_t, TAlgorithm> hash_;
|
|
};
|
|
|
|
template <typename TAlgorithm>
|
|
TAlgorithm AlgorithmsCache<TAlgorithm>::getAlgorithm(
|
|
at::IntArrayRef tensorDimensions1,
|
|
at::IntArrayRef tensorDimensions2,
|
|
int algorithmFlags,
|
|
std::function<TAlgorithm()> generatingFunc) {
|
|
int64_t seed = 0;
|
|
// Hash all of the inputs, which we wiill then use to try and look up
|
|
// a previously discovered algorithm, or fall back to generating a new one.
|
|
std::hash<int64_t> hashFn;
|
|
for (const auto num : tensorDimensions1) {
|
|
// Copied from boost::hash_combine.
|
|
// Adding 1 to differentiate between first and second vector.
|
|
seed ^= hashFn(num) + 0x9e3779b9 + (seed << 6) + (seed >> 2) + 1;
|
|
}
|
|
|
|
for (const auto num : tensorDimensions2) {
|
|
// Copied from boost::hash_combine.
|
|
seed ^= hashFn(num) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
|
|
}
|
|
|
|
// Adding 2 to differentiate from previous vectors
|
|
seed ^= hashFn(algorithmFlags) + 0x9e3779b9 + (seed << 6) + (seed >> 2) + 2;
|
|
|
|
if (seed == 0) {
|
|
return generatingFunc();
|
|
}
|
|
|
|
if (hash_.find(seed) == hash_.end()) {
|
|
TAlgorithm value = generatingFunc();
|
|
hash_[seed] = value;
|
|
}
|
|
|
|
return hash_[seed];
|
|
}
|
|
} // namespace caffe2
|
|
|
|
#endif
|