Its an async resource duh

NVIDIA · Nov 18, 2024 · 8f375b3 · 8f375b3
1 parent 6e586b8
commit 8f375b3
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 6 deletions.
diff --git a/cudax/include/cuda/experimental/__execution/env.cuh b/cudax/include/cuda/experimental/__execution/env.cuh
@@ -41,7 +41,7 @@ template <class... _Properties>
 class env_t
 {
 private:
-  using __resource   = ::cuda::experimental::mr::any_resource<_Properties...>;
+  using __resource   = ::cuda::experimental::mr::any_async_resource<_Properties...>;
   using __stream_ref = ::cuda::experimental::stream_ref;
 
   __resource __mr_           = ::cuda::experimental::mr::device_memory_resource{};

diff --git a/cudax/test/execution/env.cu b/cudax/test/execution/env.cu
@@ -10,13 +10,10 @@
 
 #include <cuda/std/type_traits>
 
-#include <cuda/experimental/__async/queries.cuh>
+#include <cuda/experimental/buffer.cuh>
 #include <cuda/experimental/execution.cuh>
+#include <cuda/experimental/memory_resource.cuh>
 
-#include "cuda/__memory_resource/properties.h"
-#include "cuda/experimental/__memory_resource/any_resource.cuh"
-#include "cuda/experimental/__memory_resource/device_memory_resource.cuh"
-#include "cuda/std/__type_traits/is_constructible.h"
 #include <catch2/catch.hpp>
 
 namespace cudax = cuda::experimental;
@@ -263,3 +260,23 @@ TEST_CASE("env_t is not constructible from a env missing querries", "[execution,
   STATIC_REQUIRE(!cuda::std::is_constructible_v<env_t, bad_env_t<true, false, true>>);
   STATIC_REQUIRE(!cuda::std::is_constructible_v<env_t, bad_env_t<true, true, false>>);
 }
+
+TEST_CASE("Can use query to construct various objects", "[execution, env]")
+{
+  SECTION("Can create an any_resource")
+  {
+    env_t env{test_resource{}};
+    cudax::mr::any_resource<cuda::mr::device_accessible> resource = env.query(cudax::get_memory_resource);
+    CHECK(resource == test_resource{});
+  }
+
+  SECTION("Can create an uninitialized_async_buffer")
+  {
+    cudax::stream stream_{};
+    env_t env{test_resource{}, stream_};
+    cudax::uninitialized_async_buffer<int, cuda::mr::device_accessible> buf{
+      env.query(cudax::get_memory_resource), env.query(cudax::get_stream), 0ull};
+    CHECK(buf.get_resource() == test_resource{});
+    CHECK(buf.get_stream() == stream_);
+  }
+}