Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
2b80806f
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,发现更多精彩内容 >>
提交
2b80806f
编写于
3月 08, 2022
作者:
M
Megvii Engine Team
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
perf(imperative/src): improve dot performance
GitOrigin-RevId: 35b5bd164ffde647125f8fa9c0ebf91195dc4f1f
上级
2f3bc2db
变更
4
隐藏空白更改
内联
并排
Showing
4 changed file
with
69 addition
and
4 deletion
+69
-4
dnn/include/megdnn/oprs/linalg.h
dnn/include/megdnn/oprs/linalg.h
+2
-1
dnn/src/common/dot.cpp
dnn/src/common/dot.cpp
+1
-1
imperative/src/impl/ops/specializations.cpp
imperative/src/impl/ops/specializations.cpp
+65
-1
src/opr/include/megbrain/opr/blas.h
src/opr/include/megbrain/opr/blas.h
+1
-1
未找到文件。
dnn/include/megdnn/oprs/linalg.h
浏览文件 @
2b80806f
...
...
@@ -150,7 +150,8 @@ public:
virtual
void
exec
(
_megdnn_tensor_in
A
,
_megdnn_tensor_in
B
,
_megdnn_tensor_out
C
,
_megdnn_workspace
workspace
)
=
0
;
void
deduce_layout
(
const
TensorLayout
&
A
,
const
TensorLayout
&
B
,
TensorLayout
&
C
);
MGE_WIN_DECLSPEC_FUC
void
deduce_layout
(
const
TensorLayout
&
A
,
const
TensorLayout
&
B
,
TensorLayout
&
C
);
virtual
size_t
get_workspace_in_bytes
(
const
TensorLayout
&
A
,
const
TensorLayout
&
B
,
const
TensorLayout
&
C
)
=
0
;
...
...
dnn/src/common/dot.cpp
浏览文件 @
2b80806f
...
...
@@ -33,7 +33,7 @@ void DotForward::check_exec(
megdnn_assert
(
workspace_in_bytes
>=
required_workspace_in_bytes
);
}
void
DotForward
::
deduce_layout
(
MGE_WIN_DECLSPEC_FUC
void
DotForward
::
deduce_layout
(
const
TensorLayout
&
A
,
const
TensorLayout
&
,
TensorLayout
&
C
)
{
C
=
TensorLayout
(
TensorShape
{
1
},
A
.
dtype
);
}
...
...
imperative/src/impl/ops/specializations.cpp
浏览文件 @
2b80806f
...
...
@@ -39,6 +39,7 @@
#include "megbrain/opr/tensor_manip.h"
#include "megbrain/opr/utility.h"
#include "../blob_manager_impl.h"
#include "../op_trait.h"
namespace
mgb
::
imperative
{
...
...
@@ -319,7 +320,70 @@ auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) {
OperatorNodeConfig
config
{
op
.
make_name
()};
return
opr
::
Dot
::
make
(
inputs
[
0
],
inputs
[
1
],
config
);
}
OP_TRAIT_REG
(
Dot
,
Dot
).
apply_on_var_node
(
apply_on_var_node
).
fallback
();
// std::shared_ptr<OpDef> make_from_op_node(cg::OperatorNodeBase* node_) {
// auto* node = &node_->cast_final_safe<opr::Dot>();
// return Dot::make(node->param());
// }
SmallVector
<
TensorPtr
>
apply_on_physical_tensor
(
const
OpDef
&
def
,
const
SmallVector
<
TensorPtr
>&
inputs
,
SmallVector
<
LogicalTensorDesc
>&
output_descs
,
const
bool
&
validated
)
{
auto
a
=
inputs
[
0
]
->
layout
();
auto
comp_node
=
inputs
[
0
]
->
comp_node
();
using
TensorND
=
megdnn
::
TensorND
;
SmallVector
<
TensorND
>
inp_tensornds
;
inp_tensornds
.
reserve
(
inputs
.
size
());
auto
dnn_opr
=
opr
::
intl
::
create_megdnn_opr
<
megdnn
::
Dot
>
(
comp_node
);
for
(
unsigned
i
=
0
;
i
<
inputs
.
size
();
++
i
)
{
auto
dnn_ten
=
inputs
[
i
]
->
dnn_tensor
();
inp_tensornds
.
push_back
(
dnn_ten
);
}
TensorLayout
oup_layout
{
inputs
[
0
]
->
dtype
()};
auto
inp1_tensor
=
inputs
[
0
]
->
dnn_tensor
();
auto
inp2_tensor
=
inputs
[
1
]
->
dnn_tensor
();
dnn_opr
->
deduce_layout
(
inp1_tensor
.
layout
,
inp2_tensor
.
layout
,
oup_layout
);
if
(
inputs
[
0
]
->
layout
().
is_empty
()
||
inputs
[
1
]
->
layout
().
is_empty
())
{
auto
fill_opr
=
opr
::
intl
::
create_megdnn_opr
<
megdnn
::
Fill
>
(
comp_node
);
DeviceTensorND
out
=
BlobManager
::
inst
()
->
alloc_workspace_with_defrag
(
comp_node
,
oup_layout
);
fill_opr
->
param
()
=
0
;
fill_opr
->
exec
(
out
.
as_megdnn
(),
{});
return
{
Tensor
::
make
(
out
)};
}
auto
wk_size
=
dnn_opr
->
get_workspace_in_bytes
(
inp_tensornds
[
0
].
layout
,
inp_tensornds
[
1
].
layout
,
output_descs
[
0
].
layout
);
DeviceTensorND
out_devtensor
=
BlobManager
::
inst
()
->
alloc_workspace_with_defrag
(
comp_node
,
oup_layout
);
TensorLayout
wk_layout
{
TensorShape
{
wk_size
},
inputs
[
0
]
->
dtype
()};
DeviceTensorND
workspace
=
BlobManager
::
inst
()
->
alloc_workspace_with_defrag
(
comp_node
,
wk_layout
);
megdnn
::
Workspace
dnn_wk
(
workspace
.
raw_ptr
(),
wk_size
);
dnn_opr
->
exec
(
inp_tensornds
[
0
],
inp_tensornds
[
1
],
out_devtensor
.
as_megdnn
(),
dnn_wk
);
return
{
Tensor
::
make
(
out_devtensor
)};
}
std
::
tuple
<
SmallVector
<
LogicalTensorDesc
>
,
bool
>
infer_output_attrs_fallible
(
const
OpDef
&
def
,
const
SmallVector
<
LogicalTensorDesc
>&
inputs
)
{
auto
&&
op_def
=
def
.
cast_final_safe
<
Dot
>
();
SmallVector
<
LogicalTensorDesc
>
dests
(
1
);
dests
[
0
].
layout
=
TensorLayout
(
TensorShape
{
1
},
inputs
[
0
].
layout
.
dtype
);
dests
[
0
].
comp_node
=
inputs
[
0
].
comp_node
;
return
{
dests
,
true
};
}
OP_TRAIT_REG
(
Dot
,
Dot
,
opr
::
Dot
)
.
apply_on_var_node
(
apply_on_var_node
)
.
infer_output_attrs_fallible
(
infer_output_attrs_fallible
)
.
apply_on_physical_tensor
(
apply_on_physical_tensor
)
.
fallback
();
}
// namespace dot
}
// namespace
...
...
src/opr/include/megbrain/opr/blas.h
浏览文件 @
2b80806f
...
...
@@ -88,7 +88,7 @@ private:
/*!
* \brief dot product of two tensors
*/
MGB_DEFINE_OPR_CLASS
(
MGB_DEFINE_OPR_CLASS
_WITH_EXPORT
(
Dot
,
cg
::
SingleCNOperatorNodeBaseT
<
mixin
::
MegDNNOprHolderImpl
<
megdnn
::
Dot
>>
)
// {
public
:
MGE_WIN_DECLSPEC_FUC
Dot
(
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录