提交 ff308e3b 编写于 作者: M Megvii Engine Team 提交者: Xu Xinran

feat(mgb/comp_node): generate uid for cuda comp node

GitOrigin-RevId: 34fa5a2fb613de80ae14034ff4337ffd2ec3f3e3
上级 32c86211
......@@ -245,6 +245,12 @@ class CudaCompNode::CompNodeImpl final: public CompNode::Impl {
throw;
});
}
uint64_t get_uid() override {
return m_uid;
}
private:
uint64_t m_uid;
};
MGB_DYN_TYPE_OBJ_FINAL_IMPL(CudaCompNode::CompNodeImpl);
......@@ -310,6 +316,17 @@ void CudaCompNodeImpl::init(
m_locator_logical = locator_logical;
m_initialized = true;
#if defined(__linux__) || defined(TARGET_OS_MAC)
FILE *fp;
fp = fopen("/dev/urandom", "r");
mgb_assert(fread(&m_uid, sizeof(m_uid), 1, fp) == 1);
fclose(fp);
#else
m_uid = std::chrono::duration_cast<std::chrono::nanoseconds>(
std::chrono::system_clock::now().time_since_epoch()
).count();
#endif
auto on_succ = [this](cudaStream_t stream) {
auto locator = m_locator;
log_comp_node_created(locator, m_locator_logical);
......
......@@ -359,6 +359,10 @@ class CompNode {
return m_impl ? m_impl->locator_logical().to_string() : "invalid";
}
uint64_t get_uid() {
return m_impl->get_uid();
}
//! get the physical locator that created this comp node
Locator locator() const {
return m_impl->locator();
......@@ -522,6 +526,10 @@ class CompNode {
virtual void add_callback(megdnn::thin_function<void()>&&);
virtual uint64_t get_uid() {
mgb_throw(MegBrainError, "get_uid is not impl yet");
};
protected:
ImplBase(free_func_t fd, free_func_t fh)
: free_device{fd}, free_host{fh} {}
......
......@@ -264,6 +264,19 @@ TEST(TestCompNodeCuda, MemNode) {
ASSERT_NE(cn00.mem_node(), cn1.mem_node());
}
TEST(TestCompNodeCuda, Uid) {
REQUIRE_GPU(2);
auto cn00 = CompNode::load("gpu0"),
cn1 = CompNode::load("gpu1"),
cn01 = CompNode::load("gpu0:0"),
cn02 = CompNode::load("gpu0:2");
ASSERT_EQ(cn00, CompNode::load("gpu0"));
ASSERT_EQ(cn00.get_uid(), cn01.get_uid());
ASSERT_NE(cn00.get_uid(), cn02.get_uid());
ASSERT_NE(cn00.get_uid(), cn1.get_uid());
}
TEST(TestCompNodeCPU, PhysicalDispatch) {
constexpr int ID = 0x2a6453e0;
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册