lightning-thunder
lightning-thunder copied to clipboard
[DTensor] Refactor to re-use prims
Investigate: Try to decrease duplication of prims - reuse the existing prims one by one to support propagating DTensor layout and providing the executors information about the supported types (Tensor, DTensor) which the executor opt in or out of. (See comment)
Note that you're converting a 2xN problem into a Nx2 problem by pushing dtensor support into prims. In the olden days, we decomposed ops to factor out common patterns, maybe this could still be a route.