[PyTorch] Fix numeric overflow caused by int-type parameters and return value...

[PyTorch] Fix numeric overflow caused by int-type parameters and return value in the roundup function (#2034)
Signed-off-by: lvdunlin <lvdunlin@xiaomi.com>
Co-authored-by: lvdunlin <lvdunlin@xiaomi.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

transformer_engine/pytorch/csrc/common.cpp

@@ -286,7 +286,7 @@ std::vector<size_t> convertShape(const NVTEShape& shape) {
   return std::vector<size_t>(shape.data, shape.data + shape.ndim);
 }
 
-int roundup(const int value, const int multiple) {
+size_t roundup(const size_t value, const size_t multiple) {
   assert(multiple > 0);
   return ((value + multiple - 1) / multiple) * multiple;
 }

transformer_engine/pytorch/csrc/common.h

@@ -417,7 +417,7 @@ void* getDataPtr(at::Tensor tensor, int offset = 0);
 
 std::vector<size_t> convertShape(const NVTEShape& shape);
 
-int roundup(const int value, const int multiple);
+size_t roundup(const size_t value, const size_t multiple);
 
 NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape);
 }  // namespace transformer_engine::pytorch