Write yield in C++ like Python (#4167)

See test for example code.

Author: zasdfgbnm

csrc/utils.h

@@ -21,6 +21,7 @@
 #include <c10/core/thread_pool.h>
 
 #include <concepts>
+#include <coroutine>
 #include <deque>
 #include <iterator>
 #include <memory>
@@ -843,4 +844,136 @@ using views::zip;
 
 #endif // C++23
 
+// Helper: turn T into reference_wrapper<U> if T is reference
+template <typename T>
+using Yielded = std::conditional_t<
+    std::is_reference_v<T>,
+    std::reference_wrapper<std::remove_reference_t<T>>,
+    T>;
+
+// Writing yield in C++20 just like Python:
+// See NVFuserTest.Generator[1-5] for usage examples
+template <typename T>
+class Generator : public std::ranges::view_interface<Generator<T>> {
+ public:
+  struct promise_type;
+  using handle_type = std::coroutine_handle<promise_type>;
+  using stored_type = Yielded<T>;
+
+  Generator(handle_type h) : coroutine_(h) {}
+  Generator(Generator&& other) noexcept : coroutine_(other.coroutine_) {
+    other.coroutine_ = nullptr;
+  }
+  Generator& operator=(Generator&& other) noexcept {
+    if (this != &other) {
+      if (coroutine_) {
+        coroutine_.destroy();
+      }
+      coroutine_ = other.coroutine_;
+      other.coroutine_ = nullptr;
+    }
+    return *this;
+  }
+  ~Generator() {
+    if (coroutine_) {
+      coroutine_.destroy();
+    }
+  }
+  Generator(const Generator&) = delete;
+  Generator& operator=(const Generator&) = delete;
+
+  struct iterator {
+    using value_type = std::remove_reference_t<T>;
+    using reference = T;
+    using difference_type = std::ptrdiff_t;
+    using iterator_category = std::input_iterator_tag;
+
+    iterator() = default;
+    explicit iterator(handle_type h) : coroutine(h) {
+      ++(*this);
+    }
+
+    reference operator*() const {
+      if constexpr (std::is_reference_v<T>) {
+        return value->get(); // unwrap reference_wrapper<T>
+      } else {
+        return *value;
+      }
+    }
+
+    iterator& operator++() {
+      coroutine.resume();
+      if (coroutine.done()) {
+        if (coroutine.promise().exception) {
+          std::rethrow_exception(coroutine.promise().exception);
+        }
+        value.reset();
+      } else {
+        value = std::ref(coroutine.promise().current_value);
+      }
+      return *this;
+    }
+
+    iterator operator++(int) {
+      auto tmp = *this;
+      ++(*this);
+      return tmp;
+    }
+    bool operator==(std::default_sentinel_t) const {
+      return !value.has_value();
+    }
+    bool operator!=(std::default_sentinel_t) const {
+      return value.has_value();
+    }
+    friend bool operator==(std::default_sentinel_t s, const iterator& it) {
+      return it == s;
+    }
+    friend bool operator!=(std::default_sentinel_t s, const iterator& it) {
+      return it != s;
+    }
+
+    handle_type coroutine = nullptr;
+    std::optional<stored_type> value;
+  };
+
+  iterator begin() const {
+    return iterator{coroutine_};
+  }
+  std::default_sentinel_t end() const {
+    return {};
+  }
+
+ private:
+  handle_type coroutine_;
+
+ public:
+  struct promise_type {
+    std::optional<stored_type> current_value;
+    std::exception_ptr exception;
+
+    auto get_return_object() {
+      return Generator{handle_type::from_promise(*this)};
+    }
+    std::suspend_always initial_suspend() {
+      return {};
+    }
+    std::suspend_always final_suspend() noexcept {
+      return {};
+    }
+    std::suspend_always yield_value(T value) {
+      if constexpr (std::is_reference_v<T>) {
+        current_value = std::ref(value); // wraps T& as reference_wrapper
+      } else {
+        current_value = std::move(value);
+      }
+      return {};
+    }
+
+    void return_void() {}
+    void unhandled_exception() {
+      exception = std::current_exception();
+    }
+  };
+};
+
 } // namespace nvfuser

tests/cpp/test_utils.cpp

@@ -18,6 +18,7 @@
 #include <scheduler/vectorize_helper.h>
 #include <tests/cpp/utils.h>
 #include <tests/cpp/validator.h>
+#include <utils.h>
 
 #include <cstdlib>
 #include <filesystem>
@@ -2072,4 +2073,101 @@ TEST_F(TestCpp23BackPort, Enumerate) {
   // bidirectional
 }
 
+namespace {
+
+// Generator that yields integers from 0 to n-1
+Generator<int> zeroToN(int n) {
+  for (int i = 0; i < n; ++i) {
+    co_yield i;
+  }
+}
+
+// Generator that yields integers from n to 2*n - 1
+Generator<int> nTo2N(int n) {
+  for (int i = n; i < 2 * n; ++i) {
+    co_yield i;
+  }
+}
+
+// Generator that yields integers from m to m + 2*n - 1
+Generator<int> mTo2NplusM(int n, int m) {
+  for (auto x : zeroToN(n)) {
+    co_yield x + m;
+  }
+  for (auto x : nTo2N(n)) {
+    co_yield x + m;
+  }
+}
+
+// Generator that yields references
+Generator<int&> items(std::vector<int>& v) {
+  for (auto& x : v) {
+    co_yield x;
+  }
+}
+
+} // namespace
+
+TEST_F(NVFuserTest, Generator1) {
+  static_assert(std::ranges::view<decltype(zeroToN(10))>);
+  std::vector<int> generated;
+  for (auto x : zeroToN(10) |
+           std::views::filter([](int x) { return x % 2 == 0; }) |
+           std::views::transform([](int x) { return x * x; })) {
+    generated.push_back(x);
+  }
+  std::vector<int> expect{0, 4, 16, 36, 64};
+  EXPECT_EQ(generated, expect);
+}
+
+TEST_F(NVFuserTest, Generator2) {
+  static_assert(std::ranges::view<decltype(mTo2NplusM(10, 10))>);
+  std::vector<int> generated;
+  for (auto x : mTo2NplusM(10, 10)) {
+    generated.push_back(x);
+  }
+  std::vector<int> expect{10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
+                          20, 21, 22, 23, 24, 25, 26, 27, 28, 29};
+  EXPECT_EQ(generated, expect);
+}
+
+TEST_F(NVFuserTest, Generator3) {
+  std::vector<int> v{0, 0, 0, 0, 0};
+  for (auto&& [i, x] : enumerate(items(v))) {
+    x = i * 10;
+  }
+  std::vector<int> expect{0, 10, 20, 30, 40};
+  EXPECT_EQ(v, expect);
+}
+
+TEST_F(NVFuserTest, Generator4) {
+  auto one2five = []() -> Generator<int> {
+    for (int i = 1; i <= 5; ++i) {
+      co_yield i;
+    }
+  };
+  std::vector<int> v;
+  for (auto x : one2five()) {
+    v.push_back(x);
+  }
+  std::vector<int> expect{1, 2, 3, 4, 5};
+  EXPECT_EQ(v, expect);
+}
+
+TEST_F(NVFuserTest, Generator5) {
+  auto excepted_exception = []() -> Generator<int> {
+    co_yield 1;
+    throw std::runtime_error("Hello, world!");
+    co_yield 2;
+  };
+  auto run_generator = [&]() {
+    for (auto x : excepted_exception()) {
+      EXPECT_EQ(x, 1);
+    }
+  };
+  EXPECT_THAT(
+      run_generator,
+      ::testing::ThrowsMessage<std::runtime_error>("Hello, world!"));
+}
+
 } // namespace nvfuser