[GitHub] [tvm] mbrookhart commented on a change in pull request #7018: [Topi] Fix GPU Dynamic Topk by Improving Dynamic Strided Slice in Topi

2020-12-03 Thread GitBox 


mbrookhart commented on a change in pull request #7018:
URL: https://github.com/apache/tvm/pull/7018#discussion_r535596794



##
File path: python/tvm/topi/cuda/sort.py
##
@@ -561,10 +561,11 @@ def topk_thrust(data, k=1, axis=-1, ret_type="both", 
is_ascend=False, dtype="int
 tag="topk_gpu",
 )
 
-if k > 0:
+if not isinstance(k, int) or k > 0:
 beg = [0] * ndim
-end = data.shape[:-1] + [k]
-out = [strided_slice(o, beg, end) for o in out]
+end = data.shape[:-1] + [k if isinstance(k, int) else 
tvm.te.size_var("dim")]
+strides = [1] * ndim
+out = [strided_slice(o, beg, end, strides) for o in out]

Review comment:
   @kevinthesun, why don't we just repeat this change in the tir topk 
above? that would fix the unit test, I think.





This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

[GitHub] [tvm] mbrookhart commented on a change in pull request #7018: [Topi] Fix GPU Dynamic Topk by Improving Dynamic Strided Slice in Topi

2020-12-02 Thread GitBox 


mbrookhart commented on a change in pull request #7018:
URL: https://github.com/apache/tvm/pull/7018#discussion_r534571294



##
File path: include/tvm/topi/transform.h
##
@@ -598,17 +598,69 @@ inline te::Tensor dynamic_strided_slice(const te::Tensor& 
x, const te::Tensor& b
  *
  * \return A Tensor whose op member is the split operation
  */
-inline Tensor strided_slice(const Tensor& x, const Array& begin, 
const Array& end,
-const Array& strides, std::string 
slice_mode = "end",
-std::string name = "T_strided_slice", std::string 
tag = kInjective) {
+inline Tensor strided_slice(const Tensor& x, const Array& begin,
+const Array& end, const Array& 
strides,
+std::string slice_mode = "end", std::string name = 
"T_strided_slice",
+std::string tag = kInjective) {
   size_t src_tensor_dim = static_cast(x->shape.size());
+  // Quick path for dynamic shape strided slice.
+  // This is for ease of use to dynamice strided slice in topi.
+  bool is_dyn = false;
+  for (size_t i = 0; i < src_tensor_dim; ++i) {
+if (!IsConstInt(x->shape[i])) {
+  is_dyn = true;
+  break;
+}
+  }
+  if (!is_dyn) {
+for (size_t i = 0; i < begin.size(); ++i) {
+  if (begin[i].defined() && !IsConstInt(begin[i])) {
+is_dyn = true;
+break;
+  }
+}
+  }
+  if (!is_dyn) {
+for (size_t i = 0; i < end.size(); ++i) {
+  if (end[i].defined() && !IsConstInt(end[i])) {
+is_dyn = true;
+break;
+  }
+}
+  }
+  if (!is_dyn) {
+for (size_t i = 0; i < strides.size(); ++i) {
+  if (strides[i].defined() && !IsConstInt(strides[i])) {
+is_dyn = true;
+break;
+  }
+}
+  }
+
+  Array out_shape;
+  if (is_dyn) {
+for (size_t i = 0; i < src_tensor_dim; ++i) {
+  out_shape.push_back(indexdiv(end[i] - begin[i], strides[i]));
+}
+return te::compute(
+out_shape,
+[&](const Array& indices) {
+  Array real_indices;
+  for (size_t i = 0; i < src_tensor_dim; ++i) {
+real_indices.push_back(indices[i] * strides[i] + begin[i]);
+  }
+  return x(real_indices);
+},
+name, tag);
+  }
+

Review comment:
   Can you move this edit into the dynamic_strided_slice function above and 
call that function here?





This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org

[GitHub] [tvm] mbrookhart commented on a change in pull request #7018: [Topi] Fix GPU Dynamic Topk by Improving Dynamic Strided Slice in Topi

2020-12-02 Thread GitBox 


mbrookhart commented on a change in pull request #7018:
URL: https://github.com/apache/tvm/pull/7018#discussion_r534568965



##
File path: include/tvm/topi/transform.h
##
@@ -598,17 +598,69 @@ inline te::Tensor dynamic_strided_slice(const te::Tensor& 
x, const te::Tensor& b
  *
  * \return A Tensor whose op member is the split operation
  */
-inline Tensor strided_slice(const Tensor& x, const Array& begin, 
const Array& end,
-const Array& strides, std::string 
slice_mode = "end",
-std::string name = "T_strided_slice", std::string 
tag = kInjective) {
+inline Tensor strided_slice(const Tensor& x, const Array& begin,
+const Array& end, const Array& 
strides,
+std::string slice_mode = "end", std::string name = 
"T_strided_slice",
+std::string tag = kInjective) {
   size_t src_tensor_dim = static_cast(x->shape.size());
+  // Quick path for dynamic shape strided slice.
+  // This is for ease of use to dynamice strided slice in topi.
+  bool is_dyn = false;
+  for (size_t i = 0; i < src_tensor_dim; ++i) {
+if (!IsConstInt(x->shape[i])) {
+  is_dyn = true;
+  break;
+}
+  }
+  if (!is_dyn) {
+for (size_t i = 0; i < begin.size(); ++i) {
+  if (begin[i].defined() && !IsConstInt(begin[i])) {
+is_dyn = true;
+break;
+  }
+}
+  }
+  if (!is_dyn) {
+for (size_t i = 0; i < end.size(); ++i) {
+  if (end[i].defined() && !IsConstInt(end[i])) {
+is_dyn = true;
+break;
+  }
+}
+  }
+  if (!is_dyn) {
+for (size_t i = 0; i < strides.size(); ++i) {
+  if (strides[i].defined() && !IsConstInt(strides[i])) {
+is_dyn = true;
+break;
+  }
+}
+  }

Review comment:
   Any chance we could do this in an outer loop over src_tensor_dim, begin, 
end, strides?





This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org