Repository navigation
Enable use of Tensor cores on cuda arch #7015
Description
Please make it possible to take advantage of Nvidia tensor cores on the cuda arch so that machine learning projects
can be integrated into taichi
it will also serve the benefit of speeding up other types of computing
as you are well aware this is a huge advantage cuda has over taichi
20x performance is nothing to scoff at when it comes to Matrix mul
Regards
Love , keep up the good work
Metadata
Metadata
Assignees
Type
Projects
- Status
Activity
bobcao3 commentedon Dec 30, 2022
Do you want the low-level access to cooperative-matrix instructions (https://developer.nvidia.com/blog/machine-learning-acceleration-vulkan-cooperative-matrices/ , and
HMMAinstructions in PTX), or do you want large matrix multiplication in the form of cuBLAS or even CUDNN integration? I doubt we will incorporate CUDNN directly because that seems like a much better fit for PyTorch. We might add the lower level instructions support because they fit better into the current programming modelzubairdroid commentedon Dec 31, 2022
Anything at this point
Just access to the the hardware will be enough
its a shame really that taichi doesnt support it
So much could be done if it did
especially machine learning
bobcao3 commentedon Dec 31, 2022
Taichi does support zero copy interop with pytotch, does that work for your case?
zubairdroid commentedon Jan 1, 2023
No need more granular control
plus it would be better as a whole if we needed to depend less on outside packages
given the nature of taichi
zubairdroid commentedon Jan 2, 2023
cuBLAS would be a better fit
So whats the plan
ailzhang commentedon Jan 28, 2023
@zubairdroid Please correct me if I'm wrong, are you requesting a cublas alternative (that uses tensor cores) written in taichi instead of interoping with cublas?
zubairdroid commentedon Feb 8, 2023
@ailzhang
Anything that allows fp16 taichi fields to use Nvidia tensor cores for multiply-accumulate while still in taichi scope
(and preferably other data types)
zubairdroid commentedon May 30, 2023
any updates
calling torch matmul from the python scope slows downthings considerably
not to mention it limits the granularity taichi has
currently its one of the main factors for people to still rely on C libraries
tf32 and fp16 speedup is nothing to scoff at
any ideas or plans for adding this feature in the future
louismhowe commentedon Jul 2, 2023
Hi, just wanted to express my interest in this feature, implemented as @bobcao3 suggests, with low-level access to warp/wave/simdgroup MMA.
My potential use case would not be for large dense MMA ops, but rather for custom kernels that fit the Taichi design philosophy, i.e., with custom data structures. Even though the widespread use of hardware MMA is merely to accelerate large dense MMA ops by reusing intermediate products in shared/LDS/threadgroup memory, this is not the only way to make good use of these instructions. The challenge is ensuring the kernel is not limited by global memory bandwidth. For example, one might use them in local iterative algorithms to free up ALU/RF resources for subsequent processing (e.g., nonlinear shaping/activation functions).
As a new Taichi user, I am not sure what the most idiomatic way or ergonomic way to do this would be - whether through
ti.simt.warp.mmaintrinsics, or automatically viati.Matrixoperations, or by some other means.Some other related ideas that may not fit with the current Taichi direction ...
mma.spfor a 2x speedup.Since I am only a new user, I do not yet feel confident in attempting to add this to Taichi myself, but I may attempt it in the future, so any suggestions regarding implementation or the best API would be welcome.
Nate711 commentedon Jan 15, 2025
Are there any updates on this? The latest GPUs A100, H100, H200 etc have relatively low fp32 flops but very high bfloat16 fp16 flops since they use tensor cores. My feeling is taichi will lose its advantage over jax and other frameworks that utilize tensor cores efficiently. Is support for tensor cores already implemented? If not, is there a timeline?
zubairdroid commentedon Apr 15, 2026
As far as i can tell
The devs have passionately ignored this feature
It would be great to be able to test genetic algorithms without throwing performance down the drain or relying on external packages