Also visit memref element type in AsmPrinter

--

PiperOrigin-RevId: 248797935

diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index f99c09f..73c74e7 100644 (file)
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -197,12 +197,17 @@ void ModuleState::visitType(Type type) {
       visitType(input);
     for (auto result : funcType.getResults())
       visitType(result);
-  } else if (auto memref = type.dyn_cast<MemRefType>()) {
+    return;
+  }
+  if (auto memref = type.dyn_cast<MemRefType>()) {
     // Visit affine maps in memref type.
     for (auto map : memref.getAffineMaps())
       recordAttributeReference(AffineMapAttr::get(map));
-  } else if (auto vecOrTensor = type.dyn_cast<ShapedType>()) {
-    visitType(vecOrTensor.getElementType());
+    // TODO(b/132735995) Remove this when MemRef is a subclass of ShapedType.
+    visitType(memref.getElementType());
+  }
+  if (auto shapedType = type.dyn_cast<ShapedType>()) {
+    visitType(shapedType.getElementType());
   }
 }