huawei
/
mindspore2022
0
0
Fork 53
Code Issues Pull Requests 29 Projects Releases Wiki Activity

代码评注赛初阶段成果提交 “你的代码写的真棒”队 #3

Open
a946079208 wants to merge 3 commits from a946079208/mindspore2022:doc_reduce_sum into master
pull from: a946079208/mindspore2022:doc_reduce_sum
merge into: huawei:master
Conversation 0 Commits 3 Files Changed 2 +84
2 changed files with 84 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
mindspore/ccsrc/plugin/device/cpu/kernel/reduce_cpu_kernel.cc
View File

@ -36,6 +36,12 @@ constexpr auto kReduceProdName = "ReduceProd";
constexpr auto kReduceAllName = "ReduceAll";
constexpr auto kReduceAnyName = "ReduceAny";
/**
提供统一的基类 Reduce，根据不同的 Reduce 类型进行特化
InitFunc 方法用于初始化相关参数，比如输入的 shape，axis，reduce_type，对应的函数等
RunFunc 执行 Kernel 的计算。
AccelerateLongVector 针对长向量做了针对性的优化
*/
template <typename T>
class ReduceCpuKernelFunc : public CpuKernelFunc {
public:
@ -65,6 +71,9 @@ class ReduceCpuKernelFunc : public CpuKernelFunc {
std::string kernel_name_;
};
/**
更新输入的参数 axis 为用户指定的需要保留规约的维度
*/
void UpdateAxis(const PrimitivePtr &prim, const CNodePtr &kernel_node, const std::string &kernel_name,
std::vector<int64_t> *axis) {
auto axis_addr = prim->GetAttr(AXIS);
@ -94,8 +103,10 @@ void ReduceCpuKernelFunc<T>::InitFunc(const CNodePtr &kernel_node) {
MS_EXCEPTION_IF_NULL(prim);
UpdateAxis(prim, kernel_node, kernel_name_, &axis_);
size_t dimension = input_shape_.size();
// 这意味着输入 axis 有效范围只能在 [-dimension, dimension-1] 之间
(void)std::transform(axis_.begin(), axis_.end(), axis_.begin(),
[dimension](const auto &a) { return a < 0 ? dimension + a : a; });
// 为了后面可以取出重复元素，先排序，这样相同的元素就会连续放置，所以可以使用 unique
sort(axis_.begin(), axis_.end());
// Delete the duplicate axis.
auto last = std::unique(axis_.begin(), axis_.end());
@ -152,6 +163,10 @@ bool ReduceCpuKernelFunc<T>::RunFunc(const std::vector<kernel::AddressPtr> &inpu
auto *input_addr = reinterpret_cast<T *>(inputs[0]->addr);
auto *output_addr = reinterpret_cast<T *>(outputs[0]->addr);
if (axis_.empty() || input_shape_.empty() || input_shape_.size() == 1) {
// 规约所有元素的条件：
// 1. 没有需要规约的 axis
// 2. 输入是一个没有 shape 的 tensor
// 3. 输入的 shape 只有一个维度，即一维向量
if (input_size < kReduceSmallVectorSize) {
// Get one ret
*output_addr = input_addr[0];
@ -162,6 +177,7 @@ bool ReduceCpuKernelFunc<T>::RunFunc(const std::vector<kernel::AddressPtr> &inpu
*output_addr /= input_size;
}
} else {
// 特化，加速长向量的情况
AccelerateLongVector(input_addr, output_addr, input_size);
}
} else {
@ -185,17 +201,29 @@ bool ReduceCpuKernelFunc<T>::RunFunc(const std::vector<kernel::AddressPtr> &inpu
++k;
}
// 以上过程的一个运算例子
// 输入的 input_shape_: [3, 2, 4, 5], axis_: [0, 2]
// 运算过程:
// i = 0, j = 0, k = 0; false[0 == 2] || false[0 != 0]; stride = 3
// i = 1, j = 1, k = 0; false[1 == 2] || true[1 != 2]; axes[0] = 1
// i = 2, j = 1, k = 1; true[1 == 2] || false[2 != 2]; stride = 12
// i = 3, j = 2, k = 1; true[2 == 2] || ?; axes[1] = 3
// 输出的 axes: [1, 3, 0, 2], stride = 12
size_t output_size = outputs[0]->size / sizeof(T);
if constexpr (std::is_same<T, float>::value) {
if (simple_execute_) {
auto task = [&](size_t start, size_t end) {
for (size_t i = start; i < end; ++i) {
// 规约 axis=1 input_shape.size=2 的情况：可以理解为二维矩阵按照行压成一维
// 因此，input_addr 需要加上 i * stride 个元素，output_addr 的第 i 个位置是规约结果
(void)ReduceSumDim2Axis1(stride, input_addr + i * stride, output_addr + i);
if (reduce_type_ == ReduceFuncType::kReduceMeanType) {
output_addr[i] /= stride;
}
}
};
// 自动并行 [0, output_size) 自动切分，然后调用 task 执行
ParallelLaunchAutoSearch(task, output_size, this, &parallel_search_info_);
return true;
}
@ -210,8 +238,10 @@ bool ReduceCpuKernelFunc<T>::RunFunc(const std::vector<kernel::AddressPtr> &inpu
auto iter = base_iter;
iter.SetPos(start * stride);
for (size_t i = start; i < end; ++i) {
// output_addr 指向第一个元素
output_addr[i] = input_addr[iter.GetPos()];
iter.GenNextPos();
// 剩下的元素调用规约函数迭代计算即可
for (size_t j = 1; j < stride; ++j) {
reduce_func_(input_addr, iter.GetPos(), &output_addr[i]);
iter.GenNextPos();
@ -226,12 +256,19 @@ bool ReduceCpuKernelFunc<T>::RunFunc(const std::vector<kernel::AddressPtr> &inpu
return true;
}
/**
AccelerateLongVector 为长向量专门优化：
使用 ParallelLaunchAutoSearch 自动并行，需要加锁避免可能存在的数据冲突
需要注意 AccelerateLongVector 只能对规约所有元素使用
在对短向量计算的时候，是单线程的；长向量是自动并行，多线程的。
*/
template <typename T>
void ReduceCpuKernelFunc<T>::AccelerateLongVector(T *input_addr, T *output_addr, size_t input_size) {
// init output_addr
*output_addr = input_addr[0];
std::mutex task_mutex;
auto task = [this, input_addr, output_addr, &task_mutex](size_t start, size_t end) {
// 跳过第 0 号元素，因为 output_addr 的初值设置为了 input_addr[0]
if (start == 0) {
++start;
}
@ -245,6 +282,8 @@ void ReduceCpuKernelFunc<T>::AccelerateLongVector(T *input_addr, T *output_addr,
++i;
}
{
// 将结果汇总到 output_addr，这里是规约了所有元素，所以 output_addr 只有一个元素
// 因此在访问同一个位置的时候，需要加锁
std::lock_guard<std::mutex> task_lock(task_mutex);
reduce_func_(&block_output, 0, output_addr);
}

45
mindspore/python/mindspore/ops/operations/math_ops.py
View File

@ -648,6 +648,11 @@ class ReduceMean(_Reduce):
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.ops as ops
>>> from mindspore import Tensor
>>>
>>> x = Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))
>>> op = ops.ReduceMean(keep_dims=True)
>>> output = op(x, 1)
@ -725,6 +730,11 @@ class ReduceSum(_Reduce):
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.ops as ops
>>> from mindspore import Tensor
>>>
>>> x = Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))
>>> op = ops.ReduceSum(keep_dims=True)
>>> output = op(x, 1)
@ -763,6 +773,16 @@ class ReduceSum(_Reduce):
[[42.]
[48.]
[54.]]]
>>> # case 5: Reduces a dimension along axis 0, 2 with list.
>>> output = op(x, [0, 2])
>>> print(output)
[[[ 72.]
[ 90.]
[108.]]]
>>> # case 6: Reduces a dimension along axis 0, 1 with tuple.
>>> output = op(x, (0, 1))
>>> print(output)
[[[45. 45. 45. 45. 45. 45.]]]
"""
@prim_attr_register
@ -810,6 +830,11 @@ class ReduceAll(_Reduce):
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.ops as ops
>>> from mindspore import Tensor
>>>
>>> x = Tensor(np.array([[True, False], [True, True]]))
>>> op = ops.ReduceAll(keep_dims=True)
>>> # case 1: Reduces a dimension by the "logicalAND" of all elements in the dimension.
@ -868,6 +893,11 @@ class ReduceAny(_Reduce):
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.ops as ops
>>> from mindspore import Tensor
>>>
>>> x = Tensor(np.array([[True, False], [True, True]]))
>>> op = ops.ReduceAny(keep_dims=True)
>>> # case 1: Reduces a dimension by the "logical OR" of all elements in the dimension.
@ -926,6 +956,11 @@ class ReduceMax(_Reduce):
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.ops as ops
>>> from mindspore import Tensor
>>>
>>> x = Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))
>>> op = ops.ReduceMax(keep_dims=True)
>>> output = op(x, 1)
@ -1012,6 +1047,11 @@ class ReduceMin(_Reduce):
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.ops as ops
>>> from mindspore import Tensor
>>>
>>> x = Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))
>>> op = ops.ReduceMin(keep_dims=True)
>>> output = op(x, 1)
@ -1134,6 +1174,11 @@ class ReduceProd(_Reduce):
``Ascend`` ``GPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.ops as ops
>>> from mindspore import Tensor
>>>
>>> x = Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))
>>> op = ops.ReduceProd(keep_dims=True)
>>> output = op(x, 1)