diff --git a/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
index b7da20c108f3e..1c31b740c9919 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DecomposeGenericByUnfoldingPermutation.cpp
@@ -8,7 +8,6 @@
 //
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
-#include <map>
 #include <utility>
 
 using namespace mlir;
@@ -125,9 +124,10 @@ computeTransposeBroadcast(AffineMap &map) {
     // refering to d3, d4. Therefore, re-base the transpose dimensions so
     // that they start from d0.
     permutation.resize(minorSize);
-    std::map<int64_t, int64_t> minorMap;
+    SmallVector<int64_t> minorMap;
+    minorMap.resize(minorSize);
     for (int64_t i = 0; i < minorSize; ++i)
-      minorMap.insert({sortedResMap[i], i});
+      minorMap[sortedResMap[i]] = i;
 
     // Re-map the dimensions.
     SmallVector<int64_t> remappedResult(minorSize);