Add string support to tile op (#14686)

### Description
Add std::string tensor type support to Tile operator


### Motivation and Context
Multiple users are hitting this missing feature:
https://github.com/microsoft/onnxruntime/issues/14511
This commit is contained in:
Ryan Hill 2023-02-16 14:59:44 -08:00 committed by GitHub
parent ae205a7924
commit 892f59b31a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 63 additions and 10 deletions

View file

@ -373,8 +373,8 @@ Do not modify directly.*
|TfIdfVectorizer|*in* X:**T**<br> *out* Y:**T1**|9+|**T** = tensor(int32), tensor(int64), tensor(string)<br/> **T1** = tensor(float)|
|ThresholdedRelu|*in* X:**T**<br> *out* Y:**T**|10+|**T** = tensor(float)|
|||[1, 9]|**T** = tensor(float)|
|Tile|*in* input:**T**<br> *in* repeats:**T1**<br> *out* output:**T**<br><br>or<br><br>*in* input:**T**<br> *in* tiles:**T**<br> *in* axis:**T**<br> *out* output:**T**|13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **T1** = tensor(int64)|
|||[6, 12]|**T** = tensor(bool), tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **T1** = tensor(int64)|
|Tile|*in* input:**T**<br> *in* repeats:**T1**<br> *out* output:**T**<br><br>or<br><br>*in* input:**T**<br> *in* tiles:**T**<br> *in* axis:**T**<br> *out* output:**T**|13+|**T** = tensor(bool), tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **T1** = tensor(int64)|
|||[6, 12]|**T** = tensor(bool), tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **T1** = tensor(int64)|
|TopK|*in* X:**T**<br> *in* K:**tensor(int64)**<br> *out* Values:**T**<br> *out* Indices:**I**<br><br>or<br><br>*in* X:**T**<br> *out* Values:**T**<br> *out* Indices:**I**|11+|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float), tensor(int32), tensor(int64)|
|||10|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float)|
|||[1, 9]|**I** = tensor(int64)<br/> **T** = tensor(double), tensor(float)|

View file

@ -34,6 +34,7 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
DataTypeImpl::GetTensorType<uint16_t>(),
DataTypeImpl::GetTensorType<uint32_t>(),
DataTypeImpl::GetTensorType<uint64_t>(),
DataTypeImpl::GetTensorType<std::string>(),
DataTypeImpl::GetTensorType<bool>()})
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()),
Tile);
@ -51,6 +52,7 @@ ONNX_CPU_OPERATOR_KERNEL(
DataTypeImpl::GetTensorType<uint16_t>(),
DataTypeImpl::GetTensorType<uint32_t>(),
DataTypeImpl::GetTensorType<uint64_t>(),
DataTypeImpl::GetTensorType<std::string>(),
DataTypeImpl::GetTensorType<bool>()})
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()),
Tile);
@ -98,6 +100,45 @@ Status TileCoreForFixedSizeTypes(const Tensor& input_tensor, Tensor& output_tens
return Status::OK();
}
Status TileCoreForStringType(const Tensor& input_tensor, Tensor& output_tensor, const int64_t* repeats, TensorAxisCounters& input_counters, const TensorPitches& output_pitches) {
const auto& input_shape = input_tensor.Shape().GetDims();
const size_t dimension_count = input_shape.size();
const auto* input = input_tensor.Data<std::string>();
auto* output = output_tensor.MutableData<std::string>();
// some helper variables that will be used along the way
size_t block_size = 0;
int64_t num_repeats = 0;
const std::string* copy = nullptr;
const int64_t innermost_dim = input_shape[dimension_count - 1];
while (input_counters) {
// Copy the input data over
block_size = SafeInt<size_t>(innermost_dim);
output = std::copy(input, input + block_size, output);
input += block_size;
// Tile data for the innermost axis
copy = output - block_size;
num_repeats = repeats[dimension_count - 1] - 1;
for (int64_t repeat = 0; repeat < num_repeats; ++repeat) {
output = std::copy(copy, copy + block_size, output);
}
// Tile data for other axes
while (input_counters.Increment()) {
block_size = onnxruntime::narrow<size_t>(output_pitches[input_counters.Axis()] * input_shape[input_counters.Axis()]);
copy = output - block_size;
num_repeats = repeats[input_counters.Axis()] - 1;
for (int64_t repeat = 0; repeat < num_repeats; ++repeat) {
output = std::copy(copy, copy + block_size, output);
}
}
}
return Status::OK();
}
namespace TileOp {
// Find the first non-1 repeat and check the input shape to the left of that dimension:
// 1) If the dim values to the left are all 1s (or don't exist), then the tiling logic is essentially copying the input buffer
@ -170,10 +211,10 @@ Status Tile::Compute(OpKernelContext* ctx) const {
// Repeat tensor has all 1s in it
if (output_shape == input_shape) {
// TODO: Handle string copies when the kernel eventually supports string type.
// For now, it shouldn't throw in the enforce as the kernel doesn't claim string support
ORT_ENFORCE(!input_tensor.IsDataType<std::string>(), "Tile doesn't support string type yet");
memcpy(output_tensor.MutableDataRaw(), input_tensor.DataRaw(), input_tensor.SizeInBytes());
if (input_tensor.IsDataType<std::string>())
std::copy(input_tensor.Data<std::string>(), input_tensor.Data<std::string>()+input_shape.Size(), output_tensor.MutableData<std::string>());
else
memcpy(output_tensor.MutableDataRaw(), input_tensor.DataRaw(), input_tensor.SizeInBytes());
return Status::OK();
}
@ -187,10 +228,7 @@ Status Tile::Compute(OpKernelContext* ctx) const {
is_batched_memcpy,
num_of_elements_per_batch,
num_of_copies_per_batch,
num_of_batch_copies)) {
// TODO: Handle string copies when the kernel eventually supports string type.
// For now, it shouldn't throw in the enforce as the kernel doesn't claim string support
ORT_ENFORCE(!input_tensor.IsDataType<std::string>(), "Tile doesn't support string type yet");
num_of_batch_copies) && !input_tensor.IsDataType<std::string>()) {
int8_t* output_data_casted = reinterpret_cast<int8_t*>(output_tensor.MutableDataRaw());
const int8_t* input_data_casted = reinterpret_cast<const int8_t*>(input_tensor.DataRaw());
@ -239,6 +277,9 @@ Status Tile::Compute(OpKernelContext* ctx) const {
static_assert(sizeof(float) == sizeof(int32_t), "Float and Int32 are of different sizes");
static_assert(sizeof(double) == sizeof(int64_t), "Double and Int64 are of different sizes");
if (input_tensor.IsDataType<std::string>())
return TileCoreForStringType(input_tensor, output_tensor, repeats, input_counters, output_pitches);
if (input_tensor.IsDataType<float>() ||
input_tensor.IsDataType<int32_t>() ||
input_tensor.IsDataType<uint32_t>())

View file

@ -16,6 +16,15 @@ std::vector<T> InputData(size_t size) {
return result;
}
template <>
std::vector<std::string> InputData<std::string>(size_t size) {
std::vector<std::string> result(size);
for (size_t i = 0; i < size; i++) {
result[i] = std::to_string(i);
}
return result;
}
template <>
std::vector<MLFloat16> InputData<MLFloat16>(size_t size) {
std::vector<MLFloat16> result(size);
@ -220,8 +229,11 @@ TEST(TensorOpTest, TileInt64Type) { RunTestWrapper<int64_t>(); }
TEST(TensorOpTest, TileUint64Type) { RunTestWrapper<uint64_t>(); }
TEST(TensorOpTest, TileStringType) { RunTestWrapper<std::string>(); }
TEST(TensorOpTest, TileBoolType) { RunTestWrapperForBool(); }
#if defined(USE_CUDA) || defined(USE_ROCM)
TEST(TensorOpTest, TileMLFloat16Type) { RunTestWrapper<MLFloat16>(); }
#endif