(func_name, A, B, value, prefetch=True)
| 110 | |
| 111 | |
| 112 | def elementwise_func(func_name, A, B, value, prefetch=True): |
| 113 | func = None |
| 114 | if A.dtype == torch.float32: |
| 115 | func = getattr(lib, f"c{func_name}_fp32", None) |
| 116 | cvalue = ct.c_float(value) |
| 117 | elif A.dtype == torch.uint8: |
| 118 | func = getattr(lib, f"c{func_name}_uint8", None) |
| 119 | cvalue = ct.c_uint8(value) |
| 120 | |
| 121 | if func is None: |
| 122 | raise NotImplementedError(f"Function not implemented: {func_name}") |
| 123 | |
| 124 | is_managed = getattr(A, "is_managed", False) |
| 125 | if is_managed and prefetch: |
| 126 | prefetch_tensor(A) |
| 127 | if B is not None: |
| 128 | prefetch_tensor(B) |
| 129 | |
| 130 | func(get_ptr(A), get_ptr(B), cvalue, ct.c_int64(A.numel())) |
| 131 | if A.is_paged or B.is_paged: |
| 132 | # paged function are fully asynchronous |
| 133 | # if we return from this function, we want to the tensor |
| 134 | # to be in the correct state, that is the final state after the |
| 135 | # operation occurred. So we synchronize. |
| 136 | if torch.cuda.is_available(): |
| 137 | torch.cuda.synchronize() |
| 138 | elif hasattr(torch, "xpu") and torch.xpu.is_available(): |
| 139 | torch.xpu.synchronize() |
| 140 | |
| 141 | |
| 142 | def fill(A, value, device=None, prefetch=True): |
no test coverage detected