MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / test_stream_mem

Function test_stream_mem

imperative/python/test/unit/core/test_stream.py:51–74  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

49
50@pytest.mark.require_ngpu(1)
51def test_stream_mem():
52 memstat = MemStat("xpux:0", "xpux:1")
53
54 x_np = np.random.randint(2 ** 10, size=(1 * 1024 * 1024,), dtype="int32")
55 unit = x_np.size * 4
56 x0 = mge.Tensor(x_np, device="xpux:0")
57
58 results = []
59 events = []
60 for i in range(100):
61 if len(events) >= 2:
62 wait_event("xpux:0", events[-2])
63 x0 = x0 + 1
64 results.append(x0.to("xpux:1", _borrow=True).sum())
65 events.append(record_event("xpux:1"))
66 del events[:-2]
67
68 y_np = x_np.sum()
69 for i, y in enumerate(results):
70 y_np += x_np.size
71 assert y_np == y.numpy()
72
73 mge._full_sync()
74 assert memstat.get_max("xpux:0") / unit < 2.1

Callers

nothing calls this directly

Calls 9

toMethod · 0.95
get_maxMethod · 0.95
wait_eventFunction · 0.90
record_eventFunction · 0.90
MemStatClass · 0.85
TensorMethod · 0.45
appendMethod · 0.45
sumMethod · 0.45
numpyMethod · 0.45

Tested by

no test coverage detected