remove thrust

Author: Awni Hannun

mlx/backend/cuda/device.h

@@ -11,7 +11,6 @@
 #include <cublasLt.h>
 #include <cuda.h>
 #include <cudnn.h>
-#include <thrust/execution_policy.h>
 
 #include <unordered_map>
 
@@ -179,11 +178,4 @@ class Device {
 Device& device(mlx::core::Device device);
 CommandEncoder& get_command_encoder(Stream s);
 
-// Return an execution policy that does not sync for result.
-// Note that not all thrust APIs support async policy, confirm before using.
-inline auto thrust_policy(cudaStream_t stream) {
-  // TODO: Connect thrust's custom allocator with mlx's allocator.
-  return thrust::cuda::par_nosync.on(stream);
-}
-
 } // namespace mlx::core::cu

mlx/backend/cuda/reduce/col_reduce.cu

@@ -8,6 +8,7 @@
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 #include <cub/block/block_load.cuh>
+#include <cub/cub.cuh>
 
 namespace mlx::core {
 

mlx/backend/cuda/utils.cpp

@@ -5,6 +5,7 @@
 #include "mlx/dtype_utils.h"
 
 #include <fmt/format.h>
+#include <cuda/cmath>
 #include <vector>
 
 namespace mlx::core {