Fix lazy adjacent_difference output selection on MUSA

Author: mingxu-2022

cub/detail/type_traits.cuh

@@ -123,6 +123,32 @@ struct adjacent_difference_output_impl<Invokable, InputT, false>
 template <typename Invokable, typename InputT>
 using adjacent_difference_output_t = typename adjacent_difference_output_impl<Invokable, InputT>::type;
 
+template <typename OutputIteratorT,
+          typename Invokable,
+          typename InputT,
+          bool = std::is_same<iterator_value_t<OutputIteratorT>, void>::value>
+struct adjacent_difference_output_select_impl
+{
+  using type = iterator_value_t<OutputIteratorT>;
+};
+
+template <typename OutputIteratorT,
+          typename Invokable,
+          typename InputT>
+struct adjacent_difference_output_select_impl<OutputIteratorT,
+                                              Invokable,
+                                              InputT,
+                                              true>
+{
+  using type = adjacent_difference_output_t<Invokable, InputT>;
+};
+
+template <typename OutputIteratorT, typename Invokable, typename InputT>
+using adjacent_difference_output_select_t =
+  typename adjacent_difference_output_select_impl<OutputIteratorT,
+                                                  Invokable,
+                                                  InputT>::type;
+
 
 } // namespace detail
 CUB_NAMESPACE_END

cub/device/dispatch/dispatch_adjacent_difference.cuh

@@ -78,9 +78,11 @@ DeviceAdjacentDifferenceDifferenceKernel(InputIteratorT input,
 
   // Prefer a concrete output value type when one is available. This avoids
   // device-lambda result introspection on MUSA while preserving write-only
-  // iterators whose value_type is void or unavailable.
-  using OutputT = detail::non_void_iterator_value_t<
-    OutputIteratorT, detail::adjacent_difference_output_t<DifferenceOpT, InputT>>;
+  // iterators whose value_type is void or unavailable. The selection must be
+  // lazy: eager alias arguments still instantiate the lambda fallback on MUSA.
+  using OutputT = detail::adjacent_difference_output_select_t<OutputIteratorT,
+                                                              DifferenceOpT,
+                                                              InputT>;
 
   using Agent = AgentDifference<ActivePolicyT,
                                 InputIteratorT,

test/test_device_adjacent_difference_lambda.cu

@@ -135,6 +135,50 @@ void TestDeviceLambdaWithoutResultType()
   }
 }
 
+void TestCapturedDeviceLambdaWithoutResultType()
+{
+  thrust::device_vector<int> rows_flat(6);
+  rows_flat[0] = 1;
+  rows_flat[1] = 2;
+  rows_flat[2] = 1;
+  rows_flat[3] = 2;
+  rows_flat[4] = 3;
+  rows_flat[5] = 4;
+
+  thrust::device_vector<int> output(3, -1);
+  auto row_ids = thrust::make_counting_iterator<long>(0);
+
+  const int *rows_ptr = thrust::raw_pointer_cast(rows_flat.data());
+  const int row_width = 2;
+
+  auto difference_op = [=] __device__(long lhs_row, long rhs_row) {
+    for (int column = 0; column < row_width; ++column)
+    {
+      const int lhs = rows_ptr[lhs_row * row_width + column];
+      const int rhs = rows_ptr[rhs_row * row_width + column];
+      if (lhs != rhs)
+      {
+        return 1;
+      }
+    }
+
+    return 0;
+  };
+
+  RunSubtractLeftCopy(row_ids, output.begin(), difference_op, 3);
+
+  thrust::host_vector<int> expected(3);
+  expected[0] = 0;
+  expected[1] = 0;
+  expected[2] = 1;
+
+  AssertEquals(output.size(), expected.size());
+  for (std::size_t i = 0; i < expected.size(); ++i)
+  {
+    AssertEquals(output[i], expected[i]);
+  }
+}
+
 void TestWriteOnlyOutputIteratorFallback()
 {
   thrust::device_vector<bool> all_differences_correct(1, true);
@@ -152,6 +196,7 @@ int main(int argc, char **argv)
   CubDebugExit(args.DeviceInit());
 
   TestDeviceLambdaWithoutResultType();
+  TestCapturedDeviceLambdaWithoutResultType();
   TestWriteOnlyOutputIteratorFallback();
 
   return 0;