From effdedf69a7eabf49378df10cbc3c3c5b2c9fa17 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Sun, 20 Oct 2024 12:09:38 -0700
Subject: [PATCH] [mlir] Avoid repeated map lookups (NFC)

---
 .../Dialect/SparseTensor/IR/SparseTensorDialect.cpp    | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 32e1dcbc2cce2..9854cfcc279b5 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -1142,16 +1142,18 @@ bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
       auto pos = dimOp.getPosition();
       if (binOp.getKind() == AffineExprKind::FloorDiv) {
         // Expect only one floordiv for each dimension.
-        if (coeffientMap.find(pos) != coeffientMap.end())
+        auto [it, inserted] = coeffientMap.try_emplace(pos);
+        if (!inserted)
           return false;
         // Record coefficient of the floordiv.
-        coeffientMap[pos] = conOp.getValue();
+        it->second = conOp.getValue();
       } else if (binOp.getKind() == AffineExprKind::Mod) {
         // Expect floordiv before mod.
-        if (coeffientMap.find(pos) == coeffientMap.end())
+        auto it = coeffientMap.find(pos);
+        if (it == coeffientMap.end())
           return false;
         // Expect mod to have the same coefficient as floordiv.
-        if (conOp.getValue() != coeffientMap[pos])
+        if (conOp.getValue() != it->second)
           return false;
         hasBlock = true;
       } else {