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:
Weixing Zhang 2020-03-23 10:51:54 -07:00 committed by GitHub
parent 57468c651c
commit 843ee346a8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 30 additions and 26 deletions

View file

@ -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;
}

View file

@ -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();

View file

@ -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

View file

@ -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();

View file

@ -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];
}

View file

@ -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();

View file

@ -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];
}