[RFC] Adding CUDA Graphs conditional node support with user-friendly APIs

I’m working on CuPy’s CUDA Graph conditional node support and considering adding a more user-friendly graph constructing API to CuPy.
I have two possible plans to achieve this goal:

- (A) “with” API
- (B) Functional API

## (A) “with” API

### pseudo code

```python
import cupy
from contextlib import contextmanager

# Add conditional node to capturing graph and
# switch to different stream to capture body graph
@contextmanager
def while_loop():
...
@contextmanager
def if_else():
...

a: cupy.ndarray = ...
b: cupy.ndarray = ...

with capture_context() as ctx:
...
with ctx.while_loop( # while loop
cond_fn=lambda: cupy.all(a == b),
cond_fn_args=() # Optional
):
... # while loop body

with ctx.if_else(cond_fn=lambda: cupy.all(a != b)) as cond:
with cond.if_():
...
with cond.else_():
# Should we consider `else` body support?
# Currently conditional node does not support `else`
# but seems to add support in the near future.
...

graph = ctx.get_graph() # get captured graph
graph.launch()
```

### Pros

- Considering the construction of CUDA graphs through stream capture, this approach provides a natural API.
- Allows for more concise writing of nested conditional nodes.

### Cons

- Defining a loop using the “with” syntax in Python may not be intuitive.
- Adding a support for an "else" body can result in an awkward or unnatural syntax.

## (B) Functional API

Is this the same approach with [torch.\_higher_order_ops.while](https://github.com/pytorch/pytorch/blob/128544399441fa0ab32f2a95237212c911da5078/torch/_higher_order_ops/while_loop.py#L61-L114)?

### pseudo code

```python
# Prepare GraphConverter class to hold graph capturing state
class GraphConverter:
def __init__(self):
self.state = ... # need to have a global state

# Function to define while loop
def while_loop(self, cond_fn):
def inner_func(body_fn, body_args=None):
... # operation to construct graph
return inner_func

def condition(self, cond_fn):
def inner_func(true_fn, false_fn=None):
...
return inner_func

def convert(self, func):
...
return func

gc = GraphConverter()
@gc.convert # A decorator to convert function to CUDA graph
def target(a: cupy.ndarray, b: cupy.ndarray):
# While
def while_fn(args):
...
gc.while_loop(cond_fn=lambda: cupy.all(a == b))(while_fn, body_args)

# You can also write as follows
@gc.while_loop(cond_fn=lambda: cupy.all(a == b))
def while_fn2():
nonlocal a, b
...

# if-then-else
def true_fn():
...
def false_fn():
...
gc.condition(cond_fn)(
true_fn=true_fn,
false_fn=false_fn # easier to add `else` support
)
# or
@gc.condition(cond_fn)
def true_fn2():
...

target(a, b) # construct and execute graph at the same time
```

### Pros

- Simplifies the addition of "else" support when the conditional node supports an "else" body.
- Prevents issues related to unintentional variable overwriting and improper resource management, as functions naturally encapsulate variable scope.

### Cons

- Writing deeply nested programs can become complex and cumbersome, although using decorators can help alleviate this issue to some extent.
- Requires the use of the "nonlocal" declaration or passing variables as arguments to access outer-scope variables.

## Current status

I began implementing the API using the functional API approach (Plan B). However, switching to Plan A shouldn't be too time-consuming since the core implementation can be shared between both approaches.

I would like to ask your thoughts on these two ideas. Please let us know your preferences, potential advantages, pitfalls, and any other insights you might have.

Start by reviewing the proposed capture_context and GraphConverter entry points, then compare the two API designs and the referenced PyTorch higher-order while_loop implementation. The issue is complete when maintainers choose an API direction and define its scope, including conditional else-body support; no implementation files or tests are identified.

Written by the indexing model from the issue text.