add a comment

apache · Nov 25, 2020 · 5943a65 · 5943a65
1 parent 91cae5f
commit 5943a65
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
@@ -699,6 +699,12 @@ def schedule_topk(outs):
 def _dyn_topk_legalize(attrs, inputs, arg_types):
     """Legalizes dyn.topk op.
 
+    On GPU, we don't directly implement a topi kernel. Instead, we combine
+    topi sort and topi strided slice to implement topk. This topi-level composition
+    doesn't work with the dynamic op. To support the dynamic op on gpu, we instead
+    legalize it into the same logic (sort + strided slice) in relay to take advantage
+    of the VM's dynamic shape capabilities
+
     Parameters
     ----------
     attrs : tvm.ir.Attrs