merge InferShapeContext and OperatorContext (#3347)

* merge InferShapeContext and OperatorContext

* OperatorBase& instead of OperatorBase*
fixstartbug
Qiao Longfei 8 years ago committed by GitHub
parent 0f3a3e9894
commit b008360b05

@ -120,10 +120,10 @@ class OperatorBase {
std::shared_ptr<std::unordered_map<std::string, int>> in_out_idxs_; std::shared_ptr<std::unordered_map<std::string, int>> in_out_idxs_;
}; };
class OperatorContext { class InferShapeContext {
public: public:
OperatorContext(const OperatorBase* op, const Scope& scope) InferShapeContext(const OperatorBase& op, const Scope& scope)
: op_(*op), scope_(scope) {} : op_(op), scope_(scope) {}
size_t InputSize() const { return op_.inputs_.size(); } size_t InputSize() const { return op_.inputs_.size(); }
@ -234,12 +234,6 @@ class OperatorContext {
const Scope& scope_; const Scope& scope_;
}; };
class InferShapeContext : public OperatorContext {
public:
InferShapeContext(const OperatorBase* op, const Scope& scope)
: OperatorContext(op, scope) {}
};
template <typename T> template <typename T>
struct EigenDeviceConverter; struct EigenDeviceConverter;
@ -255,11 +249,11 @@ struct EigenDeviceConverter<platform::GPUPlace> {
}; };
#endif #endif
class ExecutionContext : public OperatorContext { class ExecutionContext : public InferShapeContext {
public: public:
ExecutionContext(const OperatorBase* op, const Scope& scope, ExecutionContext(const OperatorBase& op, const Scope& scope,
const platform::DeviceContext* device_context) const platform::DeviceContext* device_context)
: OperatorContext(op, scope), device_context_(device_context) {} : InferShapeContext(op, scope), device_context_(device_context) {}
template <typename PlaceType, template <typename PlaceType,
typename DeviceType = typename DeviceType =
@ -311,13 +305,13 @@ class OperatorWithKernel : public OperatorBase {
std::unordered_map<OpKernelKey, std::unique_ptr<OpKernel>, OpKernelHash>; std::unordered_map<OpKernelKey, std::unique_ptr<OpKernel>, OpKernelHash>;
void InferShape(const Scope& scope) const override { void InferShape(const Scope& scope) const override {
InferShape(InferShapeContext(this, scope)); InferShape(InferShapeContext(*this, scope));
} }
void Run(const Scope& scope, void Run(const Scope& scope,
const platform::DeviceContext& dev_ctx) const final { const platform::DeviceContext& dev_ctx) const final {
auto& opKernel = AllOpKernels().at(type_).at(OpKernelKey(dev_ctx)); auto& opKernel = AllOpKernels().at(type_).at(OpKernelKey(dev_ctx));
opKernel->Compute(ExecutionContext(this, scope, &dev_ctx)); opKernel->Compute(ExecutionContext(*this, scope, &dev_ctx));
} }
static std::unordered_map<std::string /* op_type */, OpKernelMap>& static std::unordered_map<std::string /* op_type */, OpKernelMap>&

Loading…
Cancel
Save