Imported from GitHub PR openxla/xla#13813

gpu_windowed_einsum_handler pass has been re-using the empty buffer of the transformed while loop. This buffer is given by the spmd dot_handler pass. The shape of the buffer has changed from the allgathered shape of the sharded operand to the output shape of the dot which leads to a shape incompatibility error. To make the gpu handler completely safe, we will make a new element in the tuple to host the cached activation with the desired shape.
The slice index of where to write the slice into the full buffer also changes based on whether it's contracting or non-contracting dim is sharded. With the new element, we will need to determine the slice index ourselves in the handler pass.
Copybara import of the project:

--
ceeff8e5da8ecb3f382bbd8dee83e2f0c909b22d by TJ Xu <tjx@nvidia.com>:

Assign a fixed index for cached activation
Cache correct activation slice when contracting dim is sharded

--
233763b8efb4ab0045eb998b437c7b28c8f776c8 by TJ Xu <tjx@nvidia.com>:

Simplified logic in gpu einsum handler to be more generic

Merging this change closes #13813

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13813 from Tixxx:tixxx/ag_multi_fix 233763b8efb4ab0045eb998b437c7b28c8f776c8
PiperOrigin-RevId: 646461734

Imported from GitHub PR openxla/xla#13813

gpu_windowed_einsum_handler pass has been re-using the empty buffer of the transformed while loop. This buffer is given by the spmd dot_handler pass. The shape of the buffer has changed from the allgathered shape of the sharded operand to the output shape of the dot which leads to a shape incompatibility error. To make the gpu handler completely safe, we will make a new element in the tuple to host the cached activation with the desired shape.
The slice index of where to write the slice into the full buffer also changes based on whether it's contracting or non-contracting dim is sharded. With the new element, we will need to determine the slice index ourselves in the handler pass.
Copybara import of the project:

--
ceeff8e5da8ecb3f382bbd8dee83e2f0c909b22d by TJ Xu <tjx@nvidia.com>:

Assign a fixed index for cached activation
Cache correct activation slice when contracting dim is sharded

--
233763b8efb4ab0045eb998b437c7b28c8f776c8 by TJ Xu <tjx@nvidia.com>:

Simplified logic in gpu einsum handler to be more generic

Merging this change closes #13813

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13813 from Tixxx:tixxx/ag_multi_fix 233763b8efb4ab0045eb998b437c7b28c8f776c8
PiperOrigin-RevId: 646461734

Imported from GitHub PR openxla/xla#13813

gpu_windowed_einsum_handler pass has been re-using the empty buffer of the transformed while loop. This buffer is given by the spmd dot_handler pass. The shape of the buffer has changed from the allgathered shape of the sharded operand to the output shape of the dot which leads to a shape incompatibility error. To make the gpu handler completely safe, we will make a new element in the tuple to host the cached activation with the desired shape.
The slice index of where to write the slice into the full buffer also changes based on whether it's contracting or non-contracting dim is sharded. With the new element, we will need to determine the slice index ourselves in the handler pass.
Copybara import of the project:

--
ceeff8e5da8ecb3f382bbd8dee83e2f0c909b22d by TJ Xu <tjx@nvidia.com>:

Assign a fixed index for cached activation
Cache correct activation slice when contracting dim is sharded

--
233763b8efb4ab0045eb998b437c7b28c8f776c8 by TJ Xu <tjx@nvidia.com>:

Simplified logic in gpu einsum handler to be more generic

--
2220cd1a022ad519cd23ab36c31c70c9627fc76d by TJ Xu <tjx@nvidia.com>:

remove un-used variables

Merging this change closes #13813

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13813 from Tixxx:tixxx/ag_multi_fix 2220cd1a022ad519cd23ab36c31c70c9627fc76d
PiperOrigin-RevId: 646461734

Imported from GitHub PR openxla/xla#13813

gpu_windowed_einsum_handler pass has been re-using the empty buffer of the transformed while loop. This buffer is given by the spmd dot_handler pass. The shape of the buffer has changed from the allgathered shape of the sharded operand to the output shape of the dot which leads to a shape incompatibility error. To make the gpu handler completely safe, we will make a new element in the tuple to host the cached activation with the desired shape.
The slice index of where to write the slice into the full buffer also changes based on whether it's contracting or non-contracting dim is sharded. With the new element, we will need to determine the slice index ourselves in the handler pass.
Copybara import of the project:

--
ceeff8e5da8ecb3f382bbd8dee83e2f0c909b22d by TJ Xu <tjx@nvidia.com>:

Assign a fixed index for cached activation
Cache correct activation slice when contracting dim is sharded

--
233763b8efb4ab0045eb998b437c7b28c8f776c8 by TJ Xu <tjx@nvidia.com>:

Simplified logic in gpu einsum handler to be more generic

--
2220cd1a022ad519cd23ab36c31c70c9627fc76d by TJ Xu <tjx@nvidia.com>:

remove un-used variables

Merging this change closes #13813

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13813 from Tixxx:tixxx/ag_multi_fix 2220cd1a022ad519cd23ab36c31c70c9627fc76d
PiperOrigin-RevId: 646461734

Imported from GitHub PR openxla/xla#13813

gpu_windowed_einsum_handler pass has been re-using the empty buffer of the transformed while loop. This buffer is given by the spmd dot_handler pass. The shape of the buffer has changed from the allgathered shape of the sharded operand to the output shape of the dot which leads to a shape incompatibility error. To make the gpu handler completely safe, we will make a new element in the tuple to host the cached activation with the desired shape.
The slice index of where to write the slice into the full buffer also changes based on whether it's contracting or non-contracting dim is sharded. With the new element, we will need to determine the slice index ourselves in the handler pass.
Copybara import of the project:

--
ceeff8e5da8ecb3f382bbd8dee83e2f0c909b22d by TJ Xu <tjx@nvidia.com>:

Assign a fixed index for cached activation
Cache correct activation slice when contracting dim is sharded

--
233763b8efb4ab0045eb998b437c7b28c8f776c8 by TJ Xu <tjx@nvidia.com>:

Simplified logic in gpu einsum handler to be more generic

--
2220cd1a022ad519cd23ab36c31c70c9627fc76d by TJ Xu <tjx@nvidia.com>:

remove un-used variables

Merging this change closes #13813

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#13813 from Tixxx:tixxx/ag_multi_fix 2220cd1a022ad519cd23ab36c31c70c9627fc76d
PiperOrigin-RevId: 646461734

Imported from GitHub PR openxla/xla#13813

gpu_windowed_einsum_handler pass has been re-using the empty buffer of the transformed while loop. This buffer is given by the spmd dot_handler pass. The shape of the buffer has changed from the allgathered shape of the sharded operand to the output shape of the dot which leads to a shape incompatibility error. To make the gpu handler completely safe, we will make a new element in the tuple to host the cached activation with the desired shape.
The slice index of where to write the slice into the full buffer also changes based on whether it's contracting or non-contracting dim is sharded. With the new element, we will need to determine the slice index ourselves in the handler pass.
Copybara import of the project:

--
ceeff8e5da8ecb3f382bbd8dee83e2f0c909b22d by TJ Xu <tjx@nvidia.com>:

Assign a fixed index for cached activation
Cache correct activation slice when contracting dim is sharded

--
233763b8efb4ab0045eb998b437c7b28c8f776c8 by TJ Xu <tjx@nvidia.com>:

Simplified logic in gpu einsum handler to be more generic

--
2220cd1a022ad519cd23ab36c31c70c9627fc76d by TJ Xu <tjx@nvidia.com>:

remove un-used variables

Merging this change closes #13813

PiperOrigin-RevId: 646666635