mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
91797c0672
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/10946 ``` codemod -d . --extensions cc,cpp,cu,cuh,h caffe2/proto/caffe2.pb.h caffe2/proto/caffe2_pb.h ``` Reviewed By: houseroad Differential Revision: D9539945 fbshipit-source-id: 497d04720e8e7e61c05ffe1b23733d0cb774de7e
52 lines
1.6 KiB
C++
52 lines
1.6 KiB
C++
|
|
#pragma once
|
|
|
|
#include "caffe2/core/common.h"
|
|
#include "caffe2/core/transform.h"
|
|
#include "caffe2/proto/caffe2_pb.h"
|
|
#include "caffe2/utils/proto_utils.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
/**
|
|
* Common Subexpression Elimination
|
|
*
|
|
* This transforms looks for specific operators (denoted by whitelisted_ops_),
|
|
* and removes unnecessary repetition of that operator.
|
|
*
|
|
* Consider some operator of X, that reads from blob b_ written to by W.
|
|
* X_a and X_b read the output of X. However, another operator Y, is the same
|
|
* type as X, has the same arguments as X, and reads from the same input b_,
|
|
* written to by W. It's output is the same as X. Y_a, Y_b, and Y_c read from Y.
|
|
*
|
|
* Then, we can eliminate the common subexpressions X and Y, and merge them to
|
|
* Z, where X_a, X_b, Y_a, Y_b, and Y_c all read from Z.
|
|
*
|
|
*
|
|
* TODO(benz): Fix the error to not match nodes that write to external output.
|
|
*/
|
|
class CAFFE2_API CommonSubexpressionEliminationTransform : public Transform {
|
|
public:
|
|
CommonSubexpressionEliminationTransform() {
|
|
SetPatternMatchType(SORTED_WRT_EXECUTION_ORDER);
|
|
}
|
|
|
|
protected:
|
|
bool PatternRule(
|
|
const transform::Graph& g,
|
|
const std::vector<int>& subgraph,
|
|
int idx) override;
|
|
bool ValidatorRule(
|
|
const transform::Graph& g,
|
|
const std::vector<int>& subgraph) override;
|
|
bool ReplaceRule(const std::vector<int>& subgraph, transform::Graph* g_ptr)
|
|
override;
|
|
|
|
private:
|
|
bool IsWhitelisted(string op_type) {
|
|
return whitelisted_ops_.count(op_type);
|
|
}
|
|
std::set<string> whitelisted_ops_ = {"LearningRate", "FC"};
|
|
};
|
|
|
|
} // namespace caffe2
|