commit ee9b206b75dece2e7025b0d237ff6176ea927cae
author Stanley Winata <68087699+raikonenfnu@users.noreply.github.com> Thu Jan 18 16:25:19 2024 -0800
committer GitHub <noreply@github.com> Thu Jan 18 16:25:19 2024 -0800
tree a5be0cdb759aab5f8f1de07265c8837386d27f1e
parent c859e290c649b7fdda897b23cb5b99a26cd1053c

[ROCM][Ukernel] Fix index types. (#16154)

compiler/plugins/target/ROCM/builtins/ukernel/argmax_ukernel.c

diff --git a/compiler/plugins/target/ROCM/builtins/ukernel/argmax_ukernel.c b/compiler/plugins/target/ROCM/builtins/ukernel/argmax_ukernel.c
index 9c2caf3..81f4b85 100644
--- a/compiler/plugins/target/ROCM/builtins/ukernel/argmax_ukernel.c
+++ b/compiler/plugins/target/ROCM/builtins/ukernel/argmax_ukernel.c
@@ -109,7 +109,7 @@
   }
   // if there are multiple max value holder, find smallest index (argmax
   // semantics).
-  int32_t indexVal = wgMax == laneMax ? laneResult : __INT64_MAX__;
+  int64_t indexVal = wgMax == laneMax ? laneResult : __INT64_MAX__;
   laneResult = __ockl_wfred_min_i64(indexVal);
   if (laneID == 0)
     outputBuffer[output_offset] = laneResult;
@@ -191,7 +191,7 @@
   }
   // if there are multiple max value holder, find smallest index (argmax
   // semantics).
-  int32_t indexVal = wgMax == laneMax ? laneResult : __INT64_MAX__;
+  int64_t indexVal = wgMax == laneMax ? laneResult : __INT64_MAX__;
   laneResult = __ockl_wfred_min_i64(indexVal);
   if (laneID == 0)
     outputBuffer[output_offset] = laneResult;