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

ops RaggedTensorToSparse and RaggedTensorToTensor supports dynamic shape feature 

type: feature
reason: add codes to support dynamic shape for RaggedTensorToSparse and RaggedTensorToTensor.

------

Signed-off-by: wang_ziqi <wangziqi4@huawei.com>
This commit is contained in:
wang_ziqi 2023-02-17 18:48:33 +08:00
parent 7cf3697fa6
commit 37127545d4
4 changed files with 86 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

88
mindspore/ccsrc/plugin/device/cpu/kernel/ragged_tensor_to_tensor_cpu_kernel.cc
View File

@ -43,23 +43,42 @@ constexpr size_t kFirstPartitionInputIndex = 3;
}
} // namespace
void RaggedTensorToTensorCpuKernelMod::InitKernel(const CNodePtr &kernel_node) {
MS_EXCEPTION_IF_NULL(kernel_node);
kernel_name_ = common::AnfAlgo::GetCNodeName(kernel_node);
node_wpt_ = kernel_node;
row_partition_types_ = common::AnfAlgo::GetNodeAttr<std::vector<std::string>>(kernel_node, "row_partition_types");
size_t output_num = AnfAlgo::GetOutputTensorNum(kernel_node);
bool RaggedTensorToTensorCpuKernelMod::Init(const BaseOperatorPtr &base_operator,
const std::vector<KernelTensorPtr> &inputs,
const std::vector<KernelTensorPtr> &outputs) {
MS_EXCEPTION_IF_NULL(base_operator);
kernel_name_ = base_operator->name();
row_partition_types_ = GetValue<std::vector<std::string>>(base_operator->GetAttr("row_partition_types"));
shape_dtype_ = inputs[kShapeInputIndex]->GetDtype();
values_dtype_ = inputs[kValueInputIndex]->GetDtype();
size_t output_num = outputs.size();
CHECK_KERNEL_OUTPUTS_NUM(output_num, kRaggedTensorToTensorOutputsNum, kernel_name_);
return true;
}
int RaggedTensorToTensorCpuKernelMod::Resize(const BaseOperatorPtr &base_operator,
const std::vector<KernelTensorPtr> &inputs,
const std::vector<KernelTensorPtr> &outputs,
const std::map<uint32_t, tensor::TensorPtr> &inputsOnHost) {
if (int ret = KernelMod::Resize(base_operator, inputs, outputs); ret != KRET_OK) {
return ret;
}
values_shape_ = inputs[kValueInputIndex]->GetShapeVector();
default_values_shape_ = inputs[kDefaultValueInputIndex]->GetShapeVector();
output_shape_ = outputs[0]->GetShapeVector();
row_partition_shape_list_.clear();
for (int i = 0; i < SizeToLong(row_partition_types_.size()); ++i) {
row_partition_shape_list_.push_back(inputs[kFirstPartitionInputIndex + i]->GetShapeVector());
}
return KRET_OK;
}
bool RaggedTensorToTensorCpuKernelMod::Launch(const std::vector<kernel::AddressPtr> &inputs,
const std::vector<kernel::AddressPtr> &,
const std::vector<kernel::AddressPtr> &outputs) {
auto shape_dtype = AnfAlgo::GetInputDeviceDataType(node_wpt_, kShapeInputIndex);
auto values_dtype = AnfAlgo::GetInputDeviceDataType(node_wpt_, kValueInputIndex);
switch (shape_dtype) {
switch (shape_dtype_) {
case kNumberTypeInt32: {
switch (values_dtype) {
switch (values_dtype_) {
RAGGEDTENSORTOTENSOR_COMPUTE_CASE(kNumberTypeBool, int32_t, bool)
RAGGEDTENSORTOTENSOR_COMPUTE_CASE(kNumberTypeInt8, int32_t, int8_t)
RAGGEDTENSORTOTENSOR_COMPUTE_CASE(kNumberTypeInt16, int32_t, int16_t)
@ -72,12 +91,12 @@ bool RaggedTensorToTensorCpuKernelMod::Launch(const std::vector<kernel::AddressP
RAGGEDTENSORTOTENSOR_COMPUTE_CASE(kNumberTypeFloat64, int32_t, double)
default:
MS_EXCEPTION(TypeError) << "For '" << kernel_name_
<< "', unsupported values data type: " << TypeIdToType(values_dtype)->ToString();
<< "', unsupported values data type: " << TypeIdToType(values_dtype_)->ToString();
}
break;
}
case kNumberTypeInt64: {
switch (values_dtype) {
switch (values_dtype_) {
RAGGEDTENSORTOTENSOR_COMPUTE_CASE(kNumberTypeBool, int64_t, bool)
RAGGEDTENSORTOTENSOR_COMPUTE_CASE(kNumberTypeInt8, int64_t, int8_t)
RAGGEDTENSORTOTENSOR_COMPUTE_CASE(kNumberTypeInt16, int64_t, int16_t)
@ -90,13 +109,13 @@ bool RaggedTensorToTensorCpuKernelMod::Launch(const std::vector<kernel::AddressP
RAGGEDTENSORTOTENSOR_COMPUTE_CASE(kNumberTypeFloat64, int64_t, double)
default:
MS_EXCEPTION(TypeError) << "For '" << kernel_name_
<< "', unsupported values data type: " << TypeIdToType(values_dtype)->ToString();
<< "', unsupported values data type: " << TypeIdToType(values_dtype_)->ToString();
}
break;
}
default:
MS_EXCEPTION(TypeError) << "For '" << kernel_name_
<< "', unsupported shape data type: " << TypeIdToType(shape_dtype)->ToString();
<< "', unsupported shape data type: " << TypeIdToType(shape_dtype_)->ToString();
}
return true;
}
@ -108,17 +127,15 @@ void RaggedTensorToTensorCpuKernelMod::LaunchKernel(const std::vector<kernel::Ad
TYPE1 first_dimension;
GetFirstDimension<TYPE1>(inputs, &first_dimension);
std::vector<TYPE1> output_size;
auto output_shape = AnfAlgo::GetOutputDeviceShape(node_wpt_, 0);
auto values_shape = AnfAlgo::GetInputDeviceShape(node_wpt_, kValueInputIndex);
if (ragged_rank_ + values_shape.size() != output_shape.size()) {
if (ragged_rank_ + values_shape_.size() != output_shape_.size()) {
MS_LOG(EXCEPTION) << "For '" << kernel_name_
<< "', row partition size plus 'values' rank should be equal to 'shape' rank: "
<< output_shape.size() << ", but got row partition size: " << ragged_rank_
<< ", 'values' rank: " << values_shape.size();
<< output_shape_.size() << ", but got row partition size: " << ragged_rank_
<< ", 'values' rank: " << values_shape_.size();
}
output_size.reserve(output_shape.size());
for (unsigned int dim = 0; dim < output_shape.size(); dim++) {
output_size.push_back(output_shape[dim]);
output_size.reserve(output_shape_.size());
for (unsigned int dim = 0; dim < output_shape_.size(); dim++) {
output_size.push_back(output_shape_[dim]);
}
if (output_size[0] < 0) {
output_size[0] = first_dimension;
@ -150,14 +167,14 @@ void RaggedTensorToTensorCpuKernelMod::LaunchKernel(const std::vector<kernel::Ad
if (row_partition_types_[0] == "FIRST_DIM_SIZE") {
kernel::AddressPtr row_partition = inputs[i + kFirstPartitionInputIndex];
auto row_partition_ptr = reinterpret_cast<TYPE1 *>(row_partition->addr);
auto row_partition_shape = AnfAlgo::GetInputDeviceShape(node_wpt_, i + kFirstPartitionInputIndex);
auto row_partition_shape = row_partition_shape_list_[i];
Eigen::DSizes<Eigen::DenseIndex, 1> row_partition_dsize(row_partition_shape[0]);
TYPE1_flat rowET(row_partition_ptr, row_partition_dsize);
row_partition_tensor.push_back(rowET);
} else {
kernel::AddressPtr row_partition = inputs[i - 1 + kFirstPartitionInputIndex];
auto row_partition_ptr = reinterpret_cast<TYPE1 *>(row_partition->addr);
auto row_partition_shape = AnfAlgo::GetInputDeviceShape(node_wpt_, i - 1 + kFirstPartitionInputIndex);
auto row_partition_shape = row_partition_shape_list_[i - 1];
Eigen::DSizes<Eigen::DenseIndex, 1> row_partition_dsize(row_partition_shape[0]);
TYPE1_flat rowET(row_partition_ptr, row_partition_dsize);
row_partition_tensor.push_back(rowET);
@ -202,7 +219,7 @@ template <typename TYPE1>
void RaggedTensorToTensorCpuKernelMod::GetFirstDimension(const std::vector<kernel::AddressPtr> &inputs,
TYPE1 *first_dim) {
auto first_partition_tensor = inputs[kFirstPartitionInputIndex];
auto firstPartitionShape = AnfAlgo::GetInputDeviceShape(node_wpt_, kFirstPartitionInputIndex);
auto firstPartitionShape = row_partition_shape_list_[0];
std::string first_partition_type = row_partition_types_[0];
if (first_partition_type == "VALUE_ROWIDS") {
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', cannot handle 'VALUE_ROWIDS' in first dimension.";
@ -307,21 +324,19 @@ bool RaggedTensorToTensorCpuKernelMod::SetOutput(const std::vector<kernel::Addre
const std::vector<kernel::AddressPtr> &outputs,
const vector<TYPE1> &output_index) {
auto output_tensor_ptr = reinterpret_cast<TYPE2 *>(outputs[0]->addr);
auto output_tensor_shape = AnfAlgo::GetOutputDeviceShape(node_wpt_, 0);
size_t output_element_sum = SizeToLong(SizeOf(output_tensor_shape));
size_t output_element_sum = SizeToLong(SizeOf(output_shape_));
auto default_value_tensor = inputs[kDefaultValueInputIndex];
TYPE2 *default_value_pt = reinterpret_cast<TYPE2 *>(default_value_tensor->addr);
auto values_tensor = inputs[kValueInputIndex];
auto values_tensor_ptr = reinterpret_cast<TYPE2 *>(values_tensor->addr);
auto value_shape = AnfAlgo::GetInputDeviceShape(node_wpt_, kValueInputIndex);
if (value_shape.size() == 1) {
if (values_shape_.size() == 1) {
Eigen::DSizes<Eigen::DenseIndex, 1> output_tensor_dsize(output_element_sum);
TYPE2_flat outputET(output_tensor_ptr, output_tensor_dsize);
TYPE2_flat output_flat = outputET;
TYPE2 *base_output = output_flat.data();
TYPE2 default_value = default_value_pt[0];
std::fill(base_output, base_output + output_flat.size(), default_value);
Eigen::DSizes<Eigen::DenseIndex, 1> values_tensor_dsize(value_shape[0]);
Eigen::DSizes<Eigen::DenseIndex, 1> values_tensor_dsize(values_shape_[0]);
TYPE2_flat valuesET(values_tensor_ptr, values_tensor_dsize);
TYPE2_flat values_flat = valuesET;
unsigned int values_size = values_flat.size();
@ -336,13 +351,12 @@ bool RaggedTensorToTensorCpuKernelMod::SetOutput(const std::vector<kernel::Addre
}
}
} else {
auto default_values_shape = AnfAlgo::GetInputDeviceShape(node_wpt_, kDefaultValueInputIndex);
auto broadcast_shape = CPUKernelUtils::GetBroadcastShape(default_values_shape, output_tensor_shape);
if (broadcast_shape != output_tensor_shape) {
auto broadcast_shape = CPUKernelUtils::GetBroadcastShape(default_values_shape_, output_shape_);
if (broadcast_shape != output_shape_) {
MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "', unable to broadcast default_value of shape "
<< default_values_shape << " to tensor of shape " << output_tensor_shape;
<< default_values_shape_ << " to tensor of shape " << output_shape_;
}
BroadcastIterator iter(default_values_shape, output_tensor_shape, broadcast_shape);
BroadcastIterator iter(default_values_shape_, output_shape_, broadcast_shape);
TYPE2 *default_value_addr = reinterpret_cast<TYPE2 *>(inputs[kDefaultValueInputIndex]->addr);
TYPE2 *output_addr = reinterpret_cast<TYPE2 *>(outputs[0]->addr);
@ -356,7 +370,7 @@ bool RaggedTensorToTensorCpuKernelMod::SetOutput(const std::vector<kernel::Addre
TYPE2_flat output_flat = outputET;
TYPE2 *base_output = output_flat.data();
size_t values_tensor_size = SizeToLong(SizeOf(value_shape));
size_t values_tensor_size = SizeToLong(SizeOf(values_shape_));
Eigen::DSizes<Eigen::DenseIndex, 1> values_tensor_dsize(values_tensor_size);
TYPE2_flat valuesET(values_tensor_ptr, values_tensor_dsize);
TYPE2_flat values_flat = valuesET;
@ -370,7 +384,7 @@ bool RaggedTensorToTensorCpuKernelMod::SetOutput(const std::vector<kernel::Addre
}
int value_element_bytesize = value_element_size * sizeof(TYPE2);
TYPE2 *values_base = values_flat.data();
size_t values_dimsize = value_shape[0];
size_t values_dimsize = values_shape_[0];
if (values_dimsize != output_index_size) {
MS_LOG(EXCEPTION) << "For '" << kernel_name_
<< "', 'values' shape[0] must be equal to ragged tensor indices: " << output_index_size

18
mindspore/ccsrc/plugin/device/cpu/kernel/ragged_tensor_to_tensor_cpu_kernel.h
View File

@ -36,14 +36,19 @@ namespace kernel {
#define TYPE1_flat Eigen::TensorMap<Eigen::Tensor<TYPE1, 1, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned>
#define TYPE2_flat Eigen::TensorMap<Eigen::Tensor<TYPE2, 1, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned>
class RaggedTensorToTensorCpuKernelMod : public DeprecatedNativeCpuKernelMod {
class RaggedTensorToTensorCpuKernelMod : public NativeCpuKernelMod {
public:
RaggedTensorToTensorCpuKernelMod() = default;
~RaggedTensorToTensorCpuKernelMod() override = default;
void InitKernel(const CNodePtr &kernel_node) override;
bool Init(const BaseOperatorPtr &base_operator, const std::vector<KernelTensorPtr> &inputs,
const std::vector<KernelTensorPtr> &outputs) override;
bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
const std::vector<AddressPtr> &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> &inputsOnHost = std::map<uint32_t, tensor::TensorPtr>()) override;
protected:
std::vector<KernelAttr> GetOpSupport() override {
@ -73,10 +78,13 @@ class RaggedTensorToTensorCpuKernelMod : public DeprecatedNativeCpuKernelMod {
template <typename TYPE1>
void GetFirstDimension(const std::vector<kernel::AddressPtr> &inputs, TYPE1 *first_dim);
TypeId shape_dtype{kTypeUnknown};
TypeId values_dtype{kTypeUnknown};
TypeId shape_dtype_{kTypeUnknown};
TypeId values_dtype_{kTypeUnknown};
std::vector<int64_t> values_shape_;
std::vector<int64_t> output_shape_;
std::vector<int64_t> default_values_shape_;
std::vector<std::string> row_partition_types_;
CNodePtr node_wpt_;
std::vector<std::vector<int64_t>> row_partition_shape_list_;
};
} // namespace kernel
} // namespace mindspore

10
mindspore/core/ops/ragged_tensor_to_sparse.cc
View File

@ -43,6 +43,16 @@ abstract::TupleShapePtr RaggedTensorToSparseInferShape(const PrimitivePtr &primi
auto rt_dense_values_shape = input_args[kRttsInputValuesStart]->BuildShape();
auto in_values_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(rt_dense_values_shape)[kShape];
auto inputs_splits_element0_shape =
CheckAndConvertUtils::ConvertShapePtrToShapeMap(inputs_splits[0]->BuildShape())[kShape];
if (IsDynamic(inputs_splits_element0_shape) || IsDynamic(in_values_shape)) {
abstract::ShapePtr out_indices =
std::make_shared<abstract::Shape>(ShapeVector({abstract::Shape::kShapeDimAny, abstract::Shape::kShapeDimAny}));
abstract::ShapePtr out_values = std::make_shared<abstract::Shape>(ShapeVector({abstract::Shape::kShapeDimAny}));
abstract::ShapePtr out_shape = std::make_shared<abstract::Shape>(ShapeVector({abstract::Shape::kShapeDimAny}));
return std::make_shared<abstract::TupleShape>(
std::vector<abstract::BaseShapePtr>{out_indices, out_values, out_shape});
}
CheckAndConvertUtils::CheckInteger("rank of 'rt_dense_values'", SizeToLong(in_values_shape.size()), kGreaterEqual, 1,
primitive->name());
ShapeVector out_values_shape = {};

19
mindspore/core/ops/ragged_tensor_to_tensor.cc
View File

@ -31,12 +31,9 @@ BaseShapePtr RaggedTensorToTensorInferShape(const PrimitivePtr &primitive,
const std::vector<AbstractBasePtr> &input_args) {
auto prim_name = primitive->name();
auto shape_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex0]->BuildShape())[kShape];
CheckAndConvertUtils::CheckInteger("dimension of 'shape'", SizeToLong(shape_shape.size()), kEqual, 1, prim_name);
auto values_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex1]->BuildShape())[kShape];
auto default_value_shape =
CheckAndConvertUtils::ConvertShapePtrToShapeMap(input_args[kInputIndex2]->BuildShape())[kShape];
CheckAndConvertUtils::CheckInteger("dimension of 'default_value'", SizeToLong(default_value_shape.size()), kLessThan,
SizeToLong(values_shape.size()), prim_name);
auto shape_arg = input_args[kInputIndex0];
MS_EXCEPTION_IF_NULL(shape_arg);
auto output_shape = GetShapeValue(primitive, shape_arg);
@ -46,17 +43,27 @@ BaseShapePtr RaggedTensorToTensorInferShape(const PrimitivePtr &primitive,
? input_args[kInputIndex3]->cast<abstract::AbstractTuplePtr>()->elements()
: input_args[kInputIndex3]->cast<abstract::AbstractListPtr>()->elements();
auto tensors_size = tensors.size();
auto tensor0_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(tensors[0]->BuildShape())[kShape];
auto tensor0_dim = tensor0_shape.size();
const auto &row_partition_types_ptr = primitive->GetAttr("row_partition_types");
MS_EXCEPTION_IF_NULL(row_partition_types_ptr);
const auto &row_partition_types = GetValue<std::vector<std::string>>(row_partition_types_ptr);
auto types_size = row_partition_types.size();
if (IsDynamic(shape_shape) || IsDynamicRank(values_shape) || IsDynamicRank(default_value_shape) ||
IsDynamicRank(tensor0_shape)) {
return std::make_shared<abstract::Shape>(ShapeVector{abstract::Shape::kShapeRankAny});
}
CheckAndConvertUtils::CheckInteger("dimension of 'shape'", SizeToLong(shape_shape.size()), kEqual, 1, prim_name);
CheckAndConvertUtils::CheckInteger("dimension of 'default_value'", SizeToLong(default_value_shape.size()), kLessThan,
SizeToLong(values_shape.size()), prim_name);
if (tensors_size != types_size) {
MS_EXCEPTION(ValueError) << "For '" << prim_name << "', the number of row_partition_tensors must be equal to the "
<< "number of row_partition_types: " << types_size << ", but got " << tensors_size << ".";
}
if (row_partition_types[0] == "FIRST_DIM_SIZE") {
auto tensor0_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(tensors[0]->BuildShape())[kShape];
auto tensor0_dim = tensor0_shape.size();
CheckAndConvertUtils::CheckInteger("dimension of row_partition_tensors[0](for 'FIRST_DIM_SIZE')",
SizeToLong(tensor0_dim), kEqual, 0, prim_name);
if (types_size - 1 + values_rank != output_shape_rank) {
@ -66,8 +73,6 @@ BaseShapePtr RaggedTensorToTensorInferShape(const PrimitivePtr &primitive,
<< ", 'values' rank: " << values_rank << ".";
}
} else if (row_partition_types[0] == "ROW_SPLITS") {
auto tensor0_shape = CheckAndConvertUtils::ConvertShapePtrToShapeMap(tensors[0]->BuildShape())[kShape];
auto tensor0_dim = tensor0_shape.size();
CheckAndConvertUtils::CheckInteger("dimension of row_partition_tensors[0](for 'ROW_SPLITS')",
SizeToLong(tensor0_dim), kEqual, 1, prim_name);
if (types_size + values_rank != output_shape_rank) {