1d thrust scan working

python/tvm/topi/cuda/prefix_scan.py

@@ -120,19 +120,19 @@ def ceil_div(a, b):
 
 def is_thrust_available():
     """
-    Test if thrust based scan op is available.
+    Test if thrust based scan ops are available.
     """
-    return get_global_func("tvm.contrib.thrust.scan", allow_missing=True) is not None
+    return get_global_func("tvm.contrib.thrust.sum_scan", allow_missing=True) is not None
 
 
-def scan_thrust(data, axis, exclusive=True):
+def scan_thrust(data, exclusive=True):
     data_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "data_buf", data_alignment=8)
     output_buf = tvm.tir.decl_buffer(data.shape, data.dtype, "output_buf", data_alignment=8)
     return te.extern(
         [data.shape],
         [data],
         lambda ins, outs: tvm.tir.call_packed(
-            "tvm.contrib.thrust.scan", ins[0], outs[0], axis, exclusive
+            "tvm.contrib.thrust.sum_scan", ins[0], outs[0], exclusive
         ),
         dtype=[data.dtype],
         in_buffers=[data_buf],
@@ -144,15 +144,14 @@ def scan_thrust(data, axis, exclusive=True):
 
 def exclusive_scan(data, axis=-1):
     # TODO(masahi): support other binary associative operators
-    # TODO(masahi): support inclusive scan
     ndim = len(data.shape)
     if axis < 0:
         axis += ndim
     assert axis == ndim - 1, "Only support scan on the inner most axis."
 
     target = tvm.target.Target.current()
     if target and target.kind.name == "cuda" and is_thrust_available():
-        return scan_thrust(data, axis)
+        return scan_thrust(data, exclusive=True)
 
     if ndim == 1:
         data = expand_dims(data, axis=0)

src/runtime/contrib/thrust/thrust.cu

@@ -254,19 +254,25 @@ void thrust_scan(DLTensor* data,
   thrust::device_ptr<OutType> output_ptr(static_cast<OutType *>(output->data));
   const auto scan_size = data->shape[data->ndim - 1];
 
-  if (data->ndim == 1) {
+  if (data->ndim == 1 || (data->ndim == 2 && data->shape[0] == 1)) {
     if (exclusive) {
       thrust::exclusive_scan(data_ptr, data_ptr + scan_size, output_ptr);
     } else {
       thrust::inclusive_scan(data_ptr, data_ptr + scan_size, output_ptr);
     }
   } else {
+    // Use thrust segmented scan to compute scan on the inner most axis
+    // data->shape[0] * data->shape[1] * ... * data->shape[ndim - 2] scans are
+    // computed in parallel
+
+    // This is for constructing a sequence 0, 0, 0,...,1, 1, 1,...,2, 2, 2,...,
+    // without materializing the sequence vector
     auto counting_iter = thrust::counting_iterator<int>(0);
     auto key_iter = thrust::make_transform_iterator(counting_iter, [scan_size] __device__(int i) {
         return i / scan_size;
     });
-    int64_t size = 0;
-    for (int i = 0; i < data->ndim; ++i) size += data->shape[i];
+    int64_t size = 1;
+    for (int i = 0; i < data->ndim; ++i) size *= data->shape[i];
 
     if (exclusive) {
       thrust::exclusive_scan_by_key(key_iter, key_iter + size, data_ptr, output_ptr);