NFC: Use TypeSwitch to simplify existing code.

PiperOrigin-RevId: 286066371

1 files changed, 4 insertions, 8 deletions

diff --git a/mlir/lib/Transforms/Utils/Utils.cpp b/mlir/lib/Transforms/Utils/Utils.cpp
index 190b6c3155e..79a6d7a6902 100644
--- a/mlir/lib/Transforms/Utils/Utils.cpp
+++ b/mlir/lib/Transforms/Utils/Utils.cpp
@@ -22,6 +22,7 @@
 
 #include "mlir/Transforms/Utils.h"
 
+#include "mlir/ADT/TypeSwitch.h"
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/AffineStructures.h"
 #include "mlir/Analysis/Dominance.h"
@@ -47,14 +48,9 @@ static bool isMemRefDereferencingOp(Operation &op) {
 
 /// Return the AffineMapAttr associated with memory 'op' on 'memref'.
 static NamedAttribute getAffineMapAttrForMemRef(Operation *op, Value *memref) {
-  if (auto loadOp = dyn_cast<AffineLoadOp>(op))
-    return loadOp.getAffineMapAttrForMemRef(memref);
-  else if (auto storeOp = dyn_cast<AffineStoreOp>(op))
-    return storeOp.getAffineMapAttrForMemRef(memref);
-  else if (auto dmaStart = dyn_cast<AffineDmaStartOp>(op))
-    return dmaStart.getAffineMapAttrForMemRef(memref);
-  assert(isa<AffineDmaWaitOp>(op));
-  return cast<AffineDmaWaitOp>(op).getAffineMapAttrForMemRef(memref);
+  return TypeSwitch<Operation *, NamedAttribute>(op)
+      .Case<AffineDmaStartOp, AffineLoadOp, AffineStoreOp, AffineDmaWaitOp>(
+          [=](auto op) { return op.getAffineMapAttrForMemRef(memref); });
 }
 
 // Perform the replacement in `op`.