| 13 | bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so') |
| 14 | |
| 15 | def bitnet_int8xint2_linear(input0, input1, s, ws, ret): |
| 16 | out_shape = list(input0.shape) |
| 17 | out_shape[-1] = input1.shape[0] |
| 18 | |
| 19 | stream = torch.cuda.current_stream() |
| 20 | |
| 21 | M = input0.shape[0] |
| 22 | if len(out_shape) == 3: |
| 23 | M *= input0.shape[1] |
| 24 | N = input1.shape[0] |
| 25 | K = input1.shape[1] * 4 |
| 26 | |
| 27 | bitnet_lib.bitlinear_int8xint2(*[ctypes.c_void_p(input0.data_ptr()), ctypes.c_void_p(input1.data_ptr()), ctypes.c_void_p(ret.data_ptr()), ctypes.c_void_p(s.data_ptr()), ctypes.c_void_p(ws.data_ptr()), ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), ctypes.c_void_p(stream.cuda_stream)]) |
| 28 | |
| 29 | return ret |
| 30 | |
| 31 | if __name__ == '__main__': |
| 32 | test_list = [ |