[MPS] Extend addmm to integral types #160270
Conversation
Fixes #154901 [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160270
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 16 PendingAs of commit cdfaf5a with merge base 86eb65f ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| } | ||
| }); | ||
| return output; | ||
| return output; |
There was a problem hiding this comment.
Double return? Surprised this didn't error
There was a problem hiding this comment.
Yeah, I hoped some of the linters will be triggered by it, but feels like this is fine...
| constant uint3& sizes [[buffer(6)]], | ||
| uint2 tid [[thread_position_in_threadgroup]], | ||
| uint2 thread_id [[thread_position_in_grid]]) { | ||
| threadgroup T A_tile[TILE_DIM][TILE_DIM]; |
There was a problem hiding this comment.
It's ugly, but can this be rewritten as an std array too?
There was a problem hiding this comment.
Threadgroups are a bit weird(i.e. this statement affects GPU occupancy), let me give it a try in a separate PR, but make sure it would not regress the perf...
Fixes #154901 [ghstack-poisoned]
Stack from ghstack (oldest at bottom):
By adding
addmmkernel, which is a logical continuation ofmmone. The only tricking part are how alpha and beta constants are handled, which are passed asoptmath_t, i.e. that it could be, int64, int32 or floatUnified all MM flavors instantiations thru
INSTANTIATE_MM_OPSand tested thataddmmmetal kernel works as expected for floating types as well by testing it viaFixes #154901