From fa3a220eb7fcd8c48f4d34737d81b3fa3b892129 Mon Sep 17 00:00:00 2001
From: NatashaKnk <natashaknk@google.com>
Date: Fri, 9 Jun 2023 16:50:41 -0700
Subject: [PATCH] Fix bug in TopK pattern matching that sets K as 2nd dim
 rather than last. (#14041)

Test cases coming in following PR, this is just a quick fix while it's
out for review.
---
 .../StableHLO/Preprocessing/StableHLOToStableHLO.cpp          | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/StableHLOToStableHLO.cpp b/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/StableHLOToStableHLO.cpp
index 15692381bfb8..65864bd0a4e4 100644
--- a/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/StableHLOToStableHLO.cpp
+++ b/compiler/src/iree/compiler/InputConversion/StableHLO/Preprocessing/StableHLOToStableHLO.cpp
@@ -1607,7 +1607,9 @@ struct IotaSortSliceIsTopK final : OpRewritePattern<mlir::stablehlo::SortOp> {
       // Treat the first slice as inputs, the second as indices.
       if (idx == 0) {
         topV = sliceOp.getResult();
-        k = sliceOp.getLimitIndices().getValues<int64_t>()[1];
+        SmallVector<int64_t> limitIndices =
+            llvm::to_vector(sliceOp.getLimitIndices().getValues<int64_t>());
+        k = limitIndices.back();
       } else {
         topI = sliceOp.getResult();
       }