Skip to content

Commit 34e0b75

Browse files
NV-xiaoyongsNV-xiaoyongs
committed
add slangPy neural slang integarion tests
1 parent 4293931 commit 34e0b75

File tree

9 files changed

+508
-0
lines changed

9 files changed

+508
-0
lines changed

‎slangpy/tests/slangpy_tests/test_neural_bindless.py‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
2+
3+
"""
4+
Neural integration tests for bindless resource types.
5+
6+
Reviewer-requested coverage:
7+
- Bindless "pointer type" (raw pointer parameters passed via Buffer.device_address)
8+
- Bindless DescriptorHandle resources (StructuredBuffer<T>.Handle / RWStructuredBuffer<T>.Handle)
9+
"""
10+
11+
from __future__ import annotations
12+
13+
from pathlib import Path
14+
15+
import numpy as np
16+
import pytest
17+
18+
import slangpy as spy
19+
from slangpy.core.calldata import SLANG_PATH
20+
from slangpy.testing import helpers
21+
22+
23+
def _get_device_with_native_neural(device_type: spy.DeviceType) -> spy.Device:
24+
if helpers.should_skip_test_for_device(device_type):
25+
pytest.skip(f"Device type {device_type.name} not selected for this test run")
26+
27+
test_dir = Path(__file__).resolve().parent
28+
compiler_options = spy.SlangCompilerOptions(
29+
{
30+
"include_paths": [test_dir, SLANG_PATH],
31+
"debug_info": spy.SlangDebugInfoLevel.standard,
32+
"enable_experimental_features": True,
33+
}
34+
)
35+
36+
return spy.Device(
37+
type=device_type,
38+
enable_debug_layers=True,
39+
compiler_options=compiler_options,
40+
label=f"uncached-slangpy-neural-bindless-{device_type.name}",
41+
)
42+
43+
44+
# Pointer-style bindless params are supported on Vulkan. Keep this test on Vulkan only
45+
# to avoid backend-specific CUDA toolchain requirements for this integration test.
46+
POINTER_DEVICE_TYPES: list[spy.DeviceType] = [
47+
x for x in helpers.DEFAULT_DEVICE_TYPES if x in [spy.DeviceType.vulkan]
48+
]
49+
50+
51+
@pytest.mark.parametrize("device_type", POINTER_DEVICE_TYPES)
52+
def test_neural_bindless_pointer_type(device_type: spy.DeviceType) -> None:
53+
device = _get_device_with_native_neural(device_type)
54+
try:
55+
module = spy.Module(device.load_module("test_neural_bindless_pointer.slang"))
56+
57+
buf = device.create_buffer(
58+
size=4,
59+
usage=spy.BufferUsage.shader_resource,
60+
data=np.array([42], dtype=np.int32),
61+
)
62+
63+
res = int(module.read_int_ptr(buf.device_address))
64+
assert res == 42
65+
finally:
66+
device.close()
67+
68+
69+
@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
70+
def test_neural_bindless_descriptor_handle_type(device_type: spy.DeviceType) -> None:
71+
if device_type == spy.DeviceType.cuda:
72+
pytest.skip("Bindless DescriptorHandle resources not supported with CUDA yet.")
73+
74+
device = _get_device_with_native_neural(device_type)
75+
try:
76+
if not device.has_feature(spy.Feature.bindless):
77+
pytest.skip("Bindless not supported on this device.")
78+
79+
module = device.load_module("test_neural_bindless_descriptor_handle.slang")
80+
program = device.link_program(
81+
modules=[module], entry_points=[module.entry_point("compute_main")]
82+
)
83+
kernel = device.create_compute_kernel(program)
84+
85+
buffer_count = 6
86+
87+
ro_buffers: list[spy.Buffer] = []
88+
rw_buffers: list[spy.Buffer] = []
89+
for i in range(buffer_count):
90+
ro_buffers.append(
91+
device.create_buffer(
92+
size=4 * 4,
93+
usage=spy.BufferUsage.shader_resource,
94+
data=np.array([i * 10, i * 10 + 1, i * 10 + 2, i * 10 + 3], dtype=np.float32),
95+
)
96+
)
97+
rw_buffers.append(
98+
device.create_buffer(
99+
size=4 * 4,
100+
usage=spy.BufferUsage.shader_resource | spy.BufferUsage.unordered_access,
101+
data=np.zeros(4, dtype=np.float32),
102+
)
103+
)
104+
105+
buffer_info_layout = module.layout.get_type_layout(
106+
module.layout.find_type_by_name("StructuredBuffer<BufferInfo>")
107+
).element_type_layout
108+
109+
buffer_infos_buffer = device.create_buffer(
110+
size=buffer_count * buffer_info_layout.stride,
111+
usage=spy.BufferUsage.shader_resource,
112+
)
113+
results_buffer = device.create_buffer(
114+
size=buffer_count * 4,
115+
usage=spy.BufferUsage.unordered_access,
116+
)
117+
118+
c = spy.BufferCursor(buffer_info_layout, buffer_infos_buffer, load_before_write=False)
119+
for i in range(buffer_count):
120+
c[i].ro_buffer = ro_buffers[i].descriptor_handle_ro
121+
c[i].rw_buffer = rw_buffers[i].descriptor_handle_rw
122+
c[i].offset = i % 4
123+
c.apply()
124+
125+
kernel.dispatch(
126+
thread_count=[buffer_count, 1, 1],
127+
buffer_infos=buffer_infos_buffer,
128+
results=results_buffer,
129+
)
130+
131+
results = results_buffer.to_numpy().view(np.float32)
132+
expected_results = np.array(
133+
[
134+
0, # buffer 0, offset 0
135+
11, # buffer 1, offset 1
136+
22, # buffer 2, offset 2
137+
33, # buffer 3, offset 3
138+
40, # buffer 4, offset 0
139+
51, # buffer 5, offset 1
140+
],
141+
dtype=np.float32,
142+
)
143+
assert np.allclose(results, expected_results)
144+
145+
# Verify RW buffers were written.
146+
for i in range(buffer_count):
147+
rw_data = rw_buffers[i].to_numpy().view(np.float32)
148+
offset = i % 4
149+
expected_value = (i * 10 + offset) + 100.0
150+
assert np.isclose(rw_data[offset], expected_value)
151+
finally:
152+
device.close()
153+
154+
155+
if __name__ == "__main__":
156+
pytest.main([__file__, "-v", "-s"])
157+

‎slangpy/tests/slangpy_tests/test_neural_bindless_descriptor_handle.slang‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
2+
3+
// Neural integration smoke test for bindless DescriptorHandle style resources.
4+
// We import `neural` to ensure experimental module compilation works alongside bindless.
5+
6+
import neural;
7+
8+
struct BufferInfo
9+
{
10+
StructuredBuffer<float>.Handle ro_buffer;
11+
RWStructuredBuffer<float>.Handle rw_buffer;
12+
uint offset;
13+
};
14+
15+
[shader("compute")]
16+
[numthreads(1, 1, 1)]
17+
void compute_main(
18+
uint3 tid : SV_DispatchThreadID,
19+
StructuredBuffer<BufferInfo> buffer_infos,
20+
RWStructuredBuffer<float> results)
21+
{
22+
uint index = tid.x;
23+
BufferInfo info = buffer_infos[index];
24+
25+
float value = info.ro_buffer[info.offset];
26+
info.rw_buffer[info.offset] = value + 100.0;
27+
results[index] = value;
28+
}
29+

‎slangpy/tests/slangpy_tests/test_neural_bindless_pointer.slang‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
2+
3+
// Neural integration smoke test for "bindless pointer" style parameters.
4+
// This uses raw pointer parameters (passed from Python via Buffer.device_address).
5+
6+
import slangpy;
7+
import neural;
8+
9+
int read_int_ptr(int* ptr)
10+
{
11+
return ptr[0];
12+
}
13+

‎slangpy/tests/slangpy_tests/test_neural_bwd_diff_smoke.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
2+
"""
3+
Neural smoke test that actually exercises Slang autodiff (`bwd_diff(...)`).
4+
5+
Important constraints:
6+
- No dependency on sample apps under `samples/`.
7+
- No dependency on external assets (e.g. image files).
8+
9+
This uses the test-local Slang module `fflayer-bug-repro.slang` which imports the
10+
experimental `neural` module and calls `bwd_diff(loss)(DifferentialPtrPair<Storage>(...), ...)`.
11+
"""
12+
13+
from __future__ import annotations
14+
15+
from pathlib import Path
16+
17+
import numpy as np
18+
import pytest
19+
20+
import slangpy as spy
21+
from slangpy.core.calldata import SLANG_PATH
22+
from slangpy.testing import helpers
23+
24+
25+
def _get_device_with_native_neural(device_type: spy.DeviceType) -> spy.Device:
26+
if helpers.should_skip_test_for_device(device_type):
27+
pytest.skip(f"Device type {device_type.name} not selected for this test run")
28+
29+
test_dir = Path(__file__).resolve().parent
30+
compiler_options = spy.SlangCompilerOptions(
31+
{
32+
"include_paths": [test_dir, SLANG_PATH],
33+
"debug_info": spy.SlangDebugInfoLevel.standard,
34+
"enable_experimental_features": True,
35+
}
36+
)
37+
38+
return spy.Device(
39+
type=device_type,
40+
enable_debug_layers=True,
41+
compiler_options=compiler_options,
42+
label=f"uncached-slangpy-neural-bwd-diff-{device_type.name}",
43+
)
44+
45+
46+
@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
47+
def test_neural_bwd_diff_writes_param_grads(device_type: spy.DeviceType) -> None:
48+
device = _get_device_with_native_neural(device_type)
49+
try:
50+
module = spy.Module(device.load_module("fflayer-bug-repro.slang"))
51+
52+
# 2*2 weights + 2 biases = 6 floats (matches `fflayer-bug-repo.py`)
53+
params = device.create_buffer(
54+
data=np.ones((6,), dtype=np.float32),
55+
usage=spy.BufferUsage.shader_resource | spy.BufferUsage.unordered_access,
56+
)
57+
dparams = device.create_buffer(
58+
data=np.zeros((6,), dtype=np.float32),
59+
usage=spy.BufferUsage.shader_resource | spy.BufferUsage.unordered_access,
60+
)
61+
62+
module.calculate_grad(input=spy.float2(1, 1), params=params, dparams=dparams)
63+
64+
dparams_np = dparams.to_numpy().view(np.float32)
65+
assert np.any(dparams_np != 0.0)
66+
finally:
67+
device.close()
68+

‎slangpy/tests/slangpy_tests/test_neural_frontend_training.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
2+
3+
"""
4+
SlangPy integration test for neural module FFLayer (Option 2 design).
5+
6+
Tests training convergence for a simple quadratic regression task using:
7+
- FFLayer with storage passed as parameter to eval<S>()
8+
- Manual gradient computation (analytic gradients)
9+
- Simple SGD optimization
10+
11+
We fit a quadratic polynomial y = 2*x^2 - 0.5*x + 0.25 and verify convergence.
12+
"""
13+
14+
from __future__ import annotations
15+
16+
from pathlib import Path
17+
18+
import numpy as np
19+
import pytest
20+
21+
import slangpy as spy
22+
from slangpy.core.calldata import SLANG_PATH
23+
from slangpy.testing import helpers
24+
25+
26+
def _get_device_with_native_neural(device_type: spy.DeviceType) -> spy.Device:
27+
if helpers.should_skip_test_for_device(device_type):
28+
pytest.skip(f"Device type {device_type.name} not selected for this test run")
29+
30+
test_dir = Path(__file__).resolve().parent
31+
32+
# Use pre-built neural module from slang (not compiled from source)
33+
# The neural module is built as part of slang-neural-module target
34+
# Enable experimental features since neural is an experimental module
35+
compiler_options = spy.SlangCompilerOptions(
36+
{
37+
"include_paths": [test_dir, SLANG_PATH],
38+
"debug_info": spy.SlangDebugInfoLevel.standard,
39+
"enable_experimental_features": True,
40+
}
41+
)
42+
43+
return spy.Device(
44+
type=device_type,
45+
enable_debug_layers=True,
46+
compiler_options=compiler_options,
47+
label=f"uncached-slangpy-neural-frontend-{device_type.name}",
48+
)
49+
50+
51+
@pytest.mark.parametrize("device_type", helpers.DEFAULT_DEVICE_TYPES)
52+
def test_neural_frontend_training_converges(device_type: spy.DeviceType) -> None:
53+
"""
54+
Test that training converges for a simple quadratic regression task.
55+
56+
Uses FFLayer with Option 2 design (storage as parameter to eval<S>).
57+
"""
58+
device = _get_device_with_native_neural(device_type)
59+
try:
60+
module = spy.Module(device.load_module("test_neural_frontend_training.slang"))
61+
62+
param_count = int(module.get_param_count())
63+
assert param_count == 3
64+
65+
# Fit: y = 2*x^2 - 0.5*x + 0.25
66+
sample_count = 256
67+
xs = np.linspace(-1.0, 1.0, sample_count, dtype=np.float32)
68+
ys = (2.0 * xs * xs - 0.5 * xs + 0.25).astype(np.float32)
69+
70+
xs_buf = device.create_buffer(data=xs, usage=spy.BufferUsage.shader_resource)
71+
ys_buf = device.create_buffer(data=ys, usage=spy.BufferUsage.shader_resource)
72+
73+
rng = np.random.default_rng(0)
74+
params_init = (0.01 * rng.standard_normal(size=(param_count,))).astype(np.float32)
75+
76+
params = device.create_buffer(
77+
data=params_init,
78+
usage=spy.BufferUsage.shader_resource | spy.BufferUsage.unordered_access,
79+
)
80+
grads = device.create_buffer(
81+
data=np.zeros((param_count,), dtype=np.float32),
82+
usage=spy.BufferUsage.shader_resource | spy.BufferUsage.unordered_access,
83+
)
84+
85+
initial_loss = float(module.eval_loss(params, xs_buf, ys_buf, sample_count))
86+
87+
learning_rate = 0.1
88+
steps = 200
89+
for _ in range(steps):
90+
module.train_step(params, grads, xs_buf, ys_buf, sample_count, learning_rate)
91+
92+
final_loss = float(module.eval_loss(params, xs_buf, ys_buf, sample_count))
93+
94+
# Convergence: should significantly reduce MSE and reach a small absolute error.
95+
assert final_loss < initial_loss * 1e-2
96+
assert final_loss < 1e-3
97+
98+
# Parameter packing: [w0, w1, bias] for y = w0*x + w1*x^2 + bias
99+
learned = params.to_numpy().view(np.float32)[:param_count]
100+
expected = np.array([-0.5, 2.0, 0.25], dtype=np.float32)
101+
assert np.allclose(learned, expected, rtol=0.1, atol=0.1)
102+
103+
finally:
104+
device.close()
105+
106+
107+
if __name__ == "__main__":
108+
pytest.main([__file__, "-v", "-s"])

0 commit comments

Comments
 (0)