[XLA] Improve GPU memory limit handling and shape size calculation

[zhenying-liu]

This change improves accuracy of GPU memory limit calculations and provides a more maintainable interface for shape size computations.

• Change memory limit type from int64_t to uint64_t to prevent negative values and better represent memory sizes.

• Add memory space filtering to exclude host memory when calculating input/output sizes that impact GPU memory usage. This ensures we only count buffers that actually reside in device memory.

• Replace direct GetSizeOfShape() calls with ShapeSizeBytesFunction() wrapper, which provides:

  • Consistent shape size calculation across the codebase
  • Optional memory space filtering capability
  • Proper handling of dynamic shapes and their metadata

[zhenying-liu]

The two CI failures and AMD ROCm seem not related with my changes:

Error: Error: pod failed to come online with error: Error: Pod linux-x86-n2-16-xf6x6-runner-xj7vr-workflow is unhealthy with phase status Failed
Error: Process completed with exit code 1.
Error: Executing the custom container implementation failed. Please contact your self hosted runner administrator.