[flang][CodeGen] Fix use-after-free in BoxedProcedurePass

[kparzysz]

Avoid inspecting an operation that has been replaced.

This was detected by address sanitizer.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

@llvm/pr-subscribers-flang-codegen

Author: Krzysztof Parzyszek (kparzysz)

Changes

Replacing an element of an operation range while traversing the range can make the range invalid. Store the operations in a separate list, and traverse the list instead.
Additionally, avoid inspecting an operation that has been replaced.

This was detected by address sanitizer.

---

Full diff: https://github.com/llvm/llvm-project/pull/84376.diff

1 Files Affected:

• (modified) flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp (+16-3)

diff --git a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
index 4cf39716a73755..2e34b0a1b492b1 100644
--- a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
+++ b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
@@ -208,7 +208,12 @@ class BoxedProcedurePass
       mlir::IRRewriter rewriter(context);
       BoxprocTypeRewriter typeConverter(mlir::UnknownLoc::get(context));
       mlir::Dialect *firDialect = context->getLoadedDialect("fir");
-      getModule().walk([&](mlir::Operation *op) {
+      llvm::SmallVector<mlir::Operation *> operations;
+
+      getModule().walk([&](mlir::Operation *op) { operations.push_back(op); });
+
+      for (mlir::Operation *op : operations) {
+        bool opIsValid = true;
         typeConverter.setLocation(op->getLoc());
         if (auto addr = mlir::dyn_cast<BoxAddrOp>(op)) {
           mlir::Type ty = addr.getVal().getType();
@@ -220,6 +225,7 @@ class BoxedProcedurePass
             rewriter.setInsertionPoint(addr);
             rewriter.replaceOpWithNewOp<ConvertOp>(
                 addr, typeConverter.convertType(addr.getType()), addr.getVal());
+            opIsValid = false;
           } else if (typeConverter.needsConversion(resTy)) {
             rewriter.startOpModification(op);
             op->getResult(0).setType(typeConverter.convertType(resTy));
@@ -271,10 +277,12 @@ class BoxedProcedurePass
                 llvm::ArrayRef<mlir::Value>{tramp});
             rewriter.replaceOpWithNewOp<ConvertOp>(embox, toTy,
                                                    adjustCall.getResult(0));
+            opIsValid = false;
           } else {
             // Just forward the function as a pointer.
             rewriter.replaceOpWithNewOp<ConvertOp>(embox, toTy,
                                                    embox.getFunc());
+            opIsValid = false;
           }
         } else if (auto global = mlir::dyn_cast<GlobalOp>(op)) {
           auto ty = global.getType();
@@ -297,6 +305,7 @@ class BoxedProcedurePass
             rewriter.replaceOpWithNewOp<AllocaOp>(
                 mem, toTy, uniqName, bindcName, isPinned, mem.getTypeparams(),
                 mem.getShape());
+            opIsValid = false;
           }
         } else if (auto mem = mlir::dyn_cast<AllocMemOp>(op)) {
           auto ty = mem.getType();
@@ -310,6 +319,7 @@ class BoxedProcedurePass
             rewriter.replaceOpWithNewOp<AllocMemOp>(
                 mem, toTy, uniqName, bindcName, mem.getTypeparams(),
                 mem.getShape());
+            opIsValid = false;
           }
         } else if (auto coor = mlir::dyn_cast<CoordinateOp>(op)) {
           auto ty = coor.getType();
@@ -321,6 +331,7 @@ class BoxedProcedurePass
             auto toBaseTy = typeConverter.convertType(baseTy);
             rewriter.replaceOpWithNewOp<CoordinateOp>(coor, toTy, coor.getRef(),
                                                       coor.getCoor(), toBaseTy);
+            opIsValid = false;
           }
         } else if (auto index = mlir::dyn_cast<FieldIndexOp>(op)) {
           auto ty = index.getType();
@@ -332,6 +343,7 @@ class BoxedProcedurePass
             auto toOnTy = typeConverter.convertType(onTy);
             rewriter.replaceOpWithNewOp<FieldIndexOp>(
                 index, toTy, index.getFieldId(), toOnTy, index.getTypeparams());
+            opIsValid = false;
           }
         } else if (auto index = mlir::dyn_cast<LenParamIndexOp>(op)) {
           auto ty = index.getType();
@@ -343,6 +355,7 @@ class BoxedProcedurePass
             auto toOnTy = typeConverter.convertType(onTy);
             rewriter.replaceOpWithNewOp<LenParamIndexOp>(
                 index, toTy, index.getFieldId(), toOnTy, index.getTypeparams());
+            opIsValid = false;
           }
         } else if (op->getDialect() == firDialect) {
           rewriter.startOpModification(op);
@@ -354,7 +367,7 @@ class BoxedProcedurePass
           rewriter.finalizeOpModification(op);
         }
         // Ensure block arguments are updated if needed.
-        if (op->getNumRegions() != 0) {
+        if (opIsValid && op->getNumRegions() != 0) {
           rewriter.startOpModification(op);
           for (mlir::Region &region : op->getRegions())
             for (mlir::Block &block : region.getBlocks())
@@ -366,7 +379,7 @@ class BoxedProcedurePass
                 }
           rewriter.finalizeOpModification(op);
         }
-      });
+      }
     }
   }
 

[jeanPerier]

Thanks for identifying the issue

[jeanPerier]

Are you sure it is invalid to erase IR in the walk? I understand why opIsValid is needed and that it was bad to go through the region of erased ops, but I think a walk in post order is supposed to allow erasing operations:

llvm-project/mlir/include/mlir/IR/Operation.h

Line 767 in 881df55

/// (post-order by default). A callback on a block or operation is allowed to

Is there still a sanitizer issue when introducing opIsValid without creating the operations vector?

I am a bit concerned with copying the pointer of all the module operations into a SmallVector. On big Fortran applications with hundred of thousands of line in one file, the module will contain millions of operations (I doubt this go above the ~500 millions element capacity of the SmallVector), but it may be overkill for a pass that is not expected to modify a lot of IR (procedure pointers and dummy procedures are usually not all over the place).

[kparzysz]

You're right---the default walking is post-order. Let me check that again.

[kparzysz]

Updated the patch to remove the op caching.

[jeanPerier]

Thanks a lot @kparzysz for finding and fixing this issue!