Yu Yang
7 years ago
parent
6ebc6bf533
commit
fe7ed285d1
@ -1 +1,2 @@
|
||||
cc_library(var_handle SRCS var_handle.cc DEPS place)
|
||||
cc_library(op_handle_base SRCS op_handle_base.cc DEPS var_handle device_context)
|
||||
|
||||
@ -0,0 +1,84 @@
|
||||
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "paddle/fluid/framework/details/op_handle_base.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace framework {
|
||||
namespace details {
|
||||
std::string OpHandleBase::DebugString() const {
|
||||
std::stringstream ss;
|
||||
ss << "(";
|
||||
for (auto *var : inputs_) {
|
||||
ss << var->DebugString() << ", ";
|
||||
}
|
||||
ss << ") --> (";
|
||||
for (auto *var : outputs_) {
|
||||
ss << var->DebugString() << ", ";
|
||||
}
|
||||
ss << ")\n";
|
||||
return ss.str();
|
||||
}
|
||||
|
||||
OpHandleBase::~OpHandleBase() {}
|
||||
|
||||
void OpHandleBase::Run(bool use_event) {
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
if (events_.empty() && use_event) {
|
||||
for (auto &p : dev_ctx_) {
|
||||
int dev_id = boost::get<platform::CUDAPlace>(p.first).device;
|
||||
cudaSetDevice(dev_id);
|
||||
cudaEventCreateWithFlags(&events_[dev_id], cudaEventDisableTiming);
|
||||
}
|
||||
}
|
||||
#else
|
||||
PADDLE_ENFORCE(!use_event);
|
||||
#endif
|
||||
|
||||
RunImpl();
|
||||
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
if (use_event) {
|
||||
for (auto &p : dev_ctx_) {
|
||||
int dev_id = boost::get<platform::CUDAPlace>(p.first).device;
|
||||
auto stream =
|
||||
static_cast<platform::CUDADeviceContext *>(p.second)->stream();
|
||||
cudaEventRecord(events_.at(dev_id), stream);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
void OpHandleBase::Wait(platform::DeviceContext *waited_dev) {
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
if (platform::is_cpu_place(waited_dev->GetPlace()) || events_.empty()) {
|
||||
for (auto &dev_ctx : dev_ctx_) {
|
||||
dev_ctx.second->Wait();
|
||||
}
|
||||
} else {
|
||||
auto stream =
|
||||
static_cast<platform::CUDADeviceContext *>(waited_dev)->stream();
|
||||
for (auto &ev : events_) {
|
||||
PADDLE_ENFORCE(cudaStreamWaitEvent(stream, ev.second, 0));
|
||||
}
|
||||
}
|
||||
#else
|
||||
for (auto &dev_ctx : dev_ctx_) {
|
||||
dev_ctx.second->Wait();
|
||||
}
|
||||
#endif
|
||||
}
|
||||
} // namespace details
|
||||
} // namespace framework
|
||||
} // namespace paddle
|
||||
@ -0,0 +1,48 @@
|
||||
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "paddle/fluid/framework/details/var_handle.h"
|
||||
#include "paddle/fluid/platform/device_context.h"
|
||||
namespace paddle {
|
||||
namespace framework {
|
||||
namespace details {
|
||||
|
||||
struct OpHandleBase {
|
||||
std::vector<VarHandleBase *> inputs_;
|
||||
std::vector<VarHandleBase *> outputs_;
|
||||
std::unordered_map<platform::Place, platform::DeviceContext *,
|
||||
platform::PlaceHash>
|
||||
dev_ctx_;
|
||||
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
std::unordered_map<int, cudaEvent_t> events_;
|
||||
#endif
|
||||
|
||||
std::string DebugString() const;
|
||||
|
||||
virtual ~OpHandleBase();
|
||||
|
||||
void Run(bool use_event);
|
||||
|
||||
virtual void Wait(platform::DeviceContext *waited_dev);
|
||||
|
||||
protected:
|
||||
virtual void RunImpl() = 0;
|
||||
};
|
||||
|
||||
} // namespace details
|
||||
} // namespace framework
|
||||
} // namespace paddle
|
||||
Loading…
Reference in new issue