Apply suggestion from @joker-eph

Co-authored-by: Mehdi Amini <[email protected]>

Author: fabianmcg

mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp

@@ -355,34 +355,32 @@ LogicalResult PtrAddOp::inferReturnTypes(
   Type offsetType = operands[1].getType();
 
   // If neither are shaped types, result is same as base type.
-  if (!isa<ShapedType>(baseType) && !isa<ShapedType>(offsetType)) {
+  auto offTy = dyn_cast<ShapedType>(offsetType);
+  if (!offTy) {
+    // If the offset isn't shaped, the result is always the base type.
     inferredReturnTypes.push_back(baseType);
     return success();
   }
-
-  // Handle cases with shaped types.
-  if (auto baseTy = dyn_cast<ShapedType>(baseType)) {
-    // If both shaped, they must have the same shape.
-    if (auto offTy = dyn_cast<ShapedType>(offsetType)) {
-      if (offTy.getShape() != baseTy.getShape()) {
-        if (location)
-          mlir::emitError(*location) << "shapes of base and offset must match";
-        return failure();
-      }
-      // Make sure they are the same kind of shaped type.
-      if (baseType.getTypeID() != offsetType.getTypeID()) {
-        if (location)
-          mlir::emitError(*location) << "the shaped containers type must match";
-        return failure();
-      }
-    }
-    inferredReturnTypes.push_back(baseType);
-    return success();
+  auto baseTy = dyn_cast<ShapedType>(baseType);
+  if (!baseTy) {
+    // Base isn't shaped, but offset is, use the ShapedType from offset with the base pointer as element type.
+    inferredReturnTypes.push_back(offsetShapedType.clone(baseType));
+    return success();  
   }
 
-  // Base is scalar, offset is shaped.
-  auto offsetShapedType = cast<ShapedType>(offsetType);
-  inferredReturnTypes.push_back(offsetShapedType.clone(baseType));
+  // Both are shaped, their shape must match.
+  if (offTy.getShape() != baseTy.getShape()) {
+    if (location)
+      mlir::emitError(*location) << "shapes of base and offset must match";
+    return failure();
+  }
+  // Make sure they are the same kind of shaped type.
+  if (baseType.getTypeID() != offsetType.getTypeID()) {
+    if (location)
+      mlir::emitError(*location) << "the shaped containers type must match";
+    return failure();
+  }
+  inferredReturnTypes.push_back(baseType);
   return success();
 }