BucketizationUtils.h 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167
  1. #pragma once
  2. #include <ATen/core/Tensor.h>
  3. #include <ATen/native/TypeProperties.h>
  4. #include <ATen/ScalarOps.h>
  5. #ifndef AT_PER_OPERATOR_HEADERS
  6. #include <ATen/NativeFunctions.h>
  7. #else
  8. #include <ATen/ops/result_type.h>
  9. #endif
  10. namespace at {
  11. namespace native {
  12. // original values given by raw_*. If an original value is not contiguous, will make a contiguous copy to
  13. // the corresponding trimmed_* value. Additionally, if the dtypes of the boundary and input tensor do not
  14. // match, will change them to be a common super type so comparisons are done between the same types.
  15. // For any trimmed_* tensor, if its outgoing value matches what it was incoming (typically null), then the
  16. // corresponding raw_* version should be used since it was already contiguous of the right type.
  17. inline void searchsorted_maybe_trim_input_tensors(
  18. Tensor& trimmed_input,
  19. Tensor& trimmed_boundaries,
  20. Tensor& trimmed_sorter,
  21. const Tensor& raw_input,
  22. const Tensor& raw_boundaries,
  23. const Tensor& raw_sorter) {
  24. bool in_is_contiguous = raw_input.is_contiguous();
  25. bool bd_is_contiguous = raw_boundaries.is_contiguous();
  26. bool sort_is_contiguous = raw_sorter.is_contiguous();
  27. if (!in_is_contiguous) {
  28. TORCH_WARN_ONCE("torch.searchsorted(): input value tensor is non-contiguous, this will lower the performance due "
  29. "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous input value "
  30. "tensor if possible. This message will only appear once per program.");
  31. trimmed_input = raw_input.contiguous();
  32. }
  33. if (!bd_is_contiguous) {
  34. TORCH_WARN_ONCE("torch.searchsorted(): boundary tensor is non-contiguous, this will lower the performance due "
  35. "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous boundary "
  36. "tensor if possible. This message will only appear once per program.");
  37. trimmed_boundaries = raw_boundaries.contiguous();
  38. }
  39. if (!sort_is_contiguous) {
  40. TORCH_WARN_ONCE("torch.searchsorted(): sorter tensor is non-contiguous, this will lower the performance due "
  41. "to extra data copy when converting non-contiguous tensor to contiguous, please use contiguous sorter "
  42. "tensor if possible. This message will only appear once per program.");
  43. trimmed_sorter = raw_sorter.contiguous();
  44. }
  45. if (raw_input.dtype() != raw_boundaries.dtype()) {
  46. at::native::ResultTypeState state = {};
  47. state = at::native::update_result_type_state(raw_boundaries, state);
  48. state = at::native::update_result_type_state(raw_input, state);
  49. ScalarType common_stype = at::native::result_type(state);
  50. TORCH_INTERNAL_ASSERT(common_stype != ScalarType::Undefined);
  51. if (common_stype != raw_input.scalar_type()) {
  52. trimmed_input = in_is_contiguous ? raw_input.to(common_stype) : trimmed_input.to(common_stype);
  53. }
  54. if (common_stype != raw_boundaries.scalar_type()) {
  55. trimmed_boundaries = bd_is_contiguous ? raw_boundaries.to(common_stype) : trimmed_boundaries.to(common_stype);
  56. }
  57. }
  58. }
  59. /* unused but needed for internal jagged tensor class */
  60. inline void searchsorted_maybe_trim_input_tensors(
  61. Tensor& trimmed_input,
  62. Tensor& trimmed_boundaries,
  63. const Tensor& raw_input,
  64. const Tensor& raw_boundaries) {
  65. Tensor trimmed_sorter;
  66. Tensor raw_sorter;
  67. return searchsorted_maybe_trim_input_tensors(
  68. trimmed_input,
  69. trimmed_boundaries,
  70. trimmed_sorter,
  71. raw_input,
  72. raw_boundaries,
  73. raw_sorter);
  74. }
  75. inline bool searchsorted_dims_matched_before_last_dim(const Tensor& boundaries, const Tensor& input) {
  76. if (boundaries.dim() != input.dim()) {
  77. return false;
  78. }
  79. const auto& dims_bd = boundaries.sizes();
  80. const auto& dims_in = input.sizes();
  81. for (int64_t dim = 0; dim + 1 < boundaries.dim(); ++dim) {
  82. if (dims_bd[dim] != dims_in[dim]) {
  83. return false;
  84. }
  85. }
  86. return true;
  87. }
  88. inline Tensor searchsorted_scalar_tensor(const Scalar& scalar, const c10::Device& device) {
  89. auto tensor = c10::scalar_to_tensor(scalar, device);
  90. // This is to adopt the scalar promotion rules defined in native/TypeProperties.h
  91. // So we have the same type promotion rules as binary operations.
  92. tensor.unsafeGetTensorImpl()->set_wrapped_number(true);
  93. return tensor;
  94. }
  95. inline void searchsorted_pre_check(
  96. const Tensor& boundaries,
  97. const Tensor& input,
  98. const Tensor& output,
  99. const bool out_int32,
  100. const bool right,
  101. const c10::optional<c10::string_view> side_opt,
  102. const Tensor& sorter) {
  103. if (side_opt) {
  104. const c10::string_view side = *side_opt;
  105. TORCH_CHECK(side == "left" || side == "right", "torch.searchsorted(): side can only be 'left' or 'right' but ",
  106. "got ", side);
  107. // assume the user has not explicitly set (right=False, side="right")
  108. TORCH_CHECK(!right || side == "right", "torch.searchsorted(): side and right can't be set to opposites, got side "
  109. "of ", side, " while right was True");
  110. }
  111. TORCH_CHECK(boundaries.device() == input.device(), "torch.searchsorted(): boundaries and input value tensors ",
  112. "should have same device type, but got boundaries tensor device type ", boundaries.device(), " and input value ",
  113. "tensor device type ", input.device());
  114. if (sorter.defined()) {
  115. TORCH_CHECK(sorter.device() == boundaries.device(), "torch.searchsorted(): sorter and boundary tensors should ",
  116. "have same device type, but got sorter tensor device type ", sorter.device(), " and input value tensor ",
  117. "device type ", boundaries.device());
  118. TORCH_CHECK(sorter.sizes() == boundaries.sizes(), "torch.searchsorted(): boundary and sorter must have the same "
  119. "size, but got boundary tensor ", boundaries.sizes(), "and got sorter tensor ", sorter.sizes());
  120. TORCH_CHECK(sorter.scalar_type() == ScalarType::Long, "torch.searchsorted(): sorter must be a tensor of long ",
  121. "dtype but got dtype ", sorter.scalar_type());
  122. }
  123. TORCH_CHECK(input.dim() > 0 || (input.dim() == 0 && input.numel() == 1 && boundaries.dim() == 1),
  124. "torch.searchsorted(): input value can be a scalar only when boundaries tensor dimension is 1, but we got ",
  125. "boundaries tensor dim(", boundaries.dim(), ") and input value's dim(", input.dim(), ") numel(",
  126. input.numel(), ")");
  127. TORCH_CHECK(boundaries.dim() != 0, "torch.searchsorted(): boundaries tensor should have positive dimension, but ",
  128. "got 0 dimension");
  129. TORCH_CHECK(boundaries.dim() == 1 || searchsorted_dims_matched_before_last_dim(boundaries, input),
  130. "torch.searchsorted(): boundaries tensor should be 1 dimension or the first N-1 dimensions of boundaries tensor ",
  131. "and input value tensor must match, but we got boundaries tensor ", boundaries.sizes(), " and input value tensor ",
  132. input.sizes());
  133. ScalarType output_dtype = output.scalar_type();
  134. TORCH_CHECK(
  135. (output_dtype == ScalarType::Long && !out_int32) ||
  136. (output_dtype == ScalarType::Int && out_int32),
  137. "torch.searchsorted(): output tensor's dtype is wrong, it can only be Int(int32) or Long(int64) depending on ",
  138. "whether out_int32 flag is True, but we got output tensor's dtype ", output_dtype,
  139. " and out_int32 flag is ", (out_int32 ? "True" : "False"));
  140. if (out_int32) {
  141. TORCH_CHECK(boundaries.sizes().back() < INT_MAX,
  142. "torch.searchsorted(): the size of boundaries' last dimension should be less than ", INT_MAX, ", but we got ",
  143. boundaries.sizes().back());
  144. }
  145. }
  146. }}