forked from mindspore-Ecosystem/mindspore
Gitee
commit
80f788922b
|
|
@ -24,302 +24,71 @@
|
|||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
using rank::Method;
|
||||
using rank::NaOption;
|
||||
void RankCpuKernelMod::InitKernel(const CNodePtr &kernel_node) {
|
||||
MS_EXCEPTION_IF_NULL(kernel_node);
|
||||
auto input_shape = common::AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
|
||||
|
||||
static const std::map<std::string, Method> kValidMethods = {
|
||||
{"max", Method::Max}, {"min", Method::Min}, {"average", Method::Average},
|
||||
{"first", Method::First}, {"dense", Method::Dense},
|
||||
};
|
||||
auto method = common::AnfAlgo::GetNodeAttr<std::string>(kernel_node, METHOD);
|
||||
if (kValidMethods.find(method) == kValidMethods.end()) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the method must be in "
|
||||
<< Map2Str<std::map, Method>(kValidMethods) << ", but got " << method;
|
||||
}
|
||||
method_ = kValidMethods.at(method);
|
||||
|
||||
static const std::map<std::string, NaOption> kValidOptions = {
|
||||
{"keep", NaOption::Keep},
|
||||
{"top", NaOption::Top},
|
||||
{"bottom", NaOption::Bottom},
|
||||
};
|
||||
auto option = common::AnfAlgo::GetNodeAttr<std::string>(kernel_node, NA_OPTION);
|
||||
if (kValidOptions.find(option) == kValidOptions.end()) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the option must be in "
|
||||
<< Map2Str<std::map, NaOption>(kValidOptions) << ", but got " << option;
|
||||
}
|
||||
option_ = kValidOptions.at(option);
|
||||
|
||||
ascending_ = common::AnfAlgo::GetNodeAttr<bool>(kernel_node, ASCENDING);
|
||||
pct_ = common::AnfAlgo::GetNodeAttr<bool>(kernel_node, PCT);
|
||||
auto axis = common::AnfAlgo::GetNodeAttr<int64_t>(kernel_node, AXIS);
|
||||
axis_ = axis < 0 ? LongToSize(axis + SizeToLong(input_shape.size())) : LongToSize(axis);
|
||||
if (axis_ >= input_shape.size()) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_
|
||||
<< "', the evaluated axis must be smaller than the dimension of input tensor "
|
||||
<< input_shape.size() << "D, but got " << axis_;
|
||||
}
|
||||
|
||||
axisIterator_.Init(input_shape, axis_);
|
||||
SetFunc();
|
||||
|
||||
auto kernel_attr = GetKernelAttrFromNode(kernel_node);
|
||||
auto [is_match, index] = MatchKernelAttr(kernel_attr, GetOpSupport());
|
||||
namespace {
|
||||
const size_t kRankInputsNum = 1;
|
||||
const size_t kRankOutputsNum = 1;
|
||||
}; // namespace
|
||||
bool RankCpuKernelMod::Init(const BaseOperatorPtr &base_operator, const std::vector<KernelTensorPtr> &inputs,
|
||||
const std::vector<KernelTensorPtr> &outputs) {
|
||||
kernel_name_ = base_operator->name();
|
||||
auto tensor_attr = GetKernelAttrFromTensors(inputs, outputs);
|
||||
auto is_match = MatchKernelAttr(tensor_attr, GetOpSupport()).first;
|
||||
if (!is_match) {
|
||||
MS_LOG(EXCEPTION) << "Rank does not support this kernel data type: " << kernel_attr;
|
||||
}
|
||||
kernel_func_ = std::get<1>(func_list_[index]);
|
||||
const size_t kTwoIdx = 2;
|
||||
init_func_ = std::get<kTwoIdx>(func_list_[index]);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void RankCpuKernelMod::InitIOSize(const CNodePtr &kernel_node) {
|
||||
DeprecatedNativeCpuKernelMod::InitInputOutputSize(kernel_node);
|
||||
size_t element_size = axisIterator_.OuterSize() * axisIterator_.InnerSize() * axisIterator_.AxisSize();
|
||||
// id
|
||||
(void)workspace_size_list_.emplace_back((sizeof(size_t) * element_size));
|
||||
// copy element
|
||||
(void)workspace_size_list_.emplace_back((sizeof(T) * element_size));
|
||||
if constexpr (!std::is_integral_v<T>) {
|
||||
// nan flags
|
||||
(void)workspace_size_list_.emplace_back((sizeof(bool) * element_size));
|
||||
}
|
||||
}
|
||||
|
||||
void RankCpuKernelMod::SetFunc() {
|
||||
switch (method_) {
|
||||
case Method::Max: {
|
||||
func_ = [](size_t i, size_t duplicate_count, int /* culmutive_rank */, const AxisIterator &axisIterator,
|
||||
const size_t *const sort_idx, float *const output_addr) {
|
||||
for (size_t j = i - duplicate_count + 1; j < i + 1; ++j) {
|
||||
output_addr[axisIterator.GetPos(sort_idx[j])] = i + 1;
|
||||
}
|
||||
};
|
||||
} break;
|
||||
case Method::Min: {
|
||||
func_ = [](size_t i, size_t duplicate_count, int /* culmutive_rank */, const AxisIterator &axisIterator,
|
||||
const size_t *const sort_idx, float *const output_addr) {
|
||||
for (size_t j = i - duplicate_count + 1; j < i + 1; ++j) {
|
||||
output_addr[axisIterator.GetPos(sort_idx[j])] = i - duplicate_count + 2;
|
||||
}
|
||||
};
|
||||
} break;
|
||||
case Method::Average: {
|
||||
// how avg is computed directly:
|
||||
// sum = (i - duplicate_count + 1) + (i - duplicate_count + 2) +... + i
|
||||
// = duplicate_count * (2 * i - duplicate_count + 1) / 2
|
||||
// rank_sum = sum + duplicate_count = duplicate_count * (2 * i - duplicate_count + 3) / 2
|
||||
// avg = rank_sum / duplicate_count = (2 * i - duplicate_count + 3) / 2
|
||||
func_ = [](size_t i, size_t duplicate_count, int /* culmutive_rank */, const AxisIterator &axisIterator,
|
||||
const size_t *const sort_idx, float *const output_addr) {
|
||||
float avg = (2 * i - duplicate_count + 3) / 2.0;
|
||||
for (size_t j = i - duplicate_count + 1; j < i + 1; ++j) {
|
||||
output_addr[axisIterator.GetPos(sort_idx[j])] = avg;
|
||||
}
|
||||
};
|
||||
} break;
|
||||
case Method::First: {
|
||||
func_ = [](size_t i, size_t duplicate_count, int /* culmutive_rank */, const AxisIterator &axisIterator,
|
||||
const size_t *const sort_idx, float *const output_addr) {
|
||||
for (size_t j = i - duplicate_count + 1; j < i + 1; ++j) {
|
||||
output_addr[axisIterator.GetPos(sort_idx[j])] = j + 1;
|
||||
}
|
||||
};
|
||||
} break;
|
||||
case Method::Dense: {
|
||||
func_ = [](size_t i, size_t duplicate_count, int culmutive_rank, const AxisIterator &axisIterator,
|
||||
const size_t *const sort_idx, float *const output_addr) {
|
||||
for (size_t j = i - duplicate_count + 1; j < i + 1; ++j) {
|
||||
output_addr[axisIterator.GetPos(sort_idx[j])] = culmutive_rank;
|
||||
}
|
||||
};
|
||||
} break;
|
||||
case Method::MethodNotDefined:
|
||||
default:
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_
|
||||
<< "', method not init. The method must be 'max', 'min', 'average', 'first', or 'dense', "
|
||||
<< ", but got " << method_;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void RankCpuKernelMod::Launch1D(const T *input_addr, size_t *sort_idx, T *values, const AxisIterator &iter,
|
||||
float *output_addr) const {
|
||||
const size_t n = axisIterator_.AxisSize();
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
values[i] = input_addr[iter.GetPos(i)];
|
||||
MS_LOG_ERROR << "Can not match kernel based on given attr!";
|
||||
return false;
|
||||
}
|
||||
|
||||
SortIndex<T>(sort_idx, values, iter);
|
||||
|
||||
int culmutive_rank = 1;
|
||||
size_t duplicate_count = 0;
|
||||
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
duplicate_count++;
|
||||
if ((i == n - 1) || (values[sort_idx[i]] != values[sort_idx[i + 1]])) {
|
||||
func_(i, duplicate_count, culmutive_rank, iter, sort_idx, output_addr);
|
||||
culmutive_rank++;
|
||||
duplicate_count = 0;
|
||||
}
|
||||
if (Resize(base_operator, inputs, outputs) == KRET_RESIZE_FAILED) {
|
||||
MS_LOG_ERROR << "Resize failed!";
|
||||
return false;
|
||||
}
|
||||
PctConvert(output_addr, iter, culmutive_rank);
|
||||
}
|
||||
|
||||
void RankCpuKernelMod::PctConvert(float *output_addr, const AxisIterator &iter, int culmutive_rank) const {
|
||||
const size_t n = iter.AxisSize();
|
||||
if (pct_) {
|
||||
// pct calculation
|
||||
if (method_ == Method::Dense) {
|
||||
auto size = static_cast<float>(culmutive_rank - 1);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
output_addr[iter.GetPos(i)] = output_addr[iter.GetPos(i)] / size;
|
||||
}
|
||||
} else {
|
||||
auto size = static_cast<float>(n);
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
output_addr[iter.GetPos(i)] = output_addr[iter.GetPos(i)] / size;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void RankCpuKernelMod::Launch1D(const T *input_addr, size_t *sort_idx, T *values, bool *is_nan,
|
||||
const AxisIterator &iter, float *output_addr) const {
|
||||
const size_t n = iter.AxisSize();
|
||||
T nan_padding_value = GetPaddingValue<T>();
|
||||
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
const T value = input_addr[iter.GetPos(i)];
|
||||
if (std::isnan(value)) {
|
||||
values[i] = nan_padding_value;
|
||||
is_nan[i] = true;
|
||||
} else {
|
||||
values[i] = value;
|
||||
is_nan[i] = false;
|
||||
}
|
||||
}
|
||||
|
||||
SortIndex<T>(sort_idx, values, iter);
|
||||
|
||||
int culmutive_rank = 1;
|
||||
size_t duplicate_count = 0;
|
||||
size_t nans_count = 0;
|
||||
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
duplicate_count++;
|
||||
if ((i == n - 1) || std::not_equal_to<T>()(values[sort_idx[i]], values[sort_idx[i + 1]]) ||
|
||||
(is_nan[sort_idx[i]] != is_nan[sort_idx[i + 1]])) {
|
||||
if ((option_ == NaOption::Keep) && is_nan[sort_idx[i]]) {
|
||||
for (size_t j = i - duplicate_count + 1; j < i + 1; ++j) {
|
||||
output_addr[iter.GetPos(sort_idx[j])] = NAN;
|
||||
}
|
||||
} else {
|
||||
func_(i, duplicate_count, culmutive_rank, iter, sort_idx, output_addr);
|
||||
}
|
||||
if (is_nan[sort_idx[i]]) {
|
||||
nans_count = duplicate_count;
|
||||
}
|
||||
culmutive_rank++;
|
||||
duplicate_count = 0;
|
||||
}
|
||||
}
|
||||
PctConvert(output_addr, iter, culmutive_rank, nans_count);
|
||||
}
|
||||
|
||||
void RankCpuKernelMod::PctConvert(float *output_addr, const AxisIterator &iter, int culmutive_rank,
|
||||
size_t nans_count) const {
|
||||
const size_t n = iter.AxisSize();
|
||||
if (pct_) {
|
||||
// pct calculation
|
||||
if (method_ == Method::Dense) {
|
||||
auto size = static_cast<float>(culmutive_rank - 1);
|
||||
if ((option_ == NaOption::Keep && (nans_count > 0))) {
|
||||
size--;
|
||||
}
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
output_addr[iter.GetPos(i)] = output_addr[iter.GetPos(i)] / size;
|
||||
}
|
||||
} else {
|
||||
auto size = static_cast<float>(n);
|
||||
if (option_ == NaOption::Keep) {
|
||||
size -= static_cast<float>(nans_count);
|
||||
}
|
||||
for (size_t i = 0; i < n; ++i) {
|
||||
output_addr[iter.GetPos(i)] = output_addr[iter.GetPos(i)] / size;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
bool RankCpuKernelMod::LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs) {
|
||||
if (inputs.size() != 1) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the number of inputs must be 1, but got " << inputs.size();
|
||||
}
|
||||
if (outputs.size() != 1) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the number of outputs must be 1, but got " << outputs.size();
|
||||
}
|
||||
if constexpr (std::is_integral_v<T>) {
|
||||
if (workspace.size() != 2) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the number of workspaces must be 2, but got "
|
||||
<< workspace.size();
|
||||
}
|
||||
} else {
|
||||
if (workspace.size() != 3) {
|
||||
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', the number of workspaces must be 3, but got "
|
||||
<< workspace.size();
|
||||
}
|
||||
}
|
||||
auto input_addr = reinterpret_cast<T *>(inputs[0]->addr);
|
||||
auto ids_addr = reinterpret_cast<size_t *>(workspace[0]->addr);
|
||||
auto values_addr = reinterpret_cast<T *>(workspace[1]->addr);
|
||||
auto output_addr = reinterpret_cast<float *>(outputs[0]->addr);
|
||||
|
||||
auto task = [this, input_addr, ids_addr, values_addr, workspace, output_addr](size_t start, size_t end) {
|
||||
AxisIterator iter(axisIterator_);
|
||||
for (size_t index = start; index < end; index++) {
|
||||
iter.SetOffset(index);
|
||||
|
||||
size_t offset = index * iter.AxisSize();
|
||||
size_t *sort_idx = ids_addr + offset;
|
||||
T *values = values_addr + offset;
|
||||
|
||||
if constexpr (std::is_integral_v<T>) {
|
||||
Launch1D<T>(input_addr, sort_idx, values, iter, output_addr);
|
||||
} else {
|
||||
auto flags_addr = reinterpret_cast<bool *>(workspace[2]->addr);
|
||||
bool *is_nan = flags_addr + offset;
|
||||
Launch1D<T>(input_addr, sort_idx, values, is_nan, iter, output_addr);
|
||||
}
|
||||
}
|
||||
};
|
||||
ParallelLaunchAutoSearch(task, axisIterator_.OuterSize() * axisIterator_.InnerSize(), this, ¶llel_search_info_);
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<std::tuple<KernelAttr, RankCpuKernelMod::RankFunc, RankCpuKernelMod::InitFunc>>
|
||||
RankCpuKernelMod::func_list_ = {{KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
|
||||
&RankCpuKernelMod::LaunchKernel<float>, &RankCpuKernelMod::InitIOSize<float>},
|
||||
{KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeFloat32),
|
||||
&RankCpuKernelMod::LaunchKernel<double>, &RankCpuKernelMod::InitIOSize<double>},
|
||||
{KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeFloat32),
|
||||
&RankCpuKernelMod::LaunchKernel<int32_t>, &RankCpuKernelMod::InitIOSize<int32_t>},
|
||||
{KernelAttr().AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeFloat32),
|
||||
&RankCpuKernelMod::LaunchKernel<int64_t>, &RankCpuKernelMod::InitIOSize<int64_t>}};
|
||||
int RankCpuKernelMod::Resize(const BaseOperatorPtr &base_operator, const std::vector<KernelTensorPtr> &inputs,
|
||||
const std::vector<KernelTensorPtr> &outputs,
|
||||
const std::map<uint32_t, tensor::TensorPtr> &inputsOnHost) {
|
||||
MS_EXCEPTION_IF_NULL(base_operator);
|
||||
if (int ret = KernelMod::Resize(base_operator, inputs, outputs); ret != KRET_OK) {
|
||||
return ret;
|
||||
}
|
||||
input_shape_size = SizeToLong(inputs[kIndex0]->GetShapeVector().size());
|
||||
return KRET_OK;
|
||||
}
|
||||
|
||||
bool RankCpuKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &,
|
||||
const std::vector<AddressPtr> &outputs) {
|
||||
CHECK_KERNEL_INPUTS_NUM(inputs.size(), kRankInputsNum, kernel_name_);
|
||||
CHECK_KERNEL_OUTPUTS_NUM(outputs.size(), kRankOutputsNum, kernel_name_);
|
||||
auto output_data = reinterpret_cast<int64_t *>(outputs[kIndex0]->addr);
|
||||
MS_EXCEPTION_IF_NULL(output_data);
|
||||
output_data[kIndex0] = input_shape_size;
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<KernelAttr> RankCpuKernelMod::GetOpSupport() {
|
||||
std::vector<KernelAttr> support_list;
|
||||
(void)std::transform(
|
||||
func_list_.begin(), func_list_.end(), std::back_inserter(support_list),
|
||||
[](const std::tuple<KernelAttr, RankFunc, InitFunc> &tuple_item) { return std::get<0>(tuple_item); });
|
||||
return support_list;
|
||||
return {
|
||||
KernelAttr().AddInputAttr(kNumberTypeBool).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeInt).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeInt8).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeInt16).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeInt32).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeInt64).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeUInt).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeUInt8).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeUInt16).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeUInt32).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeUInt64).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeFloat).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeFloat16).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeFloat64).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeComplex).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeComplex64).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeComplex128).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeInt4).AddOutputAttr(kNumberTypeInt64),
|
||||
KernelAttr().AddInputAttr(kNumberTypeGLUInt).AddOutputAttr(kNumberTypeInt64),
|
||||
};
|
||||
}
|
||||
|
||||
MS_KERNEL_FACTORY_REG(NativeCpuKernelMod, Rank, RankCpuKernelMod);
|
||||
|
|
|
|||
|
|
@ -20,100 +20,32 @@
|
|||
#include <string>
|
||||
#include <limits>
|
||||
#include <tuple>
|
||||
#include <map>
|
||||
#include "plugin/device/cpu/kernel/cpu_kernel.h"
|
||||
#include "plugin/factory/ms_factory.h"
|
||||
#include "plugin/device/cpu/kernel/nnacl/op_base.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace kernel {
|
||||
namespace rank {
|
||||
enum Method : int {
|
||||
Average,
|
||||
Max,
|
||||
Min,
|
||||
First,
|
||||
Dense,
|
||||
MethodNotDefined,
|
||||
};
|
||||
enum NaOption : int {
|
||||
Keep,
|
||||
Top,
|
||||
Bottom,
|
||||
OptionNotDefined,
|
||||
};
|
||||
} // namespace rank
|
||||
class RankCpuKernelMod : public DeprecatedNativeCpuKernelMod {
|
||||
class RankCpuKernelMod : public NativeCpuKernelMod {
|
||||
public:
|
||||
RankCpuKernelMod() = default;
|
||||
~RankCpuKernelMod() override = default;
|
||||
|
||||
void InitKernel(const CNodePtr &kernel_node) override;
|
||||
|
||||
void SetFunc();
|
||||
|
||||
template <typename T>
|
||||
void Launch1D(const T *input_addr, size_t *sort_idx, T *values, const AxisIterator &iter, float *output_addr) const;
|
||||
template <typename T>
|
||||
void Launch1D(const T *input_addr, size_t *sort_idx, T *values, bool *is_nan, const AxisIterator &iter,
|
||||
float *output_addr) const;
|
||||
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<AddressPtr> &outputs) override {
|
||||
return kernel_func_(this, inputs, workspace, outputs);
|
||||
}
|
||||
|
||||
std::vector<KernelAttr> GetOpSupport() override;
|
||||
bool Init(const BaseOperatorPtr &base_operator, const std::vector<KernelTensorPtr> &inputs,
|
||||
const std::vector<KernelTensorPtr> &outputs) override;
|
||||
int Resize(const BaseOperatorPtr &base_operator, const std::vector<KernelTensorPtr> &inputs,
|
||||
const std::vector<KernelTensorPtr> &outputs,
|
||||
const std::map<uint32_t, tensor::TensorPtr> &others = std::map<uint32_t, tensor::TensorPtr>()) override;
|
||||
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &,
|
||||
const std::vector<AddressPtr> &outputs) override;
|
||||
|
||||
protected:
|
||||
void InitInputOutputSize(const CNodePtr &kernel_node) override { init_func_(this, kernel_node); }
|
||||
std::vector<KernelAttr> GetOpSupport() override;
|
||||
|
||||
private:
|
||||
template <typename T>
|
||||
inline void SortIndex(size_t *sort_idx, const T *values, const AxisIterator &iter) const {
|
||||
std::iota(sort_idx, sort_idx + iter.AxisSize(), 0);
|
||||
if (ascending_) {
|
||||
std::stable_sort(sort_idx, sort_idx + iter.AxisSize(),
|
||||
[values](size_t lhs, size_t rhs) { return values[lhs] < values[rhs]; });
|
||||
} else {
|
||||
std::stable_sort(sort_idx, sort_idx + iter.AxisSize(),
|
||||
[values](size_t lhs, size_t rhs) { return values[lhs] > values[rhs]; });
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
inline T GetPaddingValue() const {
|
||||
if (ascending_ != (option_ == rank::NaOption::Top)) {
|
||||
return std::numeric_limits<T>::max();
|
||||
} else {
|
||||
return std::numeric_limits<T>::min();
|
||||
}
|
||||
}
|
||||
void PctConvert(float *output_addr, const AxisIterator &iter, int culmutive_rank, size_t nans_count) const;
|
||||
void PctConvert(float *output_addr, const AxisIterator &iter, int culmutive_rank) const;
|
||||
|
||||
template <typename T>
|
||||
bool LaunchKernel(const std::vector<kernel::AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
|
||||
const std::vector<kernel::AddressPtr> &outputs);
|
||||
template <typename T>
|
||||
void InitIOSize(const CNodePtr &kernel_node);
|
||||
|
||||
using RankFunc = std::function<bool(RankCpuKernelMod *, const std::vector<kernel::AddressPtr> &,
|
||||
const std::vector<AddressPtr> &, const std::vector<kernel::AddressPtr> &)>;
|
||||
using InitFunc = std::function<void(RankCpuKernelMod *, const CNodePtr &)>;
|
||||
static std::vector<std::tuple<KernelAttr, RankFunc, InitFunc>> func_list_;
|
||||
RankFunc kernel_func_;
|
||||
InitFunc init_func_;
|
||||
|
||||
// shape info
|
||||
AxisIterator axisIterator_{};
|
||||
// parameters
|
||||
size_t axis_{0};
|
||||
rank::Method method_{rank::MethodNotDefined};
|
||||
std::function<void(size_t, size_t, int, const AxisIterator &, const size_t *const, float *const)> func_;
|
||||
rank::NaOption option_{rank::OptionNotDefined};
|
||||
bool ascending_{true};
|
||||
bool pct_{false};
|
||||
int64_t input_shape_size;
|
||||
};
|
||||
} // namespace kernel
|
||||
} // namespace mindspore
|
||||
|
||||
#endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_RANK_CPU_KERNEL_H_
|
||||
|
|
|
|||
|
|
@ -15,13 +15,65 @@
|
|||
*/
|
||||
|
||||
#include "ops/rank.h"
|
||||
#include "ops/primitive_c.h"
|
||||
#include "utils/log_adapter.h"
|
||||
#include "utils/check_convert_utils.h"
|
||||
#include "ops/op_utils.h"
|
||||
#include "mindapi/src/helper.h"
|
||||
|
||||
namespace mindspore {
|
||||
namespace ops {
|
||||
namespace {
|
||||
constexpr int64_t input_num = 1;
|
||||
} // namespace
|
||||
MIND_API_OPERATOR_IMPL(Rank, BaseOperator);
|
||||
REGISTER_PRIMITIVE_C(kNameRank, Rank);
|
||||
|
||||
class RankInfer : public abstract::OpInferBase {
|
||||
public:
|
||||
BaseShapePtr InferShape(const PrimitivePtr &primitive,
|
||||
const std::vector<AbstractBasePtr> &input_args) const override {
|
||||
MS_EXCEPTION_IF_NULL(primitive);
|
||||
CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, primitive->name());
|
||||
return abstract::kNoShape;
|
||||
}
|
||||
|
||||
TypePtr InferType(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const override {
|
||||
MS_EXCEPTION_IF_NULL(primitive);
|
||||
CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, primitive->name());
|
||||
TypePtr res = kInt64;
|
||||
return res;
|
||||
}
|
||||
|
||||
ValuePtr InferValue(const PrimitivePtr &primitive, const std::vector<AbstractBasePtr> &input_args) const override {
|
||||
MS_EXCEPTION_IF_NULL(primitive);
|
||||
auto prim_name = primitive->name();
|
||||
CheckAndConvertUtils::CheckInputArgs(input_args, kEqual, input_num, prim_name);
|
||||
auto x_type = input_args[0]->BuildType();
|
||||
MS_EXCEPTION_IF_NULL(x_type);
|
||||
if (!x_type->isa<TensorType>()) {
|
||||
MS_EXCEPTION(TypeError) << "For '" << prim_name << "', input must be a Tensor, but got: " << x_type->ToString()
|
||||
<< ".";
|
||||
}
|
||||
auto input_shape_ptr = input_args[0]->BuildShape();
|
||||
MS_EXCEPTION_IF_NULL(input_shape_ptr);
|
||||
auto shape_map = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_shape_ptr);
|
||||
auto x_shape = shape_map[kShape];
|
||||
if (IsDynamicRank(x_shape)) {
|
||||
return kAnyValue;
|
||||
}
|
||||
auto x_shape_rank = SizeToLong(x_shape.size());
|
||||
ValuePtr res = MakeValue(x_shape_rank);
|
||||
return res;
|
||||
}
|
||||
|
||||
AbstractBasePtr InferShapeAndType(const abstract::AnalysisEnginePtr &engine, const PrimitivePtr &primitive,
|
||||
const std::vector<AbstractBasePtr> &input_args) const override {
|
||||
auto type = InferType(primitive, input_args);
|
||||
auto shape = InferShape(primitive, input_args);
|
||||
auto value = InferValue(primitive, input_args);
|
||||
auto res = MakeAbstract(shape, type);
|
||||
res->set_value(value);
|
||||
return res;
|
||||
}
|
||||
};
|
||||
REGISTER_PRIMITIVE_OP_INFER_IMPL(Rank, prim::kPrimRank, RankInfer, true);
|
||||
} // namespace ops
|
||||
} // namespace mindspore
|
||||
|
|
|
|||
|
|
@ -34,8 +34,6 @@ class MIND_API Rank : public BaseOperator {
|
|||
/// \brief Init. Refer to the parameters of Python API @ref mindspore.ops.Rank for the inputs.
|
||||
void Init() const {}
|
||||
};
|
||||
MIND_API abstract::AbstractBasePtr RankInfer(const abstract::AnalysisEnginePtr &, const PrimitivePtr &primitive,
|
||||
const std::vector<abstract::AbstractBasePtr> &input_args);
|
||||
} // namespace ops
|
||||
} // namespace mindspore
|
||||
#endif // MINDSPORE_CORE_OPS_RANK_H_
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ from __future__ import absolute_import
|
|||
import operator
|
||||
|
||||
from mindspore.common import dtype as mstype
|
||||
from mindspore.common import Tensor
|
||||
from mindspore.common import Tensor, mutable
|
||||
from mindspore.ops import operations as P
|
||||
from mindspore.ops import functional as F
|
||||
from mindspore.ops.primitive import constexpr, _primexpr
|
||||
|
|
@ -1108,14 +1108,17 @@ def roll(a, shift, axis=None):
|
|||
return a
|
||||
|
||||
|
||||
@constexpr
|
||||
@_primexpr
|
||||
def _get_moved_perm(ndim, source, destination):
|
||||
"""
|
||||
Helper function for moveaxis, returns permutation after moving axes
|
||||
from source to destination.
|
||||
"""
|
||||
dest_sorted_idx = [i for i, _ in sorted(enumerate(destination), key=operator.itemgetter(1))]
|
||||
axes_orig = [i for i in range(ndim) if i not in source]
|
||||
axes_orig = mutable([], True)
|
||||
for i in range(ndim):
|
||||
if i not in source:
|
||||
axes_orig = axes_orig + [i]
|
||||
|
||||
k = 0
|
||||
m = 0
|
||||
|
|
|
|||
|
|
@ -1168,7 +1168,7 @@ class Split(Primitive):
|
|||
self.add_prim_attr('num_split', self.output_num)
|
||||
|
||||
|
||||
class Rank(PrimitiveWithInfer):
|
||||
class Rank(Primitive):
|
||||
"""
|
||||
Returns the rank of a tensor.
|
||||
|
||||
|
|
@ -1191,13 +1191,6 @@ class Rank(PrimitiveWithInfer):
|
|||
def __init__(self):
|
||||
"""Initialize Rank"""
|
||||
|
||||
def __infer__(self, x):
|
||||
validator.check_subclass("x", x['dtype'], mstype.tensor, self.name)
|
||||
out = {'shape': None,
|
||||
'dtype': None,
|
||||
'value': len(x['shape'])}
|
||||
return out
|
||||
|
||||
def __call__(self, x):
|
||||
if not isinstance(x, (Tensor, Tensor_)):
|
||||
raise TypeError("the input x must be Tensor!")
|
||||
|
|
|
|||
|
|
@ -183,6 +183,7 @@ def test_matrix_diag_part_v3_primitive_negative_k():
|
|||
np.testing.assert_allclose(result, expect)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="bug")
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.env_onecard
|
||||
@pytest.mark.platform_x86_cpu
|
||||
|
|
|
|||
|
|
@ -13,114 +13,115 @@
|
|||
# limitations under the License.
|
||||
# ============================================================================
|
||||
|
||||
from typing import List
|
||||
from random import sample
|
||||
import numpy as np
|
||||
import pytest
|
||||
from mindspore import Tensor
|
||||
import mindspore.context as context
|
||||
import mindspore.nn as nn
|
||||
from mindspore import Tensor
|
||||
from mindspore.ops import PrimitiveWithInfer, prim_attr_register
|
||||
from mindspore._checkparam import Validator as validator
|
||||
from mindspore.common import dtype as mstype
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
import mindspore.ops as ops
|
||||
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
|
||||
|
||||
|
||||
class Rank(PrimitiveWithInfer):
|
||||
"""
|
||||
Shift op frontend implementation
|
||||
"""
|
||||
|
||||
# size_t axis_{0};
|
||||
# rank::Method method_{rank::MethodNotDefined};
|
||||
# rank::NaOption option_{rank::OptionNotDefined};
|
||||
# bool ascending_{true};
|
||||
# bool pct_{false};
|
||||
@prim_attr_register
|
||||
def __init__(self, axis: int, method: str, na_option: str, ascending: bool, pct: bool):
|
||||
"""Initialize Sort"""
|
||||
self.axis = validator.check_value_type("axis", axis, [int], self.name)
|
||||
self.method = validator.check_value_type("method", method, [str], self.name)
|
||||
self.na_option = validator.check_value_type("na_option", na_option, [str], self.name)
|
||||
self.ascending = validator.check_value_type("ascending", ascending, [bool], self.name)
|
||||
self.pct = validator.check_value_type("pct", pct, [bool], self.name)
|
||||
|
||||
self.init_prim_io_names(inputs=['x'], outputs=['output'])
|
||||
|
||||
def __infer__(self, x):
|
||||
out_shapes = x['shape']
|
||||
return {
|
||||
'shape': tuple(out_shapes),
|
||||
'dtype': mstype.float32,
|
||||
'value': None
|
||||
}
|
||||
|
||||
|
||||
class RankNet(nn.Cell):
|
||||
def __init__(self, axis: int, method: str, na_option: str, ascending: bool, pct: bool):
|
||||
def __init__(self):
|
||||
super(RankNet, self).__init__()
|
||||
self.rank = Rank(axis, method, na_option, ascending, pct)
|
||||
self.rank = ops.Rank()
|
||||
|
||||
def construct(self, x):
|
||||
return self.rank(x)
|
||||
|
||||
|
||||
def pandas_rank(arr, **kwargs):
|
||||
ser = pd.DataFrame(arr)
|
||||
result = ser.rank(**kwargs)
|
||||
return result.to_numpy()
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_x86_cpu
|
||||
@pytest.mark.env_onecard
|
||||
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
|
||||
def test_rank_scalar(mode):
|
||||
"""
|
||||
Feature: ops.rank
|
||||
Description: Verify the result of rank for scalar tensor
|
||||
Expectation: success
|
||||
"""
|
||||
context.set_context(mode=mode)
|
||||
x = Tensor(np.array(1).astype(np.float32))
|
||||
expect = 0
|
||||
net = RankNet()
|
||||
output = net(x)
|
||||
assert output == expect
|
||||
|
||||
|
||||
@pytest.mark.parametrize('shape', [(10,)])
|
||||
@pytest.mark.parametrize('dtype', [np.float32, np.float64, np.int32, np.int64])
|
||||
@pytest.mark.parametrize('method', ['dense', 'first', 'max', 'min', 'average'])
|
||||
@pytest.mark.parametrize('na_option', ["keep", "top", "bottom"])
|
||||
@pytest.mark.parametrize('ascending', [True, False])
|
||||
@pytest.mark.parametrize('pct', [False, True])
|
||||
def test_rank_1d(shape: List[int], dtype, method: str, ascending: bool, pct: bool, na_option: str):
|
||||
np.random.seed(0)
|
||||
|
||||
if dtype in (np.int32, np.int64):
|
||||
arr = np.random.randint(0, 100, size=shape).astype(dtype)
|
||||
else:
|
||||
arr = np.random.random(size=shape).astype(dtype)
|
||||
arr.flat[sample(range(arr.size), int(arr.size / 10))] = np.nan
|
||||
|
||||
pd_result = pandas_rank(arr, method=method, ascending=ascending, pct=pct, na_option=na_option).flatten()
|
||||
rank = RankNet(0, method=method, ascending=ascending, pct=pct, na_option=na_option)
|
||||
mind_result = rank(Tensor(arr)).asnumpy()
|
||||
|
||||
print('arr: \n', arr, arr.dtype, arr.shape)
|
||||
print('pandas: \n', pd_result, pd_result.dtype, pd_result.shape)
|
||||
print('mind: \n', mind_result, mind_result.dtype, mind_result.shape)
|
||||
print(f'method: {method}, ascending: {ascending}, pct: {pct} na_option: {na_option}')
|
||||
assert np.allclose(pd_result, mind_result, equal_nan=True)
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_x86_cpu
|
||||
@pytest.mark.env_onecard
|
||||
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
|
||||
def test_rank_3d_tensor(mode):
|
||||
"""
|
||||
Feature: ops.rank
|
||||
Description: Verify the result of rank for 3D tensor
|
||||
Expectation: success
|
||||
"""
|
||||
context.set_context(mode=mode)
|
||||
x = Tensor(np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]).astype(np.float32))
|
||||
expect = 3
|
||||
net = RankNet()
|
||||
output = net(x)
|
||||
assert output == expect
|
||||
|
||||
|
||||
@pytest.mark.parametrize('shape', [(5, 6)])
|
||||
@pytest.mark.parametrize('dtype', [np.float32, np.float64, np.int32, np.int64])
|
||||
@pytest.mark.parametrize('method', ['dense', 'first', 'max', 'min', 'average'])
|
||||
@pytest.mark.parametrize('na_option', ["keep", "top", "bottom"])
|
||||
@pytest.mark.parametrize('axis', [0, 1])
|
||||
@pytest.mark.parametrize('ascending', [True, False])
|
||||
@pytest.mark.parametrize('pct', [False, True])
|
||||
def test_rank_2d(shape: List[int], dtype, method: str, ascending: bool, pct: bool, axis: int, na_option: str):
|
||||
np.random.seed(0)
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_x86_cpu
|
||||
@pytest.mark.env_onecard
|
||||
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
|
||||
def test_rank_4d_tensor(mode):
|
||||
"""
|
||||
Feature: ops.rank
|
||||
Description: Verify the result of rank for 4D tensor
|
||||
Expectation: success
|
||||
"""
|
||||
context.set_context(mode=mode)
|
||||
x = Tensor(np.array([[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]]).astype(np.float32))
|
||||
expect = 4
|
||||
net = RankNet()
|
||||
output = net(x)
|
||||
assert output == expect
|
||||
|
||||
if dtype in (np.int32, np.int64):
|
||||
arr = np.random.randint(0, 100, size=shape).astype(dtype)
|
||||
else:
|
||||
arr = np.random.random(size=shape).astype(dtype)
|
||||
arr.flat[sample(range(arr.size), int(arr.size / 10))] = np.nan
|
||||
|
||||
pd_result = pandas_rank(arr, method=method, ascending=ascending, pct=pct, na_option=na_option, axis=axis)
|
||||
rank = RankNet(axis=axis, method=method, ascending=ascending, pct=pct, na_option=na_option)
|
||||
mind_result = rank(Tensor(arr)).asnumpy()
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_x86_cpu
|
||||
@pytest.mark.env_onecard
|
||||
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
|
||||
def test_rank_dynamic_shape(mode):
|
||||
"""
|
||||
Feature: ops.rank
|
||||
Description: test rank with dynamic shape.
|
||||
Expectation: success
|
||||
"""
|
||||
context.set_context(mode=mode)
|
||||
x = Tensor(np.array([[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]]).astype(np.float32))
|
||||
expect = 4
|
||||
|
||||
print('arr: \n', arr, arr.dtype, arr.shape)
|
||||
print('pandas: \n', pd_result, pd_result.dtype, pd_result.shape)
|
||||
print('mind: \n', mind_result, mind_result.dtype, mind_result.shape)
|
||||
print(f'axis: {axis}, method: {method}, ascending: {ascending}, pct: {pct} na_option: {na_option}')
|
||||
assert np.allclose(pd_result, mind_result, equal_nan=True)
|
||||
net = RankNet()
|
||||
x_dyn = Tensor(shape=[None] * len(x.shape), dtype=x.dtype)
|
||||
net.set_inputs(x_dyn)
|
||||
|
||||
output = net(x)
|
||||
assert output == expect
|
||||
|
||||
|
||||
@pytest.mark.level0
|
||||
@pytest.mark.platform_x86_cpu
|
||||
@pytest.mark.env_onecard
|
||||
@pytest.mark.parametrize('mode', [context.GRAPH_MODE, context.PYNATIVE_MODE])
|
||||
def test_rank_invalid_input(mode):
|
||||
"""
|
||||
Feature: ops.rank
|
||||
Description: Test invalid input cases of rank
|
||||
Expectation: raise TypeError
|
||||
"""
|
||||
context.set_context(mode=mode)
|
||||
net = RankNet()
|
||||
with pytest.raises(TypeError):
|
||||
net(1)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
net("str")
|
||||
|
|
|
|||
Loading…
Reference in New Issue