已关闭
[Requirement|需求建议]: ascend950支持ReduceMeanWithCount算子 #1630
chenxingyu18创建于  5月21日关闭于  5月27日
chenxingyu18成员
5月21日 创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

一、背景信息 (必填)

ascend950当前不支持ReduceMeanWithCount算子,为打通功能断点,需要支持该算子

二、价值/作用 (必填)

打通950功能断点

三、设计方案 (必填)

3.1 使能方式(涉及哪些框架:如Aclnn直调、Pytorch训练等)

Aclnn直调

3.2 总体设计
3.2.1 算子支持的数据类型

fp32,fp16, bf16

3.2.2 host侧设计

使用reduce模板tiling

3.2.3 kernel侧设计

使用reduce模板,dag图如下:
struct ReduceMeanWithCountDag {
// Input 0: x
using OpCopyInX = Bind<Vec::CopyIn, Placeholder::In0>;
using CastX = Bind<Vec::Cast<PromteT, T, 0>, OpCopyInX>;

// Input 1: count (same shape as x, element-wise weight)
using OpCopyInCount    = Bind<Vec::CopyIn<T>, Placeholder::In1<T>>;
using CastCount        = Bind<Vec::Cast<PromteT, T, 0>, OpCopyInCount>;

// Input 2: count_sum (same shape as x, element-wise divisor)
using OpCopyInCountSum = Bind<Vec::CopyIn<T>, Placeholder::In2<T>>;
using CastCountSum     = Bind<Vec::Cast<PromteT, T, 0>, OpCopyInCountSum>;

// Pre-processing: weighted = x * count / count_sum
using OpMul            = Bind<Vec::Mul<PromteT>, CastX, CastCount>;
using OpDiv            = Bind<Vec::Div<PromteT>, OpMul, CastCountSum>;

// Reduce: ReduceSum on weighted data
using ReduceOp0        = Bind<Vec::ReduceSumOp<PromteT>, OpDiv>;

// Output
using Cast1            = Bind<Vec::Cast<T, PromteT, 1>, ReduceOp0>;
using OpCopyOut        = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, Cast1>;

using Outputs = Elems<OpCopyOut>;
using MemCfg  = MemOptCfg<MemLevel::LEVEL_2>;
using OpDag   = DAGSch<Outputs, void, MemCfg>;

};

3.3 支持硬件

Ascend950

3.4 算子约束限制


💡 备注(选填)

likedislike
Cchenxingyu18成员
5月21日 添加了label:requirement
Cchenxingyu18成员
5月21日 将 chenxingyu18 设为负责人
CANN-robotCANN-robot成员
5月27日 关闭了 issue
CANN-robotCANN-robot成员
5月27日 添加了label:resolved