From 0f3629cb4740a0ece3519977fd2751cccc6af5da Mon Sep 17 00:00:00 2001
From: masa <masa@pop-os.localdomain>
Date: Tue, 8 Dec 2020 19:09:11 +0900
Subject: [PATCH 01/13] add thrust stable sort

---
 python/tvm/topi/cuda/sort.py         | 25 ++++++++++++++
 src/runtime/contrib/thrust/thrust.cu | 50 ++++++++++++++++++++++++++++
 tests/python/contrib/test_sort.py    | 34 +++++++++++++++++--
 3 files changed, 107 insertions(+), 2 deletions(-)

diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
index f28d1cba096c..91e8f5e861a2 100644
--- a/python/tvm/topi/cuda/sort.py
+++ b/python/tvm/topi/cuda/sort.py
@@ -597,3 +597,28 @@ def schedule_topk(outs):
       The computation schedule for the op.
     """
     return _schedule_sort(outs)
+
+
+def stable_sort_thrust(keys, values):
+    """
+    TODO
+    """
+    keys_buf = tvm.tir.decl_buffer(keys.shape, keys.dtype, "keys_buf", data_alignment=8)
+    values_buf = tvm.tir.decl_buffer(values.shape, values.dtype, "values_buf", data_alignment=8)
+    out_bufs = [
+        tvm.tir.decl_buffer(keys.shape, keys.dtype, "keys_buf", data_alignment=8),
+        tvm.tir.decl_buffer(keys.shape, values.dtype, "values_buf", data_alignment=8),
+    ]
+    out = te.extern(
+        [keys.shape, values.shape],
+        [keys, values],
+        lambda ins, outs: tvm.tir.call_packed(
+            "tvm.contrib.thrust.stable_sort_by_key", ins[0], ins[1], outs[0], outs[1]
+        ),
+        in_buffers=[keys_buf, values_buf],
+        out_buffers=out_bufs,
+        dtype=[keys.dtype, values.dtype],
+        name="stable_sort_by_key",
+        tag="stable_sort_by_key",
+    )
+    return out[0], out[1]
diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu
index 2054db710b6d..7f9582fd564b 100644
--- a/src/runtime/contrib/thrust/thrust.cu
+++ b/src/runtime/contrib/thrust/thrust.cu
@@ -163,5 +163,55 @@ TVM_REGISTER_GLOBAL("tvm.contrib.thrust.sort")
   thrust_sort_common(input, values_out, indices_out, is_ascend, get_sort_len,
                      data_dtype, out_dtype);
 });
+
+template<typename KeyType, typename ValueType>
+void thrust_stable_sort_by_key(DLTensor* keys_in,
+                               DLTensor* values_in,
+                               DLTensor* keys_out,
+                               DLTensor* values_out) {
+  const auto size = keys_in->shape[0];
+  thrust::device_ptr<KeyType> keys_in_ptr(static_cast<KeyType *>(keys_in->data));
+  thrust::device_ptr<ValueType> values_in_ptr(static_cast<ValueType *>(values_in->data));
+  thrust::device_ptr<KeyType> keys_out_ptr(static_cast<KeyType *>(keys_out->data));
+  thrust::device_ptr<ValueType> values_out_ptr(static_cast<ValueType *>(values_out->data));
+
+  thrust::copy(keys_in_ptr, keys_in_ptr + size, keys_out_ptr);
+  thrust::copy(values_in_ptr, values_in_ptr + size, values_out_ptr);
+  thrust::stable_sort_by_key(keys_out_ptr, keys_out_ptr + size, values_out_ptr);
+}
+
+TVM_REGISTER_GLOBAL("tvm.contrib.thrust.stable_sort_by_key")
+.set_body([](TVMArgs args, TVMRetValue* ret) {
+  ICHECK_GE(args.num_args, 4);
+  DLTensor* keys_in = args[0];
+  DLTensor* values_in = args[1];
+  DLTensor* keys_out = args[2];
+  DLTensor* values_out = args[3];
+
+  auto key_dtype = DLDataType2String(keys_in->dtype);
+  auto value_dtype = DLDataType2String(values_in->dtype);
+
+  if (key_dtype == "int32") {
+    if (value_dtype == "int32") {
+      thrust_stable_sort_by_key<int, int>(keys_in, values_in, keys_out, values_out);
+    } else if (value_dtype == "float32") {
+      thrust_stable_sort_by_key<int, float>(keys_in, values_in, keys_out, values_out);
+    } else {
+      LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
+    }
+  } else if (key_dtype == "float32") {
+    if (value_dtype == "int32") {
+      thrust_stable_sort_by_key<float, int>(keys_in, values_in, keys_out, values_out);
+    } else if (value_dtype == "float32") {
+      thrust_stable_sort_by_key<float, float>(keys_in, values_in, keys_out, values_out);
+    } else {
+      LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
+    }
+  } else {
+    LOG(FATAL) << "Unsupported key dtype: " << key_dtype;
+  }
+
+});
+
 }  // namespace contrib
 }  // namespace tvm
diff --git a/tests/python/contrib/test_sort.py b/tests/python/contrib/test_sort.py
index 7bd3a9cb55b8..8ff8bf890b38 100644
--- a/tests/python/contrib/test_sort.py
+++ b/tests/python/contrib/test_sort.py
@@ -17,6 +17,7 @@
 import tvm
 import tvm.testing
 from tvm import te
+from tvm.topi.cuda import stable_sort_thrust
 import numpy as np
 
 
@@ -90,6 +91,35 @@ def test_sort_np():
     tvm.testing.assert_allclose(c.asnumpy(), np_out, rtol=1e-5)
 
 
+def test_thrust_stable_sort():
+    size = 6
+    keys = te.placeholder((size,), name="keys", dtype="int32")
+    values = te.placeholder((size,), name="values", dtype="int32")
+
+    keys_out, values_out = stable_sort_thrust(keys, values)
+
+    ctx = tvm.gpu(0)
+    target = "cuda"
+    s = te.create_schedule([keys_out.op, values_out.op])
+    f = tvm.build(s, [keys, values, keys_out, values_out], target)
+
+    keys_np = np.array([1, 4, 2, 8, 2, 7], np.int32)
+    values_np = np.random.randint(0, 10, size=(size,)).astype(np.int32)
+    keys_np_out = np.zeros(keys_np.shape, np.int32)
+    values_np_out = np.zeros(values_np.shape, np.int32)
+    a = tvm.nd.array(keys_np, ctx)
+    b = tvm.nd.array(values_np, ctx)
+    a_out = tvm.nd.array(keys_np_out, ctx)
+    b_out = tvm.nd.array(values_np_out, ctx)
+    f(a, b, a_out, b_out)
+    print(a)
+    print(b)
+    print(a_out)
+    print(b_out)
+    # tvm.testing.assert_allclose(c.asnumpy(), np_out, rtol=1e-5)
+
+
 if __name__ == "__main__":
-    test_sort()
-    test_sort_np()
+    # test_sort()
+    # test_sort_np()
+    test_thrust_stable_sort()

From 568461be1bb5f6b7d39dade6f3e0a591357ce1eb Mon Sep 17 00:00:00 2001
From: masa <masa@pop-os.localdomain>
Date: Tue, 8 Dec 2020 19:22:59 +0900
Subject: [PATCH 02/13] rename

---
 python/tvm/topi/cuda/sort.py      | 2 +-
 tests/python/contrib/test_sort.py | 8 ++++----
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
index 91e8f5e861a2..18250ae086e7 100644
--- a/python/tvm/topi/cuda/sort.py
+++ b/python/tvm/topi/cuda/sort.py
@@ -599,7 +599,7 @@ def schedule_topk(outs):
     return _schedule_sort(outs)
 
 
-def stable_sort_thrust(keys, values):
+def stable_sort_by_key_thrust(keys, values):
     """
     TODO
     """
diff --git a/tests/python/contrib/test_sort.py b/tests/python/contrib/test_sort.py
index 8ff8bf890b38..0327001e21c7 100644
--- a/tests/python/contrib/test_sort.py
+++ b/tests/python/contrib/test_sort.py
@@ -17,7 +17,7 @@
 import tvm
 import tvm.testing
 from tvm import te
-from tvm.topi.cuda import stable_sort_thrust
+from tvm.topi.cuda import stable_sort_by_key_thrust
 import numpy as np
 
 
@@ -91,12 +91,12 @@ def test_sort_np():
     tvm.testing.assert_allclose(c.asnumpy(), np_out, rtol=1e-5)
 
 
-def test_thrust_stable_sort():
+def test_thrust_stable_sort_by_key():
     size = 6
     keys = te.placeholder((size,), name="keys", dtype="int32")
     values = te.placeholder((size,), name="values", dtype="int32")
 
-    keys_out, values_out = stable_sort_thrust(keys, values)
+    keys_out, values_out = stable_sort_by_key_thrust(keys, values)
 
     ctx = tvm.gpu(0)
     target = "cuda"
@@ -122,4 +122,4 @@ def test_thrust_stable_sort():
 if __name__ == "__main__":
     # test_sort()
     # test_sort_np()
-    test_thrust_stable_sort()
+    test_thrust_stable_sort_by_key()

From 346e834ac00ce357ba5fa63c3b4300cac6c4146a Mon Sep 17 00:00:00 2001
From: masa <masa@pop-os.localdomain>
Date: Tue, 8 Dec 2020 19:40:11 +0900
Subject: [PATCH 03/13] scatter via sort working

---
 python/tvm/topi/cuda/scatter.py      | 74 +++++++++++++++++++++++++++-
 python/tvm/topi/cuda/sort.py         |  7 ++-
 src/runtime/contrib/thrust/thrust.cu |  8 +++
 3 files changed, 87 insertions(+), 2 deletions(-)

diff --git a/python/tvm/topi/cuda/scatter.py b/python/tvm/topi/cuda/scatter.py
index 89c5cd23111b..f073838f56b9 100644
--- a/python/tvm/topi/cuda/scatter.py
+++ b/python/tvm/topi/cuda/scatter.py
@@ -20,6 +20,7 @@
 from tvm import te
 from ..scatter import _verify_scatter_nd_inputs
 from .nms import atomic_add
+from .sort import stable_sort_by_key_thrust, is_thrust_available
 
 
 def ceil_div(a, b):
@@ -416,6 +417,67 @@ def gen_ir_4d(data, indices, updates, axis, out, update_func):
     return ib.get()
 
 
+def gen_scatter_1d_thrust(data, indices_sorted, updates_sorted, axis, out, _):
+    """
+    TODO
+    """
+    assert axis == 0
+    n = data.shape[0]
+
+    ib = tvm.tir.ir_builder.create()
+
+    out_ptr = ib.buffer_ptr(out)
+    data_ptr = ib.buffer_ptr(data)
+
+    max_threads = int(tvm.target.Target.current(allow_none=False).max_num_threads)
+    nthread_tx = max_threads
+
+    with ib.new_scope():
+        nthread_bx = ceil_div(n, nthread_tx)
+        tx = te.thread_axis("threadIdx.x")
+        bx = te.thread_axis("blockIdx.x")
+        ib.scope_attr(tx, "thread_extent", nthread_tx)
+        ib.scope_attr(bx, "thread_extent", nthread_bx)
+        tid = bx * nthread_tx + tx
+        with ib.if_scope(tid < n):
+            out_ptr[tid] = data_ptr[tid]
+
+    indices_ptr = ib.buffer_ptr(indices_sorted)
+    updates_ptr = ib.buffer_ptr(updates_sorted)
+
+    ni = indices_sorted.shape[0]
+
+    def do_update(ib, index, update):
+        with ib.if_scope(index < 0):
+            out_ptr[index + n] = update
+        with ib.else_scope():
+            out_ptr[index] = update
+
+    with ib.new_scope():
+        nthread_bx = ceil_div(ni, nthread_tx)
+        tx = te.thread_axis("threadIdx.x")
+        bx = te.thread_axis("blockIdx.x")
+        ib.scope_attr(tx, "thread_extent", nthread_tx)
+        ib.scope_attr(bx, "thread_extent", nthread_bx)
+        tid = bx * nthread_tx + tx
+
+        with ib.if_scope(tid == ni - 1):
+            index = indices_ptr[tid]
+            update = updates_ptr[tid]
+            do_update(ib, index, update)
+
+        with ib.else_scope():
+            with ib.if_scope(tid < ni - 1):
+                index = indices_ptr[tid]
+                index_next = indices_ptr[tid + 1]
+
+                with ib.if_scope(index != index_next):
+                    update = updates_ptr[tid]
+                    do_update(ib, index, update)
+
+    return ib.get()
+
+
 def scatter(data, indices, updates, axis=0):
     """Update data at positions defined by indices with values in updates
 
@@ -458,9 +520,18 @@ def update_func(dst_ptr, dst_index, update):
 
     out_shape = data.shape
     out_buf = tvm.tir.decl_buffer(out_shape, data.dtype, "out_buf")
+
+    in_bufs = [data]
+    if rank == 1 and is_thrust_available():
+        ir_funcs[1] = gen_scatter_1d_thrust
+        indices_sorted, updates_sorted = stable_sort_by_key_thrust(indices, updates)
+        in_bufs += [indices_sorted, updates_sorted]
+    else:
+        in_bufs += [indices, updates]
+
     out = te.extern(
         [out_shape],
-        [data, indices, updates],
+        in_bufs,
         lambda ins, outs: ir_funcs[rank](ins[0], ins[1], ins[2], axis, outs[0], update_func),
         dtype=data.dtype,
         out_buffers=[out_buf],
@@ -548,6 +619,7 @@ def gen_scatter_add_1d_atomic(data, indices, updates, axis, out, _):
     return ib.get()
 
 
+
 def scatter_add(data, indices, updates, axis=0):
     """Update data by adding values in updates at positions defined by indices
 
diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
index 18250ae086e7..f81fbf8b228b 100644
--- a/python/tvm/topi/cuda/sort.py
+++ b/python/tvm/topi/cuda/sort.py
@@ -15,9 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=invalid-name, no-member, too-many-locals, too-many-arguments, too-many-statements, singleton-comparison, unused-argument
-"""Argsort operator """
+"""Sort related operators """
 import tvm
 from tvm import te
+from tvm._ffi import get_global_func
 
 from .injective import schedule_injective_from_existing
 from ..math import identity
@@ -622,3 +623,7 @@ def stable_sort_by_key_thrust(keys, values):
         tag="stable_sort_by_key",
     )
     return out[0], out[1]
+
+
+def is_thrust_available():
+    return get_global_func("tvm.contrib.thrust.sort", allow_missing=True) is not None
diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu
index 7f9582fd564b..eb92ff22c10a 100644
--- a/src/runtime/contrib/thrust/thrust.cu
+++ b/src/runtime/contrib/thrust/thrust.cu
@@ -199,6 +199,14 @@ TVM_REGISTER_GLOBAL("tvm.contrib.thrust.stable_sort_by_key")
     } else {
       LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
     }
+  } else if (key_dtype == "int64") {
+    if (value_dtype == "int32") {
+      thrust_stable_sort_by_key<int64_t, int>(keys_in, values_in, keys_out, values_out);
+    } else if (value_dtype == "float32") {
+      thrust_stable_sort_by_key<int64_t, float>(keys_in, values_in, keys_out, values_out);
+    } else {
+      LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
+    }
   } else if (key_dtype == "float32") {
     if (value_dtype == "int32") {
       thrust_stable_sort_by_key<float, int>(keys_in, values_in, keys_out, values_out);

From be6c207fcd9a4ef30b8c8b62096e25892428e7f1 Mon Sep 17 00:00:00 2001
From: masa <masa@pop-os.localdomain>
Date: Tue, 8 Dec 2020 20:51:06 +0900
Subject: [PATCH 04/13] correctly handles negative indices

---
 cmake/modules/CUDA.cmake             |   1 +
 src/runtime/contrib/thrust/thrust.cu |  27 +++--
 tests/python/relay/test_op_level3.py | 148 +++++++++++++++------------
 3 files changed, 101 insertions(+), 75 deletions(-)

diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake
index 2583e8f3c9ca..3a0d56a7bb1e 100644
--- a/cmake/modules/CUDA.cmake
+++ b/cmake/modules/CUDA.cmake
@@ -59,6 +59,7 @@ if(USE_CUDA)
     message(STATUS "Build with Thrust support")
     cmake_minimum_required(VERSION 3.13) # to compile CUDA code
     enable_language(CUDA)
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --extended-lambda")
     file(GLOB CONTRIB_THRUST_SRC src/runtime/contrib/thrust/*.cu)
     list(APPEND RUNTIME_SRCS ${CONTRIB_THRUST_SRC})
   endif(USE_THRUST)
diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu
index eb92ff22c10a..a0dc841d16b5 100644
--- a/src/runtime/contrib/thrust/thrust.cu
+++ b/src/runtime/contrib/thrust/thrust.cu
@@ -168,15 +168,25 @@ template<typename KeyType, typename ValueType>
 void thrust_stable_sort_by_key(DLTensor* keys_in,
                                DLTensor* values_in,
                                DLTensor* keys_out,
-                               DLTensor* values_out) {
+                               DLTensor* values_out,
+                               bool for_scatter) {
   const auto size = keys_in->shape[0];
   thrust::device_ptr<KeyType> keys_in_ptr(static_cast<KeyType *>(keys_in->data));
   thrust::device_ptr<ValueType> values_in_ptr(static_cast<ValueType *>(values_in->data));
   thrust::device_ptr<KeyType> keys_out_ptr(static_cast<KeyType *>(keys_out->data));
   thrust::device_ptr<ValueType> values_out_ptr(static_cast<ValueType *>(values_out->data));
 
-  thrust::copy(keys_in_ptr, keys_in_ptr + size, keys_out_ptr);
+  if (for_scatter) {
+    auto transform_func = [size] __device__ (KeyType k) {
+      if (k < 0) return k + static_cast<KeyType>(size);
+      return k;
+    };
+    thrust::transform(keys_in_ptr, keys_in_ptr + size, keys_out_ptr, transform_func);
+  } else {
+    thrust::copy(keys_in_ptr, keys_in_ptr + size, keys_out_ptr);
+  }
   thrust::copy(values_in_ptr, values_in_ptr + size, values_out_ptr);
+
   thrust::stable_sort_by_key(keys_out_ptr, keys_out_ptr + size, values_out_ptr);
 }
 
@@ -187,31 +197,32 @@ TVM_REGISTER_GLOBAL("tvm.contrib.thrust.stable_sort_by_key")
   DLTensor* values_in = args[1];
   DLTensor* keys_out = args[2];
   DLTensor* values_out = args[3];
+  bool for_scatter = true;
 
   auto key_dtype = DLDataType2String(keys_in->dtype);
   auto value_dtype = DLDataType2String(values_in->dtype);
 
   if (key_dtype == "int32") {
     if (value_dtype == "int32") {
-      thrust_stable_sort_by_key<int, int>(keys_in, values_in, keys_out, values_out);
+      thrust_stable_sort_by_key<int, int>(keys_in, values_in, keys_out, values_out, for_scatter);
     } else if (value_dtype == "float32") {
-      thrust_stable_sort_by_key<int, float>(keys_in, values_in, keys_out, values_out);
+      thrust_stable_sort_by_key<int, float>(keys_in, values_in, keys_out, values_out, for_scatter);
     } else {
       LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
     }
   } else if (key_dtype == "int64") {
     if (value_dtype == "int32") {
-      thrust_stable_sort_by_key<int64_t, int>(keys_in, values_in, keys_out, values_out);
+      thrust_stable_sort_by_key<int64_t, int>(keys_in, values_in, keys_out, values_out, for_scatter);
     } else if (value_dtype == "float32") {
-      thrust_stable_sort_by_key<int64_t, float>(keys_in, values_in, keys_out, values_out);
+      thrust_stable_sort_by_key<int64_t, float>(keys_in, values_in, keys_out, values_out, for_scatter);
     } else {
       LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
     }
   } else if (key_dtype == "float32") {
     if (value_dtype == "int32") {
-      thrust_stable_sort_by_key<float, int>(keys_in, values_in, keys_out, values_out);
+      thrust_stable_sort_by_key<float, int>(keys_in, values_in, keys_out, values_out, for_scatter);
     } else if (value_dtype == "float32") {
-      thrust_stable_sort_by_key<float, float>(keys_in, values_in, keys_out, values_out);
+      thrust_stable_sort_by_key<float, float>(keys_in, values_in, keys_out, values_out, for_scatter);
     } else {
       LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
     }
diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index fc1929e9dc18..3dace0c4afb6 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -25,6 +25,7 @@
 from tvm.relay import create_executor, transform
 from tvm.relay.testing import check_grad, run_infer_type
 import tvm.testing
+from tvm.contrib import graph_runtime
 
 
 def test_zeros_ones():
@@ -937,10 +938,21 @@ def verify_scatter(dshape, ishape, axis=0):
         ref_res = ref_scatter(data_np, indices_np, updates_np, axis)
 
         for target, ctx in tvm.testing.enabled_targets():
-            for kind in ["graph", "debug"]:
-                intrp = relay.create_executor(kind, ctx=ctx, target=target)
-                op_res = intrp.evaluate(func)(data_np, indices_np, updates_np)
-                tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
+            for kind in ["graph"]:
+                if target == "cuda":
+                    intrp = relay.create_executor(kind, ctx=ctx, target=target)
+                    op_res = intrp.evaluate(func)(data_np, indices_np, updates_np)
+                    print(data_np)
+                    print(indices_np)
+                    print(updates_np)
+                    tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
+
+                    # mod = tvm.ir.IRModule.from_expr(func)
+                    # lib = relay.build(mod, target=target)
+                    # module = graph_runtime.GraphModule(lib["default"](ctx))
+                    # ftimer = module.module.time_evaluator("run", ctx, repeat=100, min_repeat_ms=500)
+                    # prof_res = np.array(ftimer().results) * 1e3  # convert to millisecond
+                    # print("elapsed ms:", prof_res.mean())
 
     def verify_dynamic_scatter(dshape, ishape, axis=0):
         d = relay.var("d", relay.TensorType([relay.Any() for i in range(len(dshape))], "float32"))
@@ -963,33 +975,34 @@ def verify_dynamic_scatter(dshape, ishape, axis=0):
                 op_res = intrp.evaluate()(data_np, indices_np, updates_np)
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
 
-    verify_scatter((10,), (10,), 0)
-    verify_scatter((10, 5), (10, 5), -2)
-    verify_scatter((10, 5), (10, 5), -1)
-    verify_scatter((10, 5), (3, 5), 0)
-    verify_scatter((12, 4), (7, 2), 1)
-    verify_scatter((2, 3, 4), (1, 3, 4), 0)
-    verify_scatter((2, 3, 4), (2, 1, 4), 1)
-    verify_scatter((2, 3, 4), (2, 3, 1), 2)
-    verify_scatter((4, 2, 1), (1, 1, 1), 0)
-    verify_scatter((2, 3, 4, 5), (1, 3, 4, 5), 0)
-    verify_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1)
-    verify_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2)
-    verify_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3)
-
-    verify_dynamic_scatter((10,), (10,), 0)
-    verify_dynamic_scatter((10, 5), (10, 5), -2)
-    verify_dynamic_scatter((10, 5), (10, 5), -1)
-    verify_dynamic_scatter((10, 5), (3, 5), 0)
-    verify_dynamic_scatter((12, 4), (7, 2), 1)
-    verify_dynamic_scatter((2, 3, 4), (1, 3, 4), 0)
-    verify_dynamic_scatter((2, 3, 4), (2, 1, 4), 1)
-    verify_dynamic_scatter((2, 3, 4), (2, 3, 1), 2)
-    verify_dynamic_scatter((4, 2, 1), (1, 1, 1), 0)
-    verify_dynamic_scatter((2, 3, 4, 5), (1, 3, 4, 5), 0)
-    verify_dynamic_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1)
-    verify_dynamic_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2)
-    verify_dynamic_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3)
+
+    verify_scatter((10000,), (10000,), 0)
+    # verify_scatter((10, 5), (10, 5), -2)
+    # verify_scatter((10, 5), (10, 5), -1)
+    # verify_scatter((10, 5), (3, 5), 0)
+    # verify_scatter((12, 4), (7, 2), 1)
+    # verify_scatter((2, 3, 4), (1, 3, 4), 0)
+    # verify_scatter((2, 3, 4), (2, 1, 4), 1)
+    # verify_scatter((2, 3, 4), (2, 3, 1), 2)
+    # verify_scatter((4, 2, 1), (1, 1, 1), 0)
+    # verify_scatter((2, 3, 4, 5), (1, 3, 4, 5), 0)
+    # verify_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1)
+    # verify_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2)
+    # verify_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3)
+
+    # verify_dynamic_scatter((10,), (10,), 0)
+    # verify_dynamic_scatter((10, 5), (10, 5), -2)
+    # verify_dynamic_scatter((10, 5), (10, 5), -1)
+    # verify_dynamic_scatter((10, 5), (3, 5), 0)
+    # verify_dynamic_scatter((12, 4), (7, 2), 1)
+    # verify_dynamic_scatter((2, 3, 4), (1, 3, 4), 0)
+    # verify_dynamic_scatter((2, 3, 4), (2, 1, 4), 1)
+    # verify_dynamic_scatter((2, 3, 4), (2, 3, 1), 2)
+    # verify_dynamic_scatter((4, 2, 1), (1, 1, 1), 0)
+    # verify_dynamic_scatter((2, 3, 4, 5), (1, 3, 4, 5), 0)
+    # verify_dynamic_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1)
+    # verify_dynamic_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2)
+    # verify_dynamic_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3)
 
 
 @tvm.testing.uses_gpu
@@ -1310,39 +1323,40 @@ def verify_adv_index(data_shape, index_shapes):
 
 
 if __name__ == "__main__":
-    test_cast()
-    test_zeros_ones()
-    test_unary_identity()
-    test_clip()
-    test_transpose_infer_type()
-    test_transpose()
-    test_reshape_infer_type()
-    test_reshape()
-    test_reshape_fail()
-    test_reshape_like_infer_type()
-    test_reshape_like()
-    test_take_infer_type()
-    test_take()
-    test_full_infer_type()
-    test_full()
-    test_full_like_infer_type()
-    test_full_like()
-    test_infer_type_leaky_relu()
-    test_infer_type_prelu()
-    test_squeeze()
-    test_squeeze_infer_type()
-    test_squeeze_bad_axes_infer_type()
-    test_split_infer_type()
-    test_arange()
-    test_meshgrid()
-    test_reverse()
-    test_stack()
-    test_tile()
-    test_repeat()
-    test_gather_nd()
-    test_isfinite()
-    test_isinf()
-    test_unravel_index()
-    test_sparse_to_dense()
-    test_fixed_point_multiply()
-    test_adv_index()
+    # test_cast()
+    # test_zeros_ones()
+    # test_unary_identity()
+    # test_clip()
+    # test_transpose_infer_type()
+    # test_transpose()
+    # test_reshape_infer_type()
+    # test_reshape()
+    # test_reshape_fail()
+    # test_reshape_like_infer_type()
+    # test_reshape_like()
+    # test_take_infer_type()
+    # test_take()
+    # test_full_infer_type()
+    # test_full()
+    # test_full_like_infer_type()
+    # test_full_like()
+    # test_infer_type_leaky_relu()
+    # test_infer_type_prelu()
+    # test_squeeze()
+    # test_squeeze_infer_type()
+    # test_squeeze_bad_axes_infer_type()
+    # test_split_infer_type()
+    # test_arange()
+    # test_meshgrid()
+    # test_reverse()
+    # test_stack()
+    # test_tile()
+    # test_repeat()
+    # test_gather_nd()
+    # test_isfinite()
+    # test_isinf()
+    # test_unravel_index()
+    # test_sparse_to_dense()
+    # test_fixed_point_multiply()
+    # test_adv_index()
+    test_scatter()

From 752e7849e4458bf4601bca996299c75a9f728042 Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Tue, 8 Dec 2020 21:59:33 +0900
Subject: [PATCH 05/13] clean up, add some comments

---
 python/tvm/topi/cuda/scatter.py      |  8 ++++++--
 python/tvm/topi/cuda/sort.py         |  4 ++--
 src/runtime/contrib/thrust/thrust.cu |  4 ++--
 tests/python/contrib/test_sort.py    | 24 ++++++++++++------------
 tests/python/relay/test_op_level3.py | 23 ++++++++++++-----------
 5 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/python/tvm/topi/cuda/scatter.py b/python/tvm/topi/cuda/scatter.py
index f073838f56b9..b1429ea564ea 100644
--- a/python/tvm/topi/cuda/scatter.py
+++ b/python/tvm/topi/cuda/scatter.py
@@ -471,6 +471,9 @@ def do_update(ib, index, update):
                 index = indices_ptr[tid]
                 index_next = indices_ptr[tid + 1]
 
+                # If the next neighbor in the sorted list of indices has a different index,
+                # that means thread tid is the last one to have this index.
+                # This thread can update the output.
                 with ib.if_scope(index != index_next):
                     update = updates_ptr[tid]
                     do_update(ib, index, update)
@@ -524,7 +527,9 @@ def update_func(dst_ptr, dst_index, update):
     in_bufs = [data]
     if rank == 1 and is_thrust_available():
         ir_funcs[1] = gen_scatter_1d_thrust
-        indices_sorted, updates_sorted = stable_sort_by_key_thrust(indices, updates)
+        indices_sorted, updates_sorted = stable_sort_by_key_thrust(
+            indices, updates, for_scatter=True
+        )
         in_bufs += [indices_sorted, updates_sorted]
     else:
         in_bufs += [indices, updates]
@@ -619,7 +624,6 @@ def gen_scatter_add_1d_atomic(data, indices, updates, axis, out, _):
     return ib.get()
 
 
-
 def scatter_add(data, indices, updates, axis=0):
     """Update data by adding values in updates at positions defined by indices
 
diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
index f81fbf8b228b..b6ee1f2bec21 100644
--- a/python/tvm/topi/cuda/sort.py
+++ b/python/tvm/topi/cuda/sort.py
@@ -600,7 +600,7 @@ def schedule_topk(outs):
     return _schedule_sort(outs)
 
 
-def stable_sort_by_key_thrust(keys, values):
+def stable_sort_by_key_thrust(keys, values, for_scatter=False):
     """
     TODO
     """
@@ -614,7 +614,7 @@ def stable_sort_by_key_thrust(keys, values):
         [keys.shape, values.shape],
         [keys, values],
         lambda ins, outs: tvm.tir.call_packed(
-            "tvm.contrib.thrust.stable_sort_by_key", ins[0], ins[1], outs[0], outs[1]
+            "tvm.contrib.thrust.stable_sort_by_key", ins[0], ins[1], outs[0], outs[1], for_scatter
         ),
         in_buffers=[keys_buf, values_buf],
         out_buffers=out_bufs,
diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu
index a0dc841d16b5..12d3bdcfdbc4 100644
--- a/src/runtime/contrib/thrust/thrust.cu
+++ b/src/runtime/contrib/thrust/thrust.cu
@@ -192,12 +192,12 @@ void thrust_stable_sort_by_key(DLTensor* keys_in,
 
 TVM_REGISTER_GLOBAL("tvm.contrib.thrust.stable_sort_by_key")
 .set_body([](TVMArgs args, TVMRetValue* ret) {
-  ICHECK_GE(args.num_args, 4);
+  ICHECK_GE(args.num_args, 5);
   DLTensor* keys_in = args[0];
   DLTensor* values_in = args[1];
   DLTensor* keys_out = args[2];
   DLTensor* values_out = args[3];
-  bool for_scatter = true;
+  bool for_scatter = args[4];
 
   auto key_dtype = DLDataType2String(keys_in->dtype);
   auto value_dtype = DLDataType2String(values_in->dtype);
diff --git a/tests/python/contrib/test_sort.py b/tests/python/contrib/test_sort.py
index 0327001e21c7..0a8b9b2b8e07 100644
--- a/tests/python/contrib/test_sort.py
+++ b/tests/python/contrib/test_sort.py
@@ -107,19 +107,19 @@ def test_thrust_stable_sort_by_key():
     values_np = np.random.randint(0, 10, size=(size,)).astype(np.int32)
     keys_np_out = np.zeros(keys_np.shape, np.int32)
     values_np_out = np.zeros(values_np.shape, np.int32)
-    a = tvm.nd.array(keys_np, ctx)
-    b = tvm.nd.array(values_np, ctx)
-    a_out = tvm.nd.array(keys_np_out, ctx)
-    b_out = tvm.nd.array(values_np_out, ctx)
-    f(a, b, a_out, b_out)
-    print(a)
-    print(b)
-    print(a_out)
-    print(b_out)
-    # tvm.testing.assert_allclose(c.asnumpy(), np_out, rtol=1e-5)
+    keys_in = tvm.nd.array(keys_np, ctx)
+    values_in = tvm.nd.array(values_np, ctx)
+    keys_out = tvm.nd.array(keys_np_out, ctx)
+    values_out = tvm.nd.array(values_np_out, ctx)
+    f(keys_in, values_in, keys_out, values_out)
+
+    ref_keys_out = np.sort(keys_np)
+    ref_values_out = np.array([values_np[i] for i in np.argsort(keys_np)])
+    tvm.testing.assert_allclose(keys_out.asnumpy(), ref_keys_out, rtol=1e-5)
+    tvm.testing.assert_allclose(values_out.asnumpy(), ref_values_out, rtol=1e-5)
 
 
 if __name__ == "__main__":
-    # test_sort()
-    # test_sort_np()
+    test_sort()
+    test_sort_np()
     test_thrust_stable_sort_by_key()
diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 3dace0c4afb6..c197e04d761a 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -942,17 +942,17 @@ def verify_scatter(dshape, ishape, axis=0):
                 if target == "cuda":
                     intrp = relay.create_executor(kind, ctx=ctx, target=target)
                     op_res = intrp.evaluate(func)(data_np, indices_np, updates_np)
-                    print(data_np)
-                    print(indices_np)
-                    print(updates_np)
                     tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
 
-                    # mod = tvm.ir.IRModule.from_expr(func)
-                    # lib = relay.build(mod, target=target)
-                    # module = graph_runtime.GraphModule(lib["default"](ctx))
-                    # ftimer = module.module.time_evaluator("run", ctx, repeat=100, min_repeat_ms=500)
-                    # prof_res = np.array(ftimer().results) * 1e3  # convert to millisecond
-                    # print("elapsed ms:", prof_res.mean())
+                    mod = tvm.ir.IRModule.from_expr(func)
+                    lib = relay.build(mod, target=target)
+                    module = graph_runtime.GraphModule(lib["default"](ctx))
+                    module.set_input("d", tvm.nd.array(data_np))
+                    module.set_input("i", tvm.nd.array(indices_np))
+                    module.set_input("u", tvm.nd.array(updates_np))
+                    ftimer = module.module.time_evaluator("run", ctx, repeat=100)
+                    prof_res = np.array(ftimer().results) * 1e3  # convert to millisecond
+                    print("size %d, elapsed ms: %f" % ( dshape[0], prof_res.mean()))
 
     def verify_dynamic_scatter(dshape, ishape, axis=0):
         d = relay.var("d", relay.TensorType([relay.Any() for i in range(len(dshape))], "float32"))
@@ -975,8 +975,9 @@ def verify_dynamic_scatter(dshape, ishape, axis=0):
                 op_res = intrp.evaluate()(data_np, indices_np, updates_np)
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
 
-
-    verify_scatter((10000,), (10000,), 0)
+    np.random.seed(123)
+    for size in [5000, 10000, 25000, 50000, 100000, 500000, 1000000]:
+        verify_scatter((size,), (size,), 0)
     # verify_scatter((10, 5), (10, 5), -2)
     # verify_scatter((10, 5), (10, 5), -1)
     # verify_scatter((10, 5), (3, 5), 0)

From 6447a538a7d693724d824afe577369dcdf4195ab Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Tue, 8 Dec 2020 22:19:01 +0900
Subject: [PATCH 06/13] add doc string

---
 python/tvm/topi/cuda/scatter.py | 27 +++++++++++++++++++++++++--
 python/tvm/topi/cuda/sort.py    | 31 +++++++++++++++++++++++++++++--
 2 files changed, 54 insertions(+), 4 deletions(-)

diff --git a/python/tvm/topi/cuda/scatter.py b/python/tvm/topi/cuda/scatter.py
index b1429ea564ea..68f4707eeb9b 100644
--- a/python/tvm/topi/cuda/scatter.py
+++ b/python/tvm/topi/cuda/scatter.py
@@ -418,8 +418,29 @@ def gen_ir_4d(data, indices, updates, axis, out, update_func):
 
 
 def gen_scatter_1d_thrust(data, indices_sorted, updates_sorted, axis, out, _):
-    """
-    TODO
+    """Generate scatter ir for 4d inputs
+
+    Parameters
+    ----------
+    data : tir.Tensor
+        The input data to the operator.
+
+    indices_sorted : tir.Tensor
+        The sorted index locations to update.
+
+    updates : tir.Tensor
+        The values to update, sorted by indices.
+
+    axis : int
+        The axis to scatter on
+
+    out : tir.Tensor
+        The output tensor.
+
+    Returns
+    -------
+    ret : tir
+        The computational ir.
     """
     assert axis == 0
     n = data.shape[0]
@@ -462,6 +483,7 @@ def do_update(ib, index, update):
         tid = bx * nthread_tx + tx
 
         with ib.if_scope(tid == ni - 1):
+            # The last element can always update.
             index = indices_ptr[tid]
             update = updates_ptr[tid]
             do_update(ib, index, update)
@@ -525,6 +547,7 @@ def update_func(dst_ptr, dst_index, update):
     out_buf = tvm.tir.decl_buffer(out_shape, data.dtype, "out_buf")
 
     in_bufs = [data]
+
     if rank == 1 and is_thrust_available():
         ir_funcs[1] = gen_scatter_1d_thrust
         indices_sorted, updates_sorted = stable_sort_by_key_thrust(
diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
index b6ee1f2bec21..2f7064c614f4 100644
--- a/python/tvm/topi/cuda/sort.py
+++ b/python/tvm/topi/cuda/sort.py
@@ -601,8 +601,32 @@ def schedule_topk(outs):
 
 
 def stable_sort_by_key_thrust(keys, values, for_scatter=False):
-    """
-    TODO
+    """ Sorting values with respect to keys using thrust.
+    Both keys and values will be sorted and returned.
+    Sorting is done via stable sort, so relative ordering among
+    ties are preserved.
+
+    Parameters
+    ----------
+    keys: tvm.te.Tensor
+        The 1D input keys.
+
+    values : tvm.te.Tensor,
+        The 1D input values.
+
+    for_scatter: bool, optional
+        If True, negative keys are interpreted as negative indices.
+        Before sorting, negative indices are converted to corresponding positive indices.
+        The output keys (indices) are all positive.
+        This option is introduced to optimize the scatter implementation.
+
+    Returns
+    -------
+    keys_sorted : tvm.te.Tensor
+        The sorted keys
+
+    values_sorted : tvm.te.Tensor
+        The values sorted with respect to the keys
     """
     keys_buf = tvm.tir.decl_buffer(keys.shape, keys.dtype, "keys_buf", data_alignment=8)
     values_buf = tvm.tir.decl_buffer(values.shape, values.dtype, "values_buf", data_alignment=8)
@@ -626,4 +650,7 @@ def stable_sort_by_key_thrust(keys, values, for_scatter=False):
 
 
 def is_thrust_available():
+    """
+    Test if thrust based sorting ops are available.
+    """
     return get_global_func("tvm.contrib.thrust.sort", allow_missing=True) is not None

From e16be3a57236b79e507ac58d8b4cea8ec275a4d3 Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Tue, 8 Dec 2020 22:20:27 +0900
Subject: [PATCH 07/13] remove scatter benchmark stuff

---
 tests/python/relay/test_op_level3.py | 153 ++++++++++++---------------
 1 file changed, 67 insertions(+), 86 deletions(-)

diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index c197e04d761a..82d056381666 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -25,7 +25,6 @@
 from tvm.relay import create_executor, transform
 from tvm.relay.testing import check_grad, run_infer_type
 import tvm.testing
-from tvm.contrib import graph_runtime
 
 
 def test_zeros_ones():
@@ -938,21 +937,10 @@ def verify_scatter(dshape, ishape, axis=0):
         ref_res = ref_scatter(data_np, indices_np, updates_np, axis)
 
         for target, ctx in tvm.testing.enabled_targets():
-            for kind in ["graph"]:
-                if target == "cuda":
-                    intrp = relay.create_executor(kind, ctx=ctx, target=target)
-                    op_res = intrp.evaluate(func)(data_np, indices_np, updates_np)
-                    tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
-
-                    mod = tvm.ir.IRModule.from_expr(func)
-                    lib = relay.build(mod, target=target)
-                    module = graph_runtime.GraphModule(lib["default"](ctx))
-                    module.set_input("d", tvm.nd.array(data_np))
-                    module.set_input("i", tvm.nd.array(indices_np))
-                    module.set_input("u", tvm.nd.array(updates_np))
-                    ftimer = module.module.time_evaluator("run", ctx, repeat=100)
-                    prof_res = np.array(ftimer().results) * 1e3  # convert to millisecond
-                    print("size %d, elapsed ms: %f" % ( dshape[0], prof_res.mean()))
+            for kind in ["graph", "debug"]:
+                intrp = relay.create_executor(kind, ctx=ctx, target=target)
+                op_res = intrp.evaluate(func)(data_np, indices_np, updates_np)
+                tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
 
     def verify_dynamic_scatter(dshape, ishape, axis=0):
         d = relay.var("d", relay.TensorType([relay.Any() for i in range(len(dshape))], "float32"))
@@ -975,35 +963,33 @@ def verify_dynamic_scatter(dshape, ishape, axis=0):
                 op_res = intrp.evaluate()(data_np, indices_np, updates_np)
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
 
-    np.random.seed(123)
-    for size in [5000, 10000, 25000, 50000, 100000, 500000, 1000000]:
-        verify_scatter((size,), (size,), 0)
-    # verify_scatter((10, 5), (10, 5), -2)
-    # verify_scatter((10, 5), (10, 5), -1)
-    # verify_scatter((10, 5), (3, 5), 0)
-    # verify_scatter((12, 4), (7, 2), 1)
-    # verify_scatter((2, 3, 4), (1, 3, 4), 0)
-    # verify_scatter((2, 3, 4), (2, 1, 4), 1)
-    # verify_scatter((2, 3, 4), (2, 3, 1), 2)
-    # verify_scatter((4, 2, 1), (1, 1, 1), 0)
-    # verify_scatter((2, 3, 4, 5), (1, 3, 4, 5), 0)
-    # verify_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1)
-    # verify_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2)
-    # verify_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3)
-
-    # verify_dynamic_scatter((10,), (10,), 0)
-    # verify_dynamic_scatter((10, 5), (10, 5), -2)
-    # verify_dynamic_scatter((10, 5), (10, 5), -1)
-    # verify_dynamic_scatter((10, 5), (3, 5), 0)
-    # verify_dynamic_scatter((12, 4), (7, 2), 1)
-    # verify_dynamic_scatter((2, 3, 4), (1, 3, 4), 0)
-    # verify_dynamic_scatter((2, 3, 4), (2, 1, 4), 1)
-    # verify_dynamic_scatter((2, 3, 4), (2, 3, 1), 2)
-    # verify_dynamic_scatter((4, 2, 1), (1, 1, 1), 0)
-    # verify_dynamic_scatter((2, 3, 4, 5), (1, 3, 4, 5), 0)
-    # verify_dynamic_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1)
-    # verify_dynamic_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2)
-    # verify_dynamic_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3)
+    verify_scatter((10,), (10,), 0)
+    verify_scatter((10, 5), (10, 5), -2)
+    verify_scatter((10, 5), (10, 5), -1)
+    verify_scatter((10, 5), (3, 5), 0)
+    verify_scatter((12, 4), (7, 2), 1)
+    verify_scatter((2, 3, 4), (1, 3, 4), 0)
+    verify_scatter((2, 3, 4), (2, 1, 4), 1)
+    verify_scatter((2, 3, 4), (2, 3, 1), 2)
+    verify_scatter((4, 2, 1), (1, 1, 1), 0)
+    verify_scatter((2, 3, 4, 5), (1, 3, 4, 5), 0)
+    verify_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1)
+    verify_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2)
+    verify_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3)
+
+    verify_dynamic_scatter((10,), (10,), 0)
+    verify_dynamic_scatter((10, 5), (10, 5), -2)
+    verify_dynamic_scatter((10, 5), (10, 5), -1)
+    verify_dynamic_scatter((10, 5), (3, 5), 0)
+    verify_dynamic_scatter((12, 4), (7, 2), 1)
+    verify_dynamic_scatter((2, 3, 4), (1, 3, 4), 0)
+    verify_dynamic_scatter((2, 3, 4), (2, 1, 4), 1)
+    verify_dynamic_scatter((2, 3, 4), (2, 3, 1), 2)
+    verify_dynamic_scatter((4, 2, 1), (1, 1, 1), 0)
+    verify_dynamic_scatter((2, 3, 4, 5), (1, 3, 4, 5), 0)
+    verify_dynamic_scatter((6, 3, 4, 5), (2, 3, 4, 5), 1)
+    verify_dynamic_scatter((2, 3, 8, 5), (2, 3, 1, 1), 2)
+    verify_dynamic_scatter((16, 16, 4, 5), (16, 16, 4, 5), 3)
 
 
 @tvm.testing.uses_gpu
@@ -1031,15 +1017,11 @@ def verify_scatter_add(dshape, ishape, axis=0):
         ref_res = ref_scatter_add(data_np, indices_np, updates_np, axis)
         for target, ctx in tvm.testing.enabled_targets():
             for kind in ["graph", "debug"]:
-                if target == "nvptx":
-                    # TODO(masahi): support atomic in LLVM codegen
-                    continue
                 intrp = relay.create_executor(kind, ctx=ctx, target=target)
                 op_res = intrp.evaluate(func)(data_np, indices_np, updates_np)
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
 
     verify_scatter_add((10,), (10,), 0)
-    verify_scatter_add((1000,), (1000,), 0)
     verify_scatter_add((10, 5), (10, 5), -2)
     verify_scatter_add((10, 5), (10, 5), -1)
     verify_scatter_add((10, 5), (3, 5), 0)
@@ -1324,40 +1306,39 @@ def verify_adv_index(data_shape, index_shapes):
 
 
 if __name__ == "__main__":
-    # test_cast()
-    # test_zeros_ones()
-    # test_unary_identity()
-    # test_clip()
-    # test_transpose_infer_type()
-    # test_transpose()
-    # test_reshape_infer_type()
-    # test_reshape()
-    # test_reshape_fail()
-    # test_reshape_like_infer_type()
-    # test_reshape_like()
-    # test_take_infer_type()
-    # test_take()
-    # test_full_infer_type()
-    # test_full()
-    # test_full_like_infer_type()
-    # test_full_like()
-    # test_infer_type_leaky_relu()
-    # test_infer_type_prelu()
-    # test_squeeze()
-    # test_squeeze_infer_type()
-    # test_squeeze_bad_axes_infer_type()
-    # test_split_infer_type()
-    # test_arange()
-    # test_meshgrid()
-    # test_reverse()
-    # test_stack()
-    # test_tile()
-    # test_repeat()
-    # test_gather_nd()
-    # test_isfinite()
-    # test_isinf()
-    # test_unravel_index()
-    # test_sparse_to_dense()
-    # test_fixed_point_multiply()
-    # test_adv_index()
-    test_scatter()
+    test_cast()
+    test_zeros_ones()
+    test_unary_identity()
+    test_clip()
+    test_transpose_infer_type()
+    test_transpose()
+    test_reshape_infer_type()
+    test_reshape()
+    test_reshape_fail()
+    test_reshape_like_infer_type()
+    test_reshape_like()
+    test_take_infer_type()
+    test_take()
+    test_full_infer_type()
+    test_full()
+    test_full_like_infer_type()
+    test_full_like()
+    test_infer_type_leaky_relu()
+    test_infer_type_prelu()
+    test_squeeze()
+    test_squeeze_infer_type()
+    test_squeeze_bad_axes_infer_type()
+    test_split_infer_type()
+    test_arange()
+    test_meshgrid()
+    test_reverse()
+    test_stack()
+    test_tile()
+    test_repeat()
+    test_gather_nd()
+    test_isfinite()
+    test_isinf()
+    test_unravel_index()
+    test_sparse_to_dense()
+    test_fixed_point_multiply()
+    test_adv_index()

From 23641966acc1db13af27c5634bbbde631064ea05 Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Tue, 8 Dec 2020 22:29:52 +0900
Subject: [PATCH 08/13] add more doc

---
 python/tvm/topi/cuda/scatter.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/python/tvm/topi/cuda/scatter.py b/python/tvm/topi/cuda/scatter.py
index 68f4707eeb9b..bb960ac45053 100644
--- a/python/tvm/topi/cuda/scatter.py
+++ b/python/tvm/topi/cuda/scatter.py
@@ -418,7 +418,12 @@ def gen_ir_4d(data, indices, updates, axis, out, update_func):
 
 
 def gen_scatter_1d_thrust(data, indices_sorted, updates_sorted, axis, out, _):
-    """Generate scatter ir for 4d inputs
+    """Generate scatter ir for 1d inputs, using a sorting based approach.
+    By sorting indices and comparing neighboring two indices, we can tell which
+    of elements in the indices tensor can scatter its update value into the output.
+    Sorting of indices, and sorting of updates with respect to indices, can be done
+    at the same time by thrust's sort_by_key function. It is important that sorting
+    be done in a "stable" way via stable_sort, to guaranteee deterministic output.
 
     Parameters
     ----------

From aaaba5e89a0723915a6ea3c4d22dfb612c637b93 Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Tue, 8 Dec 2020 22:34:43 +0900
Subject: [PATCH 09/13] fix typo

---
 python/tvm/topi/cuda/scatter.py      | 4 ++--
 python/tvm/topi/cuda/sort.py         | 2 +-
 tests/python/relay/test_op_level3.py | 4 ++++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/python/tvm/topi/cuda/scatter.py b/python/tvm/topi/cuda/scatter.py
index bb960ac45053..9916e2a7fa6d 100644
--- a/python/tvm/topi/cuda/scatter.py
+++ b/python/tvm/topi/cuda/scatter.py
@@ -423,7 +423,7 @@ def gen_scatter_1d_thrust(data, indices_sorted, updates_sorted, axis, out, _):
     of elements in the indices tensor can scatter its update value into the output.
     Sorting of indices, and sorting of updates with respect to indices, can be done
     at the same time by thrust's sort_by_key function. It is important that sorting
-    be done in a "stable" way via stable_sort, to guaranteee deterministic output.
+    be done in a "stable" way via stable_sort, to guarantee deterministic output.
 
     Parameters
     ----------
@@ -437,7 +437,7 @@ def gen_scatter_1d_thrust(data, indices_sorted, updates_sorted, axis, out, _):
         The values to update, sorted by indices.
 
     axis : int
-        The axis to scatter on
+        The axis to scatter on. It must be 0 for this function.
 
     out : tir.Tensor
         The output tensor.
diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
index 2f7064c614f4..3d359be46c0e 100644
--- a/python/tvm/topi/cuda/sort.py
+++ b/python/tvm/topi/cuda/sort.py
@@ -601,7 +601,7 @@ def schedule_topk(outs):
 
 
 def stable_sort_by_key_thrust(keys, values, for_scatter=False):
-    """ Sorting values with respect to keys using thrust.
+    """ Sort values with respect to keys using thrust.
     Both keys and values will be sorted and returned.
     Sorting is done via stable sort, so relative ordering among
     ties are preserved.
diff --git a/tests/python/relay/test_op_level3.py b/tests/python/relay/test_op_level3.py
index 82d056381666..fc1929e9dc18 100644
--- a/tests/python/relay/test_op_level3.py
+++ b/tests/python/relay/test_op_level3.py
@@ -1017,11 +1017,15 @@ def verify_scatter_add(dshape, ishape, axis=0):
         ref_res = ref_scatter_add(data_np, indices_np, updates_np, axis)
         for target, ctx in tvm.testing.enabled_targets():
             for kind in ["graph", "debug"]:
+                if target == "nvptx":
+                    # TODO(masahi): support atomic in LLVM codegen
+                    continue
                 intrp = relay.create_executor(kind, ctx=ctx, target=target)
                 op_res = intrp.evaluate(func)(data_np, indices_np, updates_np)
                 tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5)
 
     verify_scatter_add((10,), (10,), 0)
+    verify_scatter_add((1000,), (1000,), 0)
     verify_scatter_add((10, 5), (10, 5), -2)
     verify_scatter_add((10, 5), (10, 5), -1)
     verify_scatter_add((10, 5), (3, 5), 0)

From 0e8747747ca7bc5efc6651021ff53b8303966d76 Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Tue, 8 Dec 2020 22:39:09 +0900
Subject: [PATCH 10/13] lint fix

---
 src/runtime/contrib/thrust/thrust.cu | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu
index 12d3bdcfdbc4..a376be0b0c71 100644
--- a/src/runtime/contrib/thrust/thrust.cu
+++ b/src/runtime/contrib/thrust/thrust.cu
@@ -229,7 +229,6 @@ TVM_REGISTER_GLOBAL("tvm.contrib.thrust.stable_sort_by_key")
   } else {
     LOG(FATAL) << "Unsupported key dtype: " << key_dtype;
   }
-
 });
 
 }  // namespace contrib

From 095b157ae82384b7a734ae517d0f9bb4c37b703c Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Tue, 8 Dec 2020 22:51:59 +0900
Subject: [PATCH 11/13] silence lint

---
 src/runtime/contrib/thrust/thrust.cu | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/runtime/contrib/thrust/thrust.cu b/src/runtime/contrib/thrust/thrust.cu
index a376be0b0c71..8ccefc5ee7d2 100644
--- a/src/runtime/contrib/thrust/thrust.cu
+++ b/src/runtime/contrib/thrust/thrust.cu
@@ -177,11 +177,10 @@ void thrust_stable_sort_by_key(DLTensor* keys_in,
   thrust::device_ptr<ValueType> values_out_ptr(static_cast<ValueType *>(values_out->data));
 
   if (for_scatter) {
-    auto transform_func = [size] __device__ (KeyType k) {
+    thrust::transform(keys_in_ptr, keys_in_ptr + size, keys_out_ptr, [size] __device__(KeyType k) {
       if (k < 0) return k + static_cast<KeyType>(size);
       return k;
-    };
-    thrust::transform(keys_in_ptr, keys_in_ptr + size, keys_out_ptr, transform_func);
+    });
   } else {
     thrust::copy(keys_in_ptr, keys_in_ptr + size, keys_out_ptr);
   }
@@ -204,25 +203,31 @@ TVM_REGISTER_GLOBAL("tvm.contrib.thrust.stable_sort_by_key")
 
   if (key_dtype == "int32") {
     if (value_dtype == "int32") {
-      thrust_stable_sort_by_key<int, int>(keys_in, values_in, keys_out, values_out, for_scatter);
+      thrust_stable_sort_by_key<int, int>(keys_in, values_in, keys_out, values_out,
+                                          for_scatter);
     } else if (value_dtype == "float32") {
-      thrust_stable_sort_by_key<int, float>(keys_in, values_in, keys_out, values_out, for_scatter);
+      thrust_stable_sort_by_key<int, float>(keys_in, values_in, keys_out, values_out,
+                                            for_scatter);
     } else {
       LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
     }
   } else if (key_dtype == "int64") {
     if (value_dtype == "int32") {
-      thrust_stable_sort_by_key<int64_t, int>(keys_in, values_in, keys_out, values_out, for_scatter);
+      thrust_stable_sort_by_key<int64_t, int>(keys_in, values_in, keys_out, values_out,
+                                              for_scatter);
     } else if (value_dtype == "float32") {
-      thrust_stable_sort_by_key<int64_t, float>(keys_in, values_in, keys_out, values_out, for_scatter);
+      thrust_stable_sort_by_key<int64_t, float>(keys_in, values_in, keys_out, values_out,
+                                                for_scatter);
     } else {
       LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
     }
   } else if (key_dtype == "float32") {
     if (value_dtype == "int32") {
-      thrust_stable_sort_by_key<float, int>(keys_in, values_in, keys_out, values_out, for_scatter);
+      thrust_stable_sort_by_key<float, int>(keys_in, values_in, keys_out, values_out,
+                                            for_scatter);
     } else if (value_dtype == "float32") {
-      thrust_stable_sort_by_key<float, float>(keys_in, values_in, keys_out, values_out, for_scatter);
+      thrust_stable_sort_by_key<float, float>(keys_in, values_in, keys_out, values_out,
+                                              for_scatter);
     } else {
       LOG(FATAL) << "Unsupported value dtype: " << value_dtype;
     }

From 4a7b78725283c4ee0c0f4def8dca2201e4b6e504 Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Tue, 8 Dec 2020 22:56:12 +0900
Subject: [PATCH 12/13] fix py format

---
 python/tvm/topi/cuda/sort.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/tvm/topi/cuda/sort.py b/python/tvm/topi/cuda/sort.py
index 3d359be46c0e..0094ef1adf11 100644
--- a/python/tvm/topi/cuda/sort.py
+++ b/python/tvm/topi/cuda/sort.py
@@ -601,7 +601,7 @@ def schedule_topk(outs):
 
 
 def stable_sort_by_key_thrust(keys, values, for_scatter=False):
-    """ Sort values with respect to keys using thrust.
+    """Sort values with respect to keys using thrust.
     Both keys and values will be sorted and returned.
     Sorting is done via stable sort, so relative ordering among
     ties are preserved.

From cc513c400f86e96b738e5f3085306252ac3ded88 Mon Sep 17 00:00:00 2001
From: Masahiro Masuda <masahi129@gmail.com>
Date: Wed, 9 Dec 2020 00:01:53 +0900
Subject: [PATCH 13/13] check for thrust availablity before test

---
 tests/python/contrib/test_sort.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/python/contrib/test_sort.py b/tests/python/contrib/test_sort.py
index 0a8b9b2b8e07..9d6eb7cb3a1e 100644
--- a/tests/python/contrib/test_sort.py
+++ b/tests/python/contrib/test_sort.py
@@ -17,7 +17,7 @@
 import tvm
 import tvm.testing
 from tvm import te
-from tvm.topi.cuda import stable_sort_by_key_thrust
+from tvm.topi.cuda import stable_sort_by_key_thrust, is_thrust_available
 import numpy as np
 
 
@@ -92,6 +92,10 @@ def test_sort_np():
 
 
 def test_thrust_stable_sort_by_key():
+    if not is_thrust_available():
+        print("skip because thrust is not enabled...")
+        return
+
     size = 6
     keys = te.placeholder((size,), name="keys", dtype="int32")
     values = te.placeholder((size,), name="values", dtype="int32")