Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
b6ce02a1
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,发现更多精彩内容 >>
提交
b6ce02a1
编写于
9月 28, 2021
作者:
M
Megvii Engine Team
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
fix(subgraph): fallback back to cg if jit unsupported
GitOrigin-RevId: 853a00a4025d6e4fefa0f3ac2fe3d4cad4c3f8b9
上级
21f5a7fc
变更
4
隐藏空白更改
内联
并排
Showing
4 changed file
with
46 addition
and
19 deletion
+46
-19
imperative/python/megengine/core/tensor/utils.py
imperative/python/megengine/core/tensor/utils.py
+5
-0
imperative/python/src/ops.cpp
imperative/python/src/ops.cpp
+5
-0
imperative/src/impl/ops/utility.cpp
imperative/src/impl/ops/utility.cpp
+32
-15
src/jit/impl/nvrtc/compiler_cuda.cpp
src/jit/impl/nvrtc/compiler_cuda.cpp
+4
-4
未找到文件。
imperative/python/megengine/core/tensor/utils.py
浏览文件 @
b6ce02a1
...
...
@@ -22,6 +22,7 @@ from .._imperative_rt.core2 import (
make_shape_tuple
,
)
from
.._imperative_rt.ops
import
SubgraphBuilder
as
_SubgraphBuilder
from
.._imperative_rt.ops
import
jit_supported
from
.._wrap
import
as_device
from
..autodiff.grad
import
Function
from
..ops
import
builtin
...
...
@@ -234,6 +235,10 @@ def subgraph(
gopt_level
=
None
# disable jit and compile
jit_fusion
=
False
if
jit_fusion
and
not
jit_supported
:
jit_fusion
=
False
# jit unusable, fallback to graph compile
gopt_level
=
2
def
as_op
(
op
,
nargs
):
if
isinstance
(
op
,
str
):
assert
(
op
,
nargs
)
in
_opr_map
,
"unknown operator"
...
...
imperative/python/src/ops.cpp
浏览文件 @
b6ce02a1
...
...
@@ -652,6 +652,11 @@ void init_ops(py::module m) {
});
m
.
def
(
"set_jit_enabled"
,
&
JITFusionOp
::
set_enabled
);
bool
jit_supported
=
false
;
#if MGB_JIT
jit_supported
=
true
;
#endif
m
.
attr
(
"jit_supported"
)
=
jit_supported
;
auto
custom
=
submodule
(
m
,
"_custom"
);
init_custom
(
custom
);
...
...
imperative/src/impl/ops/utility.cpp
浏览文件 @
b6ce02a1
...
...
@@ -9,21 +9,26 @@
* "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*/
#include <
que
ue>
#include <
deq
ue>
#include "../op_trait.h"
#include "megbrain/imperative/graph_cache.h"
#include "megbrain/imperative/opr_utility.h"
#include "megbrain/imperative/ops/autogen.h"
#include "megbrain/imperative/ops/opr_attr.h"
#include "megbrain/imperative/ops/utility.h"
#include "megbrain/imperative/subgraph_detail.h"
#include "megbrain/jit/executor_opr.h"
#include "megbrain/opr/io.h"
#include "megbrain/opr/tensor_gen.h"
#include "megbrain/opr/tensor_manip.h"
#include "megbrain/opr/utility.h"
#if MGB_JIT
#include "megbrain/jit/executor_opr.h"
#endif
#include "../event_pool.h"
#include "../op_trait.h"
namespace
mgb
::
imperative
{
MGB_DYN_TYPE_OBJ_FINAL_IMPL
(
GenericPyOp
);
...
...
@@ -309,7 +314,7 @@ struct ComputingGraphHolder {
SmallVector
<
VarNode
*>
input_vars
;
SmallVector
<
VarNode
*>
output_vars
;
std
::
shared_ptr
<
DeviceMemoryAllocatorImpl
>
allocator
;
SmallVector
<
std
::
unique
_ptr
<
CompNode
::
Event
>>
events
;
SmallVector
<
std
::
shared
_ptr
<
CompNode
::
Event
>>
events
;
std
::
unique_ptr
<
cg
::
static_infer
::
StaticInferUpdater
>
updater
;
void
initialize
(
...
...
@@ -402,7 +407,7 @@ struct ComputingGraphHolder {
return
true
;
});
for
(
auto
&&
comp_node
:
comp_nodes
)
{
events
.
push_back
(
comp_node
.
create_event
(
));
events
.
push_back
(
EventPool
::
without_timer
().
alloc_shared
(
comp_node
));
events
.
back
()
->
record
();
}
}
...
...
@@ -510,7 +515,7 @@ ComputingGraphHolder<Kind>& get_computing_graph(
std
::
shared_ptr
<
OpDef
>
compiled_op
,
const
SmallVector
<
LogicalTensorDesc
>&
descs
)
{
using
ComputingGraphHolderCache
=
OpMethResultCache
<
std
::
que
ue
<
std
::
unique_ptr
<
ComputingGraphHolder
<
Kind
>>>>
;
OpMethResultCache
<
std
::
deq
ue
<
std
::
unique_ptr
<
ComputingGraphHolder
<
Kind
>>>>
;
thread_local
auto
cache
=
std
::
make_unique
<
ComputingGraphHolderCache
>
();
thread_local
size_t
nr_cg_holders
=
0
;
typename
ComputingGraphHolderCache
::
key_t
cache_key
=
{
compiled_op
,
descs
};
...
...
@@ -540,20 +545,28 @@ ComputingGraphHolder<Kind>& get_computing_graph(
}
}
if
(
holder
)
{
cg_holder_queue
.
pop
();
cg_holder_queue
.
pop
_front
();
}
}
if
(
!
holder
)
{
// create new computing graph
holder
=
std
::
make_unique
<
ComputingGraphHolder
<
Kind
>>
();
auto
&
cg_holder
=
*
holder
;
cg_holder
.
initialize
(
compiled_op
->
cast_final_safe
<
CompiledOp
>
(),
descs
);
nr_cg_holders
++
;
mgb_log_debug
(
"add new computing graph for compiled op, now %zu graphs"
,
nr_cg_holders
);
auto
create_holder
=
[
&
]
{
auto
holder
=
std
::
make_unique
<
ComputingGraphHolder
<
Kind
>>
();
auto
&
cg_holder
=
*
holder
;
cg_holder
.
initialize
(
compiled_op
->
cast_final_safe
<
CompiledOp
>
(),
descs
);
nr_cg_holders
++
;
mgb_log_debug
(
"add new computing graph for compiled op, now %zu graphs"
,
nr_cg_holders
);
return
holder
;
};
size_t
nr_graphs
=
std
::
max
(
cg_holder_queue
.
size
(),
(
size_t
)
1
);
for
(
size_t
i
=
1
;
i
<
nr_graphs
;
++
i
)
{
cg_holder_queue
.
push_front
(
create_holder
());
}
holder
=
create_holder
();
}
cg_holder_queue
.
push
(
std
::
move
(
holder
));
cg_holder_queue
.
push
_back
(
std
::
move
(
holder
));
return
*
cg_holder_queue
.
back
();
}
...
...
@@ -670,6 +683,7 @@ auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) {
// skip for dump (JITExecutor can not be dumped)
return
outputs
;
}
#if MGB_JIT
for
(
auto
&
output
:
outputs
)
{
jit
::
InternalGraphGenerator
igg
{
output
->
owner_opr
()};
std
::
vector
<
cg
::
OperatorNodeBase
*>
reverse_order
;
...
...
@@ -686,6 +700,9 @@ auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) {
auto
ig
=
igg
.
generate
();
output
=
jit
::
JITExecutor
::
make
(
ig
,
igg
.
orig_inps
()).
node
();
}
#else
mgb_assert
(
false
,
"MGB_WITH_JIT was disabled"
);
#endif
return
outputs
;
}
...
...
src/jit/impl/nvrtc/compiler_cuda.cpp
浏览文件 @
b6ce02a1
...
...
@@ -216,11 +216,11 @@ void CudaExecutable::FuncCache::compile(
ptx
=
NVRTCCompile
(
cuda_exe
->
m_source
,
major
,
minor
);
ptx_cache
=
PersistentCache
::
Blob
{
ptx
.
data
(),
ptx
.
size
()};
cache
.
put
(
cache_category
,
key
,
ptx_cache
.
val
());
mgb_log
(
"NVRTC JIT: compile %s for %d.%d: source_len=%zu ptx_len=%zu "
"time=%.3fms"
,
cuda_exe
->
m_name
.
c_str
(),
major
,
minor
,
key
.
size
,
ptx
.
size
(),
timer
.
get_msecs
());
}
mgb_log
(
"NVRTC JIT: compile %s for %d.%d: source_len=%zu ptx_len=%zu "
"time=%.3fms"
,
cuda_exe
->
m_name
.
c_str
(),
major
,
minor
,
key
.
size
,
ptx
.
size
(),
timer
.
get_msecs
());
}
void
CudaExecutable
::
FuncCache
::
exec
(
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录