@@ -330,6 +330,18 @@ std::optional<AllocationState> LatticePoint::get(mlir::Value val) const {
330330 return it->second ;
331331}
332332
333+ static mlir::Value lookThroughDeclaresAndConverts (mlir::Value value) {
334+ while (mlir::Operation *op = value.getDefiningOp ()) {
335+ if (auto declareOp = llvm::dyn_cast<fir::DeclareOp>(op))
336+ value = declareOp.getMemref ();
337+ else if (auto convertOp = llvm::dyn_cast<fir::ConvertOp>(op))
338+ value = convertOp->getOperand (0 );
339+ else
340+ return value;
341+ }
342+ return value;
343+ }
344+
333345mlir::LogicalResult AllocationAnalysis::visitOperation (
334346 mlir::Operation *op, const LatticePoint &before, LatticePoint *after) {
335347 LLVM_DEBUG (llvm::dbgs () << " StackArrays: Visiting operation: " << *op
@@ -363,10 +375,10 @@ mlir::LogicalResult AllocationAnalysis::visitOperation(
363375 mlir::Value operand = op->getOperand (0 );
364376
365377 // Note: StackArrays is scheduled in the pass pipeline after lowering hlfir
366- // to fir. Therefore, we only need to handle `fir::DeclareOp`s.
367- if ( auto declareOp =
368- llvm::dyn_cast_if_present<fir::DeclareOp>(operand. getDefiningOp ()))
369- operand = declareOp. getMemref ( );
378+ // to fir. Therefore, we only need to handle `fir::DeclareOp`s. Also look
379+ // past converts in case the pointer was changed between different pointer
380+ // types.
381+ operand = lookThroughDeclaresAndConverts (operand );
370382
371383 std::optional<AllocationState> operandState = before.get (operand);
372384 if (operandState && *operandState == AllocationState::Allocated) {
@@ -535,17 +547,12 @@ AllocMemConversion::matchAndRewrite(fir::AllocMemOp allocmem,
535547
536548 // remove freemem operations
537549 llvm::SmallVector<mlir::Operation *> erases;
538- for (mlir::Operation *user : allocmem.getOperation ()->getUsers ()) {
539- if (auto declareOp = mlir::dyn_cast_if_present<fir::DeclareOp>(user)) {
540- for (mlir::Operation *user : declareOp->getUsers ()) {
541- if (mlir::isa<fir::FreeMemOp>(user))
542- erases.push_back (user);
543- }
544- }
545-
546- if (mlir::isa<fir::FreeMemOp>(user))
547- erases.push_back (user);
548- }
550+ mlir::Operation *parent = allocmem->getParentOp ();
551+ // TODO: this shouldn't need to be re-calculated for every allocmem
552+ parent->walk ([&](fir::FreeMemOp freeOp) {
553+ if (lookThroughDeclaresAndConverts (freeOp->getOperand (0 )) == allocmem)
554+ erases.push_back (freeOp);
555+ });
549556
550557 // now we are done iterating the users, it is safe to mutate them
551558 for (mlir::Operation *erase : erases)
0 commit comments