diff --git a/llvm/lib/Target/X86/X86LowerAMXType.cpp b/llvm/lib/Target/X86/X86LowerAMXType.cpp
index 2ea78d44d67cf..eacdc8a3f639a 100644
--- a/llvm/lib/Target/X86/X86LowerAMXType.cpp
+++ b/llvm/lib/Target/X86/X86LowerAMXType.cpp
@@ -154,8 +154,8 @@ class ShapeCalculator {
 
 Value *ShapeCalculator::getRowFromCol(Instruction *II, Value *V,
                                       unsigned Granularity) {
-  if (Col2Row.count(V))
-    return Col2Row[V];
+  if (auto It = Col2Row.find(V); It != Col2Row.end())
+    return It->second;
   IRBuilder<> Builder(II);
   Value *RealRow = nullptr;
   if (isa<ConstantInt>(V))