| 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) |
no outgoing calls