mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-11 17:48:34 +00:00
Implement struct TArray and simplify code. (#3291)
* Implement operator[] for TArray and simplify the code. * fix a build error. * add a constructor with std::vector input * fix build error * update based on code review feedback Co-authored-by: Weixing Zhang <wezhan@microsoft.com>
This commit is contained in:
parent
57468c651c
commit
843ee346a8
7 changed files with 30 additions and 26 deletions
|
|
@ -39,13 +39,13 @@ __global__ void _BinaryElementWise(
|
|||
break;
|
||||
}
|
||||
int q, r;
|
||||
fdm_output_strides.data_[dim].divmod(offset, q, r);
|
||||
fdm_output_strides[dim].divmod(offset, q, r);
|
||||
if (lhs_need_compute) {
|
||||
lhs_index += static_cast<int>(lhs_padded_strides.data_[dim]) * q;
|
||||
lhs_index += static_cast<int>(lhs_padded_strides[dim]) * q;
|
||||
}
|
||||
|
||||
if (rhs_need_compute) {
|
||||
rhs_index += static_cast<int>(rhs_padded_strides.data_[dim]) * q;
|
||||
rhs_index += static_cast<int>(rhs_padded_strides[dim]) * q;
|
||||
}
|
||||
offset = r;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -81,7 +81,7 @@ struct BinaryElementwisePreparation {
|
|||
for (auto i = offset; i < out_rank; ++i) {
|
||||
// the stride for broadcast dimension is kept as 0
|
||||
if (lhs_shape.GetDims()[i - offset] != 1) {
|
||||
lhs_padded_strides.data_[i] = original_lhs_padded_strides[i];
|
||||
lhs_padded_strides[i] = original_lhs_padded_strides[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -93,7 +93,7 @@ struct BinaryElementwisePreparation {
|
|||
for (auto i = offset; i < out_rank; ++i) {
|
||||
// the stride for broadcast dimension is kept as 0
|
||||
if (rhs_shape.GetDims()[i - offset] != 1) {
|
||||
rhs_padded_strides.data_[i] = original_rhs_padded_strides[i];
|
||||
rhs_padded_strides[i] = original_rhs_padded_strides[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -101,7 +101,7 @@ struct BinaryElementwisePreparation {
|
|||
TensorPitches original_output_strides(output_shape.GetDims());
|
||||
fdm_output_strides.size_ = gsl::narrow_cast<int32_t>(out_rank);
|
||||
for (auto i = 0; i < out_rank; ++i) {
|
||||
fdm_output_strides.data_[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
|
||||
fdm_output_strides[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
|
|
|
|||
|
|
@ -48,10 +48,22 @@ struct TArray {
|
|||
ORT_ENFORCE(size <= capacity, "TArray size was set to ", size, ", exeeding the capacity limit of ", capacity);
|
||||
}
|
||||
|
||||
TArray(const std::vector<T>& vec) : TArray(static_cast<int32_t>(vec.size())) {
|
||||
memcpy(data_, vec.data(), vec.size() * sizeof(T));
|
||||
}
|
||||
|
||||
T& operator[](int32_t index) {
|
||||
return data_[index];
|
||||
}
|
||||
|
||||
__host__ __device__ __forceinline__ const T& operator[](int32_t index) const {
|
||||
return data_[index];
|
||||
}
|
||||
|
||||
static constexpr int32_t GetCapacity() { return capacity; };
|
||||
|
||||
T data_[capacity];
|
||||
int32_t size_;
|
||||
T data_[capacity];
|
||||
};
|
||||
|
||||
} // namespace cuda
|
||||
|
|
|
|||
|
|
@ -103,16 +103,8 @@ Status Slice<Tind, dynamic>::ComputeInternal(OpKernelContext* ctx) const {
|
|||
dimension_count = flattened_output_dims.size();
|
||||
}
|
||||
|
||||
TArray<int64_t> starts_buffer(gsl::narrow_cast<int32_t>(starts.size()));
|
||||
for (size_t i = 0; i < starts.size(); ++i) {
|
||||
starts_buffer.data_[i] = starts[i];
|
||||
}
|
||||
|
||||
TArray<int64_t> steps_buffer(gsl::narrow_cast<int32_t>(steps.size()));
|
||||
for (size_t i = 0; i < steps.size(); ++i) {
|
||||
steps_buffer.data_[i] = steps[i];
|
||||
}
|
||||
|
||||
TArray<int64_t> starts_buffer(starts);
|
||||
TArray<int64_t> steps_buffer(steps);
|
||||
TArray<int64_t> input_strides(gsl::narrow_cast<int32_t>(dimension_count));
|
||||
const gsl::span<int64_t> input_strides_span = gsl::make_span(input_strides.data_, input_strides.size_);
|
||||
if (p_flattened_output_dims != nullptr) {
|
||||
|
|
@ -134,8 +126,8 @@ Status Slice<Tind, dynamic>::ComputeInternal(OpKernelContext* ctx) const {
|
|||
|
||||
TensorPitches original_output_strides(p_flattened_output_dims != nullptr ? flattened_output_dims : output_dims);
|
||||
TArray<fast_divmod> output_strides(gsl::narrow_cast<int32_t>(original_output_strides.size()));
|
||||
for (size_t i = 0; i < original_output_strides.size(); ++i) {
|
||||
output_strides.data_[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
|
||||
for (int32_t i = 0; i < static_cast<int32_t>(original_output_strides.size()); ++i) {
|
||||
output_strides[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
|
||||
}
|
||||
|
||||
size_t element_size = input_tensor->DataType()->Size();
|
||||
|
|
|
|||
|
|
@ -29,11 +29,11 @@ __global__ void _SliceKernel(const int32_t dimension_count,
|
|||
break;
|
||||
}
|
||||
|
||||
output_strides.data_[dim].divmod(value, div, mod);
|
||||
input_index += (starts.data_[dim] + div * steps.data_[dim]) * input_strides.data_[dim];
|
||||
output_strides[dim].divmod(value, div, mod);
|
||||
input_index += (starts[dim] + div * steps[dim]) * input_strides[dim];
|
||||
value = mod;
|
||||
}
|
||||
input_index += starts.data_[dim] + mod * steps.data_[dim];
|
||||
input_index += starts[dim] + mod * steps[dim];
|
||||
output_data[id] = input_data[input_index];
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -97,11 +97,11 @@ Status Transpose::DoTranspose(const Transpose& kernel,
|
|||
|
||||
TArray<int64_t> input_strides(rank);
|
||||
for (auto i = 0; i < rank; i++) {
|
||||
input_strides.data_[i] = original_input_strides[permutations[i]];
|
||||
input_strides[i] = original_input_strides[permutations[i]];
|
||||
}
|
||||
TArray<fast_divmod> output_strides(rank);
|
||||
for (auto i = 0; i < rank; i++) {
|
||||
output_strides.data_[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
|
||||
output_strides[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
|
||||
}
|
||||
|
||||
size_t element_size = input.DataType()->Size();
|
||||
|
|
|
|||
|
|
@ -20,9 +20,9 @@ __global__ void _TransposeKernel(int32_t shape_rank, const TArray<int64_t> input
|
|||
break;
|
||||
}
|
||||
int out_coord, r;
|
||||
output_strides.data_[dim].divmod(output_index, out_coord, r);
|
||||
output_strides[dim].divmod(output_index, out_coord, r);
|
||||
output_index = r;
|
||||
input_index += input_strides.data_[dim] * out_coord;
|
||||
input_index += input_strides[dim] * out_coord;
|
||||
}
|
||||
output_data[id] = input_data[input_index];
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue