[dead]
Highly efficient matrix transpose in Mojo
31–40 of 70 posts
Re: Highly efficient matrix transpose in Mojo
#32Why do we ever need to transpose a matrix? Isn't it better to simply combine the transposition with whatever next operation one wishes to do with the matrix?
This could make Mojo look even better as it would ld be more compute heavy and the last step thread reduction would be less relevant.
Re: Highly efficient matrix transpose in Mojo
#33I'm not an expert in this space, but is this meaningful? I'd assume that it's more common to fuse together transposition with an operation that precedes or follows it (e.g. matmul), which should be far more efficient than materializing the entire transposition in memory if it's just an intermediate value.
Matrix transpose is a canonical example of a memory bound operation and often used to showcase optimization in a particular programming language or library. See for example the cutlass matrix transpose tutorial from Jay Shah of flash attention 3 paper: https://research.colfax-intl.com/tutorial-matrix-transpose-i...
Re: Highly efficient matrix transpose in Mojo
#34I wonder if there is a reason for not using the high level abstractions provided by Modular
Re: Highly efficient matrix transpose in Mojo
#35Does anyone use Mojo in production at all or are even hiring for Mojo?
Modular (the company behind Mojo) uses it in production. I imagine that if they have any clients then those also use Mojo in production - albeit indirectly - since all the GPU kernels used by Modular are written in Mojo.
Re: Highly efficient matrix transpose in Mojo
#36As someone said below - you'd never write just a transpose kernel - it'll be fused into something else.
Look the frontier AI companies need something other than reversing binary trees to give interview candidates
Re: Highly efficient matrix transpose in Mojo
#37I'm not an expert in this space, but is this meaningful? I'd assume that it's more common to fuse together transposition with an operation that precedes or follows it (e.g. matmul), which should be far more efficient than materializing the entire transposition in memory if it's just an intermediate value.
Matrix transpose is a canonical example of a memory bound operation and often used to showcase optimization in a particular programming language or library. See for example the cutlass matrix transpose tutorial from Jay Shah of flash attention 3 paper: https://research.colfax-intl.com/tutorial-matrix-transpose-i...
Unfortunately the issue (alluded to in the blog post you linked) is that transposes do absolutely no work but memory loads. Sure, they test that you can swizzle your accesses, but modern accelerators are all about pipelining and feeding matrix multiply units, which is considerably harder than loading from memory as fast as possible. Actually, even the Mojo post barely beats CUDA for most of its kernels, because you can hit memory bandwidth for transpose on the latest hardware using techniques from 5-10 years ago. This is definitely not true for more interesting operations.
Re: Highly efficient matrix transpose in Mojo
#38Earlier quoted context omitted.
The next operation might need the data in column major order to read it fast. So you might have to transpose first. And these maybe be concurrent stages of a processing pipeline.
Now I'm curious, how many times do you have to fully read the matrix in GPU for the total impact of reading columns to be higher than one-off actual transpose and then sequential row reads? I know it depends on lots of things, I'm after a rough estimate.
It's quite rare. Usually problems are tiled anyway and you can amortize the cost of having data in the "wrong" layout by loading coalesced in whatever is the best layout for your data and then transposing inside your tile, which gives you access to much faster memory.
Re: Highly efficient matrix transpose in Mojo
#39I wonder if there is a reason for not using the high level abstractions provided by Modular
Most interesting algorithms (e.g. with dynamic shapes, mixed computation) are typically better scheduled by hand.
Re: Highly efficient matrix transpose in Mojo
#40> This kernel archives a bandwidth of 1056.08 GB/s which is faster than the 875.46 GB/s we archived using CUDA. I believe that to be the reason because we use the PTX api for TMA transfers in Mojo.
I can't say for sure because I couldn't find the CUDA kernel but I kind of doubt this is true. You can hit memory bandwidth on Hopper without using TMA at all, which is mostly designed for accelerating asynchronous copies and reducing memory pressure. If all you are doing is a transpose you don't need any of this to go fast (though it might simplify your indexing code…?)