Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
ascend950当前不支持ReduceMeanWithCount算子,为打通功能断点,需要支持该算子
打通950功能断点
Aclnn直调
fp32,fp16, bf16
使用reduce模板tiling
使用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>;
};
Ascend950
无 💡 备注(选填)
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图如下:, Placeholder::In0>;
struct ReduceMeanWithCountDag {
// Input 0: x
using OpCopyInX = Bind<Vec::CopyIn
using CastX = Bind<Vec::Cast<PromteT, T, 0>, OpCopyInX>;
};
3.3 支持硬件
Ascend950
3.4 算子约束限制
无
💡 备注(选填)