aclnnMean/aclnnReduceSum adaption

type: feature
target: all
reason:

------

Signed-off-by: wang_ziqi <wangziqi4@huawei.com>
This commit is contained in:
wang_ziqi 2024-02-26 21:30:01 +08:00 committed by hedongdong
parent 83f322b4ba
commit 260ba5846c
47 changed files with 2314 additions and 212 deletions

View File

@ -1,6 +1,6 @@
mindspore.Tensor.mean
=====================
.. py:method:: mindspore.Tensor.mean(axis=None, keep_dims=False)
.. py:method:: mindspore.Tensor.mean(axis=None, keep_dims=False, dtype=None)
详情请参考 :func:`mindspore.ops.mean`

View File

@ -1,17 +1,18 @@
mindspore.ops.mean
==================
.. py:function:: mindspore.ops.mean(x, axis=None, keep_dims=False)
.. py:function:: mindspore.ops.mean(input, axis=None, keep_dims=False, dtype=None)
默认情况下,移除输入所有维度,返回 `x` 中所有元素的平均值。也可仅缩小指定维度 `axis` 大小至1。 `keep_dims` 控制输出和输入的维度是否相同。
默认情况下,移除输入所有维度,返回 `input` 中所有元素的平均值。也可仅缩小指定维度 `axis` 大小至1。 `keep_dims` 控制输出和输入的维度是否相同。
.. note::
Tensor类型的 `axis` 仅用作兼容旧版本,不推荐使用。
参数:
- **x** (Tensor[Number]) - 输入Tensor其数据类型为数值型。shape :math:`(N, *)` ,其中 :math:`*` 表示任意数量的附加维度。
- **axis** (Union[int, tuple(int), list(int), Tensor]) - 要减少的维度。默认值: ``None`` ,缩小所有维度。只允许常量值。假设 `x` 的秩为r取值范围[-r,r)。
- **input** (Tensor[Number]) - 输入Tensor其数据类型为数值型。shape :math:`(N, *)` ,其中 :math:`*` 表示任意数量的附加维度。
- **axis** (Union[int, tuple(int), list(int), Tensor]) - 要减少的维度。默认值: ``None`` ,缩小所有维度。只允许常量值。假设 `input` 的秩为r取值范围[-r,r)。
- **keep_dims** (bool) - 如果为 ``True`` 则保留缩小的维度大小为1。否则移除维度。默认值 ``False``
- **dtype** (:class:`mindspore.dtype`) - 期望输出Tensor的类型。默认值 ``None``
返回:
Tensor。
@ -22,7 +23,7 @@ mindspore.ops.mean
- 如果 `axis` 为一维Tensor例如取值为[1, 2],并且 `keep_dims```False`` 则输出Tensor的shape为 :math:`(x_0, x_3, ..., x_R)`
异常:
- **TypeError** - `x` 不是Tensor。
- **TypeError** - `input` 不是Tensor。
- **TypeError** - `axis` 不是以下数据类型之一int、tuple、list或Tensor。
- **TypeError** - `keep_dims` 不是bool类型。
- **ValueError** - `axis` 超出范围。

View File

@ -17,9 +17,14 @@
#include "backend/common/expander/fallback/fallback_irbuilder.h"
#include "include/common/utils/utils.h"
#include "utils/shape_utils.h"
#include "ops/op_utils.h"
namespace mindspore {
namespace expander {
namespace {
const std::set<TypeId> kIntergralSet = {kNumberTypeBool, kNumberTypeUInt8, kNumberTypeInt8, kNumberTypeInt16,
kNumberTypeInt32};
} // namespace
REG_FALLBACK_BUILDER("AddExt").SetBody(BODYFUNC(ib) {
auto x = ib->GetInput(kIndex0);
auto y = ib->GetInput(kIndex1);
@ -35,5 +40,59 @@ REG_FALLBACK_BUILDER("SubExt").SetBody(BODYFUNC(ib) {
auto alpha_tensor = ib->Cast(ib->ScalarToTensor(alpha, x->dtype()), y->dtype());
return {x - y * alpha_tensor};
});
REG_FALLBACK_BUILDER("MeanExt").SetBody(BODYFUNC(ib) {
auto input = ib->GetInput(kIndex0);
auto axis = ib->GetInput(kIndex1);
auto keep_dims = ib->GetInput(kIndex2);
auto dtype = ib->GetInput(kIndex3);
auto dtype_type = dtype->abstract()->BuildType();
MS_EXCEPTION_IF_NULL(dtype_type);
// cppcheck-suppress *
if (!dtype_type->isa<TypeNone>()) {
auto dtype_opt = ops::GetScalarValue<int64_t>(dtype->BuildValue());
MS_CHECK_VALUE(dtype_opt.has_value(), "For 'MeanExt', dtype must have valid value.");
input = ib->Cast(input, TypeIdToType(static_cast<TypeId>(dtype_opt.value())));
}
auto axis_type = axis->abstract()->BuildType();
MS_EXCEPTION_IF_NULL(axis_type);
if (axis_type->isa<TypeNone>()) {
axis = ib->Value<std::vector<int64_t>>({});
}
auto out = ib->Emit("ReduceMean", {input, axis, keep_dims});
return {out};
});
REG_FALLBACK_BUILDER("SumExt").SetBody(BODYFUNC(ib) {
auto input = ib->GetInput(kIndex0);
auto axis = ib->GetInput(kIndex1);
auto keep_dims = ib->GetInput(kIndex2);
auto dtype = ib->GetInput(kIndex3);
auto dtype_type = dtype->abstract()->BuildType();
MS_EXCEPTION_IF_NULL(dtype_type);
if (!dtype_type->isa<TypeNone>()) {
auto dtype_opt = ops::GetScalarValue<int64_t>(dtype->BuildValue());
MS_CHECK_VALUE(dtype_opt.has_value(), "For 'SumExt', dtype must have valid value.");
input = ib->Cast(input, TypeIdToType(static_cast<TypeId>(dtype_opt.value())));
} else {
auto input_type = input->dtype()->type_id();
if (kIntergralSet.find(input_type) != kIntergralSet.end()) {
input = ib->Cast(input, kInt64);
}
}
auto axis_type = axis->abstract()->BuildType();
MS_EXCEPTION_IF_NULL(axis_type);
if (axis_type->isa<TypeNone>()) {
axis = ib->Value<std::vector<int64_t>>({});
}
auto out = ib->Emit("ReduceSum", {input, axis, keep_dims, ib->Value<bool>(false)});
return {out};
});
} // namespace expander
} // namespace mindspore

View File

@ -2913,5 +2913,56 @@ REG_BPROP_BUILDER("IRFFT").SetBody(BODYFUNC(ib) {
return {grad_dout, ib->OutZeros(n), ib->OutZeros(dim), ib->OutZeros(norm)};
});
REG_BPROP_BUILDER("MeanExt").SetUnusedInputs({i0, i4}).SetBody(BODYFUNC(ib) {
auto input = ib->GetInput(kIndex0);
auto input_dtype_id = ib->GetDtypeId(input);
if (input_dtype_id == kNumberTypeComplex64 || input_dtype_id == kNumberTypeComplex128) {
MS_EXCEPTION(TypeError) << "For 'MeanExt', gradient not support for complex type currently.";
}
auto axis = ib->GetInput(kIndex1);
auto keep_dims = ib->GetInput(kIndex2);
auto dtype = ib->GetInput(kIndex3);
auto out = ib->GetInput(kIndex4);
auto dout = ib->GetInput(kIndex5);
auto axis_type = axis->abstract()->BuildType();
MS_EXCEPTION_IF_NULL(axis_type);
if (axis_type->isa<TypeNone>()) {
axis = ib->Value<std::vector<int64_t>>({});
}
auto grad = SumGrad(ib, input, axis, dout, GetValue<bool>(keep_dims->BuildValue()));
NodePtr div_shape_node;
if (IsDynamic(ib->GetShape(input)) || IsDynamic(ib->GetShape(out))) {
auto shape_out_sz = ib->DynSize(out, kFloat32);
auto div_shape = ib->DynSize(input, kFloat32) / shape_out_sz;
div_shape_node = ib->Cast(div_shape, ib->GetDtype(grad));
} else {
auto shape_out_sz = ib->GetSize(out);
if (shape_out_sz == 0) {
MS_EXCEPTION(ValueError) << "For 'MeanExt', out shape size can not be 0";
}
auto div_shape = ib->GetSize(input) / shape_out_sz;
div_shape_node = ib->Tensor(div_shape, ib->GetDtype(grad));
}
auto dx = ib->Cast(ib->RealDiv(grad, div_shape_node), ib->GetDtype(input));
return {dx, ib->OutZeros(axis), ib->OutZeros(keep_dims), ib->OutZeros(dtype)};
});
REG_BPROP_BUILDER("SumExt").SetUnusedInputs({i0, i4}).SetBody(BODYFUNC(ib) {
auto input = ib->GetInput(kIndex0);
auto axis = ib->GetInput(kIndex1);
auto keep_dims = ib->GetInput(kIndex2);
auto dtype = ib->GetInput(kIndex3);
auto dout = ib->GetInput(kIndex5);
auto axis_type = axis->abstract()->BuildType();
MS_EXCEPTION_IF_NULL(axis_type);
if (axis_type->isa<TypeNone>()) {
axis = ib->Value<std::vector<int64_t>>({});
}
auto dx = ib->Cast(SumGrad(ib, input, axis, dout, GetValue<bool>(keep_dims->BuildValue())), ib->GetDtype(input));
return {dx, ib->OutZeros(axis), ib->OutZeros(keep_dims), ib->OutZeros(dtype)};
});
REG_BPROP_BUILDERS_END
} // namespace mindspore::expander::bprop

View File

@ -38,6 +38,7 @@ BoolImmPtr ConvertBool(const py::object &obj) {
}
Int64ImmPtr ConvertInt(const py::object &obj) {
// bool is also an instance of py::int_
if (py::isinstance<py::bool_>(obj) || !py::isinstance<py::int_>(obj)) {
return nullptr;
}

View File

@ -0,0 +1,54 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "plugin/device/ascend/kernel/opapi/aclnn/mean_ext_aclnn_kernel.h"
#include <algorithm>
#include <vector>
#include <map>
#include <memory>
#include <functional>
#include "ir/tensor.h"
#include "transform/acl_ir/acl_helper.h"
#include "transform/acl_ir/op_api_convert.h"
#include "abstract/ops/primitive_infer_map.h"
namespace mindspore {
namespace kernel {
void MeanExtAscend::GetWorkSpaceInfo(const std::vector<KernelTensor *> &inputs,
const std::vector<KernelTensor *> &outputs) {
const auto axis_opt = inputs[kIndex1]->GetOptionalValueWithCheck<std::vector<int64_t>>();
if (axis_opt.has_value()) {
axis_ = axis_opt.value();
} else {
axis_ = std::vector<int64_t>{};
}
keep_dims_ = transform::ConvertKernelTensor<bool>(inputs[kIndex2]);
// Infer function has confirmed the actual dtype of output
dtype_ = outputs[kIndex0]->dtype_id();
GetWorkspaceForResize(inputs[kIndex0], axis_, keep_dims_, dtype_, outputs[kIndex0]);
}
bool MeanExtAscend::Launch(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &workspace,
const std::vector<KernelTensor *> &outputs, void *stream_ptr) {
MS_EXCEPTION_IF_NULL(stream_ptr);
ParseGenExecutor(
GEN_EXECUTOR_BOOST(op_type_, hash_id_, inputs[kIndex0], axis_, keep_dims_, dtype_, outputs[kIndex0]));
RunOp(stream_ptr, workspace);
return true;
}
MS_ACLNN_KERNEL_FACTORY_REG(MeanExt, MeanExtAscend);
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,46 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_OPAPI_ACLNN_MEAN_EXT_ACLNN_KERNEL_H_
#define MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_OPAPI_ACLNN_MEAN_EXT_ACLNN_KERNEL_H_
#include <vector>
#include <utility>
#include "ops/base_operator.h"
#include "plugin/device/ascend/kernel/opapi/aclnn_kernel_mod.h"
#include "transform/acl_ir/acl_convert.h"
namespace mindspore {
namespace kernel {
class MeanExtAscend : public AclnnKernelMod {
public:
MeanExtAscend() : AclnnKernelMod(std::move("aclnnMean")) {}
~MeanExtAscend() = default;
bool Launch(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &workspace,
const std::vector<KernelTensor *> &outputs, void *stream_ptr) override;
void GetWorkSpaceInfo(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &outputs) override;
private:
DEFINE_GET_WORKSPACE_FOR_RESIZE()
std::vector<int64_t> axis_{};
bool keep_dims_{false};
TypeId dtype_;
};
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_OPAPI_ACLNN_MEAN_EXT_ACLNN_KERNEL_H_

View File

@ -0,0 +1,54 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "plugin/device/ascend/kernel/opapi/aclnn/sum_ext_aclnn_kernel.h"
#include <algorithm>
#include <vector>
#include <map>
#include <memory>
#include <functional>
#include "ir/tensor.h"
#include "transform/acl_ir/acl_helper.h"
#include "transform/acl_ir/op_api_convert.h"
#include "abstract/ops/primitive_infer_map.h"
namespace mindspore {
namespace kernel {
void SumExtAscend::GetWorkSpaceInfo(const std::vector<KernelTensor *> &inputs,
const std::vector<KernelTensor *> &outputs) {
const auto axis_opt = inputs[kIndex1]->GetOptionalValueWithCheck<std::vector<int64_t>>();
if (axis_opt.has_value()) {
axis_ = axis_opt.value();
} else {
axis_ = std::vector<int64_t>{};
}
keep_dims_ = transform::ConvertKernelTensor<bool>(inputs[kIndex2]);
// Infer function has confirmed the actual dtype of output
dtype_ = outputs[kIndex0]->dtype_id();
GetWorkspaceForResize(inputs[kIndex0], axis_, keep_dims_, dtype_, outputs[kIndex0]);
}
bool SumExtAscend::Launch(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &workspace,
const std::vector<KernelTensor *> &outputs, void *stream_ptr) {
MS_EXCEPTION_IF_NULL(stream_ptr);
ParseGenExecutor(
GEN_EXECUTOR_BOOST(op_type_, hash_id_, inputs[kIndex0], axis_, keep_dims_, dtype_, outputs[kIndex0]));
RunOp(stream_ptr, workspace);
return true;
}
MS_ACLNN_KERNEL_FACTORY_REG(SumExt, SumExtAscend);
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,46 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_OPAPI_ACLNN_SUM_EXT_ACLNN_KERNEL_H_
#define MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_OPAPI_ACLNN_SUM_EXT_ACLNN_KERNEL_H_
#include <vector>
#include <utility>
#include "ops/base_operator.h"
#include "plugin/device/ascend/kernel/opapi/aclnn_kernel_mod.h"
#include "transform/acl_ir/acl_convert.h"
namespace mindspore {
namespace kernel {
class SumExtAscend : public AclnnKernelMod {
public:
SumExtAscend() : AclnnKernelMod(std::move("aclnnReduceSum")) {}
~SumExtAscend() = default;
bool Launch(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &workspace,
const std::vector<KernelTensor *> &outputs, void *stream_ptr) override;
void GetWorkSpaceInfo(const std::vector<KernelTensor *> &inputs, const std::vector<KernelTensor *> &outputs) override;
private:
DEFINE_GET_WORKSPACE_FOR_RESIZE()
std::vector<int64_t> axis_{};
bool keep_dims_{false};
TypeId dtype_;
};
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_OPAPI_ACLNN_SUM_EXT_ACLNN_KERNEL_H_

View File

@ -0,0 +1,59 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "plugin/device/ascend/kernel/pyboost/customize/mean_ext.h"
#include "plugin/device/ascend/hal/device/ascend_stream_manager.h"
#include "kernel/pyboost/pyboost_utils.h"
#include "plugin/device/ascend/kernel/pyboost/aclnn_utils.h"
#include "runtime/device/device_address_utils.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
tensor::TensorPtr MeanExtAscendCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype) {
OpRunner::InferOpOutput(op, input_tensor, axis, keep_dims, dtype);
std::vector<int64_t> axis_vector{};
if (axis.has_value()) {
axis_vector = ConvertValueTupleToVector<int64_t>(axis.value());
}
const auto keep_dims_imm = GetValue<bool>(keep_dims);
// Infer function has confirmed the actual dtype of output
TypeId out_dtype = op->output_abs()->GetType()->cast<TensorTypePtr>()->element()->type_id();
PyBoostUtils::PrepareOpInputs(op->device_context(), op->stream_id(), input_tensor);
PyBoostUtils::PrepareOpOutputs(op->device_context(), op->stream_id(), op->outputs());
PyBoostUtils::DispatchRun(
std::make_shared<runtime::PyBoostDeviceTask>([op, input_tensor, axis_vector, keep_dims_imm, out_dtype]() {
auto device_context = op->device_context();
PyBoostUtils::MallocOpInputs(device_context, input_tensor);
PyBoostUtils::MallocOpOutputs(device_context, op->outputs());
MS_LOG(DEBUG) << op->primitive()->name() << " Call start";
LAUNCH_ACLNN(aclnnMean, device_context, op->stream_id(), input_tensor, axis_vector, keep_dims_imm, out_dtype,
op->output(0));
MS_LOG(DEBUG) << op->primitive()->name() << " Launch end";
}));
return op->output(0);
}
} // namespace pyboost
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,35 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_
#define MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_
#include <memory>
#include <vector>
#include "ir/tensor.h"
#include "ir/value.h"
#include "kernel/pyboost/op_runner.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
tensor::TensorPtr MeanExtAscendCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype);
} // namespace pyboost
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_

View File

@ -0,0 +1,58 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "plugin/device/ascend/kernel/pyboost/customize/sum_ext.h"
#include "plugin/device/ascend/hal/device/ascend_stream_manager.h"
#include "kernel/pyboost/pyboost_utils.h"
#include "plugin/device/ascend/kernel/pyboost/aclnn_utils.h"
#include "runtime/device/device_address_utils.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
tensor::TensorPtr SumExtAscendCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype) {
OpRunner::InferOpOutput(op, input_tensor, axis, keep_dims, dtype);
std::vector<int64_t> axis_vector{};
if (axis.has_value()) {
axis_vector = ConvertValueTupleToVector<int64_t>(axis.value());
}
const auto keep_dims_imm = GetValue<bool>(keep_dims);
// Infer function has confirmed the actual dtype of output
TypeId out_dtype = op->output_abs()->GetType()->cast<TensorTypePtr>()->element()->type_id();
PyBoostUtils::PrepareOpInputs(op->device_context(), op->stream_id(), input_tensor);
PyBoostUtils::PrepareOpOutputs(op->device_context(), op->stream_id(), op->outputs());
PyBoostUtils::DispatchRun(
std::make_shared<runtime::PyBoostDeviceTask>([op, input_tensor, axis_vector, keep_dims_imm, out_dtype]() {
auto device_context = op->device_context();
PyBoostUtils::MallocOpInputs(device_context, input_tensor);
PyBoostUtils::MallocOpOutputs(device_context, op->outputs());
MS_LOG(DEBUG) << op->primitive()->name() << " Call start";
LAUNCH_ACLNN(aclnnReduceSum, device_context, op->stream_id(), input_tensor, axis_vector, keep_dims_imm, out_dtype,
op->output(0));
MS_LOG(DEBUG) << op->primitive()->name() << " Launch end";
}));
return op->output(0);
}
} // namespace pyboost
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,35 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_
#define MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_
#include <memory>
#include <vector>
#include "ir/tensor.h"
#include "ir/value.h"
#include "kernel/pyboost/op_runner.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
tensor::TensorPtr SumExtAscendCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype);
} // namespace pyboost
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_ASCEND_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_

View File

@ -17,6 +17,8 @@
#include "plugin/device/ascend/optimizer/mindir/reduce_axis_update.h"
#include <vector>
#include <memory>
#include <set>
#include <string>
#include "mindspore/core/ops/math_ops.h"
#include "include/common/utils/anfalgo.h"
@ -61,7 +63,8 @@ bool ReduceAxisUpdate::IsReduce(const BaseRef &ref) {
if (IsPrimitive(node, prim::kPrimReduceMin) || IsPrimitive(node, prim::kPrimReduceMax) ||
IsPrimitive(node, prim::kPrimReduceMean) || IsPrimitive(node, prim::kPrimReduceSum) ||
IsPrimitive(node, prim::kPrimReduceProd) || IsPrimitive(node, prim::kPrimReduceAll) ||
IsPrimitive(node, prim::kPrimReduceAny)) {
IsPrimitive(node, prim::kPrimReduceAny) || IsPrimitive(node, prim::kPrimMeanExt) ||
IsPrimitive(node, prim::kPrimSumExt)) {
return true;
}
}
@ -90,6 +93,15 @@ bool ReduceAxisUpdate::IsAxisEmpty(const ValueNodePtr &axis_node) const {
return false;
}
bool ReduceAxisUpdate::IsAxisNone(const AnfNodePtr &cnode, const ValueNodePtr &axis_node) const {
static std::set<std::string> op_name_support_none = {prim::kPrimMeanExt->name(), prim::kPrimSumExt->name()};
auto cnode_name = common::AnfAlgo::GetCNodeName(cnode);
if (op_name_support_none.find(cnode_name) == op_name_support_none.end()) {
return false;
}
return axis_node->value()->isa<None>();
}
bool ReduceAxisUpdate::IsInputScalar(const AnfNodePtr &x_node) const {
MS_EXCEPTION_IF_NULL(x_node);
auto x_shape_ptr = x_node->Shape();
@ -146,6 +158,12 @@ bool ReduceAxisUpdate::CheckMatchedDAG(const PatternMap &, const FuncGraphPtr &g
return false;
}
// If input is dynamic rank, expand axis will get wrong result.
if (IsDynamicRank(common::AnfAlgo::GetPrevNodeOutputInferShape(node, 0))) {
MS_LOG(INFO) << "The input rank of dimension of " << node->DebugString() << " is unknown.";
return false;
}
auto cnode = node->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(cnode);
@ -165,8 +183,10 @@ bool ReduceAxisUpdate::CheckMatchedDAG(const PatternMap &, const FuncGraphPtr &g
MS_LOG(INFO) << "Axis input is " << input_axis->DebugString() << ".";
auto axis_value_node = input_axis->cast<ValueNodePtr>();
if (axis_value_node == nullptr || (!IsAxisEmpty(axis_value_node) && !IsInputScalar(input_x))) {
MS_LOG(INFO) << "Axis input of node " << node->fullname_with_scope() << " is not value node or axis is not empty.";
if (axis_value_node == nullptr ||
(!(IsAxisEmpty(axis_value_node) || IsAxisNone(cnode, axis_value_node)) && !IsInputScalar(input_x))) {
MS_LOG(INFO) << "Axis input of node " << node->fullname_with_scope()
<< " is not value node or axis is not empty or none.";
return false;
} else {
MS_LOG(INFO) << "Axis of node " << node->fullname_with_scope() << " is empty.";

View File

@ -33,6 +33,7 @@ class ReduceAxisUpdate : public PatternToPatternPass {
private:
bool IsAxisEmpty(const ValueNodePtr &axis_node) const;
bool IsAxisNone(const AnfNodePtr &cnode, const ValueNodePtr &axis_node) const;
bool IsInputScalar(const AnfNodePtr &x_node) const;
static bool IsReduce(const BaseRef &ref);
};

View File

@ -0,0 +1,102 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "plugin/device/cpu/kernel/pyboost/customize/mean_ext.h"
#include "plugin/device/cpu/kernel/pyboost/auto_generate/cast.h"
#include "plugin/device/cpu/kernel/pyboost/auto_generate/mean_ext.h"
#include "kernel/pyboost/pyboost_utils.h"
#include "ops/auto_generate/gen_ops_primitive.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
namespace {
void MeanExtCPUCall(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor, const ValuePtr &axis,
const BoolImmPtr &keep_dims, const std::vector<AbstractBasePtr> &input_abs) {
PyBoostUtils::PrepareOpInputs(op->device_context(), op->stream_id(), input_tensor);
PyBoostUtils::PrepareOpOutputs(op->device_context(), op->stream_id(), op->outputs());
PyBoostUtils::DispatchRun(
std::make_shared<runtime::PyBoostDeviceTask>([op, input_tensor, axis, keep_dims, input_abs]() {
MS_LOG(DEBUG) << "For 'MeanExt', the cpu task 'ReduceMean' start";
auto device_context = op->device_context();
const auto &outputs = op->outputs();
const auto primitive = std::make_shared<Primitive>(prim::kPrimReduceMean->name());
MS_EXCEPTION_IF_NULL(primitive);
PyBoostUtils::MallocOpInputs(device_context, input_tensor);
PyBoostUtils::MallocOpOutputs(device_context, outputs);
const auto &input_address_info =
PyBoostUtils::GetAddressInfo(device_context, op->stream_id(), input_abs, input_tensor, axis, keep_dims);
const auto &output_address_info =
PyBoostUtils::GetAddressInfo(device_context, op->stream_id(), {op->output_abs()}, outputs);
PyBoostUtils::LaunchKernel(primitive, device_context, input_address_info, output_address_info);
MS_LOG(DEBUG) << "For 'MeanExt', the cpu task 'ReduceMean' end";
}));
}
} // namespace
void MeanExtCPUCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype) {
OpRunner::InferOpOutput(op, input_tensor, axis, keep_dims, dtype);
ValuePtr act_axis;
if (axis.has_value()) {
act_axis = axis.value();
} else {
act_axis = MakeValue<std::vector<int64_t>>({});
}
// Infer function has confirmed the actual dtype of output
TypeId out_dtype = op->output_abs()->GetType()->cast<TensorTypePtr>()->element()->type_id();
TensorPtr act_tensor = input_tensor;
// Call Cast before Launch ReduceMean
if (input_tensor->data_type() != out_dtype) {
MS_LOG(DEBUG) << "Call Cast cpu kernel, src dtype: " << TypeIdToString(input_tensor->data_type())
<< ", dst dtype: " << TypeIdToString(out_dtype);
act_tensor =
PyBoostUtils::CastTensor(input_tensor, out_dtype, op->device_context()->device_context_key_.device_name_);
}
// Set new input abstract for ReduceMean
std::vector<AbstractBasePtr> new_input_abs{act_tensor->ToAbstract(), act_axis->ToAbstract(), keep_dims->ToAbstract()};
// Check if dtype is matched on ReduceMean kernel
auto kernel_attr_pair =
PyBoostUtils::SelectKernel(new_input_abs, op->output_abs(), op->device_context(), prim::kPrimReduceMean->name());
if (kernel_attr_pair.first) {
MeanExtCPUCall(op, act_tensor, act_axis, keep_dims, new_input_abs);
} else {
auto &select_kernel = kernel_attr_pair.second;
auto &device_name = op->device_context()->device_context_key_.device_name_;
const auto &real_input_tensor =
PyBoostUtils::CastTensor(act_tensor, select_kernel.input_type()[0].dtype, device_name);
const auto &mean_ext_op = CREATE_PYBOOST_OP(MeanExt, device_name);
mean_ext_op->set_primitive(prim::kPrimMeanExt);
const auto out_tensor = mean_ext_op->Call(real_input_tensor, axis, keep_dims, std::nullopt);
const auto &real_output_tensor = PyBoostUtils::CastTensor(out_tensor, out_dtype, device_name);
op->set_outputs({real_output_tensor});
}
}
} // namespace pyboost
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,36 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_CPU_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_
#define MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_CPU_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_
#include <vector>
#include <memory>
#include "ir/tensor.h"
#include "ir/value.h"
#include "runtime/hardware/device_context_manager.h"
#include "kernel/pyboost/op_runner.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
void MeanExtCPUCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &x_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype);
} // namespace pyboost
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_CPU_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_

View File

@ -0,0 +1,106 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "plugin/device/cpu/kernel/pyboost/customize/sum_ext.h"
#include "plugin/device/cpu/kernel/pyboost/auto_generate/cast.h"
#include "plugin/device/cpu/kernel/pyboost/auto_generate/sum_ext.h"
#include "kernel/pyboost/pyboost_utils.h"
#include "ops/auto_generate/gen_ops_primitive.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
namespace {
void SumExtCPUCall(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor, const ValuePtr &axis,
const BoolImmPtr &keep_dims, const BoolImmPtr &skip_mode,
const std::vector<AbstractBasePtr> &input_abs) {
PyBoostUtils::PrepareOpInputs(op->device_context(), op->stream_id(), input_tensor);
PyBoostUtils::PrepareOpOutputs(op->device_context(), op->stream_id(), op->outputs());
PyBoostUtils::DispatchRun(
std::make_shared<runtime::PyBoostDeviceTask>([op, input_tensor, axis, keep_dims, skip_mode, input_abs]() {
MS_LOG(DEBUG) << "For 'SumExt', the cpu task 'ReduceSum' start";
auto device_context = op->device_context();
const auto &outputs = op->outputs();
const auto primitive = std::make_shared<Primitive>(prim::kPrimReduceSum->name());
MS_EXCEPTION_IF_NULL(primitive);
PyBoostUtils::MallocOpInputs(device_context, input_tensor);
PyBoostUtils::MallocOpOutputs(device_context, outputs);
const auto &input_address_info = PyBoostUtils::GetAddressInfo(device_context, op->stream_id(), input_abs,
input_tensor, axis, keep_dims, skip_mode);
const auto &output_address_info =
PyBoostUtils::GetAddressInfo(device_context, op->stream_id(), {op->output_abs()}, outputs);
PyBoostUtils::LaunchKernel(primitive, device_context, input_address_info, output_address_info);
MS_LOG(DEBUG) << "For 'SumExt', the cpu task 'ReduceSum' end";
}));
}
} // namespace
void SumExtCPUCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype) {
OpRunner::InferOpOutput(op, input_tensor, axis, keep_dims, dtype);
// If axis is None, Convert axis to empty tuple
ValuePtr act_axis;
if (axis.has_value()) {
act_axis = axis.value();
} else {
act_axis = MakeValue<std::vector<int64_t>>({});
}
// Infer function has confirmed the actual dtype of output
TypeId out_dtype = op->output_abs()->GetType()->cast<TensorTypePtr>()->element()->type_id();
TensorPtr act_tensor = input_tensor;
// Call Cast before Launch ReduceSum
if (input_tensor->data_type() != out_dtype) {
MS_LOG(DEBUG) << "Call Cast cpu kernel, src dtype: " << TypeIdToString(input_tensor->data_type())
<< ", dst dtype: " << TypeIdToString(out_dtype);
act_tensor =
PyBoostUtils::CastTensor(input_tensor, out_dtype, op->device_context()->device_context_key_.device_name_);
}
const auto skip_mode = std::make_shared<BoolImm>(false);
// Set new input abstract for ReduceSum
std::vector<AbstractBasePtr> new_input_abs{act_tensor->ToAbstract(), act_axis->ToAbstract(), keep_dims->ToAbstract(),
skip_mode->ToAbstract()};
// Check if dtype is matched on ReduceSum kernel
auto kernel_attr_pair =
PyBoostUtils::SelectKernel(new_input_abs, op->output_abs(), op->device_context(), prim::kPrimReduceSum->name());
if (kernel_attr_pair.first) {
SumExtCPUCall(op, act_tensor, act_axis, keep_dims, std::make_shared<BoolImm>(false), new_input_abs);
} else {
auto &select_kernel = kernel_attr_pair.second;
auto &device_name = op->device_context()->device_context_key_.device_name_;
const auto &real_input_tensor =
PyBoostUtils::CastTensor(act_tensor, select_kernel.input_type()[0].dtype, device_name);
const auto &sum_ext_op = CREATE_PYBOOST_OP(SumExt, device_name);
sum_ext_op->set_primitive(prim::kPrimSumExt);
const auto out_tensor = sum_ext_op->Call(real_input_tensor, axis, keep_dims, std::nullopt);
const auto &real_output_tensor = PyBoostUtils::CastTensor(out_tensor, out_dtype, device_name);
op->set_outputs({real_output_tensor});
}
}
} // namespace pyboost
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,36 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_CPU_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_
#define MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_CPU_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_
#include <vector>
#include <memory>
#include "ir/tensor.h"
#include "ir/value.h"
#include "runtime/hardware/device_context_manager.h"
#include "kernel/pyboost/op_runner.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
void SumExtCPUCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &x_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype);
} // namespace pyboost
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_CPU_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_

View File

@ -0,0 +1,90 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "plugin/device/gpu/kernel/pyboost/customize/mean_ext.h"
#include "plugin/device/gpu/kernel/pyboost/auto_generate/cast.h"
#include "runtime/hardware/device_context_manager.h"
#include "plugin/device/gpu/hal/device/gpu_device_manager.h"
#include "ops/auto_generate/gen_ops_primitive.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
namespace {
void MeanExtGPUCall(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor, const ValuePtr &axis,
const BoolImmPtr &keep_dims) {
PyBoostUtils::PrepareOpInputs(op->device_context(), op->stream_id(), input_tensor);
PyBoostUtils::PrepareOpOutputs(op->device_context(), op->stream_id(), op->outputs());
PyBoostUtils::DispatchRun(std::make_shared<runtime::PyBoostDeviceTask>([op, input_tensor, axis, keep_dims]() {
MS_LOG(DEBUG) << "For 'MeanExt', the gpu task 'ReduceMean' start";
auto device_context = op->device_context();
const auto &outputs = op->outputs();
const auto primitive = std::make_shared<Primitive>(prim::kPrimReduceMean->name());
MS_EXCEPTION_IF_NULL(primitive);
PyBoostUtils::MallocOpInputs(device_context, input_tensor);
PyBoostUtils::MallocOpOutputs(device_context, outputs);
// Set new Abstract for ReduceMean
std::vector<AbstractBasePtr> input_abs{input_tensor->ToAbstract(), axis->ToAbstract(), keep_dims->ToAbstract()};
const auto &input_address_info =
PyBoostUtils::GetAddressInfo(device_context, op->stream_id(), input_abs, input_tensor, axis, keep_dims);
const auto &output_address_info =
PyBoostUtils::GetAddressInfo(device_context, op->stream_id(), {op->output_abs()}, outputs);
auto stream = device::gpu::GPUDeviceManager::GetInstance().GetStream(op->stream_id());
PyBoostUtils::LaunchKernel(primitive, device_context, input_address_info, output_address_info, stream);
static auto sync = MsContext::GetInstance()->get_param<bool>(MS_CTX_ENABLE_PYNATIVE_SYNCHRONIZE);
if (sync && !device_context->device_res_manager_->SyncAllStreams()) {
MS_LOG(EXCEPTION) << "SyncStream failed for op " << primitive->name();
}
MS_LOG(DEBUG) << "For 'MeanExt', the gpu task 'ReduceMean' end";
}));
}
} // namespace
void MeanExtGPUCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype) {
OpRunner::InferOpOutput(op, input_tensor, axis, keep_dims, dtype);
// If axis is None, Convert axis to empty tuple
ValuePtr act_axis;
if (axis.has_value()) {
act_axis = axis.value();
} else {
act_axis = MakeValue<std::vector<int64_t>>({});
}
// Infer function has confirmed the actual dtype of output
TypeId out_dtype = op->output_abs()->GetType()->cast<TensorTypePtr>()->element()->type_id();
TensorPtr act_tensor = input_tensor;
// Call Cast before Launch ReduceMean
if (input_tensor->data_type() != out_dtype) {
MS_LOG(DEBUG) << "Call Cast gpu kernel, src dtype: " << TypeIdToString(input_tensor->data_type())
<< ", dst dtype: " << TypeIdToString(out_dtype);
const auto &cast_op = CREATE_PYBOOST_OP(Cast, op->device_context()->device_context_key_.device_name_);
cast_op->set_primitive(prim::kPrimCast);
act_tensor = cast_op->Call(input_tensor, std::make_shared<Int64Imm>(out_dtype));
}
MeanExtGPUCall(op, act_tensor, act_axis, keep_dims);
}
} // namespace pyboost
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,36 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_
#define MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_
#include <vector>
#include <memory>
#include "ir/tensor.h"
#include "ir/value.h"
#include "runtime/hardware/device_context_manager.h"
#include "kernel/pyboost/op_runner.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
void MeanExtGPUCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype);
} // namespace pyboost
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_PYBOOST_CUSTOMIZE_MEAN_EXT_H_

View File

@ -0,0 +1,92 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "plugin/device/gpu/kernel/pyboost/customize/sum_ext.h"
#include "plugin/device/gpu/kernel/pyboost/auto_generate/cast.h"
#include "runtime/hardware/device_context_manager.h"
#include "plugin/device/gpu/hal/device/gpu_device_manager.h"
#include "ops/auto_generate/gen_ops_primitive.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
namespace {
void SumExtGPUCall(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor, const ValuePtr &axis,
const BoolImmPtr &keep_dims, const BoolImmPtr &skip_mode) {
PyBoostUtils::PrepareOpInputs(op->device_context(), op->stream_id(), input_tensor);
PyBoostUtils::PrepareOpOutputs(op->device_context(), op->stream_id(), op->outputs());
PyBoostUtils::DispatchRun(
std::make_shared<runtime::PyBoostDeviceTask>([op, input_tensor, axis, keep_dims, skip_mode]() {
MS_LOG(DEBUG) << "For 'SumExt', the gpu task 'ReduceSum' start";
auto device_context = op->device_context();
const auto &outputs = op->outputs();
const auto primitive = std::make_shared<Primitive>(prim::kPrimReduceSum->name());
MS_EXCEPTION_IF_NULL(primitive);
PyBoostUtils::MallocOpInputs(device_context, input_tensor);
PyBoostUtils::MallocOpOutputs(device_context, outputs);
// Set new Abstract for ReduceSum
std::vector<AbstractBasePtr> input_abs{input_tensor->ToAbstract(), axis->ToAbstract(), keep_dims->ToAbstract(),
skip_mode->ToAbstract()};
const auto &input_address_info = PyBoostUtils::GetAddressInfo(device_context, op->stream_id(), input_abs,
input_tensor, axis, keep_dims, skip_mode);
const auto &output_address_info =
PyBoostUtils::GetAddressInfo(device_context, op->stream_id(), {op->output_abs()}, outputs);
auto stream = device::gpu::GPUDeviceManager::GetInstance().GetStream(op->stream_id());
PyBoostUtils::LaunchKernel(primitive, device_context, input_address_info, output_address_info, stream);
static auto sync = MsContext::GetInstance()->get_param<bool>(MS_CTX_ENABLE_PYNATIVE_SYNCHRONIZE);
if (sync && !device_context->device_res_manager_->SyncAllStreams()) {
MS_LOG(EXCEPTION) << "SyncStream failed for op " << primitive->name();
}
MS_LOG(DEBUG) << "For 'SumExt', the gpu task 'ReduceSum' end";
}));
}
} // namespace
void SumExtGPUCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype) {
OpRunner::InferOpOutput(op, input_tensor, axis, keep_dims, dtype);
// If axis is None, Convert axis to empty tuple
ValuePtr act_axis;
if (axis.has_value()) {
act_axis = axis.value();
} else {
act_axis = MakeValue<std::vector<int64_t>>({});
}
// Infer function has confirmed the actual dtype of output
TypeId out_dtype = op->output_abs()->GetType()->cast<TensorTypePtr>()->element()->type_id();
TensorPtr act_tensor = input_tensor;
// Call Cast before Launch ReduceSum
if (input_tensor->data_type() != out_dtype) {
MS_LOG(DEBUG) << "Call Cast gpu kernel, src dtype: " << TypeIdToString(input_tensor->data_type())
<< ", dst dtype: " << TypeIdToString(out_dtype);
const auto &cast_op = CREATE_PYBOOST_OP(Cast, op->device_context()->device_context_key_.device_name_);
cast_op->set_primitive(prim::kPrimCast);
act_tensor = cast_op->Call(input_tensor, std::make_shared<Int64Imm>(out_dtype));
}
SumExtGPUCall(op, act_tensor, act_axis, keep_dims, std::make_shared<BoolImm>(false));
}
} // namespace pyboost
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,36 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_
#define MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_
#include <vector>
#include <memory>
#include "ir/tensor.h"
#include "ir/value.h"
#include "runtime/hardware/device_context_manager.h"
#include "kernel/pyboost/op_runner.h"
namespace mindspore {
namespace kernel {
namespace pyboost {
void SumExtGPUCustomize(const std::shared_ptr<OpRunner> &op, const TensorPtr &input_tensor,
const std::optional<ValueTuplePtr> &axis, const BoolImmPtr &keep_dims,
const std::optional<Int64ImmPtr> &dtype);
} // namespace pyboost
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_MINDSPORE_CCSRC_PLUGIN_DEVICE_GPU_KERNEL_PYBOOST_CUSTOMIZE_SUM_EXT_H_

View File

@ -0,0 +1,77 @@
mean_ext:
description: |
Reduces all dimension of a tensor by averaging all elements in the dimension, by default.
And reduce a dimension of `input` along the specified `axis`. `keep_dims`
determines whether the dimensions of the output and input are the same.
Args:
input (Tensor[Number]): The input tensor. The dtype of the tensor to be reduced is number.
:math:`(N, *)` where :math:`*` means, any number of additional dimensions.
axis (Union[int, tuple(int), list(int)]): The dimensions to reduce. Default: ``None`` , reduce all dimensions.
Only constant value is allowed. Assume the rank of `input` is r, and the value range is [-r,r).
keep_dims (bool): If ``True``, keep these reduced dimensions and the length is 1.
If ``False``, don't keep these dimensions. Default: ``False`` .
dtype (:class:`mindspore.dtype`): The desired data type of returned Tensor. Default: ``None``.
Returns:
Tensor, has the same data type as input tensor.
- If `axis` is None, and `keep_dims` is False,
the output is a 0-D tensor representing the product of all elements in the input tensor.
- If `axis` is int, set as 1, and `keep_dims` is False,
the shape of output is :math:`(x_0, x_2, ..., x_R)`.
- If `axis` is tuple(int), set as (1, 2), and `keep_dims` is ``False`` ,
the shape of output is :math:`(x_0, x_3, ..., x_R)`.
Raises:
TypeError: If `x` is not a Tensor.
TypeError: If `axis` is not one of the following: int, tuple or list.
TypeError: If `keep_dims` is not a bool.
ValueError: If `axis` is out of range.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import mindspore
>>> import numpy as np
>>> from mindspore import Tensor, ops
>>> x = Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))
>>> output = ops.mean(x, 1, keep_dims=True)
>>> result = output.shape
>>> print(result)
(3, 1, 5, 6)
>>> # case 1: Reduces a dimension by averaging all elements in the dimension.
>>> x = Tensor(np.array([[[2, 2, 2, 2, 2, 2], [2, 2, 2, 2, 2, 2], [2, 2, 2, 2, 2, 2]],
... [[4, 4, 4, 4, 4, 4], [5, 5, 5, 5, 5, 5], [6, 6, 6, 6, 6, 6]],
... [[6, 6, 6, 6, 6, 6], [8, 8, 8, 8, 8, 8], [10, 10, 10, 10, 10, 10]]]),
... mindspore.float32)
>>> output = ops.mean(x)
>>> print(output)
5.0
>>> print(output.shape)
()
>>> # case 2: Reduces a dimension along the axis 0
>>> output = ops.mean(x, 0, True)
>>> print(output)
[[[4. 4. 4. 4. 4. 4.]
[5. 5. 5. 5. 5. 5.]
[6. 6. 6. 6. 6. 6.]]]
>>> # case 3: Reduces a dimension along the axis 1
>>> output = ops.mean(x, 1, True)
>>> print(output)
[[[2. 2. 2. 2. 2. 2.]]
[[5. 5. 5. 5. 5. 5.]]
[[8. 8. 8. 8. 8. 8.]]]
>>> # case 4: Reduces a dimension along the axis 2
>>> output = ops.mean(x, 2, True)
>>> print(output)
[[[ 2.]
[ 2.]
[ 2.]]
[[ 4.]
[ 5.]
[ 6.]]
[[ 6.]
[ 8.]
[10.]]]

View File

@ -0,0 +1,54 @@
sum_ext:
description: |
Calculate sum of Tensor elements over a given dim.
Args:
input (Tensor): The input tensor.
dim (Union[None, int, tuple(int), list(int)]): Dimensions along which a sum is performed.
If ``None``, sum all the elements of the input tensor.
If the `dim` is a tuple or list of ints, a sum is performed on all the dimensions specified in the tuple.
Must be in the range :math:`[-input.ndim, input.ndim)` . Default: ``None``.
keepdim (bool): Whether the output tensor has dim retained or not.
If ``True``, keep these reduced dimensions and the length is 1.
If ``False``, don't keep these dimensions. Default: ``False``.
dtype (:class:`mindspore.dtype`): The desired data type of returned Tensor. Default: ``None``.
Returns:
A Tensor, sum of elements over a given dim in `input`.
Raises:
TypeError: If `input` is not a Tensor.
TypeError: If `dim` is not an int, tulpe(int), list(int) or None.
ValueError: If `dim` is not in the range :math:`[-input.ndim, input.ndim)` .
TypeError: If `keepdim` is not a bool.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import mindspore
>>> import numpy as np
>>> from mindspore import Tensor, ops
>>> from mindspore import dtype as mstype
>>> x = Tensor(np.array([[[1, 1, 1, 1, 1, 1], [2, 2, 2, 2, 2, 2], [3, 3, 3, 3, 3, 3]],
... [[4, 4, 4, 4, 4, 4], [5, 5, 5, 5, 5, 5], [6, 6, 6, 6, 6, 6]],
... [[7, 7, 7, 7, 7, 7], [8, 8, 8, 8, 8, 8], [9, 9, 9, 9, 9, 9]]]), mstype.float32)
>>> out = ops.sum(x)
>>> print(out)
270.0
>>> out = ops.sum(x, dim=2)
>>> print(out)
[[ 6. 12. 18.]
[24. 30. 36.]
[42. 48. 54.]]
>>> out = ops.sum(x, dim=2, keepdim=True)
>>> print(out)
[[[ 6.]
[12.]
[18.]]
[[24.]
[30.]
[36.]]
[[42.]
[48.]
[54.]]]

View File

@ -0,0 +1,26 @@
#operator mean_ext
mean_ext:
args:
input:
dtype: tensor
axis:
dtype: tuple[int]
type_cast: int, list[int], tensor
default: None
keep_dims:
dtype: bool
default: False
dtype:
dtype: TypeId
arg_handler: dtype_to_type_id
default: None
returns:
output:
dtype: tensor
dispatch:
enable: True
Ascend: MeanExtAscend
CPU: MeanExtCPU
GPU: MeanExtGPU
function:
name: mean

View File

@ -0,0 +1,24 @@
#operator sum_ext
sum_ext:
args:
input:
dtype: tensor
dim:
dtype: tuple[int]
type_cast: int, list[int]
default: None
keepdim:
dtype: bool
default: False
dtype:
dtype: TypeId
arg_handler: dtype_to_type_id
default: None
returns:
output:
dtype: tensor
dispatch:
enable: True
Ascend: SumExtAscend
CPU: SumExtCPU
GPU: SumExtGPU

View File

@ -0,0 +1,32 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "ops/ops_frontend_func_impl.h"
namespace mindspore::ops {
namespace {
constexpr auto kMeanExt = "MeanExt";
} // namespace
class MeanExtFrontendFuncImpl : public OpFrontendFuncImpl {
public:
ValuePtr InferValue(const PrimitivePtr &, const std::vector<AbstractBasePtr> &input_args) const override {
return InferValueCallback::GetInstance().CallPyInferValue(kMeanExt, input_args);
}
};
REGISTER_PRIMITIVE_FUNCTION_FRONTEND_FUNC_IMPL(kMeanExt, MeanExtFrontendFuncImpl);
} // namespace mindspore::ops

View File

@ -0,0 +1,32 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "ops/ops_frontend_func_impl.h"
namespace mindspore::ops {
namespace {
constexpr auto kSumExt = "SumExt";
} // namespace
class SumExtFrontendFuncImpl : public OpFrontendFuncImpl {
public:
ValuePtr InferValue(const PrimitivePtr &, const std::vector<AbstractBasePtr> &input_args) const override {
return InferValueCallback::GetInstance().CallPyInferValue(kSumExt, input_args);
}
};
REGISTER_PRIMITIVE_FUNCTION_FRONTEND_FUNC_IMPL(kSumExt, SumExtFrontendFuncImpl);
} // namespace mindspore::ops

View File

@ -1,30 +0,0 @@
/**
* Copyright 2023 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "ops/ops_func_impl/mean.h"
#include "ops/op_utils.h"
namespace mindspore {
namespace ops {
BaseShapePtr MeanFuncImpl::InferShape(const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) const {
return ReduceExtInferShape(primitive, input_args);
}
TypePtr MeanFuncImpl::InferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const {
return ReduceExtInferType(primitive, input_args);
}
} // namespace ops
} // namespace mindspore

View File

@ -0,0 +1,38 @@
/**
* Copyright 2023 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "ops/ops_func_impl/mean_ext.h"
#include <memory>
#include "ops/ops_func_impl/reduce_arithmetic.h"
#include "ops/op_utils.h"
namespace mindspore {
namespace ops {
BaseShapePtr MeanExtFuncImpl::InferShape(const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) const {
return ReduceExtandInferShape(primitive, input_args);
}
TypePtr MeanExtFuncImpl::InferType(const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) const {
if (input_args[kIndex3]->GetType()->isa<TypeNone>()) {
return input_args[kIndex0]->GetType()->Clone();
}
auto dtype_opt = GetScalarValue<int64_t>(input_args[kIndex3]->GetValue());
MS_CHECK_VALUE(dtype_opt.has_value(), primitive->name() + " error: dtype input should has valid value.");
return std::make_shared<TensorType>(TypeIdToType(static_cast<TypeId>(dtype_opt.value())));
}
} // namespace ops
} // namespace mindspore

View File

@ -14,17 +14,15 @@
* limitations under the License.
*/
#ifndef MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_SUM_H_
#define MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_SUM_H_
#ifndef MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_MEAN_EXT_H_
#define MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_MEAN_EXT_H_
#include <memory>
#include <vector>
#include "mindapi/base/types.h"
#include "ops/ops_func_impl/op_func_impl.h"
namespace mindspore {
namespace ops {
class MIND_API SumFuncImpl : public OpFuncImpl {
class MIND_API MeanExtFuncImpl : public OpFuncImpl {
public:
BaseShapePtr InferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const override;
TypePtr InferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const override;
@ -32,4 +30,4 @@ class MIND_API SumFuncImpl : public OpFuncImpl {
} // namespace ops
} // namespace mindspore
#endif // MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_SUM_H_
#endif // MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_MEAN_EXT_H_

View File

@ -86,5 +86,68 @@ BaseShapePtr ReduceInferShape(const PrimitivePtr &primitive, const std::vector<A
}
return std::make_shared<abstract::Shape>(out_shape);
}
BaseShapePtr ReduceExtandInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
auto keep_dims_value = input_args[kInputIndex2]->GetValue();
auto keep_dims_opt = GetScalarValue<bool>(keep_dims_value);
if (MS_UNLIKELY(!keep_dims_opt.has_value())) {
return std::make_shared<abstract::Shape>(ShapeVector({abstract::Shape::kShapeRankAny}));
}
auto keep_dims = keep_dims_opt.value();
auto x_shape = input_args[kInputIndex0]->GetShape()->GetShapeVector();
// If axis is None
if (input_args[kInputIndex1]->GetType()->isa<TypeNone>()) {
return keep_dims
? std::make_shared<abstract::Shape>(IsDynamicRank(x_shape) ? x_shape : ShapeVector(x_shape.size(), 1))
: std::make_shared<abstract::Shape>(ShapeVector({}));
}
auto axis_array_opt = GetArrayValue<int64_t>(input_args[kInputIndex1]);
if (axis_array_opt.has_value()) {
// If axis is empty tuple and keep_dims is False, return a zero-dimensional Tensor
if (axis_array_opt->size() == 0 && !keep_dims) {
return std::make_shared<abstract::Shape>(ShapeVector({}));
}
}
if (IsDynamicRank(x_shape)) {
return std::make_shared<abstract::Shape>(x_shape);
}
if (!axis_array_opt.has_value()) {
// axis is dynamic.
return keep_dims ? std::make_shared<abstract::Shape>(ShapeVector(x_shape.size(), -1))
: std::make_shared<abstract::Shape>(ShapeVector({abstract::Shape::kShapeRankAny}));
}
auto x_shape_size = x_shape.size();
auto axis_array = axis_array_opt.value();
// All values of the axis are known.
if (!axis_array.HasUnknownValue()) {
std::vector<int64_t> axis_vec = axis_array.ToVector();
std::vector<int64_t> real_axis_vec;
(void)std::transform(
axis_vec.begin(), axis_vec.end(), std::back_inserter(real_axis_vec),
[&x_shape_size, &primitive](const int64_t &axis) { return CalRealAixs(axis, x_shape_size, primitive); });
auto out_shape = ReduceFuncCalShapeInferImpl(primitive, x_shape, real_axis_vec, keep_dims);
return std::make_shared<abstract::Shape>(out_shape);
}
// If the axis has unknown value, the reduction position will be any of the input dimensions.
if (!keep_dims) {
MS_CHECK_VALUE(x_shape.size() >= axis_array_opt->size(),
CheckAndConvertUtils::FormatCheckInRangeMsg("axis size", axis_array_opt->size(), kIncludeLeft,
{0, x_shape.size()}, primitive));
return std::make_shared<abstract::Shape>(ShapeVector(x_shape.size() - axis_array_opt->size(), -1));
}
auto out_shape = ShapeVector(x_shape.size(), -1);
for (size_t i = 0; i < axis_array.size(); ++i) {
if (!axis_array.IsValueUnknown(i)) {
auto axis = CalRealAixs(axis_array[i], x_shape_size, primitive);
out_shape[axis] = 1;
}
}
return std::make_shared<abstract::Shape>(out_shape);
}
} // namespace ops
} // namespace mindspore

View File

@ -24,6 +24,8 @@
namespace mindspore {
namespace ops {
BaseShapePtr ReduceInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args);
BaseShapePtr ReduceExtandInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args);
int64_t CalRealAixs(const int64_t &axis, const size_t &x_shape_size, const PrimitivePtr &primitive);
} // namespace ops
} // namespace mindspore
#endif // MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_REDUCE_REDUCE_ARITHMETIC_H_

View File

@ -1,30 +0,0 @@
/**
* Copyright 2023 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "ops/ops_func_impl/sum.h"
#include "ops/op_utils.h"
namespace mindspore {
namespace ops {
BaseShapePtr SumFuncImpl::InferShape(const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) const {
return ReduceExtInferShape(primitive, input_args);
}
TypePtr SumFuncImpl::InferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const {
return ReduceExtInferType(primitive, input_args);
}
} // namespace ops
} // namespace mindspore

View File

@ -0,0 +1,48 @@
/**
* Copyright 2023 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "ops/ops_func_impl/sum_ext.h"
#include <set>
#include "ops/ops_func_impl/reduce_arithmetic.h"
#include "ops/op_utils.h"
namespace mindspore {
namespace ops {
BaseShapePtr SumExtFuncImpl::InferShape(const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) const {
return ReduceExtandInferShape(primitive, input_args);
}
TypePtr SumExtFuncImpl::InferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const {
TypeId type_id;
if (input_args[kIndex3]->GetType()->isa<TypeNone>()) {
auto tensor_type = input_args[kIndex0]->GetType()->cast<TensorTypePtr>();
MS_EXCEPTION_IF_NULL(tensor_type);
type_id = tensor_type->element()->type_id();
static std::set<TypeId> intergral_set = {kNumberTypeBool, kNumberTypeUInt8, kNumberTypeInt8, kNumberTypeInt16,
kNumberTypeInt32};
if (intergral_set.find(type_id) != intergral_set.end()) {
type_id = kNumberTypeInt64;
}
} else {
auto dtype_opt = GetScalarValue<int64_t>(input_args[kIndex3]->GetValue());
MS_CHECK_VALUE(dtype_opt.has_value(), primitive->name() + " error: dtype input should has valid value.");
type_id = static_cast<TypeId>(dtype_opt.value());
}
return std::make_shared<TensorType>(TypeIdToType(type_id));
}
} // namespace ops
} // namespace mindspore

View File

@ -14,8 +14,8 @@
* limitations under the License.
*/
#ifndef MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_MEAN_H_
#define MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_MEAN_H_
#ifndef MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_SUM_EXT_H_
#define MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_SUM_EXT_H_
#include <memory>
#include <vector>
@ -24,7 +24,7 @@
namespace mindspore {
namespace ops {
class MIND_API MeanFuncImpl : public OpFuncImpl {
class MIND_API SumExtFuncImpl : public OpFuncImpl {
public:
BaseShapePtr InferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const override;
TypePtr InferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const override;
@ -32,4 +32,4 @@ class MIND_API MeanFuncImpl : public OpFuncImpl {
} // namespace ops
} // namespace mindspore
#endif // MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_MEAN_H_
#endif // MINDSPORE_CORE_OPS_OPS_FUNC_IMPL_SUM_EXT_H_

View File

@ -70,7 +70,7 @@ itemsize_map = {mstype.bool_: 1, mstype.int8: 1, mstype.uint8: 1,
nan_tensor = Tensor(float('nan'), dtype=mstype.float32)
def mean(x, axis=None, keep_dims=False):
def mean(x, axis=None, keep_dims=False, dtype=None):
"""
Reduces a dimension of a tensor by averaging all elements in the dimension.
@ -78,6 +78,7 @@ def mean(x, axis=None, keep_dims=False):
axis (Union[None, int, tuple(int), list(int)]): Dimensions of reduction,
when axis is None or empty tuple, reduce all dimensions. Default: ().
keep_dims (bool): Whether to keep the reduced dimensions. Default: False.
dtype (:class:`mindspore.dtype`): The desired data type of returned Tensor. Default: ``None``.
Returns:
Tensor, has the same data type as input tensor.
@ -93,7 +94,7 @@ def mean(x, axis=None, keep_dims=False):
>>> print(output)
2.0
"""
return F.mean(x, axis, keep_dims)
return F.mean(x, axis, keep_dims, dtype)
def ndimension(x):

View File

@ -1848,11 +1848,11 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
"""
return tensor_operator_registry.get('log2')(self)
def mean(self, axis=None, keep_dims=False):
def mean(self, axis=None, keep_dims=False, dtype=None):
"""
For details, please refer to :func:`mindspore.ops.mean`.
"""
return tensor_operator_registry.get('mean')(self, axis, keep_dims)
return tensor_operator_registry.get('mean')(self, axis, keep_dims, dtype)
def amin(self, axis=None, keepdims=False, *, initial=None, where=None):
"""
@ -3355,14 +3355,9 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
>>> print(input_x.sum(axis=1))
[10. 35.]
"""
if initial is not None and not isinstance(initial, (int, float, bool)):
raise TypeError(f"For Tensor.sum, initial must be int, float or bool, but got {type(initial)}.")
res = tensor_operator_registry.get("sum")(self, axis, keepdims)
if initial is not None:
res += initial
if dtype is not None:
res = res.astype(dtype)
return res
if initial is None:
return tensor_operator_registry.get("sum")(self, axis, keepdims, dtype=dtype)
return tensor_operator_registry.get("sum")(self, axis, keepdims, dtype=dtype) + initial
def sum_to_size(self, *size):
r"""

View File

@ -41,7 +41,7 @@ from mindspore.ops.auto_generate import (minimum, maximum, mul, sin, sinc, sinh,
matrix_exp, sqrt, rsqrt, square, trace, nextafter, abs, acos, acosh, angle,
asin, asinh, atan, atan2, atanh, ceil, equal, erf, erfc, erfinv, exp, expm1,
floor, floor_divide, floor_mod, gcd, greater, greater_equal, less, less_equal,
log, log1p, neg, not_equal, pow, round, isfinite, argmax)
log, log1p, neg, not_equal, pow, round, isfinite, argmax, mean, sum_ext)
from mindspore.nn import layer
from mindspore._checkparam import check_is_number
from mindspore import _checkparam as validator
@ -6636,94 +6636,6 @@ def amax(input, axis=None, keepdims=False, *, initial=None, where=None):
return _get_cache_prim(P.ReduceMax)(keepdims)(input, axis)
def mean(x, axis=None, keep_dims=False):
r"""
Reduces all dimension of a tensor by averaging all elements in the dimension, by default.
And reduce a dimension of `input` along the specified `axis`. `keep_dims`
determines whether the dimensions of the output and input are the same.
Note:
The `axis` with tensor type is only used for compatibility with older versions and is not recommended.
Args:
x (Tensor[Number]): The input tensor. The dtype of the tensor to be reduced is number.
:math:`(N, *)` where :math:`*` means, any number of additional dimensions.
axis (Union[int, tuple(int), list(int), Tensor]): The dimensions to reduce. Default: ``None`` ,
reduce all dimensions. Only constant value is allowed. Assume the rank of `input` is r,
and the value range is [-r,r).
keep_dims (bool): If ``True`` , keep these reduced dimensions and the length is 1.
If ``False`` , don't keep these dimensions. Default: ``False`` .
Returns:
Tensor, has the same data type as input tensor.
- If `axis` is ``None`` , and `keep_dims` is ``False`` ,
the output is a 0-D tensor representing the product of all elements in the input tensor.
- If `axis` is int, set as 1, and `keep_dims` is ``False`` ,
the shape of output is :math:`(x_0, x_2, ..., x_R)`.
- If `axis` is tuple(int), set as (1, 2), and `keep_dims` is ``False`` ,
the shape of output is :math:`(x_0, x_3, ..., x_R)`.
- If `axis` is 1-D Tensor, set as [1, 2], and `keep_dims` is ``False`` ,
the shape of output is :math:`(x_0, x_3, ..., x_R)`.
Raises:
TypeError: If `x` is not a Tensor.
TypeError: If `axis` is not one of the following: int, tuple, list or Tensor.
TypeError: If `keep_dims` is not a bool.
ValueError: If `axis` is out of range.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import mindspore
>>> import numpy as np
>>> from mindspore import Tensor, ops
>>> x = Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))
>>> output = ops.mean(x, 1, keep_dims=True)
>>> result = output.shape
>>> print(result)
(3, 1, 5, 6)
>>> # case 1: Reduces a dimension by averaging all elements in the dimension.
>>> x = Tensor(np.array([[[2, 2, 2, 2, 2, 2], [2, 2, 2, 2, 2, 2], [2, 2, 2, 2, 2, 2]],
... [[4, 4, 4, 4, 4, 4], [5, 5, 5, 5, 5, 5], [6, 6, 6, 6, 6, 6]],
... [[6, 6, 6, 6, 6, 6], [8, 8, 8, 8, 8, 8], [10, 10, 10, 10, 10, 10]]]),
... mindspore.float32)
>>> output = ops.mean(x)
>>> print(output)
5.0
>>> print(output.shape)
()
>>> # case 2: Reduces a dimension along the axis 0
>>> output = ops.mean(x, 0, True)
>>> print(output)
[[[4. 4. 4. 4. 4. 4.]
[5. 5. 5. 5. 5. 5.]
[6. 6. 6. 6. 6. 6.]]]
>>> # case 3: Reduces a dimension along the axis 1
>>> output = ops.mean(x, 1, True)
>>> print(output)
[[[2. 2. 2. 2. 2. 2.]]
[[5. 5. 5. 5. 5. 5.]]
[[8. 8. 8. 8. 8. 8.]]]
>>> # case 4: Reduces a dimension along the axis 2
>>> output = ops.mean(x, 2, True)
>>> print(output)
[[[ 2.]
[ 2.]
[ 2.]]
[[ 4.]
[ 5.]
[ 6.]]
[[ 6.]
[ 8.]
[10.]]]
"""
if axis is None:
axis = ()
return _get_cache_prim(P.ReduceMean)(keep_dims)(x, axis)
def prod(input, axis=None, keep_dims=False):
r"""
Reduces a dimension of a tensor by multiplying all elements in the dimension, by default. And also can
@ -10055,23 +9967,7 @@ def sum(input, dim=None, keepdim=False, *, dtype=None):
[48.]
[54.]]]
"""
if not isinstance(input, Tensor):
raise TypeError(f"For 'sum', 'input' must be Tensor, but got{type(input)}")
if dim is not None and not isinstance(dim, (int, tuple, list)):
raise TypeError(f"For 'sum', 'dim' must be int, tuple(int), list(int) or None, but got {type(dim)}")
if not isinstance(keepdim, bool):
raise TypeError(f"For 'sum', 'keepdim' must be bool, but got {type(keepdim)}")
if input.dtype == mstype.bool_:
input = input.astype(mstype.int64)
if dtype is not None:
input = input.astype(dtype)
reduce_sum = _get_cache_prim(P.ReduceSum)(keep_dims=keepdim)
if dim is not None:
out = reduce_sum(input, dim)
else:
out = reduce_sum(input)
return out
return sum_ext(input, dim, keepdim, dtype)
def tanhshrink(input):

View File

@ -3003,7 +3003,7 @@ def bidense(input1, input2, weight, bias=None):
output = output.transpose(2, 0, 1) * input2
output = output.sum(2).swapaxes(0, 1)
if bias is not None:
output = bias_add_(output, bias)
output = bias_add_(output.astype(bias.dtype), bias)
if len(input1_shape) != 2:
output_shape = input1_shape[:-1] + (-1,)
output = output.reshape(output_shape)

View File

@ -1248,6 +1248,35 @@ def _infer_value_for_Reduce(input_x, axis, keep_dims, prim_name):
return value
def _infer_value_for_ReduceExtand(input_x, axis, keep_dims, dtype, prim_name):
"""Infer value for Common ReduceExtand op."""
value = None
if input_x is not None:
prim_map = {
'MeanExt': np.mean,
'SumExt': np.sum,
}
np_reduce_extand_func = prim_map.get(prim_name, None)
if np_reduce_extand_func is not None:
value = input_x.asnumpy()
if isinstance(axis, int):
pass
elif axis:
axis = tuple(set(axis))
else:
axis = tuple(range(len(value.shape)))
if dtype is not None:
np_dtype = mstype.dtype_to_nptype(typing.type_id_to_type(dtype))
value = np_reduce_extand_func(value, axis, dtype=np_dtype, keepdims=keep_dims)
else:
value = np_reduce_extand_func(value, axis, keepdims=keep_dims)
value = np.array(value)
value = Tensor(value)
return value
def infer_value_for_Cast(x, dst_type_enum):
"""Infer value for Cast op."""
if x is None:
@ -1304,6 +1333,16 @@ def infer_value_for_ReduceAny(input_x, axis, keep_dims):
return _infer_value_for_Reduce(input_x, axis, keep_dims, 'ReduceAny')
def infer_value_for_MeanExt(input_x, axis, keep_dims, dtype):
"""Infer value for MeanExt op."""
return _infer_value_for_ReduceExtand(input_x, axis, keep_dims, dtype, 'MeanExt')
def infer_value_for_SumExt(input_x, axis, keep_dims, dtype):
"""Infer value for SumExt op."""
return _infer_value_for_ReduceExtand(input_x, axis, keep_dims, dtype, 'SumExt')
def infer_value_for_Diag(input_x):
"""Infer value for Diag op."""
if input_x is None:

View File

@ -187,3 +187,22 @@ def test_fallback_sum_with_x_unsupported_operand_type_error_2():
with pytest.raises(TypeError) as ex:
foo()
assert "unsupported operand type" in str(ex.value)
@pytest.mark.level0
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.platform_x86_ascend_training
@pytest.mark.env_onecard
def test_fallback_sum_with_x_tensor_n_default_2():
"""
Feature: JIT Fallback
Description: Test sum() in graph mode with input x tensor and input n default.
Expectation: No exception.
"""
@jit
def foo():
x = sum(Tensor([[1, 1], [2, 2]]))
return x
out = foo()
assert np.allclose(out.asnumpy(), np.array([3, 3]))

View File

@ -0,0 +1,486 @@
# Copyright 2024 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
import pytest
import numpy as np
import mindspore as ms
import mindspore.ops.function as F
import mindspore.common.dtype as mstype
from mindspore import ops
from mindspore import Tensor
from mindspore.ops.function import mean, prod
from tests.st.utils import test_utils
from tests.st.ops.dynamic_shape.test_op_utils import TEST_OP
def generate_random_input(shape, dtype):
return np.random.randn(*shape).astype(dtype)
def generate_expect_forward_output(name, x, axis=None, keep_dims=False, dtype=None):
if name == "mean":
return np.mean(x, axis=axis, dtype=dtype, keepdims=keep_dims)
if name == "prod":
return np.prod(x, axis=axis, dtype=dtype, keepdims=keep_dims)
if name == "sum":
return np.sum(x, axis=axis, dtype=dtype, keepdims=keep_dims)
return None
def mean_func(x, axis=None, keep_dims=False, dtype=None):
return mean(x, axis, keep_dims, dtype)
def sum_func(x, axis=None, keep_dims=False, dtype=None):
return F.sum(x, axis, keep_dims, dtype=dtype)
def prod_func(x, axis=None, keep_dims=False, dtype=None):
return prod(x, axis, keep_dims, dtype)
@test_utils.run_with_cell
def mean_forward_func(x, axis=None, keep_dims=False, dtype=None):
return mean_func(x, axis, keep_dims, dtype)
@test_utils.run_with_cell
def sum_forward_func(x, axis=None, keep_dims=False, dtype=None):
return sum_func(x, axis, keep_dims, dtype=dtype)
@test_utils.run_with_cell
def prod_forward_func(x, axis=None, keep_dims=False, dtype=None):
return prod(x, axis, keep_dims, dtype)
@test_utils.run_with_cell
def mean_backward_func(x, axis=None, keep_dims=False, dtype=None):
return ops.grad(mean_forward_func, (0))(x, axis, keep_dims, dtype)
@test_utils.run_with_cell
def sum_backward_func(x, axis=None, keep_dims=False, dtype=None):
return ops.grad(sum_forward_func, (0))(x, axis, keep_dims, dtype)
@test_utils.run_with_cell
def prod_backward_func(x, axis=None, keep_dims=False, dtype=None):
return ops.grad(prod_forward_func, (0))(x, axis, keep_dims, dtype)
@pytest.mark.level0
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize('keep_dims', [False, True])
@pytest.mark.parametrize('in_dtype', [mstype.float16])
@pytest.mark.parametrize('out_dtype', [mstype.float32])
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
def test_mean_normal(keep_dims, in_dtype, out_dtype, context_mode):
"""
Feature: pyboost function.
Description: test function mean forward and backward on ascend with different datatype.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
axis = (0, -1)
x = generate_random_input((2, 3, 4, 5), mstype.dtype_to_nptype(in_dtype))
output = mean_forward_func(ms.Tensor(x), axis, keep_dims, out_dtype)
expect = generate_expect_forward_output("mean", x, axis, keep_dims, mstype.dtype_to_nptype(out_dtype))
np.testing.assert_equal(output.dtype, out_dtype)
np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
axis = (0, -1)
x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(mstype.dtype_to_nptype(in_dtype))
grads = mean_backward_func(ms.Tensor(x), axis, False, out_dtype)
expect = np.full((2, 3, 4), 1 / (2 * 4), mstype.dtype_to_nptype(in_dtype))
np.testing.assert_equal(grads.dtype, in_dtype)
np.testing.assert_allclose(grads.asnumpy(), expect, rtol=1e-3)
@pytest.mark.level0
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
def test_mean_default(context_mode):
"""
Feature: pyboost function.
Description: test function mean forward and backward on ascend with default args.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
x = generate_random_input((2, 3, 4, 5), np.float32)
output = mean_forward_func(ms.Tensor(x))
expect = generate_expect_forward_output("mean", x)
np.testing.assert_equal(output.dtype, mstype.float32)
np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
x1 = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.float32)
grads = mean_backward_func(ms.Tensor(x1))
expect = np.full((2, 3, 4), 1 / (2 * 3 * 4), np.float32)
np.testing.assert_equal(grads.dtype, mstype.float32)
np.testing.assert_allclose(grads.asnumpy(), expect, rtol=1e-3)
@pytest.mark.level1
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
def test_mean_dynamic(context_mode):
"""
Feature: pyboost function.
Description: test function mean with dynamic shape and rank.
Expectation: expect correct result.
"""
input1 = Tensor(generate_random_input((2, 3, 4, 5), np.float32))
axis1 = (0, -1)
input2 = Tensor(generate_random_input((3, 3, 4, 4), np.float32))
axis2 = (0, -1)
TEST_OP(mean_func, [[input1, axis1], [input2, axis2]], mode=context_mode, grad=True)
input3 = Tensor(generate_random_input((3, 4, 5), np.float16))
axis3 = ()
keep_dims3 = False
dtype3 = mstype.float32
input4 = Tensor(generate_random_input((3, 4), np.float16))
axis4 = ()
keep_dims4 = False
dtype4 = mstype.float32
TEST_OP(mean_func, [[input3, axis3, keep_dims3, dtype3], [input4, axis4, keep_dims4, dtype4]],
nontensor_dynamic_type='None', mode=context_mode, grad=True, test_resize=False)
input5 = Tensor(generate_random_input((2, 3, 4), np.float32))
input6 = Tensor(generate_random_input((2, 3), np.float32))
TEST_OP(mean_func, [[input5], [input6]], mode=context_mode, grad=True)
@pytest.mark.level0
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize('keep_dims', [False, True])
@pytest.mark.parametrize('in_dtype', [mstype.float16])
@pytest.mark.parametrize('out_dtype', [mstype.float32])
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
def test_sum_normal(keep_dims, in_dtype, out_dtype, context_mode):
"""
Feature: pyboost function.
Description: test function sum forward on ascend with different datatype.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
axis = (0, -1)
x = generate_random_input((2, 3, 4, 5), mstype.dtype_to_nptype(in_dtype))
output = sum_forward_func(ms.Tensor(x), axis, keep_dims, out_dtype)
expect = generate_expect_forward_output("sum", x, axis, keep_dims, mstype.dtype_to_nptype(out_dtype))
np.testing.assert_equal(output.dtype, out_dtype)
np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
axis = (0, -1)
x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(mstype.dtype_to_nptype(in_dtype))
grads = sum_backward_func(ms.Tensor(x), axis, False, out_dtype)
expect = np.ones((2, 3, 4), mstype.dtype_to_nptype(in_dtype))
np.testing.assert_equal(grads.dtype, in_dtype)
np.testing.assert_allclose(grads.asnumpy(), expect, rtol=1e-3)
@pytest.mark.level0
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
def test_sum_default(context_mode):
"""
Feature: pyboost function.
Description: test function sum on ascend with default args.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
x = generate_random_input((2, 3, 4, 5), np.float32)
output = sum_forward_func(ms.Tensor(x))
expect = generate_expect_forward_output("sum", x)
np.testing.assert_equal(output.dtype, mstype.float32)
np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
x1 = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(np.float32)
grads = sum_backward_func(ms.Tensor(x1))
expect = np.ones((2, 3, 4), np.float32)
np.testing.assert_equal(grads.dtype, mstype.float32)
np.testing.assert_allclose(grads.asnumpy(), expect, rtol=1e-3)
@pytest.mark.level1
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
def test_sum_dynamic(context_mode):
"""
Feature: pyboost function.
Description: test function sum with dynamic shape and rank.
Expectation: expect correct result.
"""
input1 = Tensor(generate_random_input((2, 3, 4, 5), np.float32))
axis1 = (0, -1)
input2 = Tensor(generate_random_input((3, 3, 4, 4), np.float32))
axis2 = (0, -1)
TEST_OP(sum_func, [[input1, axis1], [input2, axis2]], mode=context_mode, grad=True)
input3 = Tensor(generate_random_input((3, 4, 5), np.float32))
axis3 = ()
keep_dims3 = False
dtype3 = mstype.int32
input4 = Tensor(generate_random_input((3, 4), np.float32))
axis4 = ()
keep_dims4 = False
dtype4 = mstype.int64
TEST_OP(sum_func, [[input3, axis3, keep_dims3, dtype3], [input4, axis4, keep_dims4, dtype4]],
nontensor_dynamic_type='None', mode=context_mode, grad=True, test_resize=False)
input5 = Tensor(generate_random_input((2, 3, 4), np.float32))
input6 = Tensor(generate_random_input((2, 3), np.float32))
TEST_OP(sum_func, [[input5], [input6]], mode=context_mode, grad=True)
@pytest.mark.level1
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize('axis', [(-1), ()])
@pytest.mark.parametrize('in_dtype', [mstype.float16])
@pytest.mark.parametrize('out_dtype', [mstype.float32, mstype.int8, mstype.uint8, mstype.complex128])
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
def test_sum_vaild_dtype(axis, in_dtype, out_dtype, context_mode):
"""
Feature: pyboost function.
Description: test function sum forward on ascend with different datatype.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
x = generate_random_input((2, 3, 4), mstype.dtype_to_nptype(in_dtype))
output = sum_forward_func(ms.Tensor(x), axis, False, out_dtype)
np.testing.assert_equal(output.dtype, out_dtype)
x1 = generate_random_input((3, 4, 5), mstype.dtype_to_nptype(in_dtype))
grads = sum_backward_func(ms.Tensor(x1), axis, False, out_dtype)
np.testing.assert_equal(grads.dtype, in_dtype)
@pytest.mark.level1
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize('axis', [(-1), ()])
@pytest.mark.parametrize('in_dtype', [mstype.bool_, mstype.int8, mstype.int16, mstype.int32, mstype.uint8])
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
def test_sum_default_dtype(axis, in_dtype, context_mode):
"""
Feature: pyboost function.
Description: test function sum forward on ascend with different datatype.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
x = generate_random_input((2, 3, 4), mstype.dtype_to_nptype(in_dtype))
output = sum_forward_func(ms.Tensor(x), axis, False, None)
np.testing.assert_equal(output.dtype, mstype.int64)
x1 = generate_random_input((3, 4, 5), mstype.dtype_to_nptype(in_dtype))
grads = sum_backward_func(ms.Tensor(x1), axis, False, None)
np.testing.assert_equal(grads.dtype, in_dtype)
@pytest.mark.level0
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize('keep_dims', [False, True])
@pytest.mark.parametrize('in_dtype', [mstype.float16])
@pytest.mark.parametrize('out_dtype', [mstype.float32])
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
@pytest.mark.skip(reason="No support yet")
def test_prod_normal(keep_dims, in_dtype, out_dtype, context_mode):
"""
Feature: pyboost function.
Description: test function prod forward on ascend with different datatype.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
axis = 0
x = generate_random_input((2, 3, 4, 5), mstype.dtype_to_nptype(in_dtype))
output = prod_forward_func(ms.Tensor(x), axis, keep_dims, out_dtype)
expect = generate_expect_forward_output("prod", x, axis, keep_dims, mstype.dtype_to_nptype(out_dtype))
np.testing.assert_equal(output.dtype, out_dtype)
np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
axis = -1
x = np.arange(2 * 3 * 4).reshape(2, 3, 4).astype(mstype.dtype_to_nptype(in_dtype))
grads = prod_backward_func(ms.Tensor(x), axis, False, out_dtype)
expect = np.array([[[6.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
[2.1000e+02, 1.6800e+02, 1.4000e+02, 1.2000e+02],
[9.9000e+02, 8.8000e+02, 7.9200e+02, 7.2000e+02]],
[[2.7300e+03, 2.5200e+03, 2.3400e+03, 2.1840e+03],
[5.8140e+03, 5.4720e+03, 5.1680e+03, 4.8960e+03],
[1.0626e+04, 1.0120e+04, 9.6600e+03, 9.2400e+03]]]).astype(mstype.dtype_to_nptype(in_dtype))
np.testing.assert_equal(grads.dtype, in_dtype)
np.testing.assert_allclose(grads.asnumpy(), expect, rtol=1e-3)
@pytest.mark.level1
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize('axis', [0, -1])
@pytest.mark.parametrize('keep_dims', [False, True])
@pytest.mark.parametrize('in_dtype', [mstype.float16])
@pytest.mark.parametrize('out_dtype', [mstype.float32])
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
@pytest.mark.skip(reason="No support yet")
def test_prod_normal_1d(axis, keep_dims, in_dtype, out_dtype, context_mode):
"""
Feature: pyboost function.
Description: test function prod forward on ascend with different datatype.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
x = np.random.randn(5).astype(mstype.dtype_to_nptype(in_dtype))
output = prod_forward_func(ms.Tensor(x), axis, keep_dims, out_dtype)
expect = generate_expect_forward_output("prod", x, axis, keep_dims, mstype.dtype_to_nptype(out_dtype))
np.testing.assert_equal(output.dtype, out_dtype)
np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-2)
@pytest.mark.level0
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
@pytest.mark.skip(reason="No support yet")
def test_prod_default(context_mode):
"""
Feature: pyboost function.
Description: test function prod on ascend with default args.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
x = generate_random_input((2, 3, 4, 5), np.float32)
output = prod_forward_func(ms.Tensor(x))
expect = generate_expect_forward_output("prod", x)
np.testing.assert_equal(output.dtype, mstype.float32)
np.testing.assert_allclose(output.asnumpy(), expect, rtol=1e-3)
x1 = np.arange(2 * 3).reshape(2, 3).astype(np.float32)
grads = prod_backward_func(ms.Tensor(x1))
expect = np.array([[120, 0, 0], [0, 0, 0]]).astype(np.float32)
np.testing.assert_equal(grads.dtype, mstype.float32)
np.testing.assert_allclose(grads.asnumpy(), expect, rtol=1e-3)
@pytest.mark.level1
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
@pytest.mark.skip(reason="No support yet")
def test_prod_dynamic(context_mode):
"""
Feature: pyboost function.
Description: test function prod with dynamic shape and rank.
Expectation: expect correct result.
"""
input1 = Tensor(generate_random_input((2, 3, 4, 5), np.float32))
axis1 = -1
input2 = Tensor(generate_random_input((3, 3, 4, 4), np.float32))
axis2 = -1
TEST_OP(prod_func, [[input1, axis1], [input2, axis2]], mode=context_mode, grad=True)
input3 = Tensor(generate_random_input((3, 4, 5), np.float32))
axis3 = 0
keep_dims3 = False
dtype3 = mstype.int32
input4 = Tensor(generate_random_input((3, 4), np.float32))
axis4 = 0
keep_dims4 = False
dtype4 = mstype.int64
TEST_OP(prod_func, [[input3, axis3, keep_dims3, dtype3], [input4, axis4, keep_dims4, dtype4]],
nontensor_dynamic_type='None', mode=context_mode, grad=True, test_resize=False)
input5 = Tensor(generate_random_input((2, 3, 4), np.float32))
input6 = Tensor(generate_random_input((2, 3), np.float32))
TEST_OP(prod_func, [[input5], [input6]], mode=context_mode, grad=True)
@pytest.mark.level1
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize('axis', [-1, None])
@pytest.mark.parametrize('in_dtype', [mstype.float16])
@pytest.mark.parametrize('out_dtype', [mstype.float32, mstype.int8, mstype.uint8])
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
@pytest.mark.skip(reason="No support yet")
def test_prod_vaild_dtype(axis, in_dtype, out_dtype, context_mode):
"""
Feature: pyboost function.
Description: test function prod forward on ascend with different datatype.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
x = generate_random_input((2, 3, 4), mstype.dtype_to_nptype(in_dtype))
output = prod_forward_func(ms.Tensor(x), axis, False, out_dtype)
np.testing.assert_equal(output.dtype, out_dtype)
x1 = generate_random_input((3, 4, 5), mstype.dtype_to_nptype(in_dtype))
grads = prod_backward_func(ms.Tensor(x1), axis, False, out_dtype)
np.testing.assert_equal(grads.dtype, in_dtype)
@pytest.mark.level1
@pytest.mark.env_onecard
@pytest.mark.platform_x86_cpu
@pytest.mark.platform_x86_gpu_training
@pytest.mark.platform_arm_ascend_training
@pytest.mark.parametrize('axis', [-1, None])
@pytest.mark.parametrize('in_dtype', [mstype.int8, mstype.int16, mstype.int32, mstype.uint8])
@pytest.mark.parametrize("context_mode", [ms.GRAPH_MODE, ms.PYNATIVE_MODE])
@pytest.mark.skip(reason="No support yet")
def test_prod_default_dtype(axis, in_dtype, context_mode):
"""
Feature: pyboost function.
Description: test function prod forward on ascend with different datatype.
Expectation: expect correct result.
"""
ms.context.set_context(mode=context_mode)
x = generate_random_input((2, 3, 4), mstype.dtype_to_nptype(in_dtype))
output = prod_forward_func(ms.Tensor(x), axis, False, None)
np.testing.assert_equal(output.dtype, mstype.int64)
x1 = generate_random_input((3, 4, 5), mstype.dtype_to_nptype(in_dtype))
grads = prod_backward_func(ms.Tensor(x1), axis, False, None)
np.testing.assert_equal(grads.dtype, in_dtype)

View File

@ -0,0 +1,271 @@
/**
* Copyright 2024 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include <memory>
#include "common/common_test.h"
#include "ops/ops_func_impl/mean_ext.h"
#include "ops/ops_func_impl/sum_ext.h"
// #include "ops/ops_func_impl/prod_ext.h"
#include "ops/auto_generate/gen_ops_name.h"
#include "ops/test_ops.h"
#include "ops/test_ops_cmp_utils.h"
#include "ops/test_value_utils.h"
#include "abstract/abstract_value.h"
#include "abstract/ops/primitive_infer_map.h"
namespace mindspore {
namespace ops {
namespace {
struct ReduceExtandParams {
ShapeVector input_shape;
TypePtr input_type;
ShapeVector output_shape;
TypePtr output_type;
AbstractBasePtr axis;
AbstractBasePtr keep_dims;
AbstractBasePtr dtype;
};
static auto value_any = std::make_shared<abstract::AbstractScalar>(kValueAny, kTypeAny);
static auto value_none = std::make_shared<abstract::AbstractScalar>(kValueAny, kTypeNone);
static auto keep_dims_true = std::make_shared<BoolImm>(true)->ToAbstract();
static auto keep_dims_false = std::make_shared<BoolImm>(false)->ToAbstract();
static auto dtype_float64 = std::make_shared<Int64Imm>(kNumberTypeFloat64)->ToAbstract();
static auto dtype_int32 = std::make_shared<Int64Imm>(kNumberTypeInt32)->ToAbstract();
static auto dtype_int16 = std::make_shared<Int64Imm>(kNumberTypeInt16)->ToAbstract();
static auto dtype_int8 = std::make_shared<Int64Imm>(kNumberTypeInt8)->ToAbstract();
static auto dtype_uint8 = std::make_shared<Int64Imm>(kNumberTypeUInt8)->ToAbstract();
static auto dtype_bool = std::make_shared<Int64Imm>(kNumberTypeBool)->ToAbstract();
AbstractBasePtr CreateInt(const int &value) {
return CreatePyInt(value)->ToAbstract();
}
AbstractBasePtr CreateIntTuple(const std::vector<NumberContainer> &value) {
return CreatePyIntTuple(value)->ToAbstract();
}
template <typename T>
tensor::TensorPtr CreateTensor(const ShapeVector &shape, const TypeId &dtype, std::vector<T> value) {
void *data_ptr = &value[0];
auto tensor = std::make_shared<tensor::Tensor>(dtype, shape, data_ptr, dtype);
return tensor;
}
static std::map<std::string, OpFuncImplPtr> reduce_extand_func_impl = {
{kNameMeanExt, std::make_shared<MeanExtFuncImpl>()},
{kNameSumExt, std::make_shared<SumExtFuncImpl>()},
// {kNameProdExt, std::make_shared<ProdExtFuncImpl>()},
};
} // namespace
class TestReduceExtand : public TestOps, public testing::WithParamInterface<std::tuple<const char *, ReduceExtandParams>> {};
TEST_P(TestReduceExtand, dyn_shape) {
const auto &op_name = std::get<0>(GetParam());
const auto &param = std::get<1>(GetParam());
ASSERT_TRUE(reduce_extand_func_impl.find(op_name) != reduce_extand_func_impl.end());
auto op_impl = reduce_extand_func_impl[op_name];
ASSERT_NE(op_impl, nullptr);
auto prim = std::make_shared<Primitive>(op_name);
ASSERT_NE(prim, nullptr);
auto input = std::make_shared<abstract::AbstractTensor>(param.input_type, param.input_shape);
ASSERT_NE(input, nullptr);
auto input_args = std::vector<AbstractBasePtr>{input, param.axis, param.keep_dims, param.dtype};
auto expect_shape = std::make_shared<abstract::TensorShape>(param.output_shape);
ASSERT_NE(expect_shape, nullptr);
auto expect_type = std::make_shared<TensorType>(param.output_type);
ASSERT_NE(expect_type, nullptr);
auto out_shape = op_impl->InferShape(prim, input_args);
auto out_type = op_impl->InferType(prim, input_args);
ShapeCompare(out_shape, expect_shape);
TypeCompare(out_type, expect_type);
}
auto ReduceExtandTestCase = testing::ValuesIn(
{ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 1, 4}, kFloat32, CreateIntTuple({1}), keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 4}, kFloat32, CreateIntTuple({1}), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {1, 1, 4}, kFloat32, CreateIntTuple({0, 1}), keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {4}, kFloat32, CreateIntTuple({0, 1}), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 3, 1}, kFloat32, CreateIntTuple({-1}), keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 4}, kFloat32, CreateIntTuple({-2}), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 1, 1}, kFloat32, CreateIntTuple({-1, -2}), keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {4}, kFloat32, CreateIntTuple({-2, -3}), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-1, 1, -1}, kFloat32, CreateIntTuple({kValueAny, 1}), keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-1}, kFloat32, CreateIntTuple({kValueAny, 1}), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-1, -1, -1}, kFloat32, CreateIntTuple({kValueAny, kValueAny}), keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-1}, kFloat32, CreateIntTuple({kValueAny, kValueAny}), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-1, -1, -1}, kFloat32, value_any, keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-2}, kFloat32, value_any, keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {}, kFloat32, CreateIntTuple({}), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-2}, kFloat32, CreateIntTuple({1}), value_any, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-2}, kFloat32, CreateIntTuple({1, 2}), value_any, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, 1, 4}, kFloat32, CreateIntTuple({1}), keep_dims_true, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, 4}, kFloat32, CreateIntTuple({1}), keep_dims_false, value_none},
ReduceExtandParams{{-1, 3, 4}, kFloat32, {1, 3, 1}, kFloat32, CreateIntTuple({0, 2}), keep_dims_true, value_none},
ReduceExtandParams{{-1, 3, 4}, kFloat32, {3}, kFloat32, CreateIntTuple({0, 2}), keep_dims_false, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, 1, -1}, kFloat32, CreateIntTuple({kValueAny, 1}), keep_dims_true, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1}, kFloat32, CreateIntTuple({kValueAny, 1}), keep_dims_false, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, -1, -1}, kFloat32, CreateIntTuple({kValueAny, kValueAny}), keep_dims_true, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1}, kFloat32, CreateIntTuple({kValueAny, kValueAny}), keep_dims_false, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, -1, -1}, kFloat32, value_any, keep_dims_true, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-2}, kFloat32, value_any, keep_dims_false, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {}, kFloat32, CreateIntTuple({}), keep_dims_false, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-2}, kFloat32, CreateIntTuple({1}), value_any, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-2}, kFloat32, CreateIntTuple({1, 2}), value_any, value_none},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat32, CreateIntTuple({1}), keep_dims_true, value_none},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat32, CreateIntTuple({0, 2}), keep_dims_false, value_none},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat32, CreateIntTuple({kValueAny, 1}), keep_dims_true, value_none},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat32, value_any, keep_dims_true, value_none},
ReduceExtandParams{{-2}, kFloat32, {}, kFloat32, CreateIntTuple({}), keep_dims_false, value_none},
ReduceExtandParams{{-1, -1, -1}, kFloat32, {1, 1, 1}, kFloat64, value_none, keep_dims_true, dtype_float64},
ReduceExtandParams{{-1, -1, -1}, kFloat32, {}, kFloat64, value_none, keep_dims_false, dtype_float64},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat64, value_none, keep_dims_true, dtype_float64},
ReduceExtandParams{{-2}, kFloat32, {}, kFloat64, value_none, keep_dims_false, dtype_float64},
ReduceExtandParams{{-1, -1, -1}, kFloat32, {-2}, kFloat64, value_none, value_any, dtype_float64},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat64, value_none, value_any, dtype_float64},
ReduceExtandParams{{}, kFloat32, {-2}, kFloat32, CreateIntTuple({0}), value_any, value_none},
ReduceExtandParams{{}, kFloat32, {}, kFloat32, CreateIntTuple({0}), keep_dims_true, value_none},
ReduceExtandParams{{}, kFloat32, {}, kFloat32, CreateIntTuple({0}), keep_dims_false, value_none},
ReduceExtandParams{{}, kFloat32, {-2}, kFloat32, value_any, value_any, value_none},
ReduceExtandParams{{}, kFloat32, {}, kFloat32, value_any, keep_dims_true, value_none},
ReduceExtandParams{{}, kFloat32, {-2}, kFloat32, value_any, keep_dims_false, value_none}});
auto ReduceExtandTestCase_ProdExt = testing::ValuesIn(
{ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 1, 4}, kFloat32, CreateInt(1), keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 4}, kFloat32, CreateInt(1), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 3, 1}, kFloat32, CreateInt(-1), keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {2, 4}, kFloat32, CreateInt(-2), keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-1, -1, -1}, kFloat32, value_any, keep_dims_true, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-1, -1}, kFloat32, value_any, keep_dims_false, value_none},
ReduceExtandParams{{2, 3, 4}, kFloat32, {-2}, kFloat32, CreateInt(1), value_any, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, 1, 4}, kFloat32, CreateInt(1), keep_dims_true, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, 4}, kFloat32, CreateInt(1), keep_dims_false, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, -1, -1}, kFloat32, value_any, keep_dims_true, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-1, -1}, kFloat32, value_any, keep_dims_false, value_none},
ReduceExtandParams{{-1, -1, 4}, kFloat32, {-2}, kFloat32, CreateInt(1), value_any, value_none},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat32, CreateInt(1), keep_dims_true, value_none},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat32, value_any, keep_dims_true, value_none},
ReduceExtandParams{{-1, -1, -1}, kFloat32, {}, kFloat64, value_none, keep_dims_false, dtype_float64},
ReduceExtandParams{{-2}, kFloat32, {}, kFloat64, value_none, keep_dims_false, dtype_float64},
ReduceExtandParams{{-1, -1, -1}, kFloat32, {-2}, kFloat64, value_none, value_any, dtype_float64},
ReduceExtandParams{{-2}, kFloat32, {-2}, kFloat64, value_none, value_any, dtype_float64},
ReduceExtandParams{{}, kFloat32, {-2}, kFloat32, CreateInt(0), value_any, value_none},
ReduceExtandParams{{}, kFloat32, {}, kFloat32, CreateInt(0), keep_dims_true, value_none},
ReduceExtandParams{{}, kFloat32, {}, kFloat32, CreateInt(0), keep_dims_false, value_none},
ReduceExtandParams{{}, kFloat32, {-2}, kFloat32, value_any, value_any, value_none},
ReduceExtandParams{{}, kFloat32, {}, kFloat32, value_any, keep_dims_true, value_none},
ReduceExtandParams{{}, kFloat32, {}, kFloat32, value_any, keep_dims_false, value_none}});
auto ReduceExtandTestCase_ExtraDtype = testing::ValuesIn(
{ReduceExtandParams{{}, kFloat32, {-2}, kFloat32, value_any, value_any, value_none},
ReduceExtandParams{{}, kComplex64, {-2}, kComplex64, value_any, value_any, value_none},
ReduceExtandParams{{}, kInt32, {-2}, kInt64, value_any, value_any, value_none},
ReduceExtandParams{{}, kInt16, {-2}, kInt64, value_any, value_any, value_none},
ReduceExtandParams{{}, kInt8, {-2}, kInt64, value_any, value_any, value_none},
ReduceExtandParams{{}, kUInt8, {-2}, kInt64, value_any, value_any, value_none},
ReduceExtandParams{{}, kBool, {-2}, kInt64, value_any, value_any, value_none},
ReduceExtandParams{{}, kInt32, {-2}, kBool, value_any, value_any, dtype_bool},
ReduceExtandParams{{}, kInt16, {-2}, kUInt8, value_any, value_any, dtype_uint8},
ReduceExtandParams{{}, kInt8, {-2}, kBool, value_any, value_any, dtype_bool},
ReduceExtandParams{{}, kUInt8, {-2}, kInt16, value_any, value_any, dtype_int16},
ReduceExtandParams{{}, kBool, {-2}, kInt32, value_any, value_any, dtype_int32}});
INSTANTIATE_TEST_CASE_P(TestMeanExtGroup, TestReduceExtand,
testing::Combine(testing::ValuesIn({kNameMeanExt}), ReduceExtandTestCase));
INSTANTIATE_TEST_CASE_P(TestSumExtGroup, TestReduceExtand,
testing::Combine(testing::ValuesIn({kNameSumExt}), ReduceExtandTestCase));
INSTANTIATE_TEST_CASE_P(TestSumExtGroup_ExtraDtype, TestReduceExtand,
testing::Combine(testing::ValuesIn({kNameSumExt}), ReduceExtandTestCase_ExtraDtype));
// INSTANTIATE_TEST_CASE_P(TestProdExtGroup, TestReduceExtand,
// testing::Combine(testing::ValuesIn({kNameProdExt}), ReduceExtandTestCase_ProdExt));
// INSTANTIATE_TEST_CASE_P(TestProdExtGroup_ExtraDtype, TestReduceExtand,
// testing::Combine(testing::ValuesIn({kNameProdExt}), ReduceExtandTestCase_ExtraDtype));
struct ReduceExtandInferValueParams {
tensor::TensorPtr input;
AbstractBasePtr axis;
AbstractBasePtr keep_dims;
AbstractBasePtr dtype;
tensor::TensorPtr out;
};
class TestReduceExtandInferValue : public TestOps, public testing::WithParamInterface<std::tuple<const char *, ReduceExtandInferValueParams>> {};
TEST_P(TestReduceExtandInferValue, dyn_shape_infer_value) {
const auto &op_name = std::get<0>(GetParam());
const auto &param = std::get<1>(GetParam());
auto primitive = std::make_shared<Primitive>(op_name);
ASSERT_NE(primitive, nullptr);
ASSERT_NE(param.input, nullptr);
auto input = param.input->ToAbstract();
ASSERT_NE(input, nullptr);
auto input_args = abstract::AbstractBasePtrList{input, param.axis, param.keep_dims, param.dtype};
auto value_opt = abstract::InferValueByFuncImpl(primitive, input_args);
if (!value_opt.has_value()) {
MS_LOG(ERROR) << op_name << " have no infer value implement!";
ASSERT_TRUE(false);
}
auto infer_out = value_opt.value();
if (infer_out == nullptr) {
MS_LOG(ERROR) << op_name << " can not infer value with inputs: " << input_args;
ASSERT_TRUE(false);
}
auto infer_tensor = infer_out->cast<tensor::TensorPtr>();
ASSERT_NE(infer_tensor, nullptr);
ASSERT_TRUE(infer_tensor->ValueEqual(*param.out));
}
auto ReduceExtandInferValueTestCase_MeanExt = testing::ValuesIn(
{ReduceExtandInferValueParams{
CreateTensor<float>(ShapeVector{2, 2}, kNumberTypeFloat32, std::vector<float>{2, 3, 4, 5}),
CreateIntTuple({0}), keep_dims_false, dtype_float64,
CreateTensor<double>(ShapeVector{2}, kNumberTypeFloat64, std::vector<double>{3, 4})},
ReduceExtandInferValueParams{
CreateTensor<float>(ShapeVector{2, 2}, kNumberTypeFloat32, std::vector<float>{2, 3, 4, 5}),
value_none, keep_dims_false, dtype_float64,
CreateTensor<double>(ShapeVector{}, kNumberTypeFloat64, std::vector<double>{3.5})},
ReduceExtandInferValueParams{
CreateTensor<float>(ShapeVector{2, 2}, kNumberTypeFloat32, std::vector<float>{2, 3, 4, 5}),
CreateIntTuple({0}), keep_dims_true, value_none,
CreateTensor<float>(ShapeVector{1, 2}, kNumberTypeFloat32, std::vector<float>{3, 4})}});
auto ReduceExtandInferValueTestCase_SumExt = testing::ValuesIn(
{ReduceExtandInferValueParams{
CreateTensor<float>(ShapeVector{2, 2}, kNumberTypeFloat32, std::vector<float>{2, 3, 4, 5}),
CreateIntTuple({0}), keep_dims_false, dtype_float64,
CreateTensor<double>(ShapeVector{2}, kNumberTypeFloat64, std::vector<double>{6, 8})},
ReduceExtandInferValueParams{
CreateTensor<float>(ShapeVector{2, 2}, kNumberTypeFloat32, std::vector<float>{2, 3, 4, 5}),
value_none, keep_dims_false, dtype_float64,
CreateTensor<double>(ShapeVector{}, kNumberTypeFloat64, std::vector<double>{14})},
ReduceExtandInferValueParams{
CreateTensor<float>(ShapeVector{2, 2}, kNumberTypeFloat32, std::vector<float>{2, 3, 4, 5}),
CreateIntTuple({0}), keep_dims_true, value_none,
CreateTensor<float>(ShapeVector{1, 2}, kNumberTypeFloat32, std::vector<float>{6, 8})}});
INSTANTIATE_TEST_CASE_P(TestMeanExtInferValueGroup, TestReduceExtandInferValue,
testing::Combine(testing::ValuesIn({kNameMeanExt}), ReduceExtandInferValueTestCase_MeanExt));
INSTANTIATE_TEST_CASE_P(TestSumExtInferValueGroup, TestReduceExtandInferValue,
testing::Combine(testing::ValuesIn({kNameSumExt}), ReduceExtandInferValueTestCase_SumExt));
} // namespace ops
} // namespace mindspore

View File

@ -0,0 +1,21 @@
# Copyright 2024 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""setup for pytest"""
import mindspore.context as context
# pylint: disable=unused-argument
def setup_module(module):
context.set_context(mode=context.GRAPH_MODE)

View File

@ -75,20 +75,6 @@ def test_fallback_sum_with_x_tensor_n_default():
assert out.asnumpy() == 6
def test_fallback_sum_with_x_tensor_n_default_2():
"""
Feature: JIT Fallback
Description: Test sum() in graph mode with input x tensor and input n default.
Expectation: No exception.
"""
@jit
def foo():
x = sum(Tensor([[1, 1], [2, 2]]))
return x
out = foo()
assert np.allclose(out.asnumpy(), np.array([3, 3]))
def test_fallback_sum_with_x_numpy_array_n_default_2():
"""
Feature: JIT Fallback