(func_name, A, B, value, prefetch=True)
| 153 | |
| 154 | |
| 155 | def elementwise_func(func_name, A, B, value, prefetch=True): |
| 156 | func = None |
| 157 | if A.dtype == torch.float32: |
| 158 | func = getattr(lib, f"c{func_name}_fp32", None) |
| 159 | cvalue = ct.c_float(value) |
| 160 | elif A.dtype == torch.uint8: |
| 161 | func = getattr(lib, f"c{func_name}_uint8", None) |
| 162 | cvalue = ct.c_uint8(value) |
| 163 | |
| 164 | if func is None: |
| 165 | raise NotImplementedError(f"Function not implemented: {func_name}") |
| 166 | |
| 167 | is_managed = getattr(A, "is_managed", False) |
| 168 | if is_managed and prefetch: |
| 169 | prefetch_tensor(A) |
| 170 | if B is not None: |
| 171 | prefetch_tensor(B) |
| 172 | |
| 173 | func(get_ptr(A), get_ptr(B), cvalue, ct.c_int64(A.numel())) |
| 174 | if A.is_paged or B.is_paged: |
| 175 | # paged function are fully asynchronous |
| 176 | # if we return from this function, we want to the tensor |
| 177 | # to be in the correct state, that is the final state after the |
| 178 | # operation occurred. So we synchronize. |
| 179 | torch.cuda.synchronize() |
| 180 | |
| 181 | |
| 182 | def fill(A, value, device=None, prefetch=True): |
no test coverage detected
searching dependent graphs…