()
| 143 | ) |
| 144 | @build_and_clean("matmul_scale.cpp", "matmul_scale.cu") |
| 145 | def test_gpu_func(): |
| 146 | class MatMulScale(Function): |
| 147 | def __init__(self, scale): |
| 148 | super().__init__() |
| 149 | self.scale = scale |
| 150 | |
| 151 | def forward(self, lhs, rhs): |
| 152 | op = custom.MatMulScaleForward(scale=self.scale) |
| 153 | self.lhs = lhs |
| 154 | self.rhs = rhs |
| 155 | return apply(op, lhs, rhs)[0] |
| 156 | |
| 157 | def backward(self, ograd): |
| 158 | op = custom.MatMulScaleBackward(scale=self.scale) |
| 159 | return apply(op, ograd, self.lhs, self.rhs) |
| 160 | |
| 161 | def gen_matmul_data(seed, m, k, n, low=-0.5, high=0.5, dtype=np.float32): |
| 162 | rng = np.random.RandomState(seed=seed) |
| 163 | lhs_np = rng.uniform(low=low, high=high, size=(m, k)).astype(dtype) |
| 164 | rhs_np = rng.uniform(low=low, high=high, size=(k, n)).astype(dtype) |
| 165 | ograd_np = rng.uniform(low=low, high=high, size=(m, n)).astype(dtype) |
| 166 | scale = rng.uniform(low=0.1, high=0.9, size=(1)).astype(np.float32)[0] |
| 167 | |
| 168 | return lhs_np, rhs_np, ograd_np, scale |
| 169 | |
| 170 | def builtin_func(lhs, rhs, scale): |
| 171 | out = F.matmul(lhs, rhs) * scale |
| 172 | return out |
| 173 | |
| 174 | def test_matmul_scale(m=1, k=1, n=1, seed=2021): |
| 175 | lhs_np, rhs_np, _, scale = gen_matmul_data(seed, m, k, n) |
| 176 | custom_lhs, custom_rhs = Tensor(lhs_np), Tensor(rhs_np) |
| 177 | builtin_lhs, builtin_rhs = Tensor(lhs_np), Tensor(rhs_np) |
| 178 | |
| 179 | custom_func = MatMulScale(scale=scale) |
| 180 | custom_out = custom_func(custom_lhs, custom_rhs) |
| 181 | builtin_out = builtin_func(builtin_lhs, builtin_rhs, scale) |
| 182 | |
| 183 | np.testing.assert_allclose(custom_out, builtin_out, rtol=1e-3, atol=1e-5) |
| 184 | |
| 185 | def test_matmul_scale_trace(m=1, k=1, n=1, seed=2021): |
| 186 | @jit.trace(capture_as_const=True) |
| 187 | def func_dumper(lhs, rhs, *, net): |
| 188 | return net(lhs, rhs) |
| 189 | |
| 190 | lhs_np, rhs_np, _, scale = gen_matmul_data(seed, m, k, n) |
| 191 | lhs_tensor, rhs_tensor = Tensor(lhs_np), Tensor(rhs_np) |
| 192 | func = MatMulScale(scale=scale) |
| 193 | real = func_dumper(lhs_tensor, rhs_tensor, net=func) |
| 194 | real = func_dumper(lhs_tensor, rhs_tensor, net=func) |
| 195 | |
| 196 | ref = builtin_func(Tensor(lhs_np), Tensor(rhs_np), scale) |
| 197 | np.testing.assert_allclose(real.numpy(), ref.numpy(), rtol=1e-3, atol=1e-5) |
| 198 | |
| 199 | test_matmul_scale(128, 256, 64, 2028) |
| 200 | test_matmul_scale(64, 32, 16, 2029) |
| 201 | |
| 202 | test_matmul_scale_trace(64, 32, 16, 2030) |
nothing calls this directly
no test coverage detected