Fix memory_coloring pass when MIGRAPHX_NSTREAMS > 2 (#3757)

I noticed allocation segments reaching close to uint64 max value, which clearly would throw an out of memory error when trying to allocate on the GPU. This happened when MIGRAPHX_NSTREAMS was set to 2 or greater and the model somehow was large enough to trigger it.

Changing from auto to size_t seems to fix the issue.

After some further investigation, the segment start value is the culprit which becomes int32. I've updated n as size_t for clarity.

src/memory_coloring.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -177,9 +177,9 @@ struct allocation_segment
     {
         assert(ins->get_shape().bytes() > 0);
         // Compute alignment
-        auto n = 1 + (ins->get_shape().bytes() - 1) / alignment;
+        std::size_t n = 1 + (ins->get_shape().bytes() - 1) / alignment;
         assert(n > 0);
-        auto start = 0;
+        std::size_t start = 0;
         // Insert at end if it cant fit at the begining
         if(segments.empty() or segments.begin()->first <= n)
         {

test/memory_coloring_test.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2025 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -3805,4 +3805,18 @@ TEST_CASE(test_tuple)
     CHECK(is_disjoint({a1, a2}));
 }
 
+TEST_CASE(test_large_offsets)
+{
+    migraphx::module m;
+
+    auto a1 = add_alloc(m, {migraphx::shape::float_type, {10000000000}});
+    auto m1 = m.add_instruction(pass_op{}, a1);
+    auto a2 = add_alloc(m, {migraphx::shape::float_type, {10000000000}});
+    m.add_instruction(pass_op{}, a2, m1);
+    run_pass(m);
+    CHECK(m.get_parameter_shape("scratch").bytes() == 80000000000);
+    CHECK(no_allocate(m));
+    CHECK(is_disjoint({a1, a2}));
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }