Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
MegEngine 天元
MegEngine
提交
b7c9361f
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,发现更多精彩内容 >>
提交
b7c9361f
编写于
3月 17, 2022
作者:
M
Megvii Engine Team
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
perf(mge/functional): add infer_output_attrs_fallible for some ops
GitOrigin-RevId: 33ae4b18e9038469170cde4ecc63427247d40f4b
上级
a4327c4d
变更
1
隐藏空白更改
内联
并排
Showing
1 changed file
with
84 addition
and
0 deletion
+84
-0
imperative/src/impl/ops/specializations.cpp
imperative/src/impl/ops/specializations.cpp
+84
-0
未找到文件。
imperative/src/impl/ops/specializations.cpp
浏览文件 @
b7c9361f
...
...
@@ -61,6 +61,44 @@ auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) {
return
opr
::
Dimshuffle
::
make
(
inputs
[
0
],
ds
.
pattern
,
0UL
,
config
);
}
std
::
tuple
<
SmallVector
<
LogicalTensorDesc
>
,
bool
>
infer_output_attrs_fallible
(
const
OpDef
&
def
,
const
SmallVector
<
LogicalTensorDesc
>&
inputs
)
{
auto
&&
ds
=
static_cast
<
const
Dimshuffle
&>
(
def
);
mgb_assert
(
ds
.
pattern
.
size
()
<=
TensorShape
::
MAX_NDIM
,
"Dimshuffle pattern exceeds max length of %zd"
,
TensorShape
::
MAX_NDIM
);
size_t
nr_inp
=
inputs
.
size
();
mgb_assert
(
nr_inp
==
1
,
"Dimshuffle expects 1 inputs; got %lu actually"
,
nr_inp
);
auto
&&
src
=
inputs
[
0
];
TensorShape
out_shape
;
if
(
src
.
layout
.
ndim
==
0
)
{
return
{{{
TensorLayout
(
out_shape
,
src
.
layout
.
dtype
),
src
.
comp_node
}},
false
};
}
size_t
pattern_ndim
=
*
std
::
max_element
(
ds
.
pattern
.
begin
(),
ds
.
pattern
.
end
())
+
1
;
mgb_assert
(
src
.
layout
.
ndim
==
pattern_ndim
,
"input ndim mismatch for Dimshuffle: expect=%zd actual=%zd"
,
pattern_ndim
,
src
.
layout
.
ndim
);
size_t
idx
=
0
;
bool
input_used
[
TensorLayout
::
MAX_NDIM
]
=
{
0
};
for
(
auto
i
:
ds
.
pattern
)
{
if
(
i
<
0
)
{
out_shape
[
idx
]
=
1
;
}
else
{
input_used
[
i
]
=
true
;
out_shape
[
idx
]
=
src
.
layout
.
shape
[
i
];
}
++
idx
;
}
for
(
size_t
i
=
0
;
i
<
pattern_ndim
;
++
i
)
{
mgb_assert
(
input_used
[
i
]
||
src
.
layout
.
shape
[
i
]
==
1
,
"non-1 dim discarded in Dimshuffle: ishp=%s dim=%zd"
,
src
.
layout
.
megdnn
::
TensorShape
::
to_string
().
c_str
(),
i
);
}
return
{{{
TensorLayout
(
out_shape
,
src
.
layout
.
dtype
),
src
.
comp_node
}},
true
};
}
SmallVector
<
TensorPtr
>
apply_on_physical_tensor
(
const
OpDef
&
def
,
const
SmallVector
<
TensorPtr
>&
inputs
,
SmallVector
<
LogicalTensorDesc
>&
output_descs
,
const
bool
&
validated
)
{
...
...
@@ -110,6 +148,7 @@ OP_TRAIT_REG(Dimshuffle, Dimshuffle, opr::Dimshuffle)
.
make_from_op_node
(
make_from_op_node
)
.
apply_on_var_node
(
apply_on_var_node
)
.
apply_on_physical_tensor
(
apply_on_physical_tensor
)
.
infer_output_attrs_fallible
(
infer_output_attrs_fallible
)
.
fallback
();
}
// namespace dimshuffle
}
// namespace
...
...
@@ -127,6 +166,22 @@ auto apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) {
return
opr
::
AxisAddRemove
::
make
(
inputs
[
0
],
param
,
config
);
}
std
::
tuple
<
SmallVector
<
LogicalTensorDesc
>
,
bool
>
infer_output_attrs_fallible
(
const
OpDef
&
def
,
const
SmallVector
<
LogicalTensorDesc
>&
inputs
)
{
auto
&&
op_def
=
def
.
cast_final_safe
<
AddAxis
>
();
size_t
nr_inp
=
inputs
.
size
();
mgb_assert
(
nr_inp
==
1
,
"AddAxis expects 1 inputs; got %lu actually"
,
nr_inp
);
auto
&&
src
=
inputs
[
0
];
auto
olayout
=
src
.
layout
;
if
(
src
.
layout
.
ndim
==
0
)
{
return
{{{
TensorLayout
(
src
.
layout
.
dtype
),
src
.
comp_node
}},
false
};
}
for
(
auto
&&
i
:
op_def
.
axis
)
{
olayout
.
add_axis_cont_inplace
(
i
);
}
return
{{{
olayout
,
src
.
comp_node
}},
true
};
}
SmallVector
<
TensorPtr
>
apply_on_physical_tensor
(
const
OpDef
&
def
,
const
SmallVector
<
TensorPtr
>&
inputs
,
SmallVector
<
LogicalTensorDesc
>&
output_descs
,
const
bool
&
validated
)
{
...
...
@@ -145,6 +200,7 @@ SmallVector<TensorPtr> apply_on_physical_tensor(
OP_TRAIT_REG
(
AddAxis
,
AddAxis
)
.
apply_on_var_node
(
apply_on_var_node
)
.
apply_on_physical_tensor
(
apply_on_physical_tensor
)
.
infer_output_attrs_fallible
(
infer_output_attrs_fallible
)
.
fallback
();
}
// namespace add_axis
}
// namespace
...
...
@@ -188,9 +244,37 @@ SmallVector<TensorPtr> apply_on_physical_tensor(
return
{
Tensor
::
make
(
src
->
blob
(),
src
->
offset
(),
tlayout
)};
}
std
::
tuple
<
SmallVector
<
LogicalTensorDesc
>
,
bool
>
infer_output_attrs_fallible
(
const
OpDef
&
def
,
const
SmallVector
<
LogicalTensorDesc
>&
inputs
)
{
auto
&&
op_def
=
def
.
cast_final_safe
<
RemoveAxis
>
();
size_t
nr_inp
=
inputs
.
size
();
mgb_assert
(
nr_inp
==
1
,
"RemoveAxis expects 1 inputs; got %lu actually"
,
nr_inp
);
auto
&&
src
=
inputs
[
0
];
auto
olayout
=
src
.
layout
;
if
(
src
.
layout
.
ndim
==
0
)
{
return
{{{
TensorLayout
(
src
.
layout
.
dtype
),
src
.
comp_node
}},
false
};
}
for
(
auto
&&
i
:
op_def
.
axis
)
{
if
(
olayout
.
ndim
==
1
)
{
mgb_assert
(
olayout
.
shape
[
0
]
==
1
&&
i
==
0
,
"can not remove axis %u from tensor of shape=%s"
,
i
,
olayout
.
megdnn
::
TensorShape
::
to_string
().
c_str
());
}
else
{
mgb_assert
(
i
<
olayout
.
ndim
&&
olayout
.
shape
[
i
]
==
1
,
"can not remove axis %u from tensor of shape=%s"
,
i
,
olayout
.
megdnn
::
TensorShape
::
to_string
().
c_str
());
olayout
.
remove_axis_inplace
(
i
);
}
}
return
{{{
olayout
,
src
.
comp_node
}},
true
};
}
OP_TRAIT_REG
(
RemoveAxis
,
RemoveAxis
)
.
apply_on_var_node
(
apply_on_var_node
)
.
apply_on_physical_tensor
(
apply_on_physical_tensor
)
.
infer_output_attrs_fallible
(
infer_output_attrs_fallible
)
.
fallback
();
}
// namespace remove_axis
}
// namespace
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录