ThreadLocalStream in C++ (#3405)

tests/scheduler_tests.cpp

@@ -104,6 +104,28 @@ TEST_CASE("test new stream in threads") {
   }
 }
 
+TEST_CASE("test thread local stream") {
+  auto s = new_thread_local_stream(default_device());
+  int result = sum(arange(10, s)).item<int>();
+
+  std::atomic<int> finished = 0;
+  std::vector<std::thread> threads;
+  int num_threads = 4;
+  for (int i = 0; i < 4; ++i) {
+    threads.emplace_back([&]() {
+      int r = sum(arange(10, s)).item<int>();
+      CHECK_EQ(result, r);
+      finished += 1;
+      clear_streams();
+    });
+  }
+
+  for (auto& t : threads) {
+    t.join();
+  }
+  CHECK_EQ(finished, num_threads);
+}
+
 TEST_CASE("test get streams") {
   // Initialize default CPU stream before querying
   default_stream(Device::cpu);