Add TT, TN, NT, NN tests for HopperMultipleMatmulScheduler #3310
Conversation
rdspring1
changed the title
|
I created TN test with |
rdspring1
force-pushed
the
hopper_matmul_tests
branch
from
d8bc1a6 to
7c8f375
Compare
rdspring1
changed the title
| auto tv0 = makeContigConcreteTensor({-1, -1, 1}, dtype); // A [M, K, b] | ||
| auto tv1 = makeContigConcreteTensor({1, -1, -1}, dtype); // B [b, K, N] |
There was a problem hiding this comment.
In this case the MmaOp input order is MKN, and the output gets reordered with a root->logical reordering to MNK. In the TN case there's no such reordering because the logical order of inputs is MNK. Note that we can also have allocation domains set on the inputs. Maybe we could parametrize all the combinations i.e. the orders of the allocation and logical domains of the inputs?
There was a problem hiding this comment.
Maybe we could parametrize all the combinations i.e. the orders of the allocation and logical domains of the inputs?
When the allocation and logical domains are different, would the input operand be not contiguous, concrete tensors?
There was a problem hiding this comment.
They can be contiguous, concrete, and have permuted allocation domain. Contiguity is with respect to allocation domain, so e.g. a tensor of logical shape [5, 7] and stride [7, 1] is contiguous, but so is one with logical shape [5, 7] and stride [1, 5]. The latter would correspond to having a swapped allocation domain in nvFuser.
|
!test |
This PR modifies `schedulePrologues` to use TMA loads to move mma operands to shared memory. Stacked on #3324 and #3310. ## Details 1. Input operands are loaded into shared memory via `CpAsyncBulkTensorTile` LoadStoreOp. 2. Replace `LdMatrix` operation with basic set. 3. Modified `scheduleOperandSmemStores` to apply swizzling to avoid bank conflicts. 4. Refactor `swizzleSharedMemory` by moving the analysis component to a separate function named `analyzeSwizzleSharedMemory`. 5. Create `tmaSwizzleSharedMemory` function that uses `analyzeSwizzleSharedMemory` and then finds the appropriate tma swizzle format. 6. Disable loop rotation. There is an issue with tma loads and circular buffering. Not sure if loop rotation is required for hopper matmul. 7. Expect hopper matmul tests to give incorrect results.
This PR creates four tests for the
HopperMultiMatmulScheduler. Each tests covers a different matmul layout - TT, TN, NT, and NN where the input arguments are already broadcasted.