Fix ExpandMemoryOpPass don't work properly

core/iwasm/compilation/aot_llvm_extra.cpp

@@ -82,43 +82,40 @@ class ExpandMemoryOpPass : public PassInfoMixin<ExpandMemoryOpPass>
 PreservedAnalyses
 ExpandMemoryOpPass::run(Function &F, FunctionAnalysisManager &AM)
 {
-    Intrinsic::ID ID = F.getIntrinsicID();
-    bool Changed = false;
-
-    for (auto I = F.user_begin(), E = F.user_end(); I != E;) {
-        Instruction *Inst = cast<Instruction>(*I);
-        ++I;
-
-        switch (ID) {
-            case Intrinsic::memcpy:
-            {
-                auto *Memcpy = cast<MemCpyInst>(Inst);
-                Function *ParentFunc = Memcpy->getParent()->getParent();
-                const TargetTransformInfo &TTI =
-                    AM.getResult<TargetIRAnalysis>(*ParentFunc);
-                expandMemCpyAsLoop(Memcpy, TTI);
-                Memcpy->eraseFromParent();
-                Changed = true;
-                break;
+    SmallVector<MemIntrinsic *, 16> MemCalls;
+
+    /* Iterate over all instructions in the function, looking for memcpy,
+     * memmove, and memset.  When we find one, expand it into a loop. */
+
+    for (auto &BB : F) {
+        for (auto &Inst : BB) {
+            if (auto *Memcpy = dyn_cast_or_null<MemCpyInst>(&Inst)) {
+                MemCalls.push_back(Memcpy);
             }
-            case Intrinsic::memmove:
-            {
-                auto *Memmove = cast<MemMoveInst>(Inst);
-                expandMemMoveAsLoop(Memmove);
-                Memmove->eraseFromParent();
-                Changed = true;
-                break;
+            else if (auto *Memmove = dyn_cast_or_null<MemMoveInst>(&Inst)) {
+                MemCalls.push_back(Memmove);
             }
-            case Intrinsic::memset:
-            {
-                auto *Memset = cast<MemSetInst>(Inst);
-                expandMemSetAsLoop(Memset);
-                Memset->eraseFromParent();
-                Changed = true;
-                break;
+            else if (auto *Memset = dyn_cast_or_null<MemSetInst>(&Inst)) {
+                MemCalls.push_back(Memset);
             }
-            default:
-                break;
+        }
+    }
+
+    for (MemIntrinsic *MemCall : MemCalls) {
+        if (MemCpyInst *Memcpy = dyn_cast<MemCpyInst>(MemCall)) {
+            Function *ParentFunc = Memcpy->getParent()->getParent();
+            const TargetTransformInfo &TTI =
+                AM.getResult<TargetIRAnalysis>(*ParentFunc);
+            expandMemCpyAsLoop(Memcpy, TTI);
+            Memcpy->eraseFromParent();
+        }
+        else if (MemMoveInst *Memmove = dyn_cast<MemMoveInst>(MemCall)) {
+            expandMemMoveAsLoop(Memmove);
+            Memmove->eraseFromParent();
+        }
+        else if (MemSetInst *Memset = dyn_cast<MemSetInst>(MemCall)) {
+            expandMemSetAsLoop(Memset);
+            Memset->eraseFromParent();
         }
     }
 
@@ -297,13 +294,6 @@ aot_apply_llvm_new_pass_manager(AOTCompContext *comp_ctx, LLVMModuleRef module)
         FPM.addPass(SLPVectorizerPass());
         FPM.addPass(LoadStoreVectorizerPass());
 
-        /* Run specific passes for AOT indirect mode in last since general
-           optimization may create some intrinsic function calls like
-           llvm.memset, so let's remove these function calls here. */
-        if (comp_ctx->is_indirect_mode) {
-            FPM.addPass(ExpandMemoryOpPass());
-        }
-
         if (comp_ctx->enable_llvm_pgo || comp_ctx->use_prof_file) {
             /* LICM pass: loop invariant code motion, attempting to remove
                as much code from the body of a loop as possible. Experiments
@@ -341,6 +331,15 @@ aot_apply_llvm_new_pass_manager(AOTCompContext *comp_ctx, LLVMModuleRef module)
         else {
             MPM.addPass(PB.buildPerModuleDefaultPipeline(OL));
         }
+
+        /* Run specific passes for AOT indirect mode in last since general
+            optimization may create some intrinsic function calls like
+            llvm.memset, so let's remove these function calls here. */
+        if (comp_ctx->is_indirect_mode) {
+            FunctionPassManager FPM1;
+            FPM1.addPass(ExpandMemoryOpPass());
+            MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM1)));
+        }
     }
 
     MPM.run(*M, MAM);