Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
e6943017
MegEngine
项目概览
MegEngine 天元
/
MegEngine
12 个月 前同步成功
通知
393
Star
4704
Fork
582
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
MegEngine
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
e6943017
编写于
6月 08, 2022
作者:
M
Megvii Engine Team
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
fix(ops/jit): skip lookup include path when nvcc executable not found
GitOrigin-RevId: f5e7dce1c52fe099de8e90684bbe5b0566a5115a
上级
02c1a0c3
变更
2
显示空白变更内容
内联
并排
Showing
2 changed file
with
73 addition
and
11 deletion
+73
-11
imperative/python/test/unit/core/test_subgraph.py
imperative/python/test/unit/core/test_subgraph.py
+57
-0
src/core/impl/utils/cuda_helper.cpp
src/core/impl/utils/cuda_helper.cpp
+16
-11
未找到文件。
imperative/python/test/unit/core/test_subgraph.py
浏览文件 @
e6943017
import
functools
import
os
import
platform
import
subprocess
import
sys
import
numpy
as
np
import
pytest
...
...
@@ -152,3 +156,56 @@ def test_subgraph_jit_backward():
y1
=
x1
*
x1
y2
=
mul
(
x2
,
x2
)
gm
.
backward
(
y2
)
@
pytest
.
mark
.
skipif
(
platform
.
system
()
!=
"Linux"
,
reason
=
"jit fusion is only available on Linux"
,
)
def
test_subgraph_jit
():
prog
=
"""
import megengine
import numpy as np
from megengine.core.tensor.utils import subgraph_fn
# 3 * 4 * 5 > MEGDNN_MAX_NDIM
x_np = np.random.rand(3, 4, 5).astype("float32")
x1 = megengine.Tensor(x_np)
x2 = megengine.Tensor(x_np)
@subgraph_fn(
"Mul",
dtype=x1.dtype,
device=x1.device,
nr_inputs=2,
gopt_level=None,
jit_fusion=True,
custom_grad=True,
)
def mul(inputs, f, c):
x, y = inputs[0:2]
z = f("*", x, y)
(dz,) = yield (z,)
dx = f("*", dz, y)
dy = f("*", dz, x)
yield (dx, dy)
y, = mul(x1, x2)
# ensure execution
y.numpy()
"""
env
=
dict
(
os
.
environ
)
if
"PATH"
in
env
:
# remove nvcc from environ["PATH"]
path
=
env
[
"PATH"
]
paths
=
path
.
split
(
os
.
pathsep
)
paths
=
[
path
for
path
in
paths
if
not
(
os
.
path
.
isdir
(
path
)
and
"nvcc"
in
os
.
listdir
(
path
))
]
path
=
os
.
pathsep
.
join
(
paths
)
env
[
"PATH"
]
=
path
# previous program may be stored in persistent cache
env
[
"MGE_FASTRUN_CACHE_TYPE"
]
=
"MEMORY"
subprocess
.
check_call
([
sys
.
executable
,
"-c"
,
prog
],
env
=
env
)
src/core/impl/utils/cuda_helper.cpp
浏览文件 @
e6943017
...
...
@@ -17,6 +17,8 @@ using namespace mgb;
#include <unistd.h>
#endif
namespace
{
#ifndef PATH_MAX
#define PATH_MAX 4096
#endif
...
...
@@ -32,6 +34,7 @@ using namespace mgb;
#define PATH_SPLITER '\\'
#define ENV_PATH "Path"
#define NVCC_EXE "nvcc.exe"
void
*
dlopen
(
const
char
*
file
,
int
)
{
return
static_cast
<
void
*>
(
LoadLibrary
(
file
));
}
...
...
@@ -108,15 +111,15 @@ std::string find_file_in_envs_with_intmd(
std
::
string
get_nvcc_root_path
()
{
auto
nvcc_root_path
=
find_file_in_envs_with_intmd
({
ENV_PATH
},
NVCC_EXE
);
if
(
nvcc_root_path
.
empty
())
{
mgb_throw
(
MegBrainError
,
"nvcc not found. Add your nvcc to your environment Path"
);
return
{};
}
else
{
auto
idx
=
nvcc_root_path
.
rfind
(
PATH_SPLITER
);
return
nvcc_root_path
.
substr
(
0
,
idx
+
1
);
}
}
}
// namespace
std
::
vector
<
std
::
string
>
mgb
::
get_cuda_include_path
()
{
#if MGB_CUDA
std
::
vector
<
std
::
string
>
paths
;
...
...
@@ -131,6 +134,7 @@ std::vector<std::string> mgb::get_cuda_include_path() {
// 2. use nvcc path
auto
nvcc_path
=
get_nvcc_root_path
();
if
(
!
nvcc_path
.
empty
())
{
auto
cudart_header_path
=
nvcc_path
+
".."
+
PATH_SPLITER
+
"include"
+
PATH_SPLITER
+
"cuda_runtime.h"
;
//! double check path_to_nvcc/../include/cuda_runtime.h exists
...
...
@@ -139,6 +143,7 @@ std::vector<std::string> mgb::get_cuda_include_path() {
paths
.
emplace_back
(
nvcc_path
+
".."
);
paths
.
emplace_back
(
nvcc_path
+
".."
+
PATH_SPLITER
+
"include"
);
}
}
// 3. use libcudart.so library path
char
cuda_lib_path
[
PATH_MAX
];
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录