tensorcore support - #9283
Draft
abadams wants to merge 76 commits into
Draft
Conversation
The highest we could target was 8.6, so on anything newer we emitted sm_86 PTX and left it to the driver to JIT. Now Ada, Hopper and both Blackwells can be named, and host target detection picks them up. It makes no measurable difference to the tensor core matmul, which is what prompted this: the kernel only uses instructions that have been available since 7.0, so there's nothing for a newer target to do. cuBLAS on the same device runs an sm_80 kernel for the same reason. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The comment above halide_target_feature_t lists three places to keep in sync when adding a feature, and PyEnums.cpp was the one I missed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The cuda runtime capped kernels at 64 registers per thread when loading a module, with an HL_CUDA_JIT_MAX_REGISTERS environment variable as an escape hatch. Capping registers trades spilling against occupancy, and ptxas has more information about the kernel and the device than we do. Two apps set the escape hatch to 256 to compensate. Both are removed here; on an sm_86 device the 256 cap is now slightly slower than letting the driver decide. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
get_subtile partitions accesses to a tile-memory allocation between the distinct sub-tiles it holds, and is_load_of_multiramp digs a load out from under the casts, broadcasts and lane permutations that can wrap it. Neither is specific to AMX, so move them next to the MultiRamp machinery they are built on. This drops 126 lines from ExtractTileOperations. Supporting this, is_multiramp learns to see through a shuffle of a single vector when the shuffle is a reshaping rather than a gather: either a transpose, or any permutation of a one-dimensional multiramp whose lane indices are themselves a multiramp. MultiRamp::transpose applies the former to a multiramp directly, splitting a dim in two where the transpose falls inside one. Routing the AMX operands through is_load_of_multiramp lets that pass match loads it previously missed, and lets it check the operand types against the type actually loaded from memory rather than the type of the expression wrapping it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
is_load_of_multiramp stripped any number of casts off a load and returned only the Load underneath, so a caller had no way to see that the values had been cast. ExtractTileOperations then typed its tile registers from the load, and the signedness of those types picks which of the four integer tdpb instructions runs. A pipeline multiplying uint8 buffers reinterpreted as int8 compiled to tdpbuud - an unsigned multiply for a signed algorithm. Peel at most one cast, so the element type of the original Expr and the type of the returned Load together say whether the values were cast and to what, and have the AMX integer path reject a mismatch. The float path still reads the load's type, which is what it wants: the bf16 to f32 widening cast is expected there, and tile_load takes the bf16. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A transpose's shuffle mask is itself a multiramp of constants - transposing n lanes with c columns gives a mask of shape (n/c, c) and strides (c, 1) - so the general reshaping-shuffle case in is_multiramp already covers it. Drop the separate is_transpose case and MultiRamp::transpose, which existed only to serve it and had no other caller. That leaves one rule for shuffles of a single vector: it is a reshaping rather than a gather if the mask is a multiramp of constants, and we can say what it reshapes to when the shuffled vector is one-dimensional. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Dividing a vector by a broadcast only folds when the two are vectorized the same way, which they aren't when the numerator is a nested vector. Add rules for the two shapes that come up: a broadcast numerator, where the division can be pushed inwards until the two line up, and a ramp whose lanes are each repeated by an inner broadcast, where repeating a lane doesn't change the set of values so the existing first-and-last-lane argument still applies. While here, replace the can_prove predicate on the rule being generalized with a structural one. Matching the base as a multiple of the denominator is enough to know the quotient is uniform, given the ramp doesn't span far enough to reach the next multiple. The structural form also generalizes it: the stride may be any non-negative constant, the base's multiplier need only be a multiple of the denominator rather than equal to it, and the base may be affine rather than linear. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A store into shared memory whose value is a plain load from global memory is emitted as an asynchronous copy on sm_80 and later, which moves the data without routing it through registers. The copies issued in a producer are waited for at the end of it. This takes the shared-memory tensor core matmul from 31.4 to 39.5 TFlop/s at 2048^3 on an RTX 5060 Ti. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
If nothing between a barrier and the next one at the same level loads what was stored before it, the later barrier can do the job of both. The fence types of the elided barrier are added to the one that remains. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A Func stored in GPUSharedAsync goes in GPU shared memory, but is written by an asynchronous copy instruction that moves the data straight from global memory without routing it through registers. Asking for the memory type is a promise that every store to it is a copy the hardware can make that way, so a store that isn't reports what the copy engine requires and how a schedule usually satisfies it, rather than quietly falling back to a load and a store. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
AMXTile was never added, and GPUSharedAsync is new. List them in the same order as the enum in Expr.h so that a missing one is easier to spot. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Cover each of the three copy widths the hardware supports at several element sizes, plus the shapes a staged input takes: a two-dimensional tile, a padded stride, two inputs staged into one kernel, and a Func::in wrapper. The error test checks that breaking each constraint produces a user error saying which one, rather than crashing or quietly falling back to a load and a store. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The peephole fired on any store into shared memory that matched, so a Func in plain GPUShared got an asynchronous copy too, and there was no way to ask for the synchronous version. That made the memory type only control whether a store that didn't match was an error, rather than whether the copy was asynchronous at all. Require the destination to be GPUSharedAsync, so that GPUShared and GPUSharedAsync are the two ways to ask for the two lowerings, and a schedule can compare them. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A store whose value wasn't a plain load was always reported as the Func not being a copy, but that also caught two cases where it is a copy and something else is wrong. A source read with a stride is broken into a shuffle of dense loads before it reaches here, and a source computed elsewhere in the kernel is a load from the wrong place. Report those as themselves. Also test the predicated case, which needs an align_storage to keep the destination aligned, or the alignment requirement is what fails first. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The density of the two ends of the copy is checked by strided_ramp_base, whose default stride of one is doing the work. That is easy to misread as extracting an address, so say so where it is called. A source read with a stride never reaches that check, because it is broken into a shuffle of dense loads first and fails the earlier test for being a plain load. Storing the staged Func in the opposite order to the one it is read in is what reaches it, so test that. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This is the GPU counterpart to the existing AMXTile support. Scheduling a matmul accumulator with .store_in(MemoryType::WMMAAccumulator) makes a vanilla Halide matrix multiply compile to wmma instructions. The new pass, extract_wmma_operations, recognizes the three operations that a tensor core accumulator supports - initialization from zero or from a matrix in memory, accumulation of a matrix multiply, and copying the result out - and rewrites them as intrinsics. Anything else is an error, because the layout of a fragment across the registers of a warp isn't architecturally specified, so those instructions are the only way in or out. Nothing in the schedule says the accumulator is spread over a warp, so the pass also introduces the loop over the 32 lanes. The fragment intrinsics are pure functions of a matrix value and the lane. The matrix is named by a Load of the whole of it, with the lanes in row-major order regardless of how it's laid out in memory, so that the passes that track uses of an allocation see the access and its true footprint. Its layout in memory and the distance between its rows or columns are recovered from the strides of that Load's index. Copying an accumulator out is a predicated store of the whole matrix, where the predicate says which entries this lane holds. Supported: float16 operands, float32 or float16 accumulators, all three tile shapes, row- and column-major operands and results, several accumulator fragments per warp, several warps per block, and operands staged through shared memory. Along the way: - is_multiramp learns to see through the lane permutations the simplifier and flatten_nested_ramps introduce, so those passes are free to rewrite the accesses and the backend puts them back. - The simplifier lifts broadcasts out of pure elementwise calls, so a lane-uniform value stays recognizable as one. - The subtile partitioning that AMX was doing moves to MultiRamp, shared between the two passes. - The CUDA runtime no longer caps registers per thread at 64. That was costing the tensor core matmul about 1.7x, and cost the two apps that set HL_CUDA_JIT_MAX_REGISTERS to work around it. apps/tensorcore_matmul reaches 40 TFlop/s on an RTX 5060 Ti, against 8 for the best non-tensor-core schedule and 50 for cuBLAS. apps/tensorcore_resize ports the block-sparse resampling algorithm, in which a resize becomes a dense matrix multiply. Its cudaonly schedule works; the tensorcore one is blocked on a simplifier gap documented in its README. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
An accumulator scheduled at block level, with the reduction loop above the loop over warps, is the shape a matmul needs in order to stage its operand panels into shared memory once per block and have every warp reuse them. Three things were in the way: - The index of an access to the accumulator then depends on which warp is doing it. That dependence selects between the per-thread copies of the allocation rather than between subtiles within one, so it's substituted away before working out which subtile an access refers to. - Halide sees an accumulator outside the thread loops as shared between threads, so it keeps the atomic node around the update. Codegen then scalarizes the store, one lane at a time, and the wmma intrinsics fall apart. An accumulator is per-thread register storage, so there's nothing to race with, and the pass drops the atomic. - The loop over warps was being given the innermost thread dimension, which the loop over lanes needs. Inside an accumulator allocation, a thread loop is one dimension further out than it looks. Also tune apps/tensorcore_matmul, which now unrolls a couple of reduction steps to get more operand loads in flight. Together with the register cap removal this takes it from 20.7 to 26.1 TFlop/s at 1024, 32.6 to 38.8 at 2048, and 40.2 to 39.6 at 4096, against cuBLAS at 43.3, 49.6 and 50.0. Staging the operands through shared memory is now expressible, and there's a test for it, but it isn't yet a win: it peaks at 37.4 TFlop/s against 41.2 for loading the operands straight from global memory and letting L2 do the reuse. The barrier between staging and computing serializes the two, and recovering that needs double buffering, which Halide expresses with async and ring_buffer. Those lower to Fork nodes and semaphores, which have no meaning in device code, so that's a separate piece of work. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…e app The tile shape and the set of sub-tiles were tracked across the whole pass rather than per allocation, so a pipeline with two accumulators of different shapes was rejected. apps/tensorcore_resize has one of each, and now works: on an RTX 5060 Ti it downsamples a 3840x2160 image by 4x in 0.246 ms against 0.359 ms for the cuda-only schedule. Both of that app's schedules round the output up to a multiple of the tile size, so the runner gives them an output buffer of that size. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The memory type now covers all three matrices of a multiply, with the role inferred from use: an allocation accumulated into by a matrix multiply is the accumulator, and one read as an operand is that operand. An operand staged this way is loaded into fragment registers where the schedule says to compute it, and reused by every multiply that reads it, which is a hoist that nothing downstream can do when the loop isn't unrolled. The pass now tracks a scope of fragments across two passes over the whole statement rather than one allocation at a time, so fragments can nest and several multiplies can be in flight at once. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two things were costing the tensor core matmul instructions it didn't need to spend. Each producer waited for its own asynchronous copies, so a block staging two operand panels paid the latency of the first before it had even issued the second. Wait instead at the barrier or the load that needs the data, so every copy a block issues can be in flight at once. Fragments also round-tripped through their allocations as Halide vectors. NVPTX holds a wide float vector in pairs of registers, so every tensor core instruction unpacked its accumulator into eight registers and packed the result back up again - about 117 instructions per reduction step per warp, which was most of the instruction stream. Read and write fragments a register at a time instead, so they never become vectors. At 2048^3 on an RTX 5060 Ti this takes the best tensor core matmul from 40.0 to 46.0 TFlop/s. cuBLAS is 49.6. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Without the annotation it has to assume the largest block the hardware supports and allocate registers for that. Measured no difference on the tensor core matmul, whose register count is pinned by its live accumulators, but it is information the backend should be passing on. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
With fragments no longer packed into and out of vectors, a wider tile per warp pays off. At 1024^3 / 2048^3 / 4096^3 on an RTX 5060 Ti the tensor core schedule goes from 27.1 / 39.3 / 40.3 to 31.3 / 45.2 / 46.2 TFlop/s, against cuBLAS at 43.3 / 49.6 / 50.0. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Of the strides that keep the rows aligned enough for an asynchronous copy, the one in use was among the worse ones: it left 2.2M shared load bank conflicts against 23K for the next multiple of eight up. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The tensor cores multiply brain floats into single precision and eight-bit integers into 32-bit ones, as well as halves. The pieces that were specific to halves: Fragment sizes were two constants. Only the accumulator one is universal - an accumulator tile holds eight elements per lane whatever its type and shape. Operand fragments hold exactly their share of the matrix, except half precision ones, which are a fixed sixteen elements per lane whatever the shape, with the hardware replicating across lanes when the shape holds fewer. So the operand size is now derived from the shape and type rather than assumed. The element type in the intrinsic names was guessed from the bit width, which would have silently emitted a half precision instruction for brain floats. It is now an explicit table that errors on anything unsupported, and a second table gives the operand and accumulator combinations that have an instruction. The intrinsics take a fragment as 32-bit registers. Halide's half vectors match that signature directly, but brain floats and bytes have to be bitcast to i32. A fragment that is a single register comes back as that register rather than as a struct holding one of them. wmma_matmul covers the new types across every shape, layout, and staging option it already covered halves for. Its operands are now small integers, which every type here represents exactly and whose dot products stay exact in the accumulator, so it checks for equality rather than to a tolerance. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The shapes were tuned for half precision operands accumulated into single precision, and the other types were borrowing them. Sweeping sixty shapes at each size and operand type gives a shape per pair, and the operand type turns out to matter as much as the size. Bytes want a tall block spread over four warps, where the 16-bit types want a wide one over fewer: byte operands make the loads cheap enough to pay for a much larger accumulator. That is worth 45% at 1024 and 43% at 4096 over the shapes they were borrowing. Brain floats pick out exactly the same shapes as halves at every size, so they share a row. Also builds the half-into-half variant, which the app could not previously ask for. Its operands are sparse zeros and ones so that the dot products stay under 2048, the largest integer half precision represents exactly, which keeps the check exact. The comment now records all seven measured configurations at all three sizes rather than a couple of ratios. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
cublas does brain floats, eight-bit integers, and a half accumulator too, so the table can compare like with like at every pair of types rather than leaving the new ones without a reference. Brain floats land where halves do, within a few percent of cublas, and the half accumulator is ahead of it at 2048. Eight bit is the weak one: 55% of cublas at 1024 and 62% at 4096. cublas is not using the wmma instructions there - eight-bit operands have an mma shape with twice the reduction depth per instruction, which this schedule cannot reach. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The staging vector was eight elements, which is the widest asynchronous copy for a 16-bit operand but half of it for a byte. Making it sixteen bytes rather than eight elements is the same element-versus-byte confusion the padding had. With that fixed, and with the staging depth swept as well - bytes fit twice the reduction depth in the same shared memory, and mostly want to - the eight-bit variant goes from 61 to 63 TFlop/s at 1024 and from 81 to 90 at 4096. That leaves it at 58% to 69% of cublas rather than 55% to 62%. The gap is still real, but it was not the wmma instructions, as the comment claimed before this: it was a schedule tuned for 16-bit operands. The table is now in GFlop/s with a peak column, which shows how far each row is from what the hardware could do rather than only from cublas. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The previous comment measured the schedules against a peak computed from the maximum clock and an assumed rate per SM per clock. Both inputs were wrong: the part averages 2817 MHz while benchmarking rather than its 3090 maximum, and wmma multiplies bytes at the same rate it multiplies halves into halves, not twice it. Issuing back-to-back wmma instructions out of registers with no memory traffic gives the real ceilings: 51541, 99626 and 100650 GOP/s. Against those the schedules reach 95%, 87% and 89%, and the half-into-half one matches cublas exactly. That is a much better account of them than 39% of a peak they could never reach. It also settles the eight-bit gap, which an earlier commit message got wrong in both directions. cublas is 29% past the wmma ceiling, so it cannot be using wmma. The mma instructions reach 188355 GOP/s at the same shape, 1.87x. Both do the same 8192 ops per instruction - the earlier claim that the deeper reduction was more work per instruction was wrong - so what differs is purely how fast the two families issue. The gap is structural after all, but not for the reason first given. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
It reaches 28% of what the cuda cores can do where the tensor core schedules reach 87% to 95% of their instructions, and the difference is one missing scheduling primitive rather than an old schedule. Sharing a staged panel across a block needs the reduction chunk loop above the thread loops and the accumulator below them, but an accumulator that outlives the chunk loop lands at block level, where a Register allocation is sized for the whole block tile and spills. The tensor core schedules only get around it because a WMMAFragment allocation at block level is already per-lane. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Four had gone stale as the generator grew: the header still said only half precision reached the tensor cores, tiles_x said the best block gets smaller as the matrices grow, the block shape comment said bytes want four warps when they want two, and the tensor core schedule quoted ratios against cublas from before the retune. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Main factored the test for which allocations stay in registers into allocation_goes_to_registers, so the tensor core fragment case moves into that helper rather than the caller.
Running the filter five times inside the lambda and synchronizing once at the end amortizes the per-launch synchronization just as well as the hand-rolled sampling loop did, and leaves the sampling to the helper. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The operands are small integers so the results can be checked exactly, which also makes them very compressible - three quarters of them are zero. Measuring against dense random operands instead moves nothing by more than a couple of percent, on either side of the comparison. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The two sides had their own allocations, so they were multiplying different data - Halide's small integers against cublas's memset pattern. Fill a Halide buffer, let running the filter copy it down, and pass halide_cuda_get_device_ptr of it to cublas, so both multiply the same numbers. That drops the separate cudaMalloc and memset entirely. Doing that made the two implementations pair up naturally, so each pair of types is now one function that runs the filter, checks it, times it, and times cublas beside it. The runner is 64 lines shorter. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The runner was carrying its own element get and set for buffers of a type with no C++ equivalent, which is really just bfloat16. Halide already has usable float16_t and bfloat16_t, and Float16.h already declares halide_type_of for both, so the buffers can be plain typed ones and the dispatch goes away. The cost is linking libHalide into a program that otherwise only needs the generated code, which a TODO records. Float16.h depends on nothing but HalideRuntime.h and its implementation touches nothing in the compiler, so making those definitions inline and shipping the header - or moving the types to the runtime - would remove the need. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Its beta was one, so it read the output and added to it, which the filters here do not do. Zero makes the two do the same work, and is worth up to ten percent to cublas at the widths where the output is large relative to the operands. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ives MemoryType::AMXTile and MemoryType::WMMAFragment both describe storage for a matrix tile in whatever form a target's matrix unit keeps one, so merge them into MemoryType::Tile. AMXTile remains as a deprecated alias. WMMAFragment was never released, so it is simply removed. Add tile_init, tile_load, tile_store and tile_matmul to Stage and Func. They are sugar over the existing directives: each reorders the dimensions that make up a tile to be innermost and vectorizes them, tile_matmul additionally marking the stage atomic so that a reduction dimension can be vectorized. Everything but tile_store also sets the memory type, since those are the ones that produce a tile. The order of a tile's dimensions is fixed by the instruction rather than free to schedule, so imposing it here lets schedules drop it: reorder only permutes the dimensions it names among the positions they already hold, so a file-local helper names every dimension to move the tile ones inwards. The same four directives now describe both targets' matrix units, lowering to tilezero/tileloadd/tdpbf16ps/tilestored on x86 and to the corresponding wmma operations on CUDA. Rename test/correctness/tiled_matmul.cpp to amx_matmul.cpp to say which matrix unit it covers, and register wmma_matmul.cpp in the correctness CMakeLists, where it was missing. In apps/cuda_mat_mul, express the paired splits as tile calls, drop the reorder arguments the directives now supply, and give the tensor core dimensions the names mmx, mmy and mmr so that Vars are no longer spelled like RVars. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Its schedule needs adjusting for the tile directives, so it will return in a later PR. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Own header first, then internal headers, then system headers, with a blank line between each group. clang-format sorts within blank-line-separated blocks but does not reorder the blocks, so these had survived reformatting. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…dams/wmma_intrinsics
wmma_matmul's two consumer schedules and most of tiled_matmul_errors were still spelled out as reorder, atomic and vectorize. In tiled_matmul_errors that includes the schedule_matmul helper, which is the valid scaffolding nine scenarios share, so it should read the way such a schedule is meant to be written. scenario_no_matmul keeps the desugared spelling. It puts a value that is not a matrix multiply into tile memory on purpose, which is not something the directives can express. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
# Conflicts: # src/FuseGPUThreadLoops.cpp
…dams/wmma_intrinsics
The value of each of these depends on which lane of the warp it runs in, but nothing in their arguments did, so the dependence was invisible in the IR. wmma_lane_owns was the worst case: its arguments were three constants, so only its impurity stopped it being lifted out of the loop over lanes, and it would have been reasonable to make it pure. The lane goes last, so the argument indices the backend uses are unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
abadams
marked this pull request as draft
August 5, 2026 17:41
# Conflicts: # test/correctness/CMakeLists.txt
Load::make lost the default for is_streaming when the short forms were added, so the call that left it off no longer compiles. The short form is what it wanted anyway: an unpredicated load from an internal buffer. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit 3fe9e54)
Both lost the default for is_streaming when the short forms were added, so the calls that left it off no longer compile. Most of these are unpredicated accesses to internal buffers, which is what the short form means. The load of a matrix a tensor core instruction takes a fragment from has a real image and parameter, so it keeps the long form and passes is_streaming explicitly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…adams/wmma_intrinsics
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds support for targeting tensor cores via wmma instructions, unifying that path with our amx support. It adds three new intrinsic families for wmma. They are all pure functions of values that do not touch memory:
wmma_matrix_to_fragment_*: Take a large vector that represents an entire matrix and extract just the lanes that a particular warp lane owns in nvidia's opaque representation. This is effectively a shuffle node, but by an unknown mask.
wmma_fragment_to_matrix_d: Take an wmma fragment owned by a given lane and expand it back into a full matrix, with only the entries known to this lane specified. This is not very useful by itself, but it pairs with the next intrinsic:
wmma_lane_owns: A matrix-sized boolean vector that returns whether or not a given gpu lane knows that matrix entry. Storing a value resulting from wmma_fragment_to_matrix_d with a predicate wmma_lane_owns is an accurate representation of a wmma tile store instructions and is pattern-matched to the same.
The intrinsics also take a bunch of matrix-tile-shape args.
The way these are reached from the front-end is the same as amx support: matrix ops are represented as nested vectorization. However, this PR adds some syntactic sugar for it:
With these changes you can get pretty close to cublas performance with reasonable-looking schedules. See the changes to apps/cuda_mat_mul.
Opening as a draft, to get CI coverage while still getting underlying functionality merged. Some of the changes in this branch are already in separate PRs or are going to be carved off into separate PRs.