(
A: Tensor,
B: Tensor,
out: Optional[torch.Tensor] = None,
transposed_A=False,
transposed_B=False,
)
| 1336 | |
| 1337 | @deprecated("This function is deprecated and will be removed in a future release.", category=FutureWarning) |
| 1338 | def igemm( |
| 1339 | A: Tensor, |
| 1340 | B: Tensor, |
| 1341 | out: Optional[torch.Tensor] = None, |
| 1342 | transposed_A=False, |
| 1343 | transposed_B=False, |
| 1344 | ): |
| 1345 | sout = check_matmul(A, B, out, transposed_A, transposed_B) |
| 1346 | if out is None: |
| 1347 | out = torch.zeros(size=sout, dtype=torch.int32, device=A.device) |
| 1348 | if len(A.shape) == 3 and len(B.shape) == 3: |
| 1349 | if A.shape[0] == B.shape[0] and A.shape[2] == B.shape[1]: |
| 1350 | return batched_igemm(A, B, out) |
| 1351 | |
| 1352 | sA = A.shape |
| 1353 | sB = B.shape |
| 1354 | if transposed_A and len(sA) == 2: |
| 1355 | sA = (sA[1], sA[0]) |
| 1356 | elif transposed_A and len(sA) == 3: |
| 1357 | sA = (sA[0], sA[2], sA[0]) |
| 1358 | if transposed_B and len(sB) == 2: |
| 1359 | sB = (sB[1], sB[0]) |
| 1360 | elif transposed_B and len(sB) == 3: |
| 1361 | sB = (sB[0], sB[2], sB[0]) |
| 1362 | # this is a mess: cuBLAS expect column major, but PyTorch is row major. |
| 1363 | # So to perform the matrix multiplication, we have to treat A, B, and C matrices |
| 1364 | # (transpose of row major is column major) |
| 1365 | # This means we compute B^T A^T = C^T and we explicitly switch the dimensions of each of these |
| 1366 | |
| 1367 | # matrices in the input arguments for cuBLAS |
| 1368 | # column major: A @ B = C: [m, k] @ [k, n] = [m, n] |
| 1369 | # row major: B^T @ A^T = C^T: [m, k] @ [k, n] = [m, n] |
| 1370 | # column major with row major layout: B^T @ A^T = C^T: [k, m] @ [n, k] = [n, m] |
| 1371 | if len(sB) == 2: |
| 1372 | if B.stride()[0] == B.shape[1]: |
| 1373 | transposed_B = False |
| 1374 | elif B.stride()[1] == B.shape[0]: |
| 1375 | transposed_B = True |
| 1376 | if len(A.shape) == 2: |
| 1377 | if A.stride()[0] == A.shape[1]: |
| 1378 | transposed_A = False |
| 1379 | elif A.stride()[1] == A.shape[0]: |
| 1380 | transposed_A = True |
| 1381 | else: |
| 1382 | if A.stride()[1] == A.shape[2]: |
| 1383 | transposed_A = False |
| 1384 | elif A.stride()[2] == A.shape[1]: |
| 1385 | transposed_A = True |
| 1386 | |
| 1387 | if len(sA) == 2: |
| 1388 | n = sA[0] |
| 1389 | ldb = A.stride()[1 if transposed_A else 0] |
| 1390 | elif len(sA) == 3 and len(sB) == 2: |
| 1391 | n = sA[0] * sA[1] |
| 1392 | ldb = sA[2] |
| 1393 | |
| 1394 | m = sB[1] |
| 1395 | k = sB[0] |
nothing calls this directly
no test coverage detected