char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
3e037f313
char/aliceai-t5-moe
1#pragma once 2// This file contains primitives that expose the tensor core PTX instructions for CUDA code. 3// The primitives can be used in a similar way as the nvcuda::wmma interface but with a well-defined memory layout. 4// The documentation for the PTX instructions can be found under: 5// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-multiply-accumulate-operation-using-mma-instruction 6// 7// Like with nvcuda::wmma there are three types of matrix tiles: A, B, and C with A @ B = C. 8// A is a row-major matrix with shape M x K. 9// B is a column-major matrix with shape K x N. 10// C is a column-major matrix with shape M x N. 11// A, B, and C are represented using the same fundamental data type: a row-major matrix with I rows and J columns. 12// Note that J is measured in physical 32 bit elements instead of logical elements. 13// The methods get_i and get_j can be used to get the physical 32 bit index of the lth element of a thread within a tile. 14// All matrix tiles have ne physical 32 bit elements per warp. 15// 16// As described in the PTX documentation, all pointers for load_ldmatrix must be to shared memory and aligned to 16 bytes. 17// The API in this file also assumes that the pointers for load_generic are aligned to 16 bytes, unaligned pointers are considered undefined behavior. 18 19#include "common.cuh" 20 21// On Volta each warp is doing 4 8x8 mma operations in parallel. 22// The basic memory layout for a 32x8 output tile is to stack 4 input tiles in I direction and to mirror the B tile. 23// However, the i indices in this file are by default permuted to simplify the index calculations. 24// #define GGML_CUDA_MMA_NO_VOLTA_PERM 25 26#if CUDART_VERSION >= 11080 27 28static __device__ __forceinline__ int ggml_cuda_movmatrix(const int x) { 29 int ret = 0; 30 31#ifdef TURING_MMA_AVAILABLE 32 asm("movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;" 33 : "=r"(ret) : "r"(x)); 34#else 35 GGML_UNUSED(x); 36 NO_DEVICE_CODE; 37#endif // defined(TURING_MMA_AVAILABLE) 38 return ret; 39} 40 41#else 42 43static __device__ __forceinline__ int ggml_cuda_movmatrix(const int x) { 44 // Imagine transposing row-major matrix to column-major matrix. 45 const int src_i_low = 2 * (threadIdx.x % 4); 46 const int src_i_high = src_i_low + 1; 47 const int src_j = threadIdx.x / 4; 48 49 const int src_laneid_low = src_i_low * 4 + src_j / 2; 50 const int src_laneid_high = src_i_high * 4 + src_j / 2; 51 52 const int shift_low = ((src_j + 0) % 2) * 16; 53 const int shift_high = ((src_j + 1) % 2) * 16; 54 55 const int ret_low = (__shfl_sync(0xFFFFFFFF, x, src_laneid_low, WARP_SIZE) >> shift_low) & 0x0000FFFF; 56 const int ret_high = (__shfl_sync(0xFFFFFFFF, x, src_laneid_high, WARP_SIZE) << shift_high) & 0xFFFF0000; 57 58 return ret_low | ret_high; 59} 60 61#endif // CUDART_VERSION >= 11080 62 63static __device__ __forceinline__ half2 ggml_cuda_movmatrix(const half2 x) { 64 half2 ret; 65 *((int *) &ret) = ggml_cuda_movmatrix(*((const int *) &x)); 66 return ret; 67} 68 69namespace ggml_cuda_mma { 70 71 // Some architectures like Volta or CDNA3 perform multiple matrix multiplications per warp in parallel, 72 // effectively the warp is being split into subgroups of threads that each perform a single mma instruction. 73 // In those cases the data can be split in different ways across the warp. 74 enum data_layout { 75 // By default the data uses the I direction as its major dimension and the J direction as its minor dimension. 76 // For the A/C matrices this means I major == row major, J major == column major. 77 // For the B matrix this means I major == column major, J major == row major. 78 // MIRRORED == Each data value is held exactly once per thread subgroup. 79 DATA_LAYOUT_I_MAJOR = 0, // Always used for Turing, Ampere, Ada Lovelace, consumer Blackwell, matrix A&B for RDNA4 and CDNA. 80 DATA_LAYOUT_J_MAJOR = 10, // Matrix C for CDNA and RDNA4, int and float matrix C for RDNA3. 81 DATA_LAYOUT_I_MAJOR_MIRRORED = 20, // Volta, matrix A&B for RDNA3. 82 DATA_LAYOUT_J_MAJOR_MIRRORED = 30, 83 DATA_LAYOUT_I_MAJOR_SCRAMBLED = 40, // Scrambled matrix C for faster transposition (RDNA4/CDNA), convert to float to unscramble. 84 }; 85 // Implemented mma combinations are: 86 // - (I_MAJOR, I_MAJOR) -> I_MAJOR 87 // - (I_MAJOR, I_MAJOR_MIRRORED) -> I_MAJOR 88 // - (I_MAJOR, J_MAJOR_MIRRORED) -> I_MAJOR 89 90 static constexpr __device__ data_layout get_input_data_layout() { 91#if defined(RDNA3) || defined(VOLTA_MMA_AVAILABLE) 92 return DATA_LAYOUT_I_MAJOR_MIRRORED; 93#else 94 return DATA_LAYOUT_I_MAJOR; 95#endif // defined(RDNA3) || defined(VOLTA_MMA_AVAILABLE) 96 } 97 98 template <int I_, int J_, typename T, data_layout ds_=DATA_LAYOUT_I_MAJOR> 99 struct tile {}; 100 101 template <int I_, int J_, typename T> 102 struct tile<I_, J_, T, DATA_LAYOUT_I_MAJOR> { 103 static constexpr int I = I_; 104 static constexpr int J = J_; 105 static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; 106 107#if defined(AMD_MFMA_AVAILABLE) 108 static constexpr int ne = I * J / 64; 109 T x[ne] = {0}; 110 111 static constexpr __device__ bool supported() { 112 if (I == 16 && J == 8) return true; 113 if (I == 32 && J == 4) return true; 114 if (I == 16 && J == 16) return true; 115 if (I == 32 && J == 32) return true; 116 return false; 117 } 118 119 static __device__ __forceinline__ int get_i(const int l) { 120 if constexpr (I == 16 && J == 4) { 121 return threadIdx.x % 16; 122 } else if constexpr (I == 16 && J == 8) { 123 return threadIdx.x % 16; 124 } else if constexpr (I == 32 && J == 4) { 125 return threadIdx.x % 32; 126 } else if constexpr (I == 16 && J == 16) { 127 return threadIdx.x % 16; 128 } else if constexpr (I == 32 && J == 32) { 129 return threadIdx.x % 32; 130 } else { 131 NO_DEVICE_CODE; 132 return -1; 133 } 134 } 135 136 static __device__ __forceinline__ int get_j(const int l) { 137 if constexpr (I == 16 && J == 4) { 138 return threadIdx.x / 16; 139 } else if constexpr (I == 16 && J == 8) { 140 return 2 * (threadIdx.x / 16) + l; 141 } else if constexpr (I == 32 && J == 4) { 142 return 2 * (threadIdx.x / 32) + l; 143 } else if constexpr (I == 16 && J == 16) { 144 return 4 * (threadIdx.x / 16) + l; 145 } else if constexpr (I == 32 && J == 32) { 146 return 4 * (threadIdx.x / 32) + 8 * (l / 4) + (l % 4); 147 } else { 148 NO_DEVICE_CODE; 149 return -1; 150 } 151 } 152#elif defined(VOLTA_MMA_AVAILABLE) 153 static constexpr int ne = I * J / 32; 154 T x[ne] = {0}; 155 156 static constexpr __device__ bool supported() { 157 if (I == 32 && J == 8) return true; 158 return false; 159 } 160 161 static __device__ __forceinline__ int get_i(const int l) { 162 if constexpr (I == 32 && J == 8) { 163#ifdef GGML_CUDA_MMA_NO_VOLTA_PERM 164 return (((threadIdx.x % 16) / 4) * 8) + ((threadIdx.x / 16) * 4) + (l & 2) + (threadIdx.x % 2); 165#else 166 return (l & 2) + (threadIdx.x & ~2); 167#endif // GGML_CUDA_MMA_NO_VOLTA_PERM 168 } else { 169 NO_DEVICE_CODE; 170 return -1; 171 } 172 } 173 174 static __device__ __forceinline__ int get_j(const int l) { 175 if constexpr (I == 32 && J == 8) { 176 return (threadIdx.x & 2) + (l & (4 + 1)); 177 } else { 178 NO_DEVICE_CODE; 179 return -1; 180 } 181 } 182#elif defined(AMD_WMMA_AVAILABLE) 183 static constexpr int ne = I * J / 32; 184 T x[ne] = {0}; 185 186 static constexpr __device__ bool supported() { 187 if (I == 16 && J == 16) return true; 188 if (I == 16 && J == 8) return true; 189 if (I == 16 && J == 4) return true; 190 return false; 191 } 192 193 static __device__ __forceinline__ int get_i(const int l) { 194 if constexpr (supported()) { 195 return threadIdx.x % 16; 196 } else { 197 NO_DEVICE_CODE; 198 return -1; 199 } 200 } 201 202 static __device__ __forceinline__ int get_j(const int l) { 203 if constexpr (I == 16 && J == 16) { 204#if defined(RDNA3) 205 if constexpr (std::is_same_v<T, float> || std::is_same_v<T, int>) { 206 // matrix C 207 return 2 * l + (threadIdx.x / 16); 208 } else { 209 // matrix A&B 210 return l; 211 } 212#else 213 // matrix C is the transposed matrix A&B on RDNA4 214 return ne * (threadIdx.x / 16) + l; 215#endif // defined(RDNA3) 216 } else if constexpr (I == 16 && J == 8) { 217 // mmq input for RDNA4 218 return ne * (threadIdx.x / 16) + l; 219 } else if constexpr (I == 16 && J == 4) { 220 return ne * (threadIdx.x / 16) + l; 221 } else { 222 NO_DEVICE_CODE; 223 return -1; 224 } 225 } 226#else 227 static constexpr int ne = I * J / 32; 228 T x[ne] = {0}; 229 230 static constexpr __device__ bool supported() { 231 if (I == 8 && J == 4) return true; 232 if (I == 8 && J == 8) return true; 233 if (I == 16 && J == 8) return true; 234 if (I == 16 && J == 16) return true; 235 if (I == 32 && J == 8) return true; 236 return false; 237 } 238 239 static __device__ __forceinline__ int get_i(const int l) { 240 if constexpr (I == 8 && J == 4) { 241 return threadIdx.x / 4; 242 } else if constexpr (I == 8 && J == 8) { 243 return threadIdx.x / 4; 244 } else if constexpr (I == 16 && J == 8) { 245 return ((l / 2) * 8) + (threadIdx.x / 4); 246 } else if constexpr (I == 16 && J == 16) { 247 return (((l / 2) % 2) * 8) + (threadIdx.x / 4); 248 } else if constexpr (I == 32 && J == 8) { 249 return tile<16, 8, T>::get_i(l); // Memory layout simply repeated with same pattern in i direction. 250 } else { 251 NO_DEVICE_CODE; 252 return -1; 253 } 254 } 255 256 static __device__ __forceinline__ int get_j(const int l) { 257 if constexpr (I == 8 && J == 4) { 258 return threadIdx.x % 4; 259 } else if constexpr (I == 8 && J == 8) { 260 return (l * 4) + (threadIdx.x % 4); 261 } else if constexpr (I == 16 && J == 8) { 262 return ((threadIdx.x % 4) * 2) + (l % 2); 263 } else if constexpr (I == 16 && J == 16) { 264 return ((l / 4) * 8) + ((threadIdx.x % 4) * 2) + (l % 2); 265 } else if constexpr (I == 32 && J == 8) { 266 return tile<16, 8, T>::get_j(l); // Memory layout simply repeated with same pattern in i direction. 267 } else { 268 NO_DEVICE_CODE; 269 return -1; 270 } 271 } 272#endif // defined(GGML_USE_HIP) 273 }; 274 275 template <int I_, int J_> 276 struct tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR> { 277 static constexpr int I = I_; 278 static constexpr int J = J_; 279 static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; 280 281#if defined(VOLTA_MMA_AVAILABLE) 282 static constexpr int ne = I * J / WARP_SIZE; 283 half2 x[ne] = {{0.0f, 0.0f}}; 284 285 static constexpr __device__ bool supported() { 286 if (I == 32 && J == 4) return true; 287 return false; 288 } 289 290 static __device__ __forceinline__ int get_i(const int l) { 291 if constexpr (I == 32 && J == 4) { 292#ifdef GGML_CUDA_MMA_NO_VOLTA_PERM 293 return (((threadIdx.x % 16) / 4) * 8) + ((threadIdx.x / 16) * 4) + (threadIdx.x % 4); 294#else 295 return threadIdx.x; 296#endif // GGML_CUDA_MMA_NO_VOLTA_PERM 297 } else { 298 NO_DEVICE_CODE; 299 return -1; 300 } 301 } 302 303 static __device__ __forceinline__ int get_j(const int l) { 304 if constexpr (I == 32 && J == 4) { 305 return l; 306 } else { 307 NO_DEVICE_CODE; 308 return -1; 309 } 310 } 311#elif defined(AMD_WMMA_AVAILABLE) 312 static constexpr int ne = I * J / 32; 313 half2 x[ne] = {{0.0f, 0.0f}}; 314 315 static constexpr __device__ bool supported() { 316 if (I == 16 && J == 8) return true; 317 if (I == 16 && J == 16) return true; 318 if (I == 32 && J == 8) return true; 319 return false; 320 } 321 322 static __device__ __forceinline__ int get_i(const int l) { 323 if constexpr (I == 16 && J == 8) { 324 return threadIdx.x % 16; 325 } else if constexpr (I == 16 && J == 16) { 326 return threadIdx.x % 16; 327 } else if constexpr (I == 32 && J == 8) { 328 return (threadIdx.x % 16) * 2 + l / (ne/2); 329 } else { 330 NO_DEVICE_CODE; 331 return -1; 332 } 333 } 334 335 static __device__ __forceinline__ int get_j(const int l) { 336 if constexpr (I == 16 && J == 8) { 337 return (threadIdx.x / 16) * ne + l; 338 } else if constexpr (I == 16 && J == 16) { 339#ifdef RDNA3 340 return l*2 + (threadIdx.x / 16); 341#else 342 return (threadIdx.x / 16) * ne + l; 343#endif // RDNA3 344 } else if constexpr (I == 32 && J == 8) { 345 return (threadIdx.x / 16) * (ne/2) + l % (ne/2); 346 } else { 347 NO_DEVICE_CODE; 348 return -1; 349 } 350 } 351#elif defined(AMD_MFMA_AVAILABLE) 352 static constexpr int ne = I * J / 64; 353 half2 x[ne] = {{0.0f, 0.0f}}; 354 355 static constexpr __device__ bool supported() { 356 if (I == 16 && J == 8) return true; 357 if (I == 16 && J == 16) return true; 358 if (I == 32 && J == 8) return true; 359 return false; 360 } 361 362 static __device__ __forceinline__ int get_i(const int l) { 363 if constexpr (I == 16 && J == 8) { 364 return threadIdx.x % 16; 365 } else if constexpr (I == 16 && J == 16) { 366 return threadIdx.x % 16; 367 } else if constexpr (I == 32 && J == 8) { 368 return (threadIdx.x % 16) * 2 + l / (ne/2); 369 } else { 370 NO_DEVICE_CODE; 371 return -1; 372 } 373 } 374 375 static __device__ __forceinline__ int get_j(const int l) { 376 if constexpr (I == 16 && J == 8) { 377 return (threadIdx.x / 16) * ne + l; 378 } else if constexpr (I == 16 && J == 16) { 379 return (threadIdx.x / 16) * ne + l; 380 } else if constexpr (I == 32 && J == 8) { 381 return (threadIdx.x / 16) * (ne/2) + l % (ne/2); 382 } else { 383 NO_DEVICE_CODE; 384 return -1; 385 } 386 } 387#else 388 static constexpr int ne = I * J / WARP_SIZE; 389 half2 x[ne] = {{0.0f, 0.0f}}; 390 391 static constexpr __device__ bool supported() { 392 if (I == 8 && J == 4) return true; 393 if (I == 8 && J == 8) return true; 394 if (I == 16 && J == 8) return true; 395 if (I == 16 && J == 16) return true; 396 if (I == 32 && J == 8) return true; 397 return false; 398 } 399 400 static __device__ __forceinline__ int get_i(const int l) { 401 if constexpr (I == 8 && J == 8) { 402 return threadIdx.x / 4; 403 } else if constexpr (I == 16 && J == 4) { 404 return (l * 8) + (threadIdx.x / 4); 405 } else if constexpr (I == 16 && J == 8) { 406 return ((l % 2) * 8) + (threadIdx.x / 4); 407 } else if constexpr (I == 32 && J == 8) { 408 return ((l / 4) * 16) + ((l % 2) * 8) + (threadIdx.x / 4); 409 } else { 410 NO_DEVICE_CODE; 411 return -1; 412 } 413 } 414 415 static __device__ __forceinline__ int get_j(const int l) { 416 if constexpr (I == 8 && J == 8) { 417 return (l * 4) + (threadIdx.x % 4); 418 } else if constexpr (I == 16 && J == 4) { 419 return threadIdx.x % 4; 420 } else if constexpr (I == 16 && J == 8) { 421 return ((l / 2) * 4) + (threadIdx.x % 4); 422 } else if constexpr (I == 32 && J == 8) { 423 return ((l & 2) * 2) + (threadIdx.x % 4); 424 } else { 425 NO_DEVICE_CODE; 426 return -1; 427 } 428 } 429#endif // defined(VOLTA_MMA_AVAILABLE) 430 }; 431 432 template <int I_, int J_> 433 struct tile<I_, J_, nv_bfloat162, DATA_LAYOUT_I_MAJOR> { 434 static constexpr int I = I_; 435 static constexpr int J = J_; 436 static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; 437 438#if defined(AMD_WMMA_AVAILABLE) 439 static constexpr int ne = tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::ne; 440 nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; 441 442 static constexpr __device__ bool supported() { 443 return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::supported(); 444 } 445 446 static __device__ __forceinline__ int get_i(const int l) { 447 return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_i(l); 448 } 449 450 static __device__ __forceinline__ int get_j(const int l) { 451 return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_j(l); 452 } 453#elif defined(AMD_MFMA_AVAILABLE) 454 static constexpr int ne = tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::ne; 455 nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; 456 457 static constexpr __device__ bool supported() { 458 return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::supported(); 459 } 460 461 static __device__ __forceinline__ int get_i(const int l) { 462 return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_i(l); 463 } 464 465 static __device__ __forceinline__ int get_j(const int l) { 466 return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_j(l); 467 } 468#else 469 static constexpr int ne = I * J / WARP_SIZE; 470 nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; 471 472 static constexpr __device__ bool supported() { 473 if (I == 8 && J == 8) return true; 474 if (I == 16 && J == 4) return true; 475 if (I == 16 && J == 8) return true; 476 return false; 477 } 478 479 static __device__ __forceinline__ int get_i(const int l) { 480 if constexpr (I == 8 && J == 8) { 481 return threadIdx.x / 4; 482 } else if constexpr (I == 16 && J == 4) { 483 return (l * 8) + (threadIdx.x / 4); 484 } else if constexpr (I == 16 && J == 8) { 485 return ((l % 2) * 8) + (threadIdx.x / 4); 486 } else { 487 NO_DEVICE_CODE; 488 return -1; 489 } 490 } 491 492 static __device__ __forceinline__ int get_j(const int l) { 493 if constexpr (I == 8 && J == 8) { 494 return (l * 4) + (threadIdx.x % 4); 495 } else if constexpr (I == 16 && J == 4) { 496 return threadIdx.x % 4; 497 } else if constexpr (I == 16 && J == 8) { 498 return ((l / 2) * 4) + (threadIdx.x % 4); 499 } else { 500 NO_DEVICE_CODE; 501 return -1; 502 } 503 } 504#endif // defined(AMD_WMMA_AVAILABLE) 505 }; 506 507 template <int I_, int J_, typename T> 508 struct tile<I_, J_, T, DATA_LAYOUT_J_MAJOR> { 509 static constexpr int I = I_; 510 static constexpr int J = J_; 511 static constexpr data_layout dl = DATA_LAYOUT_J_MAJOR; 512 513 static constexpr int ne = tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::ne; 514 T x[ne] = {0}; 515 516 static constexpr __device__ bool supported() { 517 return tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::supported(); 518 } 519 520 static __device__ __forceinline__ int get_i(const int l) { 521 return tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::get_j(l); 522 } 523 524 static __device__ __forceinline__ int get_j(const int l) { 525 return tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::get_i(l); 526 } 527 }; 528 529 template <int I_, int J_, typename T> 530 struct tile<I_, J_, T, DATA_LAYOUT_I_MAJOR_MIRRORED> { 531 static constexpr int I = I_; 532 static constexpr int J = J_; 533 static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; 534 535 // RDNA3 536 static constexpr int ne = I * J / 32 * 2; 537 538 T x[ne] = {0}; 539 540 static constexpr __device__ bool supported() { 541 if (I == 16 && J == 16) return true; 542 if (I == 16 && J == 8) return true; 543 if (I == 16 && J == 4) return true; 544 if (I == 32 && J == 8) return true; 545 return false; 546 } 547 548 static __device__ __forceinline__ int get_i(const int l) { 549 if constexpr (I == 16) { 550 return threadIdx.x % 16; 551 } else if constexpr (I == 32) { 552 return (threadIdx.x % 16) * 2 + l / (ne/2); 553 } else { 554 NO_DEVICE_CODE; 555 return -1; 556 } 557 } 558 559 static __device__ __forceinline__ int get_j(const int l) { 560 if constexpr (I == 16) { 561 return l; 562 } else if constexpr (I == 32) { 563 return l % (ne/2); 564 } else { 565 NO_DEVICE_CODE; 566 return -1; 567 } 568 } 569 }; 570 571 template <int I_, int J_> 572 struct tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> { 573 static constexpr int I = I_; 574 static constexpr int J = J_; 575 static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; 576#if defined(RDNA3) 577 static constexpr int ne = tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::ne; 578 579 half2 x[ne] = {{0.0f, 0.0f}}; 580 581 static constexpr __device__ bool supported() { 582 return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::supported(); 583 } 584 585 static __device__ __forceinline__ int get_i(const int l) { 586 return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_i(l); 587 } 588 589 static __device__ __forceinline__ int get_j(const int l) { 590 return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_j(l); 591 } 592#else // Volta 593 static constexpr int ne = I * J / (WARP_SIZE/4); 594 595 half2 x[ne] = {{0.0f, 0.0f}}; 596 597 static constexpr __device__ bool supported() { 598 if (I == 8 && J == 4) return true; 599 return false; 600 } 601 602 static __device__ __forceinline__ int get_i(const int /*l*/) { 603 if constexpr (I == 8 && J == 4) { 604 return ((threadIdx.x / 16) * 4) + (threadIdx.x % 4); 605 } else { 606 NO_DEVICE_CODE; 607 return -1; 608 } 609 } 610 611 static __device__ __forceinline__ int get_j(const int l) { 612 if constexpr (I == 8 && J == 4) { 613 return l; 614 } else { 615 NO_DEVICE_CODE; 616 return -1; 617 } 618 } 619#endif // defined(RDNA3) 620 }; 621 622 template <int I_, int J_> 623 struct tile<I_, J_, nv_bfloat162, DATA_LAYOUT_I_MAJOR_MIRRORED> { 624 static constexpr int I = I_; 625 static constexpr int J = J_; 626 static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; 627 static constexpr int ne = tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::ne; 628 629 nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; 630 631 static constexpr __device__ bool supported() { 632 return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::supported(); 633 } 634 635 static __device__ __forceinline__ int get_i(const int l) { 636 return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_i(l); 637 } 638 639 static __device__ __forceinline__ int get_j(const int l) { 640 return tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::get_j(l); 641 } 642 }; 643 644 template <int I_, int J_> 645 struct tile<I_, J_, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> { 646 static constexpr int I = I_; 647 static constexpr int J = J_; 648 static constexpr data_layout dl = DATA_LAYOUT_J_MAJOR_MIRRORED; 649 static constexpr int ne = I * J / (WARP_SIZE/4); 650 651 half2 x[ne] = {{0.0f, 0.0f}}; 652 653 static constexpr __device__ bool supported() { 654 if (I == 8 && J == 4) return true; 655 return false; 656 } 657 658 static __device__ __forceinline__ int get_i(const int l) { 659 if constexpr (I == 8 && J == 4) { 660 return ((l / 2) * 4) + (threadIdx.x % 4); 661 } else { 662 NO_DEVICE_CODE; 663 return -1; 664 } 665 } 666 667 static __device__ __forceinline__ int get_j(const int l) { 668 if constexpr (I == 8 && J == 4) { 669 return ((threadIdx.x / 16) * 2) + (l % 2); 670 } else { 671 NO_DEVICE_CODE; 672 return -1; 673 } 674 } 675 }; 676 677 template <int I_, int J_> 678 struct tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> { 679 static constexpr int I = I_; 680 static constexpr int J = J_; 681 static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_SCRAMBLED; 682 683 static constexpr int ne = I * J / ggml_cuda_get_physical_warp_size(); 684 half2 x[ne] = {{0.0f, 0.0f}}; 685 686 static constexpr __device__ bool supported() { 687 if (I == 16 && J == 16) return true; 688 return false; 689 } 690 691 static __device__ __forceinline__ int get_i(const int l) { 692 return tile<I_, J_, half2, DATA_LAYOUT_I_MAJOR>::get_i(l); 693 } 694 }; 695 696 static __device__ __forceinline__ tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> unscramble(const tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & t) { 697#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)) 698 tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> ret; 699#pragma unroll 700 for (int l0 = 0; l0 < t.ne/2; ++l0) { 701 ret.x[2*l0 + 0] = __lows2half2(t.x[l0], t.x[l0 + t.ne/2]); 702 ret.x[2*l0 + 1] = __highs2half2(t.x[l0], t.x[l0 + t.ne/2]); 703 } 704 return ret; 705#else 706 NO_DEVICE_CODE; 707 GGML_UNUSED(t); 708#endif // defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)) 709 } 710 711#if defined(TURING_MMA_AVAILABLE) 712 template <int I, int J> 713 static __device__ __forceinline__ tile<I, J/2, half2> get_half2(const tile<I, J, float> & tile_float) { 714 tile<I, J/2, half2> ret; 715#pragma unroll 716 for (int l0 = 0; l0 < tile_float.ne; l0 += 2) { 717 ret.x[l0/2] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); 718 } 719 return ret; 720 } 721 722 static __device__ __forceinline__ tile<8, 8, half2> get_transposed(const tile<16, 4, half2> & t) { 723 tile<8, 8, half2> ret; 724 ret.x[0] = ggml_cuda_movmatrix(t.x[0]); 725 ret.x[1] = ggml_cuda_movmatrix(t.x[1]); 726 727 return ret; 728 } 729#elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3) 730 static __device__ __forceinline__ tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> get_half2( 731 const tile<16, 16, float, DATA_LAYOUT_I_MAJOR> & tile_float) { 732 tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> ret; 733#pragma unroll 734 for (int l = 0; l < tile_float.ne; ++l) { 735 float tmp[2]; 736 int i = threadIdx.x / 16; 737 tmp[i] = tile_float.x[l]; 738 i ^= 1; 739 tmp[i] = __shfl_xor_sync(0xFFFFFFFF, tile_float.x[l], 16, WARP_SIZE); 740 ret.x[l] = make_half2(tmp[0], tmp[1]); 741 } 742 return ret; 743 } 744#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) 745 template <int I, int J> 746 static __device__ __forceinline__ tile<I, J/2, half2> get_half2(const tile<I, J, float> & tile_float) { 747 tile<I, J/2, half2> ret; 748#pragma unroll 749 for (int l0 = 0; l0 < tile_float.ne; l0 += 2) { 750 ret.x[l0/2] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); 751 } 752 return ret; 753 } 754 755 static __device__ __forceinline__ tile<8, 8, half2> get_transposed(const tile<16, 4, half2> & t) { 756 NO_DEVICE_CODE; 757 return tile<8, 8, half2>{}; 758 } 759#else // Volta 760 template <int I, int J> 761 static __device__ __forceinline__ tile<I, J/2, half2> get_half2(const tile<I, J, float> & tile_float) { 762 tile<I, J/2, half2> ret; 763#pragma unroll 764 for (int l0 = 0; l0 < tile_float.ne; l0 += 4) { 765 ret.x[l0/2 + 0] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); 766 ret.x[l0/2 + 1] = make_half2(tile_float.x[l0 + 2], tile_float.x[l0 + 3]); 767 768 // On Volta FP16 and FP32 tiles have a different memory layout, 769 // for the conversion threads with an offset of 2 need to exchange half their values: 770 ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)] = __shfl_xor_sync( 771 0xFFFFFFFF, ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)], 2, WARP_SIZE); 772 } 773 return ret; 774 } 775#endif // defined(TURING_MMA_AVAILABLE) 776 777 template <int I, int J, typename T, data_layout dl> 778 static __device__ __forceinline__ void load_generic(tile<I, J, T, dl> & t, const T * __restrict__ xs0, const int stride) { 779#pragma unroll 780 for (int l = 0; l < t.ne; ++l) { 781 t.x[l] = xs0[t.get_i(l)*stride + t.get_j(l)]; 782 } 783 } 784 785 template <typename T> 786 static __device__ __forceinline__ void load_ldmatrix( 787 tile<8, 8, T> & t, const T * __restrict__ xs0, const int stride) { 788#ifdef TURING_MMA_AVAILABLE 789 int * xi = (int *) t.x; 790 const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + ((threadIdx.x / t.I) * (t.J / 2)) % t.J; 791 asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];" 792 : "=r"(xi[0]), "=r"(xi[1]) 793 : "l"(xs)); 794#else 795 GGML_UNUSED_VARS(t, xs0, stride); 796 NO_DEVICE_CODE; 797#endif // TURING_MMA_AVAILABLE 798 } 799 800 template <typename T, data_layout dl> 801 static __device__ __forceinline__ void load_ldmatrix( 802 tile<16, 4, T, dl> & t, const T * __restrict__ xs0, const int stride) { 803#ifdef TURING_MMA_AVAILABLE 804 int * xi = (int *) t.x; 805 const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride; 806 asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];" 807 : "=r"(xi[0]), "=r"(xi[1]) 808 : "l"(xs)); 809#elif defined(AMD_WMMA_AVAILABLE) 810#ifdef RDNA3 811 static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout"); 812 static_assert(sizeof(t.x) == 16, "bad ne"); 813 ggml_cuda_memcpy_1<8>(t.x + 0, xs0 + t.get_i(0)*stride + 0); 814 ggml_cuda_memcpy_1<8>(t.x + 2, xs0 + t.get_i(0)*stride + 2); 815#else 816 static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); 817 static_assert(sizeof(t.x) == 8, "bad ne"); 818 ggml_cuda_memcpy_1<8>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0)); 819#endif // RDNA3 820#elif defined(AMD_MFMA_AVAILABLE) 821 static_assert(sizeof(t.x) == 4, "bad ne"); 822 ggml_cuda_memcpy_1<4>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0)); 823#else 824 GGML_UNUSED_VARS(t, xs0, stride); 825 NO_DEVICE_CODE; 826#endif // TURING_MMA_AVAILABLE 827 } 828 829 template <typename T, data_layout dl> 830 static __device__ __forceinline__ void load_ldmatrix( 831 tile<16, 8, T, dl> & t, const T * __restrict__ xs0, const int stride) { 832#if defined(TURING_MMA_AVAILABLE) 833 int * xi = (int * ) t.x; 834 const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2); 835 asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" 836 : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) 837 : "l"(xs)); 838#elif defined(VOLTA_MMA_AVAILABLE) 839 ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + 0, xs0 + t.get_i(0)*stride + 0); 840 ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + 4, xs0 + t.get_i(4)*stride + 4); 841#elif defined(AMD_WMMA_AVAILABLE) 842#ifdef RDNA3 843 static_assert(dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout"); 844 static_assert(sizeof(t.x) == 32, "bad ne"); 845 ggml_cuda_memcpy_1<16>(t.x + 0, xs0 + t.get_i(0)*stride + 0); 846 ggml_cuda_memcpy_1<16>(t.x + 4, xs0 + t.get_i(0)*stride + 4); 847#else 848 static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); 849 static_assert(sizeof(t.x) == 16, "bad ne"); 850 ggml_cuda_memcpy_1<16>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0)); 851#endif // RDNA3 852#elif defined(AMD_MFMA_AVAILABLE) 853 static_assert(sizeof(t.x) == 8, "bad ne"); 854 ggml_cuda_memcpy_1<8>(t.x, xs0 + t.get_i(0)*stride + t.get_j(0)); 855#else 856 GGML_UNUSED_VARS(t, xs0, stride); 857 NO_DEVICE_CODE; 858#endif // TURING_MMA_AVAILABLE 859 } 860 861 static __device__ __forceinline__ void load_ldmatrix( 862 tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) { 863 ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride); 864 } 865 866 static __device__ __forceinline__ void load_ldmatrix( 867 tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) { 868#pragma unroll 869 for (int l0 = 0; l0 < t.ne; l0 += 2) { 870 ggml_cuda_memcpy_1<2*sizeof(half2)>(t.x + l0, xs0 + t.get_i(l0)*stride + t.get_j(l0)); 871 } 872 } 873 874 static __device__ __forceinline__ void load_ldmatrix( 875 tile<32, 4, half2> & t, const half2 * __restrict__ xs0, const int stride) { 876#if defined(VOLTA_MMA_AVAILABLE) 877 ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride); 878#else 879 GGML_UNUSED_VARS(t, xs0, stride); 880 NO_DEVICE_CODE; 881#endif // defined(VOLTA_MMA_AVAILABLE) 882 } 883 884 template <int I, typename T, data_layout dl> 885 static __device__ __forceinline__ void load_ldmatrix_trans( 886 tile<I, 8, T, dl> & t, const T * __restrict__ xs0, const int stride) { 887#ifdef TURING_MMA_AVAILABLE 888 static_assert(I == 16, "bad tile width"); 889 static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout"); 890 int * xi = (int *) t.x; 891 const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2); 892 asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];" 893 : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3]) 894 : "l"(xs)); 895#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) 896 static_assert(dl == DATA_LAYOUT_I_MAJOR || dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout"); 897 if constexpr (I == 32) { 898#pragma unroll 899 for (int l0 = 0; l0 < t.ne/2; ++l0) { 900 const half2 tmp0 = xs0[(2*t.get_j(l0) + 0)*stride + t.get_i(l0)/2]; 901 const half2 tmp1 = xs0[(2*t.get_j(l0) + 1)*stride + t.get_i(l0)/2]; 902 903 t.x[l0] = __lows2half2(tmp0, tmp1); 904 t.x[l0 + t.ne/2] = __highs2half2(tmp0, tmp1); 905 } 906 } else { 907 half * xh = (half *) t.x; 908#pragma unroll 909 for (int l = 0; l < t.ne; ++l) { 910 xh[2*l + 0] = ((const half *) xs0)[(2*t.get_j(l) + 0)*(2*stride) + t.get_i(l)]; 911 xh[2*l + 1] = ((const half *) xs0)[(2*t.get_j(l) + 1)*(2*stride) + t.get_i(l)]; 912 } 913 } 914#else 915 GGML_UNUSED_VARS(t, xs0, stride); 916 NO_DEVICE_CODE; 917#endif // TURING_MMA_AVAILABLE 918 } 919 920 static __device__ __forceinline__ void mma( 921 tile<16, 8, int> & D, const tile<16, 4, int> & A, const tile<8, 4, int> & B) { 922#ifdef TURING_MMA_AVAILABLE 923#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 924 asm("mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" 925 : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3]) 926 : "r"(A.x[0]), "r"(A.x[1]), "r"(B.x[0])); 927#else 928 // On Turing m16n8k16 mma is not available, use 2x m8n8k16 mma instead: 929 asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" 930 : "+r"(D.x[0]), "+r"(D.x[1]) 931 : "r"(A.x[0]), "r"(B.x[0])); 932 asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" 933 : "+r"(D.x[2]), "+r"(D.x[3]) 934 : "r"(A.x[1]), "r"(B.x[0])); 935#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 936#else 937 GGML_UNUSED_VARS(D, A, B); 938 NO_DEVICE_CODE; 939#endif // TURING_MMA_AVAILABLE 940 } 941 942 static __device__ __forceinline__ void mma( 943 tile<16, 8, int> & D, const tile<16, 8, int> & A, const tile<8, 8, int> & B) { 944#ifdef TURING_MMA_AVAILABLE 945#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 946 asm("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" 947 : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3]) 948 : "r"(A.x[0]), "r"(A.x[1]), "r"(A.x[2]), "r"(A.x[3]), "r"(B.x[0]), "r"(B.x[1])); 949#else 950 // On Turing m16n8k32 mma is not available, use 4x m8n8k16 mma instead: 951 asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" 952 : "+r"(D.x[0]), "+r"(D.x[1]) 953 : "r"(A.x[0]), "r"(B.x[0])); 954 asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" 955 : "+r"(D.x[2]), "+r"(D.x[3]) 956 : "r"(A.x[1]), "r"(B.x[0])); 957 asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" 958 : "+r"(D.x[0]), "+r"(D.x[1]) 959 : "r"(A.x[2]), "r"(B.x[1])); 960 asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" 961 : "+r"(D.x[2]), "+r"(D.x[3]) 962 : "r"(A.x[3]), "r"(B.x[1])); 963#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 964#else 965 GGML_UNUSED_VARS(D, A, B); 966 NO_DEVICE_CODE; 967#endif // TURING_MMA_AVAILABLE 968 } 969 970 static __device__ __forceinline__ void mma( 971 tile<16, 4, half2> & D, const tile<16, 8, half2> & A, const tile<8, 8, half2> & B) { 972#ifdef TURING_MMA_AVAILABLE 973 const int * Axi = (const int *) A.x; 974 const int * Bxi = (const int *) B.x; 975 int * Dxi = (int *) D.x; 976#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 977 asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" 978 : "+r"(Dxi[0]), "+r"(Dxi[1]) 979 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); 980#else 981 // On Turing m16n8k16 mma is not available, use 2x m8n8k8 mma instead: 982 asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" 983 : "+r"(Dxi[0]), "+r"(Dxi[1]) 984 : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); 985 asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" 986 : "+r"(Dxi[0]), "+r"(Dxi[1]) 987 : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1])); 988#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 989#else 990 GGML_UNUSED_VARS(D, A, B); 991 NO_DEVICE_CODE; 992#endif // TURING_MMA_AVAILABLE 993 } 994 995 static __device__ __forceinline__ void mma( 996 tile<16, 8, half2> & D, const tile<16, 8, half2> & A, const tile<16, 8, half2> & B) { 997#ifdef TURING_MMA_AVAILABLE 998 const int * Axi = (const int *) A.x; 999 const int * Bxi = (const int *) B.x; 1000 int * Dxi = (int *) D.x; 1001#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 1002 asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" 1003 : "+r"(Dxi[0]), "+r"(Dxi[1]) 1004 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[2])); 1005 asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" 1006 : "+r"(Dxi[2]), "+r"(Dxi[3]) 1007 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]), "r"(Bxi[3])); 1008#else 1009 // On Turing m16n8k16 mma is not available, use 4x m8n8k8 mma instead: 1010 asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" 1011 : "+r"(Dxi[0]), "+r"(Dxi[1]) 1012 : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); 1013 asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" 1014 : "+r"(Dxi[0]), "+r"(Dxi[1]) 1015 : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2])); 1016 asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" 1017 : "+r"(Dxi[2]), "+r"(Dxi[3]) 1018 : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[1])); 1019 asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" 1020 : "+r"(Dxi[2]), "+r"(Dxi[3]) 1021 : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3])); 1022#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 1023#elif defined(AMD_WMMA_AVAILABLE) 1024#if defined(RDNA4) 1025 using halfx8_t = __attribute__((ext_vector_type(8))) _Float16; 1026 halfx8_t& acc_frag = reinterpret_cast<halfx8_t&>(D.x[0]); 1027 const halfx8_t& a_frag = reinterpret_cast<const halfx8_t&>(A.x[0]); 1028 const halfx8_t& b_frag = reinterpret_cast<const halfx8_t&>(B.x[0]); 1029 acc_frag = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag); 1030#else 1031 GGML_UNUSED_VARS(D, A, B); 1032 NO_DEVICE_CODE; 1033#endif // defined(RDNA4) 1034#elif defined(AMD_MFMA_AVAILABLE) 1035 // MFMA: FP16 input, FP32 accumulate, convert back to half2. 1036 using halfx4_t = __attribute__((ext_vector_type(4))) _Float16; 1037 using floatx4_t = __attribute__((ext_vector_type(4))) float; 1038 1039 // Convert existing half2 accumulator to float for MFMA: 1040 floatx4_t acc_f32; 1041 { 1042 const halfx4_t acc_h = reinterpret_cast<const halfx4_t&>(D.x[0]); 1043#pragma unroll 1044 for (int i = 0; i < 4; ++i) { 1045 acc_f32[i] = (float)acc_h[i]; 1046 } 1047 } 1048 1049 const halfx4_t& a_frag = reinterpret_cast<const halfx4_t&>(A.x[0]); 1050 const halfx4_t& b_frag = reinterpret_cast<const halfx4_t&>(B.x[0]); 1051 acc_f32 = __builtin_amdgcn_mfma_f32_16x16x16f16(a_frag, b_frag, acc_f32, 0, 0, 0); 1052 1053 // Convert back to half2: 1054 { 1055 halfx4_t result_h; 1056#pragma unroll 1057 for (int i = 0; i < 4; ++i) { 1058 result_h[i] = (_Float16)acc_f32[i]; 1059 } 1060 reinterpret_cast<halfx4_t&>(D.x[0]) = result_h; 1061 } 1062#else 1063 GGML_UNUSED_VARS(D, A, B); 1064 NO_DEVICE_CODE; 1065#endif // TURING_MMA_AVAILABLE 1066 } 1067 1068 static __device__ __forceinline__ void mma( 1069 tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & D, const tile<32, 8, half2, DATA_LAYOUT_I_MAJOR> & A, 1070 const tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> & B) { 1071#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)) 1072 tile<16, 8, half2> * D16 = (tile<16, 8, half2> *) &D; 1073 const tile<16, 8, half2> * A16 = (const tile<16, 8, half2> *) &A; 1074 mma(D16[0], A16[0], B); 1075 mma(D16[1], A16[1], B); 1076#else 1077 GGML_UNUSED_VARS(D, A, B); 1078 NO_DEVICE_CODE; 1079#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) && defined(RDNA4) 1080 } 1081 1082 template <data_layout dl_ab, data_layout dl_d> 1083 static __device__ __forceinline__ void mma( 1084 tile<16, 8, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<8, 8, float, dl_ab> & B) { 1085#ifdef AMPERE_MMA_AVAILABLE 1086 const int * Axi = (const int *) A.x; 1087 const int * Bxi = (const int *) B.x; 1088 int * Dxi = (int *) D.x; 1089 asm("mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" 1090 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1091 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); 1092#else 1093 GGML_UNUSED_VARS(D, A, B); 1094 NO_DEVICE_CODE; 1095#endif // AMPERE_MMA_AVAILABLE 1096 } 1097 1098 template <data_layout dl_ab, data_layout dl_d> 1099 static __device__ __forceinline__ void mma( 1100 tile<16, 16, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<16, 8, float, dl_ab> & B) { 1101#ifdef AMD_MFMA_AVAILABLE 1102 using floatx4_t = __attribute__((ext_vector_type(4))) float; 1103 floatx4_t& acc_frag = reinterpret_cast<floatx4_t&>(D.x[0]); 1104#if defined(CDNA3) 1105 using floatx2_t = __attribute__((ext_vector_type(2))) float; 1106 const floatx2_t& a_frag = reinterpret_cast<const floatx2_t&>(A.x[0]); 1107 const floatx2_t& b_frag = reinterpret_cast<const floatx2_t&>(B.x[0]); 1108 acc_frag = __builtin_amdgcn_mfma_f32_16x16x8_xf32(a_frag, b_frag, acc_frag, 0, 0, 0); 1109#elif defined(CDNA4) || defined(CDNA2) || defined(CDNA1) 1110 // CDNA4 (gfx950) does not support xf32 MFMA, use f32 path like CDNA2/CDNA1 1111#pragma unroll 1112 for (int i = 0; i < 2; ++i) { 1113 acc_frag = __builtin_amdgcn_mfma_f32_16x16x4f32(A.x[i], B.x[i], acc_frag, 0, 0, 0); 1114 } 1115#else 1116 GGML_UNUSED_VARS(D, A, B); 1117 NO_DEVICE_CODE; 1118#endif // defined(CDNA3) 1119#else 1120 GGML_UNUSED_VARS(D, A, B); 1121 NO_DEVICE_CODE; 1122#endif // AMD_MFMA_AVAILABLE 1123 } 1124 1125 template <ggml_type type> 1126 static __device__ __forceinline__ void mma_block_scaled_fp4(tile<16, 8, float> & D, 1127 const tile<16, 8, int> & A, 1128 const tile<8, 8, int> & B, 1129 uint32_t a_scale, 1130 uint32_t b_scale) { 1131#ifdef BLACKWELL_MMA_AVAILABLE 1132 const int * Axi = (const int *) A.x; 1133 const int * Bxi = (const int *) B.x; 1134 float * Dxi = (float *) D.x; 1135 1136 if constexpr (type == GGML_TYPE_MXFP4) { 1137 asm volatile( 1138 "mma.sync.aligned.kind::mxf4.block_scale.scale_vec::2X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue8m0 " 1139 "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3}, " 1140 "%10, {0, 0}, %11, {0, 0};" 1141 : "+f"(Dxi[0]), "+f"(Dxi[1]), "+f"(Dxi[2]), "+f"(Dxi[3]) 1142 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]), "r"(a_scale), "r"(b_scale)); 1143 } else { 1144 asm volatile( 1145 "mma.sync.aligned.kind::mxf4nvf4.block_scale.scale_vec::4X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue4m3 " 1146 "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3}, " 1147 "%10, {0, 0}, %11, {0, 0};" 1148 : "+f"(Dxi[0]), "+f"(Dxi[1]), "+f"(Dxi[2]), "+f"(Dxi[3]) 1149 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]), "r"(a_scale), "r"(b_scale)); 1150 } 1151#else 1152 GGML_UNUSED_VARS(D, A, B, a_scale, b_scale); 1153#endif // BLACKWELL_MMA_AVAILABLE 1154 } 1155 1156 static __device__ __forceinline__ void mma( 1157 tile<16, 8, float> & D, const tile<16, 8, half2> & A, const tile<8, 8, half2> & B) { 1158#ifdef TURING_MMA_AVAILABLE 1159 const int * Axi = (const int *) A.x; 1160 const int * Bxi = (const int *) B.x; 1161 int * Dxi = (int *) D.x; 1162#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 1163 asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" 1164 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1165 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); 1166#else 1167 // On Turing m16n8k16 mma is not available, use 2x m8n8k8 mma instead: 1168 asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" 1169 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1170 : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); 1171 asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" 1172 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1173 : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1])); 1174#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 1175#else 1176 GGML_UNUSED_VARS(D, A, B); 1177 NO_DEVICE_CODE; 1178#endif // TURING_MMA_AVAILABLE 1179 } 1180 1181 static __device__ __forceinline__ void mma( 1182 tile<16, 8, float> & D, const tile<16, 8, nv_bfloat162> & A, const tile<8, 8, nv_bfloat162> & B) { 1183#ifdef AMPERE_MMA_AVAILABLE 1184 const int * Axi = (const int *) A.x; 1185 const int * Bxi = (const int *) B.x; 1186 int * Dxi = (int *) D.x; 1187 asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" 1188 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1189 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); 1190#else 1191 GGML_UNUSED_VARS(D, A, B); 1192 NO_DEVICE_CODE; 1193#endif // AMPERE_MMA_AVAILABLE 1194 } 1195 1196 template <data_layout dl_ab, data_layout dl_d> 1197 static __device__ __forceinline__ void mma( 1198 tile<16, 16, float, dl_d> & D, const tile<16, 8, half2, dl_ab> & A, const tile<16, 8, half2, dl_ab> & B) { 1199#ifdef TURING_MMA_AVAILABLE 1200 const int * Axi = (const int *) A.x; 1201 const int * Bxi = (const int *) B.x; 1202 int * Dxi = (int *) D.x; 1203#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 1204 asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" 1205 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1206 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[2])); 1207 asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" 1208 : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) 1209 : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]), "r"(Bxi[3])); 1210#else 1211 // On Turing m16n8k16 mma is not available, use 4x m8n8k8 mma instead: 1212 asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" 1213 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1214 : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); 1215 asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" 1216 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1217 : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2])); 1218 asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" 1219 : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) 1220 : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[1])); 1221 asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" 1222 : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) 1223 : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3])); 1224#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE 1225#elif defined(AMD_WMMA_AVAILABLE) 1226#if defined(RDNA4) 1227 using halfx8_t = __attribute__((ext_vector_type(8))) _Float16; 1228 using floatx8_t = __attribute__((ext_vector_type(8))) float; 1229 floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]); 1230 const halfx8_t& a_frag = reinterpret_cast<const halfx8_t&>(A.x[0]); 1231 const halfx8_t& b_frag = reinterpret_cast<const halfx8_t&>(B.x[0]); 1232 acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag); 1233#elif defined(RDNA3) 1234 using halfx16_t = __attribute__((ext_vector_type(16))) _Float16; 1235 using floatx8_t = __attribute__((ext_vector_type(8))) float; 1236 floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]); 1237 const halfx16_t& a_frag = reinterpret_cast<const halfx16_t&>(A.x[0]); 1238 const halfx16_t& b_frag = reinterpret_cast<const halfx16_t&>(B.x[0]); 1239 acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a_frag, b_frag, acc_frag); 1240#else 1241 GGML_UNUSED_VARS(D, A, B); 1242 NO_DEVICE_CODE; 1243#endif // RDNA4 1244#elif defined(AMD_MFMA_AVAILABLE) 1245 using halfx4_t = __attribute__((ext_vector_type(4))) _Float16; 1246 using floatx4_t = __attribute__((ext_vector_type(4))) float; 1247 floatx4_t& acc_frag = reinterpret_cast<floatx4_t&>(D.x[0]); 1248 const halfx4_t& a_frag = reinterpret_cast<const halfx4_t&>(A.x[0]); 1249 const halfx4_t& b_frag = reinterpret_cast<const halfx4_t&>(B.x[0]); 1250 acc_frag = __builtin_amdgcn_mfma_f32_16x16x16f16(a_frag, b_frag, acc_frag, 0, 0, 0); 1251#else 1252 GGML_UNUSED_VARS(D, A, B); 1253 NO_DEVICE_CODE; 1254#endif // TURING_MMA_AVAILABLE 1255 } 1256 1257 template <data_layout dl_ab, data_layout dl_d> 1258 static __device__ __forceinline__ void mma( 1259 tile<16, 16, float, dl_d> & D, const tile<16, 8, nv_bfloat162, dl_ab> & A, const tile<16, 8, nv_bfloat162, dl_ab> & B) { 1260#if defined(AMD_WMMA_AVAILABLE) 1261#if defined(RDNA4) 1262 using bf16x8_t = __attribute__((ext_vector_type(8))) __bf16; 1263 using floatx8_t = __attribute__((ext_vector_type(8))) float; 1264 floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]); 1265 const bf16x8_t& a_frag = reinterpret_cast<const bf16x8_t&>(A.x[0]); 1266 const bf16x8_t& b_frag = reinterpret_cast<const bf16x8_t&>(B.x[0]); 1267 acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(a_frag, b_frag, acc_frag); 1268#elif defined(RDNA3) 1269 using bf16x16_t = __attribute__((ext_vector_type(16))) __bf16; 1270 using floatx8_t = __attribute__((ext_vector_type(8))) float; 1271 floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]); 1272 const bf16x16_t& a_frag = reinterpret_cast<const bf16x16_t&>(A.x[0]); 1273 const bf16x16_t& b_frag = reinterpret_cast<const bf16x16_t&>(B.x[0]); 1274 acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a_frag, b_frag, acc_frag); 1275#else 1276 GGML_UNUSED_VARS(D, A, B); 1277 NO_DEVICE_CODE; 1278#endif // defined(RDNA4) 1279#elif defined(AMD_MFMA_AVAILABLE) 1280 using floatx4_t = __attribute__((ext_vector_type(4))) float; 1281 floatx4_t& acc_frag = reinterpret_cast<floatx4_t&>(D.x[0]); 1282#if defined(CDNA4) || defined(CDNA3) || defined(CDNA2) 1283 using bf16x4_t = __attribute__((ext_vector_type(4))) __bf16; 1284 const bf16x4_t& a_frag = reinterpret_cast<const bf16x4_t&>(A.x[0]); 1285 const bf16x4_t& b_frag = reinterpret_cast<const bf16x4_t&>(B.x[0]); 1286 acc_frag = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a_frag, b_frag, acc_frag, 0, 0, 0); 1287#elif defined(CDNA1) 1288#pragma unroll 1289 for (int i = 0; i < 2; ++i) { 1290 using bf16x2_t = __attribute__((ext_vector_type(2))) __bf16; 1291 const bf16x2_t& a_frag = reinterpret_cast<const bf16x2_t&>(A.x[i]); 1292 const bf16x2_t& b_frag = reinterpret_cast<const bf16x2_t&>(B.x[i]); 1293 acc_frag = __builtin_amdgcn_mfma_f32_16x16x8bf16(a_frag, b_frag, acc_frag, 0, 0, 0); 1294 } 1295#else 1296 GGML_UNUSED_VARS(D, A, B); 1297 NO_DEVICE_CODE; 1298#endif // defined(CDNA3) || defined(CDNA2) 1299#else 1300 GGML_UNUSED_VARS(D, A, B); 1301 NO_DEVICE_CODE; 1302#endif // defined(AMD_WMMA_AVAILABLE) 1303 } 1304 1305 template <data_layout dl_d, data_layout dl_ab> 1306 static __device__ __forceinline__ void mma( 1307 tile<16, 16, int, dl_d> & D, const tile<16, 8, int, dl_ab> & A, const tile<16, 8, int, dl_ab> & B) { 1308#if defined(AMD_MFMA_AVAILABLE) 1309 using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; 1310 int32x4_t * acc = (int32x4_t *) D.x; 1311#if defined(CDNA4) || defined(CDNA3) 1312 acc[0] = __builtin_amdgcn_mfma_i32_16x16x32_i8(((int64_t *) A.x)[0], ((int64_t *) B.x)[0], acc[0], 0, 0, 0); 1313#elif defined(CDNA2) || defined(CDNA1) 1314 acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[0], B.x[0], acc[0], 0, 0, 0); 1315 acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[1], B.x[1], acc[0], 0, 0, 0); 1316#endif // defined(CDNA4) || defined(CDNA3) 1317#elif defined(AMD_WMMA_AVAILABLE) 1318 using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int; 1319 int32x8_t * acc = (int32x8_t *) D.x; 1320#if defined(RDNA4) 1321 using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int; 1322 int32x2_t * a_vec = (int32x2_t *) A.x; 1323 int32x2_t * b_vec = (int32x2_t *) B.x; 1324 acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[0], true, b_vec[0], acc[0], true); 1325 acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[1], true, b_vec[1], acc[0], true); 1326#elif defined(RDNA3) 1327 using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; 1328 int32x4_t * a_vec = (int32x4_t *) A.x; 1329 int32x4_t * b_vec = (int32x4_t *) B.x; 1330 acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], true); 1331 acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[1], true, b_vec[1], acc[0], true); 1332#endif // RDNA4 1333#else 1334 GGML_UNUSED_VARS(D, A, B); 1335 NO_DEVICE_CODE; 1336#endif // AMD_MFMA_AVAILABLE 1337 } 1338 1339 static __device__ __forceinline__ void mma( 1340 tile<32, 32, int> & D, const tile<32, 4, int> & A, const tile<32, 4, int> & B) { 1341#if defined(AMD_MFMA_AVAILABLE) 1342 using int32x16_t = __attribute__((__vector_size__(16 * sizeof(int)))) int; 1343 int32x16_t * acc = (int32x16_t *) D.x; 1344#if defined(CDNA4) || defined(CDNA3) 1345 acc[0] = __builtin_amdgcn_mfma_i32_32x32x16_i8(((int64_t *) A.x)[0], ((int64_t *) B.x)[0], acc[0], 0, 0, 0); 1346#elif defined(CDNA2) || defined(CDNA1) 1347 acc[0] = __builtin_amdgcn_mfma_i32_32x32x8i8(A.x[0], B.x[0], acc[0], 0, 0, 0); 1348 acc[0] = __builtin_amdgcn_mfma_i32_32x32x8i8(A.x[1], B.x[1], acc[0], 0, 0, 0); 1349#endif // defined(CDNA4) || defined(CDNA3) 1350 1351#else 1352 GGML_UNUSED_VARS(D, A, B); 1353 NO_DEVICE_CODE; 1354#endif // AMD_MFMA_AVAILABLE 1355 } 1356 1357 template <typename T1, typename T2, int J, int K> 1358 static __device__ __forceinline__ void mma( 1359 tile<32, J, T1> & D, const tile<32, K, T2> & A, const tile<J, K, T2> & B) { 1360 tile <16, J, T1> * D16 = reinterpret_cast< tile<16, J, T1> *>(&D); 1361 const tile<16, K, T2> * A16 = reinterpret_cast<const tile<16, K, T2> *>(&A); 1362 mma(D16[0], A16[0], B); 1363 mma(D16[1], A16[1], B); 1364 } 1365 1366 static __device__ __forceinline__ void mma( 1367 tile<32, 8, float> & D, const tile<32, 4, half2> & A, const tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) { 1368#if defined(VOLTA_MMA_AVAILABLE) 1369 const int * Axi = (const int *) A.x; 1370 const int * Bxi = (const int *) B.x; 1371 int * Dxi = (int *) D.x; 1372 asm("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " 1373 "{%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9}, {%10, %11}, {%0, %1, %2, %3, %4, %5, %6, %7};" 1374 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]), "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) 1375 : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]), "r"(Bxi[1])); 1376 asm("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " 1377 "{%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9}, {%10, %11}, {%0, %1, %2, %3, %4, %5, %6, %7};" 1378 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]), "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) 1379 : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]), "r"(Bxi[3])); 1380#else 1381 GGML_UNUSED_VARS(D, A, B); 1382 NO_DEVICE_CODE; 1383#endif // defined(VOLTA_MMA_AVAILABLE) 1384 } 1385 1386 static __device__ __forceinline__ void mma( 1387 tile<32, 4, half2> & D, const tile<32, 4, half2> & A, const tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & B) { 1388#if defined(VOLTA_MMA_AVAILABLE) 1389 const int * Axi = (const int *) A.x; 1390 const int * Bxi = (const int *) B.x; 1391 int * Dxi = (int *) D.x; 1392 asm("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 " 1393 "{%0, %1, %2, %3}, {%4, %5}, {%6, %7}, {%0, %1, %2, %3};" 1394 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1395 : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]), "r"(Bxi[1])); 1396 asm("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 " 1397 "{%0, %1, %2, %3}, {%4, %5}, {%6, %7}, {%0, %1, %2, %3};" 1398 : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) 1399 : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]), "r"(Bxi[3])); 1400#else 1401 GGML_UNUSED_VARS(D, A, B); 1402 NO_DEVICE_CODE; 1403#endif // defined(VOLTA_MMA_AVAILABLE) 1404 } 1405 1406 static __device__ __forceinline__ void mma( 1407 tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> & D, const tile<32, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & A, 1408 const tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) { 1409#if defined(AMD_WMMA_AVAILABLE) && defined(RDNA3) 1410 using halfx16_t = __attribute__((ext_vector_type(16))) _Float16; 1411 halfx16_t * xD = (halfx16_t *) D.x; 1412 const halfx16_t * xA = (const halfx16_t *) A.x; 1413 const halfx16_t * xB = (const halfx16_t *) B.x; 1414 xD[0] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(xA[0], xB[0], xD[0], /*opsel =*/ 0); 1415 xD[0] = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(xA[1], xB[0], xD[0], /*opsel =*/ 1); 1416#else 1417 GGML_UNUSED_VARS(D, A, B); 1418 NO_DEVICE_CODE; 1419#endif // TURING_MMA_AVAILABLE 1420 } 1421 1422 template <data_layout dl_d, data_layout dl_ab> 1423 static __device__ __forceinline__ void mma( 1424 tile<16, 16, int, dl_d> & D, const tile<16, 4, int, dl_ab> & A, const tile<16, 4, int, dl_ab> & B) { 1425#if defined(AMD_MFMA_AVAILABLE) 1426 using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; 1427 int32x4_t * acc = (int32x4_t *) D.x; 1428#if defined(CDNA4) || defined(CDNA3) 1429 const int64_t xA = uint32_t(A.x[0]); 1430 const int64_t xB = uint32_t(B.x[0]); 1431 acc[0] = __builtin_amdgcn_mfma_i32_16x16x32_i8(xA, xB, acc[0], 0, 0, 0); 1432#elif defined(CDNA2) || defined(CDNA1) 1433 acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[0], B.x[0], acc[0], 0, 0, 0); 1434#endif // defined(CDNA4) || defined(CDNA3) 1435#elif defined(AMD_WMMA_AVAILABLE) 1436 using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int; 1437 int32x8_t * acc = (int32x8_t *) D.x; 1438#if defined(RDNA4) 1439 using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int; 1440 int32x2_t * a_vec = (int32x2_t *) A.x; 1441 int32x2_t * b_vec = (int32x2_t *) B.x; 1442 acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[0], true, b_vec[0], acc[0], false); 1443#elif defined(RDNA3) 1444 using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; 1445 int32x4_t * a_vec = (int32x4_t *) A.x; 1446 int32x4_t * b_vec = (int32x4_t *) B.x; 1447 acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], false); 1448#endif // RDNA4 1449#else 1450 GGML_UNUSED(D); 1451 GGML_UNUSED(A); 1452 GGML_UNUSED(B); 1453 NO_DEVICE_CODE; 1454#endif // AMD_WMMA_AVAILABLE 1455 } 1456}