Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!35475 fix scatter max gpu kernel and code check 

Merge pull request !35475 from polyhedral/max
This commit is contained in:
i-robot 2022-06-07 02:00:23 +00:00 committed by Gitee
parent c3bfbf2c3c bd7278de51
commit 7b5287ffc7
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
7 changed files with 27 additions and 217 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/ccsrc/plugin/device/gpu/kernel/cuda_impl/cuda_ops/scatter_functor_impl.cu
View File

@ -69,7 +69,7 @@ __global__ void ScatterMaxKernel(S size_limit, const size_t inner_size, const si
continue;
}
const size_t current_pos = indices[index] * inner_size + offset;
input[current_pos] = updates[pos] > input[current_pos] ? updates[pos] : input[current_pos];
MsAtomicMax(&input[current_pos], updates[pos]);
}
}

8
mindspore/core/ops/scatter_arithmetic.cc
View File

@ -19,6 +19,8 @@
#include <set>
#include <string>
#include "ops/scatter_update.h"
#include "ops/scatter_min.h"
#include "ops/scatter_max.h"
#include "abstract/ops/primitive_infer_map.h"
#include "ops/op_utils.h"
#include "utils/check_convert_utils.h"
@ -88,5 +90,11 @@ AbstractBasePtr ScatterArithmeticInfer(const abstract::AnalysisEnginePtr &, cons
MIND_API_OPERATOR_IMPL(ScatterUpdate, BaseOperator);
REGISTER_PRIMITIVE_EVAL_IMPL(ScatterUpdate, prim::kPrimScatterUpdate, ScatterArithmeticInfer, nullptr, true);
MIND_API_OPERATOR_IMPL(ScatterMin, BaseOperator);
REGISTER_PRIMITIVE_EVAL_IMPL(ScatterMin, prim::kPrimScatterMin, ScatterArithmeticInfer, nullptr, true);
MIND_API_OPERATOR_IMPL(ScatterMax, BaseOperator);
REGISTER_PRIMITIVE_EVAL_IMPL(ScatterMax, prim::kPrimScatterMax, ScatterArithmeticInfer, nullptr, true);
} // namespace ops
} // namespace mindspore

94
mindspore/core/ops/scatter_max.cc
View File

@ -1,94 +0,0 @@
/**
* Copyright 2022 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/scatter_max.h"
#include <set>
#include <map>
#include <string>
#include "abstract/ops/primitive_infer_map.h"
#include "ops/op_utils.h"
#include "utils/check_convert_utils.h"
#include "mindapi/src/helper.h"
namespace mindspore {
namespace ops {
namespace {
abstract::ShapePtr ScatterMaxInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
BaseShapePtr input_x_shape_ptr = input_args[kInputIndex0]->BuildShape();
MS_EXCEPTION_IF_NULL(input_x_shape_ptr);
BaseShapePtr indices_shape_ptr = input_args[kInputIndex1]->BuildShape();
MS_EXCEPTION_IF_NULL(indices_shape_ptr);
BaseShapePtr updates_shape_ptr = input_args[kInputIndex2]->BuildShape();
MS_EXCEPTION_IF_NULL(updates_shape_ptr);
if (input_x_shape_ptr->IsDynamic()) {
MS_EXCEPTION(ValueError) << "For " << primitive->name() << ", "
<< "the 'input_x' does not support dynamic shape, but got the shape of 'input_x' is "
<< input_x_shape_ptr->ToString();
}
if (indices_shape_ptr->IsDynamic() || updates_shape_ptr->IsDynamic()) {
return input_x_shape_ptr->cast<abstract::ShapePtr>();
}
std::vector<int64_t> input_x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_x_shape_ptr)[kShape];
std::vector<int64_t> indices_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(indices_shape_ptr)[kShape];
std::vector<int64_t> updates_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(updates_shape_ptr)[kShape];
std::vector<int64_t> check_update_shape(indices_shape);
for (int64_t i = 1; i < SizeToLong(input_x_shape.size()); ++i) {
check_update_shape.push_back(input_x_shape[i]);
}
if (updates_shape != check_update_shape) {
MS_EXCEPTION(ValueError) << "For " << primitive->name() << ", "
<< "updates_shape = indices_shape + x_shape[1:], but got x_shape: "
<< input_x_shape_ptr->ToString() << ", indices_shape: " << indices_shape_ptr->ToString()
<< ", updates_shape: " << updates_shape_ptr->ToString() << ".";
}
auto output_shape = input_args[kInputIndex0]->BuildShape()->cast<abstract::ShapePtr>();
return output_shape;
}
TypePtr ScatterMaxInferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
auto input_x_type_ptr = input_args[kInputIndex0]->BuildType();
auto indiecs_type_ptr = input_args[kInputIndex1]->BuildType();
auto updates_type_ptr = input_args[kInputIndex2]->BuildType();
auto prim_name = primitive->name();
const std::set<TypePtr> indices_types = {kInt32, kInt64};
const std::set<TypePtr> valid_types = {kInt32, kInt64, kFloat16, kFloat32, kFloat64};
(void)CheckAndConvertUtils::CheckTensorTypeValid("indices type", indiecs_type_ptr, indices_types, prim_name);
(void)CheckAndConvertUtils::CheckTensorTypeValid("input_x type", input_x_type_ptr, valid_types, prim_name);
(void)CheckAndConvertUtils::CheckTensorTypeValid("updates type", updates_type_ptr, valid_types, prim_name);
std::map<std::string, TypePtr> type_dict;
type_dict.emplace("input_x", input_x_type_ptr);
type_dict.emplace("updates", updates_type_ptr);
return CheckAndConvertUtils::CheckTensorTypeSame(type_dict, common_valid_types, prim_name);
}
} // namespace
MIND_API_OPERATOR_IMPL(ScatterMax, BaseOperator);
AbstractBasePtr ScatterMaxInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) {
MS_EXCEPTION_IF_NULL(primitive);
const int64_t input_num = 3;
(void)CheckAndConvertUtils::CheckInputArgs(input_args, kGreaterEqual, input_num, primitive->name());
auto infer_type = ScatterMaxInferType(primitive, input_args);
auto infer_shape = ScatterMaxInferShape(primitive, input_args);
return abstract::MakeAbstract(infer_shape, infer_type);
}
REGISTER_PRIMITIVE_EVAL_IMPL(ScatterMax, prim::kPrimScatterMax, ScatterMaxInfer, nullptr, true);
} // namespace ops
} // namespace mindspore

3
mindspore/core/ops/scatter_max.h
View File

@ -30,9 +30,6 @@ class MIND_API ScatterMax : public BaseOperator {
/// \brief Constructor.
ScatterMax() : BaseOperator(kNameScatterMax) { InitIOName({"input_x", "indices", "updates"}, {"output"}); }
};
abstract::AbstractBasePtr ScatterMaxInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<abstract::AbstractBasePtr> &input_args);
} // namespace ops
} // namespace mindspore

94
mindspore/core/ops/scatter_min.cc
View File

@ -1,94 +0,0 @@
/**
* Copyright 2022 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/scatter_min.h"
#include <set>
#include <map>
#include <string>
#include "abstract/ops/primitive_infer_map.h"
#include "ops/op_utils.h"
#include "utils/check_convert_utils.h"
#include "mindapi/src/helper.h"
namespace mindspore {
namespace ops {
namespace {
abstract::ShapePtr ScatterMinInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
BaseShapePtr input_x_shape_ptr = input_args[kInputIndex0]->BuildShape();
MS_EXCEPTION_IF_NULL(input_x_shape_ptr);
BaseShapePtr indices_shape_ptr = input_args[kInputIndex1]->BuildShape();
MS_EXCEPTION_IF_NULL(indices_shape_ptr);
BaseShapePtr updates_shape_ptr = input_args[kInputIndex2]->BuildShape();
MS_EXCEPTION_IF_NULL(updates_shape_ptr);
if (input_x_shape_ptr->IsDynamic()) {
MS_EXCEPTION(ValueError) << "For " << primitive->name() << ", "
<< "the 'input_x' does not support dynamic shape, but got the shape of 'input_x' is "
<< input_x_shape_ptr->ToString();
}
if (indices_shape_ptr->IsDynamic() || updates_shape_ptr->IsDynamic()) {
return input_x_shape_ptr->cast<abstract::ShapePtr>();
}
std::vector<int64_t> input_x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_x_shape_ptr)[kShape];
std::vector<int64_t> indices_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(indices_shape_ptr)[kShape];
std::vector<int64_t> updates_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(updates_shape_ptr)[kShape];
std::vector<int64_t> check_update_shape(indices_shape);
for (int64_t i = 1; i < SizeToLong(input_x_shape.size()); ++i) {
check_update_shape.push_back(input_x_shape[i]);
}
if (updates_shape != check_update_shape) {
MS_EXCEPTION(ValueError) << "For " << primitive->name() << ", "
<< "updates_shape = indices_shape + x_shape[1:], but got x_shape: "
<< input_x_shape_ptr->ToString() << ", indices_shape: " << indices_shape_ptr->ToString()
<< ", updates_shape: " << updates_shape_ptr->ToString() << ".";
}
auto output_shape = input_args[kInputIndex0]->BuildShape()->cast<abstract::ShapePtr>();
return output_shape;
}
TypePtr ScatterMinInferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
auto input_x_type_ptr = input_args[kInputIndex0]->BuildType();
auto indiecs_type_ptr = input_args[kInputIndex1]->BuildType();
auto updates_type_ptr = input_args[kInputIndex2]->BuildType();
auto prim_name = primitive->name();
const std::set<TypePtr> indices_types = {kInt32, kInt64};
const std::set<TypePtr> valid_types = {kInt32, kInt64, kFloat16, kFloat32, kFloat64};
(void)CheckAndConvertUtils::CheckTensorTypeValid("indices type", indiecs_type_ptr, indices_types, prim_name);
(void)CheckAndConvertUtils::CheckTensorTypeValid("input_x type", input_x_type_ptr, valid_types, prim_name);
(void)CheckAndConvertUtils::CheckTensorTypeValid("updates type", updates_type_ptr, valid_types, prim_name);
std::map<std::string, TypePtr> type_dict;
type_dict.emplace("input_x", input_x_type_ptr);
type_dict.emplace("updates", updates_type_ptr);
return CheckAndConvertUtils::CheckTensorTypeSame(type_dict, common_valid_types, prim_name);
}
} // namespace
MIND_API_OPERATOR_IMPL(ScatterMin, BaseOperator);
AbstractBasePtr ScatterMinInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) {
MS_EXCEPTION_IF_NULL(primitive);
const int64_t input_num = 3;
(void)CheckAndConvertUtils::CheckInputArgs(input_args, kGreaterEqual, input_num, primitive->name());
auto infer_type = ScatterMinInferType(primitive, input_args);
auto infer_shape = ScatterMinInferShape(primitive, input_args);
return abstract::MakeAbstract(infer_shape, infer_type);
}
REGISTER_PRIMITIVE_EVAL_IMPL(ScatterMin, prim::kPrimScatterMin, ScatterMinInfer, nullptr, true);
} // namespace ops
} // namespace mindspore

3
mindspore/core/ops/scatter_min.h
View File

@ -30,9 +30,6 @@ class MIND_API ScatterMin : public BaseOperator {
/// \brief Constructor.
ScatterMin() : BaseOperator(kNameScatterMin) { InitIOName({"input_x", "indices", "updates"}, {"output"}); }
};
abstract::AbstractBasePtr ScatterMinInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<abstract::AbstractBasePtr> &input_args);
} // namespace ops
} // namespace mindspore

40
tests/st/ops/gpu/test_scatter_func_op.py
View File

@ -150,7 +150,7 @@ def test_scatter_func_small_float32():
# max
output = scatter_func_net("max", inputx, indices, updates)
expected = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]])
expected = np.array([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]])
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
# min
@ -189,7 +189,7 @@ def test_scatter_func_input_updated():
# max
net = TestScatterFuncNet("max", lock, inputx, indices, updates)
net()
expected = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]])
expected = np.array([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]])
np.testing.assert_array_almost_equal(net.inputx.asnumpy(), expected)
# min
@ -455,8 +455,8 @@ def test_scatter_func_input_less_than_1_float32():
output = scatter_func_net("max", inputx, indices, updates)
expected = np.array(
[
[37.0, 38.0, 39.0],
[34.0, 35.0, 66.0],
[55.0, 56.0, 57.0],
[64.0, 65.0, 66.0],
[67.0, 68.0, 69.0],
],
dtype=np.float32,
@ -494,7 +494,7 @@ def test_scatter_func_float16():
# max
output = scatter_func_net("max", inputx, indices, updates)
expected = np.array([[6.0, 1.0, 2.0], [3.0, 4.0, 5.0]])
expected = np.array([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]])
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
# min
@ -546,8 +546,8 @@ def test_scatter_func_large_float16():
# max
output = scatter_func_net("max", inputx, indices, updates)
expected = np.array([
[[63.0, 64.0, 65.0, 66.0], [67.0, 68.0, 69.0, 70.0], [71.0, 72.0, 73.0, 74.0]],
[[99.0, 100.0, 101.0, 102.0], [103.0, 104.0, 105.0, 106.0], [95.0, 96.0, 97.0, 98.0]],
[[75.0, 76.0, 77.0, 78.0], [79.0, 80.0, 81.0, 82.0], [83.0, 84.0, 85.0, 86.0]],
[[99.0, 100.0, 101.0, 102.0], [103.0, 104.0, 105.0, 106.0], [107.0, 108.0, 109.0, 110.0]],
])
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
@ -595,8 +595,8 @@ def test_scatter_func_disordered_float16():
expected = np.array(
[
[95.0, 96.0, 97.0, 98.0],
[67.0, 68.0, 69.0, 70.0],
[99.0, 100.0, 101.0, 102.0],
[79.0, 80.0, 81.0, 82.0],
[107.0, 108.0, 109.0, 110.0],
]
)
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
@ -650,8 +650,8 @@ def test_scatter_func_large_int32():
# max
output = scatter_func_net("max", inputx, indices, updates)
expected = np.array([
[[63.0, 64.0, 65.0, 66.0], [67.0, 68.0, 69.0, 70.0], [71.0, 72.0, 73.0, 74.0]],
[[99.0, 100.0, 101.0, 102.0], [103.0, 104.0, 105.0, 106.0], [95.0, 96.0, 97.0, 98.0]],
[[75.0, 76.0, 77.0, 78.0], [79.0, 80.0, 81.0, 82.0], [83.0, 84.0, 85.0, 86.0]],
[[99.0, 100.0, 101.0, 102.0], [103.0, 104.0, 105.0, 106.0], [107.0, 108.0, 109.0, 110.0]],
])
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
@ -699,8 +699,8 @@ def test_scatter_func_disordered_int32():
expected = np.array(
[
[95.0, 96.0, 97.0, 98.0],
[67.0, 68.0, 69.0, 70.0],
[99.0, 100.0, 101.0, 102.0],
[79.0, 80.0, 81.0, 82.0],
[107.0, 108.0, 109.0, 110.0],
]
)
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
@ -747,7 +747,7 @@ def test_scatter_func_disordered_dynamic_int32():
# max
output = scatter_func_d_net("max", inputx, indices, updates)
expected = np.array(
[[95.0, 96.0, 97.0, 98.0], [67.0, 68.0, 69.0, 70.0], [99.0, 100.0, 101.0, 102.0]]
[[95.0, 96.0, 97.0, 98.0], [79.0, 80.0, 81.0, 82.0], [107.0, 108.0, 109.0, 110.0]]
).astype(np.int32)
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
@ -839,7 +839,7 @@ def test_scatter_func_input_less_than_1_dynamic_float32():
# update
output = scatter_func_d_net("update", inputx, indices, updates)
expected = np.array(
[[37.0, 38.0, 39.0], [34.0, 35.0, 66.0], [67.0, 68.0, 69.0],], dtype=np.float32,
[[37.0, 38.0, 39.0], [34.0, 35.0, 66.0], [67.0, 68.0, 69.0]], dtype=np.float32,
)
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
@ -870,7 +870,7 @@ def test_scatter_func_input_less_than_1_dynamic_float32():
# max
output = scatter_func_d_net("max", inputx, indices, updates)
expected = np.array(
[[37.0, 38.0, 39.0], [34.0, 35.0, 66.0], [67.0, 68.0, 69.0],], dtype=np.float32,
[[55.0, 56.0, 57.0], [64.0, 65.0, 66.0], [67.0, 68.0, 69.0]], dtype=np.float32,
)
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
@ -921,8 +921,8 @@ def test_scatter_func_dynamic_two_inputs():
output_1, output_2 = scatter_func_d2_net(
"max", inputx, indices_1, updates_1, indices_2, updates_2
)
expected_1 = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]])
expected_2 = np.array([[17.0, 16.0, 15.0], [11.0, 10.0, 9.0]])
expected_1 = np.array([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]])
expected_2 = np.array([[17.0, 16.0, 15.0], [11.0, 10.0, 11.0]])
np.testing.assert_array_almost_equal(output_1.asnumpy(), expected_1)
np.testing.assert_array_almost_equal(output_2.asnumpy(), expected_2)
@ -988,7 +988,3 @@ def test_scatter_func_updates_vmap():
output = VmapNet(ScatterFuncVmapNet("min"), inputx, (0, None, 0), 0)(indices, updates)
expected = np.array([[0.1, 0.1, 2.2], [1.2, 1.3, 5.5]]).astype(np.float32)
np.testing.assert_array_almost_equal(output.asnumpy(), expected)
if __name__ == "__main__":
test_scatter_func_indices_vmap()
test_scatter_func_updates_vmap()