NFC: Extract launch dimensions computation from TritonWrapper.
This is another prefactoring to make Triton fusions compatible with FusionInterface and partially fused HLO. For the former, we need the LaunchDimension computation to be a separate function. For the latter, we change the launch dimension function signatures to no longer take an HloComputation, because we don't yet have that during fusion (at least not a complete one). For now, this change is a no-op, since we do not yet have any boundary functions for non-fusion ops. PiperOrigin-RevId: 565015567
Showing
想要评论请 注册 或 登录