Change Redzone space limit for XLA GPU

[kaixih]

This PR changes how the redzone space limit is set in the XLA gpu conv algorithm picker.

1. It sets the numeric max of int64 for the input/output allocator. So, we can have consistent behavior with the gemm picker.
2. It allows the adjustment of the space limit for the scratch allocator via an env var. So, users can adjust it via XLA_FLAGS=--xla_gpu_redzone_scratch_max_megabytes=6144.

cc. @nluehr

[cheshire]

Hi @kaixih , could you provide some more context on what is the desired goal?

[kaixih]

@cheshire Sure. Basically, we found that the max space limit of redzone allocator for the XLA conv is set to be 4GB, which is insufficient for some models that expect large input/output tensors. In addition, we also noticed that this limit is not adjustable during runtime. So, compared to the XLA gemm, which has already set the limit of the input/output redzone allocator to the numeric max of int, we think it might be appropriate to follow it for the XLA conv. Moreover, we introduced a new env var to control the scratch redzone allocator max limit as well in case it needs to be adjusted.

[kaixih]

@cheshire Can you help check what blocks the merging? It seems some "Google internal checks" failed. Thanks.

[akuegel]

There was a merge conflict in xla.proto, you used the same tag as was already used in a recent change. I fixed that and got your PR merged.

[gbaned]

Seems auto-merge is not happening but the changes are merged into master now, so we can close this. Thank you for the PR.