[ROCm] std::clamp work-around for hip-clang compiler (pytorch#127812)

Fixes pytorch#127666.

Other std math functions are replaced with those in the global namespace during hipify. HIP does not claim to support every function in the C++ standard library. std::clamp is not yet supported and we have been relying on the std implementation. For Fedora 40 + gcc 14, a host-side assert is used which is not supported. Work-around this by replacing std::clamp with min and max.  Using #ifndef USE_ROCM to differentiate between CUDA using std::clamp and the ROCm replacement broke Windows builds. The replacement generates the same PTX as std::clamp, so using the replacement unconditionally. The replacement generates the same PTX as std::clamp. See https://godbolt.org/z/Wde9KW3v4 for a sample.

Original patch comes from @lamikr. Modified to improve efficiency.

lamikr/rocm_sdk_builder#37

Co-authored-by: Nikita Shulga <[email protected]>
Pull Request resolved: pytorch#127812
Approved by: https://github.com/hongxiayang, https://github.com/malfet

Author: 2 people

aten/src/ATen/native/cuda/IndexKernel.cu

@@ -259,7 +259,14 @@ void index_put_kernel_quantized_cuda(TensorIterator& iter, const IntArrayRef ind
 
     gpu_index_kernel(iter, index_size, index_stride, [inv_scale, zero_point, qmin, qmax]C10_DEVICE(char* const out_data, const char* const in_data, const int64_t offset) {
       int64_t qvalue = static_cast<int64_t>(zero_point + nearbyintf(*(float*)in_data * inv_scale));
-      qvalue = std::clamp(qvalue, qmin, qmax);
+      // See https://github.com/pytorch/pytorch/issues/127666
+      // and https://github.com/pytorch/pytorch/issues/128253.
+      // hip-clang std::clamp __glibcxx_assert_fail host function when building on Fedora40/gcc14.
+      // The following replaces std::clamp(qvalue, qmin, qmax) and is a viable solution for
+      // both CUDA and ROCm since std::clamp and this replacement generates the same PTX.
+      // Using #ifdef USE_ROCM to differentiate caused Windows build failures.
+      // The replacement should generate the same PTX as std::clamp. See https://godbolt.org/z/Wde9KW3v4
+      qvalue = (qvalue < qmin) ? qmin : (qmax < qvalue) ? qmax : qvalue;
       *(scalar_t*)(out_data + offset) = static_cast<scalar_t>(qvalue);
     });
   });