MCPcopy Create free account
hub / github.com/NVIDIA/cutlass / test_binary_ops

Function test_binary_ops

test/examples/CuTeDSL/test_math.py:95–109  ·  view source on GitHub ↗
()

Source from the content-addressed store, hash-verified

93
94
95def test_binary_ops():
96 mag_inp = torch.tensor([3.5, -2.0, 0.0, 1.0], device="cuda", dtype=torch.float32)
97 sign_inp = torch.tensor([-1.0, 1.0, -1.0, 1.0], device="cuda", dtype=torch.float32)
98 expected = torch.tensor([-3.5, 2.0, -0.0, 1.0], device="cuda", dtype=torch.float32)
99 out = torch.zeros(4, device="cuda", dtype=torch.float32)
100 mag_inp_cute = from_dlpack(mag_inp)
101 sign_inp_cute = from_dlpack(sign_inp)
102 out_cute = from_dlpack(out)
103
104 args = (mag_inp_cute, sign_inp_cute, out_cute)
105
106 cute.compile(_binary_ops_host, *args)(*args)
107 torch.cuda.synchronize()
108
109 assert torch.equal(out, expected)

Callers

nothing calls this directly

Calls 5

from_dlpackFunction · 0.90
equalMethod · 0.80
tensorMethod · 0.45
compileMethod · 0.45
synchronizeMethod · 0.45

Tested by

no test coverage detected