Skip to content

Commit 862a1a2

Browse files
Medyan-NaserMedyan-Naser
committed
add test case to show autodiff for WARP is broken
1 parent 8dd97c3 commit 862a1a2

1 file changed

Lines changed: 250 additions & 0 deletions

File tree

‎examples/cfd/test_stepper_autodiff.py‎

Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,250 @@
1+
"""
2+
Test: XLB Stepper Autodiff - JAX vs Warp Comparison
3+
4+
This script tests whether gradients propagate through the XLB stepper
5+
for both JAX and Warp backends. It performs identical tests on both
6+
backends and compares the results side-by-side.
7+
8+
Expected result: JAX works, Warp does not (stepper lacks adjoint kernels).
9+
10+
Usage:
11+
python examples/cfd/test_stepper_autodiff.py
12+
"""
13+
import numpy as np
14+
15+
print()
16+
print("=" * 70)
17+
print("XLB STEPPER AUTODIFF TEST")
18+
print("=" * 70)
19+
print()
20+
print("This test checks if gradients propagate through the LBM stepper.")
21+
print("We run the SAME test on both JAX and Warp backends and compare.")
22+
print()
23+
24+
# =============================================================================
25+
# SETUP BOTH BACKENDS
26+
# =============================================================================
27+
28+
import warp as wp
29+
wp.init()
30+
31+
import xlb
32+
from xlb.compute_backend import ComputeBackend
33+
from xlb.precision_policy import PrecisionPolicy
34+
from xlb.grid import grid_factory
35+
from xlb.operator.stepper import IncompressibleNavierStokesStepper
36+
from xlb.operator.boundary_condition import FullwayBounceBackBC
37+
from xlb.operator.macroscopic import Macroscopic
38+
import xlb.velocity_set
39+
40+
# Common parameters
41+
grid_shape = (32, 32)
42+
omega = 1.8
43+
precision_policy = PrecisionPolicy.FP32FP32
44+
45+
print("-" * 70)
46+
print("TEST CONFIGURATION")
47+
print("-" * 70)
48+
print(f" Grid shape: {grid_shape}")
49+
print(f" Omega: {omega}")
50+
print(f" Precision: FP32FP32")
51+
print(f" Boundary: FullwayBounceBackBC (walls)")
52+
print(f" Collision: BGK")
53+
print(f" Test: Forward 1 step -> Compute rho -> MSE Loss -> Backward")
54+
print()
55+
56+
# =============================================================================
57+
# WARP BACKEND SETUP
58+
# =============================================================================
59+
60+
warp_velocity_set = xlb.velocity_set.D2Q9(
61+
precision_policy=precision_policy,
62+
compute_backend=ComputeBackend.WARP,
63+
)
64+
65+
xlb.init(
66+
velocity_set=warp_velocity_set,
67+
default_backend=ComputeBackend.WARP,
68+
default_precision_policy=precision_policy,
69+
)
70+
71+
warp_grid = grid_factory(grid_shape, compute_backend=ComputeBackend.WARP)
72+
73+
box = warp_grid.bounding_box_indices()
74+
walls = [box["bottom"][i] + box["top"][i] + box["left"][i] + box["right"][i] for i in range(warp_velocity_set.d)]
75+
walls = np.unique(np.array(walls), axis=-1).tolist()
76+
warp_bc = FullwayBounceBackBC(indices=walls)
77+
78+
warp_stepper = IncompressibleNavierStokesStepper(
79+
grid=warp_grid,
80+
boundary_conditions=[warp_bc],
81+
collision_type="BGK",
82+
)
83+
84+
warp_f_0, warp_f_1, warp_bc_mask, warp_missing_mask = warp_stepper.prepare_fields()
85+
86+
warp_macro = Macroscopic(
87+
velocity_set=warp_velocity_set,
88+
precision_policy=precision_policy,
89+
compute_backend=ComputeBackend.WARP,
90+
)
91+
92+
q = warp_velocity_set.q
93+
shape_4d = (*grid_shape, 1)
94+
95+
@wp.kernel
96+
def warp_loss_kernel(rho: wp.array4d(dtype=wp.float32), loss: wp.array(dtype=wp.float32)):
97+
i, j, k = wp.tid()
98+
wp.atomic_add(loss, 0, rho[0, i, j, k] ** 2.0)
99+
100+
# =============================================================================
101+
# JAX BACKEND SETUP
102+
# =============================================================================
103+
104+
import jax
105+
import jax.numpy as jnp
106+
from jax import value_and_grad
107+
108+
jax_velocity_set = xlb.velocity_set.D2Q9(
109+
precision_policy=precision_policy,
110+
compute_backend=ComputeBackend.JAX,
111+
)
112+
113+
xlb.init(
114+
velocity_set=jax_velocity_set,
115+
default_backend=ComputeBackend.JAX,
116+
default_precision_policy=precision_policy,
117+
)
118+
119+
jax_grid = grid_factory(grid_shape, compute_backend=ComputeBackend.JAX)
120+
121+
jax_box = jax_grid.bounding_box_indices()
122+
jax_walls = [jax_box["bottom"][i] + jax_box["top"][i] + jax_box["left"][i] + jax_box["right"][i] for i in range(jax_velocity_set.d)]
123+
jax_walls = np.unique(np.array(jax_walls), axis=-1).tolist()
124+
jax_bc = FullwayBounceBackBC(indices=jax_walls)
125+
126+
jax_stepper = IncompressibleNavierStokesStepper(
127+
grid=jax_grid,
128+
boundary_conditions=[jax_bc],
129+
collision_type="BGK",
130+
)
131+
132+
jax_f_0, jax_f_1, jax_bc_mask, jax_missing_mask = jax_stepper.prepare_fields()
133+
134+
# =============================================================================
135+
# RUN TESTS
136+
# =============================================================================
137+
138+
# --- WARP TEST ---
139+
f_in_warp = wp.zeros((q, *shape_4d), dtype=wp.float32, requires_grad=True)
140+
f_out_warp = wp.zeros((q, *shape_4d), dtype=wp.float32, requires_grad=True)
141+
rho_warp = wp.zeros((1, *shape_4d), dtype=wp.float32, requires_grad=True)
142+
u_warp = wp.zeros((2, *shape_4d), dtype=wp.float32, requires_grad=True)
143+
loss_warp = wp.zeros((1,), dtype=wp.float32, requires_grad=True)
144+
wp.copy(f_in_warp, warp_f_0)
145+
146+
with wp.Tape() as tape:
147+
f_out_warp, f_in_warp = warp_stepper(f_in_warp, f_out_warp, warp_bc_mask, warp_missing_mask, omega, 0)
148+
rho_warp, u_warp = warp_macro(f_out_warp, rho_warp, u_warp)
149+
wp.launch(warp_loss_kernel, inputs=[rho_warp], outputs=[loss_warp], dim=rho_warp.shape[1:])
150+
151+
warp_loss_val = float(loss_warp.numpy()[0])
152+
loss_warp.grad.fill_(1.0)
153+
tape.backward()
154+
155+
warp_f_in_grad = f_in_warp.grad.numpy() if f_in_warp.grad is not None else np.zeros_like(warp_f_0.numpy())
156+
warp_f_out_grad = f_out_warp.grad.numpy() if f_out_warp.grad is not None else np.zeros_like(warp_f_0.numpy())
157+
warp_rho_grad = rho_warp.grad.numpy() if rho_warp.grad is not None else np.zeros((1, *shape_4d))
158+
159+
warp_f_in_grad_norm = float(np.linalg.norm(warp_f_in_grad))
160+
warp_f_out_grad_norm = float(np.linalg.norm(warp_f_out_grad))
161+
warp_rho_grad_norm = float(np.linalg.norm(warp_rho_grad))
162+
163+
# --- JAX TEST ---
164+
def jax_forward_and_loss(f_in):
165+
f_out, _ = jax_stepper(f_in, jax_f_1, jax_bc_mask, jax_missing_mask, omega, 0)
166+
rho = jnp.sum(f_out, axis=0)
167+
return jnp.sum(rho ** 2)
168+
169+
jax_loss_val, jax_grad = value_and_grad(jax_forward_and_loss)(jax_f_0)
170+
jax_loss_val = float(jax_loss_val)
171+
jax_grad_norm = float(jnp.linalg.norm(jax_grad))
172+
173+
# =============================================================================
174+
# SIDE-BY-SIDE RESULTS
175+
# =============================================================================
176+
177+
print("=" * 70)
178+
print("RESULTS: SIDE-BY-SIDE COMPARISON")
179+
print("=" * 70)
180+
print()
181+
print(f"{'Metric':<35} {'WARP':<15} {'JAX':<15}")
182+
print("-" * 65)
183+
print(f"{'Loss value':<35} {warp_loss_val:<15.4f} {jax_loss_val:<15.4f}")
184+
print(f"{'d(Loss)/d(f_input) gradient norm':<35} {warp_f_in_grad_norm:<15.4f} {jax_grad_norm:<15.4f}")
185+
print()
186+
187+
print("-" * 70)
188+
print("GRADIENT FLOW ANALYSIS (Warp only - to debug where gradients stop)")
189+
print("-" * 70)
190+
print()
191+
print(" In Warp, we can check gradients at each stage of the computation:")
192+
print()
193+
print(f" 1. loss.grad (set manually) : 1.0 (seed)")
194+
print(f" 2. d(loss)/d(rho) gradient norm : {warp_rho_grad_norm:.4f}")
195+
print(f" 3. d(loss)/d(f_out) gradient norm : {warp_f_out_grad_norm:.4f}")
196+
print(f" 4. d(loss)/d(f_in) gradient norm : {warp_f_in_grad_norm:.4f} <-- THIS IS THE PROBLEM")
197+
print()
198+
print(" Gradient flows: loss -> rho -> f_out (through Macroscopic) ✓")
199+
print(" Gradient STOPS: f_out -> f_in (through Stepper) ✗")
200+
print()
201+
202+
print("=" * 70)
203+
print("DIAGNOSIS")
204+
print("=" * 70)
205+
print()
206+
207+
if warp_f_in_grad_norm == 0 and jax_grad_norm > 0:
208+
print(" ISSUE CONFIRMED: Warp stepper does not propagate gradients.")
209+
print()
210+
print(" WHY THIS HAPPENS:")
211+
print(" -----------------")
212+
print(" Warp's autodiff (wp.Tape) requires either:")
213+
print(" a) Automatic adjoint generation (works for simple kernels), or")
214+
print(" b) Manual @wp.func_grad adjoint implementations")
215+
print()
216+
print(" XLB's stepper kernel (nse_stepper.py) has characteristics that")
217+
print(" PREVENT automatic adjoint generation:")
218+
print(" - Early returns: 'if _boundary_id == wp.uint8(255): return'")
219+
print(" - Integer conditionals and mask operations")
220+
print(" - Complex nested @wp.func calls without adjoints")
221+
print()
222+
print(" The Macroscopic operator DOES work because it's a simple")
223+
print(" summation kernel that Warp can auto-differentiate.")
224+
print()
225+
print(" JAX WORKS because it uses source-code transformation (not tape)")
226+
print(" which can differentiate through any Python/JAX code automatically.")
227+
print()
228+
print(" TO FIX (requires XLB core changes):")
229+
print(" ------------------------------------")
230+
print(" Add @wp.func_grad adjoint implementations for:")
231+
print(" - xlb/operator/collision/bgk.py: warp_functional()")
232+
print(" - xlb/operator/stream/stream.py: warp_functional()")
233+
print(" - xlb/operator/equilibrium/*.py: warp_functional()")
234+
print()
235+
print(" RECOMMENDATION:")
236+
print(" ----------------")
237+
print(" Use JAX backend for differentiable LBM applications until")
238+
print(" Warp adjoint kernels are implemented in XLB.")
239+
else:
240+
print(" Unexpected result - please investigate.")
241+
242+
print()
243+
print("=" * 70)
244+
print("SUMMARY")
245+
print("=" * 70)
246+
print()
247+
print(f" WARP: Loss={warp_loss_val:.2f}, Gradient={warp_f_in_grad_norm:.2f} --> {'BROKEN' if warp_f_in_grad_norm == 0 else 'OK'}")
248+
print(f" JAX: Loss={jax_loss_val:.2f}, Gradient={jax_grad_norm:.2f} --> {'OK' if jax_grad_norm > 0 else 'BROKEN'}")
249+
print()
250+
print("=" * 70)

0 commit comments

Comments
 (0)