diff --git a/mindspore/ccsrc/plugin/device/cpu/kernel/reduce_cpu_kernel.cc b/mindspore/ccsrc/plugin/device/cpu/kernel/reduce_cpu_kernel.cc index ab1feff4d5a..560f216817d 100644 --- a/mindspore/ccsrc/plugin/device/cpu/kernel/reduce_cpu_kernel.cc +++ b/mindspore/ccsrc/plugin/device/cpu/kernel/reduce_cpu_kernel.cc @@ -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 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 *axis) { auto axis_addr = prim->GetAttr(AXIS); @@ -94,8 +103,10 @@ void ReduceCpuKernelFunc::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::RunFunc(const std::vector &inpu auto *input_addr = reinterpret_cast(inputs[0]->addr); auto *output_addr = reinterpret_cast(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::RunFunc(const std::vector &inpu *output_addr /= input_size; } } else { + // 特化，加速长向量的情况 AccelerateLongVector(input_addr, output_addr, input_size); } } else { @@ -185,17 +201,29 @@ bool ReduceCpuKernelFunc::RunFunc(const std::vector &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::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, ¶llel_search_info_); return true; } @@ -210,8 +238,10 @@ bool ReduceCpuKernelFunc::RunFunc(const std::vector &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::RunFunc(const std::vector &inpu return true; } +/** +AccelerateLongVector 为长向量专门优化： +使用 ParallelLaunchAutoSearch 自动并行，需要加锁避免可能存在的数据冲突 +需要注意 AccelerateLongVector 只能对规约所有元素使用 +在对短向量计算的时候，是单线程的；长向量是自动并行，多线程的。 +*/ template void ReduceCpuKernelFunc::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::AccelerateLongVector(T *input_addr, T *output_addr, ++i; } { + // 将结果汇总到 output_addr，这里是规约了所有元素，所以 output_addr 只有一个元素 + // 因此在访问同一个位置的时候，需要加锁 std::lock_guard task_lock(task_mutex); reduce_func_(&block_output, 0, output_addr); } diff --git a/mindspore/python/mindspore/ops/operations/math_ops.py b/mindspore/python/mindspore/ops/operations/math_ops.py index addaf14fca4..0a8669c4999 100644 --- a/mindspore/python/mindspore/ops/operations/math_ops.py +++ b/mindspore/python/mindspore/ops/operations/math_ops.py @@ -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)