Add LowerBound

This commit is contained in:
bubtltk 2021-09-15 13:57:18 +08:00
parent a6faf2160e
commit e514fc3feb
10 changed files with 440 additions and 1 deletions

View File

@ -0,0 +1,76 @@
/**
* Copyright 2021 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 "backend/kernel_compiler/cpu/lower_bound_cpu_kernel.h"
#include "runtime/device/cpu/cpu_device_address.h"
namespace {
size_t kDataSizeThreshold_ = 4 * 1024;
}
namespace mindspore {
namespace kernel {
template <typename I, typename O>
void LowerBoundCPUKernel<I, O>::InitKernel(const CNodePtr &kernel_node) {
sorted_x_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 0);
values_shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 1);
output_shape_ = AnfAlgo::GetOutputDeviceShape(kernel_node, 0);
size_t size_exp = 2;
if (sorted_x_shape_.size() != values_shape_.size() || sorted_x_shape_.size() != size_exp ||
sorted_x_shape_[0] != values_shape_[0]) {
MS_LOG(EXCEPTION) << "The shape of input is invalid.";
}
sorted_x_num_ = sorted_x_shape_[0] * sorted_x_shape_[1];
values_num_ = values_shape_[0] * values_shape_[1];
output_num_ = output_shape_[0] * output_shape_[1];
if (values_num_ != output_num_) {
MS_LOG(EXCEPTION) << "Infer the shape of output error.";
}
}
template <typename I, typename O>
bool LowerBoundCPUKernel<I, O>::Launch(const std::vector<kernel::AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &,
const std::vector<kernel::AddressPtr> &outputs) {
auto sorted_x_data_addr = reinterpret_cast<I *>(inputs[0]->addr);
auto values_data_addr = reinterpret_cast<I *>(inputs[1]->addr);
auto output_data_addr = reinterpret_cast<O *>(outputs[0]->addr);
size_t sorted_x_data_column = sorted_x_shape_[1];
size_t values_data_column = values_shape_[1];
auto task = [&](size_t start, size_t end) {
for (size_t i = 0; i < values_num_; i++) {
size_t seq_row = i / values_data_column;
size_t low = seq_row * sorted_x_data_column;
size_t up = (seq_row + 1) * sorted_x_data_column - 1;
while (low <= up) {
size_t mid = (low + up) / 2;
if (values_data_addr[i] <= sorted_x_data_addr[mid]) {
up = mid - 1;
} else {
low = mid + 1;
}
}
output_data_addr[i] = static_cast<O>(low - seq_row * sorted_x_data_column);
}
};
if (values_num_ * sizeof(I) < kDataSizeThreshold_) {
task(0, values_num_);
} else {
CPUKernelUtils::ParallelFor(task, values_num_);
}
return true;
}
} // namespace kernel
} // namespace mindspore

View File

@ -0,0 +1,135 @@
/**
* Copyright 2021 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_BACKEND_KERNEL_COMPILER_CPU_LOWER_BOUND_CPU_KERNEL_H_
#define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_LOWER_BOUND_CPU_KERNEL_H_
#include <vector>
#include "backend/kernel_compiler/cpu/cpu_kernel.h"
#include "backend/kernel_compiler/cpu/cpu_kernel_factory.h"
namespace mindspore {
namespace kernel {
template <typename I, typename O>
class LowerBoundCPUKernel : public CPUKernel {
public:
LowerBoundCPUKernel() = default;
~LowerBoundCPUKernel() override = default;
void InitKernel(const CNodePtr &kernel_node) override;
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
const std::vector<AddressPtr> &outputs) override;
private:
std::vector<size_t> sorted_x_shape_;
std::vector<size_t> values_shape_;
std::vector<size_t> output_shape_;
size_t sorted_x_num_;
size_t values_num_;
size_t output_num_;
};
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeFloat16).AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, float16, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, float, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeFloat64).AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, double, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound, KernelAttr().AddInputAttr(kNumberTypeInt8).AddInputAttr(kNumberTypeInt8).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, int8_t, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeInt16).AddInputAttr(kNumberTypeInt16).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, int16_t, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeInt32).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, int32_t, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeInt64).AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, int64_t, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeUInt8).AddInputAttr(kNumberTypeUInt8).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, uint8_t, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeUInt16).AddInputAttr(kNumberTypeUInt16).AddOutputAttr(kNumberTypeInt32),
LowerBoundCPUKernel, uint16_t, int32_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeFloat16).AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, float16, int64_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeFloat32).AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, float, int64_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeFloat64).AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, double, int64_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound, KernelAttr().AddInputAttr(kNumberTypeInt8).AddInputAttr(kNumberTypeInt8).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, int8_t, int64_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeInt16).AddInputAttr(kNumberTypeInt16).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, int16_t, int64_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeInt32).AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, int32_t, int64_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeInt64).AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, int64_t, int64_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeUInt8).AddInputAttr(kNumberTypeUInt8).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, uint8_t, int64_t);
MS_REG_CPU_KERNEL_T_S(
LowerBound,
KernelAttr().AddInputAttr(kNumberTypeUInt16).AddInputAttr(kNumberTypeUInt16).AddOutputAttr(kNumberTypeInt64),
LowerBoundCPUKernel, uint16_t, int64_t);
} // namespace kernel
} // namespace mindspore
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_LOWER_BOUND_CPU_KERNEL_H_

View File

@ -113,6 +113,7 @@ constexpr auto kSplitV = "SplitV";
constexpr auto kDynamicBroadcastTo = "DynamicBroadcastTo";
constexpr auto kReshape = "Reshape";
constexpr auto kLstsq = "Lstsq";
constexpr auto kLowerBound = "LowerBound";
// NN
constexpr auto kCTCLoss = "CTCLoss";
@ -309,6 +310,7 @@ inline const PrimitivePtr kPrimImag = std::make_shared<Primitive>(kImag);
inline const PrimitivePtr kPrimConj = std::make_shared<Primitive>(kConj);
inline const PrimitivePtr kPrimExtractVolumePatches = std::make_shared<Primitive>("ExtractVolumePatches");
inline const PrimitivePtr kPrimLstsq = std::make_shared<Primitive>(kLstsq);
inline const PrimitivePtr kPrimLowerBound = std::make_shared<Primitive>(kLowerBound);
// NN
inline const PrimitivePtr kPrimCeLU = std::make_shared<Primitive>("CeLU");

View File

@ -0,0 +1,71 @@
/**
* Copyright 2021 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/lower_bound.h"
namespace mindspore {
namespace ops {
namespace {
abstract::ShapePtr LowerBoundInferShape(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
auto x_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[0]->BuildShape())[kShape];
auto values_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[1]->BuildShape())[kShape];
size_t size_exp = 2;
if (x_shape.size() != size_exp) {
MS_EXCEPTION(ValueError) << "The rank of sorted_x need to be equal to 2, but got " << values_shape.size();
}
if (values_shape.size() != size_exp) {
MS_EXCEPTION(ValueError) << "The rank of values need to be equal to 2, but got " << values_shape.size();
}
if (x_shape[0] != values_shape[0]) {
MS_EXCEPTION(ValueError) << "The shape of values is " << input_args[1]->BuildShape()->ToString()
<< ", but the shape of sorted_x is " << input_args[0]->BuildShape()->ToString()
<< ". The first dimension of the shape of sorted_x must be equal to that of values.";
}
return std::make_shared<abstract::Shape>(values_shape);
}
TypePtr LowerBoundInferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) {
std::map<std::string, TypePtr> input_types;
std::set<TypePtr> input_valid_types = {kFloat16, kFloat32, kFloat64, kInt8, kInt16, kInt32, kInt64, kUInt8, kUInt16};
TypePtr sorted_x_type = input_args[0]->BuildType();
TypePtr values_type = input_args[1]->BuildType();
(void)input_types.emplace("sorted_x", sorted_x_type);
(void)input_types.emplace("values", values_type);
(void)CheckAndConvertUtils::CheckTensorTypeSame(input_types, input_valid_types, primitive->name());
auto dtype_attr = primitive->GetAttr("out_type");
auto out_type = dtype_attr->cast<TypePtr>();
auto out_type_id = out_type->type_id();
MS_EXCEPTION_IF_NULL(out_type);
if (out_type_id != kInt32->type_id() && out_type_id != kInt64->type_id()) {
MS_EXCEPTION(TypeError) << "out_type must be int32 or int64.";
}
return out_type;
}
} // namespace
AbstractBasePtr LowerBoundInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) {
MS_EXCEPTION_IF_NULL(primitive);
const int64_t input_num = 2;
CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, primitive->name());
auto infer_type = LowerBoundInferType(primitive, input_args);
auto infer_shape = LowerBoundInferShape(primitive, input_args);
return abstract::MakeAbstract(infer_shape, infer_type);
}
REGISTER_PRIMITIVE_EVAL_IMPL(LowerBound, prim::kPrimLowerBound, LowerBoundInfer, nullptr, true);
} // namespace ops
} // namespace mindspore

View File

@ -0,0 +1,46 @@
/**
* Copyright 2021 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_CORE_OPS_LOWER_BOUND_H_
#define MINDSPORE_CORE_OPS_LOWER_BOUND_H_
#include <map>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "ops/primitive_c.h"
#include "ops/op_utils.h"
#include "abstract/abstract_value.h"
#include "abstract/primitive_infer_map.h"
#include "utils/check_convert_utils.h"
namespace mindspore {
namespace ops {
constexpr auto kNameLowerBound = "LowerBound";
class LowerBound : public PrimitiveC {
public:
LowerBound() : PrimitiveC(kNameLowerBound) { InitIOName({"sorted_x", "values"}, {"y"}); }
~LowerBound() = default;
MS_DECLARE_PARENT(LowerBound, PrimitiveC);
};
AbstractBasePtr LowerBoundInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args);
using PrimLowerBound = std::shared_ptr<LowerBound>;
} // namespace ops
} // namespace mindspore
#endif // MINDSPORE_CORE_OPS_LOWER_BOUND_H_

View File

@ -93,3 +93,4 @@ from .resize_bilinear_grad import _resize_bilinear_grad_aicpu
from .scatter_elements import _scatter_elements_aicpu
from .non_max_suppression import _non_max_suppression_aicpu
from .square import _square_aicpu
from .lower_bound import _lower_bound_aicpu

View File

@ -0,0 +1,47 @@
# Copyright 2021 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.
# ============================================================================
"""LowerBound op"""
from mindspore.ops.op_info_register import op_info_register, AiCPURegOp, DataType
lower_bound_op_info = AiCPURegOp("LowerBound") \
.fusion_type("OPAQUE") \
.input(0, "sorted_x", "required") \
.input(1, "values", "required") \
.output(0, "y", "required") \
.dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.I32_Default) \
.dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.I32_Default) \
.dtype_format(DataType.F64_Default, DataType.F64_Default, DataType.I32_Default) \
.dtype_format(DataType.I8_Default, DataType.I8_Default, DataType.I32_Default) \
.dtype_format(DataType.I16_Default, DataType.I16_Default, DataType.I32_Default) \
.dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default) \
.dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I32_Default) \
.dtype_format(DataType.U8_Default, DataType.U8_Default, DataType.I32_Default) \
.dtype_format(DataType.U16_Default, DataType.U16_Default, DataType.I32_Default) \
.dtype_format(DataType.F16_Default, DataType.F16_Default, DataType.I64_Default) \
.dtype_format(DataType.F32_Default, DataType.F32_Default, DataType.I64_Default) \
.dtype_format(DataType.F64_Default, DataType.F64_Default, DataType.I64_Default) \
.dtype_format(DataType.I8_Default, DataType.I8_Default, DataType.I64_Default) \
.dtype_format(DataType.I16_Default, DataType.I16_Default, DataType.I64_Default) \
.dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I64_Default) \
.dtype_format(DataType.I64_Default, DataType.I64_Default, DataType.I64_Default) \
.dtype_format(DataType.U8_Default, DataType.U8_Default, DataType.I64_Default) \
.dtype_format(DataType.U16_Default, DataType.U16_Default, DataType.I64_Default) \
.get_op_info()
@op_info_register(lower_bound_op_info)
def _lower_bound_aicpu():
"""LowerBound aicpu register"""
return

View File

@ -34,7 +34,8 @@ from .array_ops import (Argmax, Argmin, Cast, Concat, Pack, Stack, Unpack, Unsta
UnsortedSegmentProd, UnsortedSegmentSum, SpaceToDepth, DepthToSpace, SpaceToBatch,
BatchToSpace, SpaceToBatchND, BatchToSpaceND, BroadcastTo, InplaceUpdate, ReverseSequence,
EmbeddingLookup, Unique, GatherD, Identity, Range, MaskedFill, MaskedSelect, SearchSorted,
TensorScatterMax, TensorScatterMin, TensorScatterSub, ScatterElements, ExtractVolumePatches)
TensorScatterMax, TensorScatterMin, TensorScatterSub, ScatterElements, ExtractVolumePatches,
LowerBound)
from .comm_ops import (AllGather, AllReduce, NeighborExchange, NeighborExchangeV2, AlltoAll, _AllSwap, ReduceScatter,
Broadcast,
_MirrorOperator, _MirrorMiniStepOperator, _MiniStepAllGather, ReduceOp, _VirtualDataset,
@ -233,6 +234,7 @@ __all__ = [
'Lerp',
'Less',
'LessEqual',
'LowerBound',
'RealDiv',
'Div',
'DivNoNan',

View File

@ -6833,3 +6833,56 @@ class Lstsq(Primitive):
validator.check_type_name("l2_regularizer", l2_regularizer, 0.0, self.name)
self.fast = fast
self.l2_regularizer = l2_regularizer
class LowerBound(Primitive):
"""
Returns a tensor that contains the index for finding the lower bound of the value
of the input values element in the input sorted_x.
Args:
out_type (:class:`mindspore.dtype`): An optional data type of `mindspore.dtype.int32` and
`mindspore.dtype.int64`. Default: `mindspore.dtype.int32`.
Inputs:
- **sorted_x** (Tensor) - The input tensor whose dtype is real number and the data of each row must be sorted
in ascending order. The rank must be 2.
- **values** (Tensor) - The input tensor whose dtype is the same as `sorted_x` and the first dimension of the
shape of `values` must be equal to that of `sorted_x`. The rank must be 2.
Outputs:
Tensor, whose dtype is determined by `out_type` and whose shape is the same as that of `values`.
Raises:
TypeError: If `sorted_x` is not a Tensor.
TypeError: If `values` is not a Tensor.
TypeError: If `out_type` is invalid.
TypeError: If the type of `sorted_x` is not the same as that of `values`.
ValueError: If rank of the `sorted_x` is not equal to 2.
ValueError: If rank of the `values` is not equal to 2.
ValueError: If the first dimension of the shape of `sorted_x` is not equal to that of `values`.
Supported Platforms:
``CPU``
Examples:
>>> import mindspore
>>> import numpy as np
>>> from mindspore import Tensor
>>> import mindspore.ops as ops
>>> lowerbound = ops.LowerBound(out_type = mindspore.int32)
>>> sorted_x = Tensor(np.arange(12).reshape(3, 4).astype(np.int8))
>>> values = Tensor(np.array([[3], [4], [8]]).astype(np.int8))
>>> output = lowerbound(sorted_x, values)
>>> print(output)
[[3]
[0]
[0]]
"""
@prim_attr_register
def __init__(self, out_type=mstype.int32):
"""Initialize LowerBound"""
valid_values = (mstype.int32, mstype.int64)
validator.check_type_name("out_type", out_type, valid_values, self.name)
self.init_prim_io_names(inputs=['sorted_x', 'values'], outputs=['y'])

View File

@ -2761,6 +2761,12 @@ test_case_array_ops = [
'desc_inputs': [Tensor(np.array([[2, 1, 5], [3, 5, 1], [1, 1, 1]]), mstype.float32),
Tensor(np.array([[10, 5], [15, 8], [7, 4]]), mstype.float32)],
'skip': ['backward']}),
('LowerBound', {
'block': P.LowerBound(),
'desc_inputs': [Tensor(np.arange(20).reshape(4, 5), mstype.int8),
Tensor([[3], [5], [7], [8]], mstype.int8)],
'skip': ['backward'],
}),
]
test_case_image_ops = [