From cf64728581f047f5e2ab7dd64a5efdfff7a3f762 Mon Sep 17 00:00:00 2001
From: Huang Qi <huangqi3@xiaomi.com>
Date: Fri, 28 Jul 2023 14:31:03 +0800
Subject: [PATCH] Fix ExpandMemoryOpPass don't work properly

---
 core/iwasm/compilation/aot_llvm_extra.cpp | 44 +++++++++--------------
 1 file changed, 16 insertions(+), 28 deletions(-)

diff --git a/core/iwasm/compilation/aot_llvm_extra.cpp b/core/iwasm/compilation/aot_llvm_extra.cpp
index 1ccb2f5490..5771884e33 100644
--- a/core/iwasm/compilation/aot_llvm_extra.cpp
+++ b/core/iwasm/compilation/aot_llvm_extra.cpp
@@ -82,43 +82,29 @@ class ExpandMemoryOpPass : public PassInfoMixin<ExpandMemoryOpPass>
 PreservedAnalyses
 ExpandMemoryOpPass::run(Function &F, FunctionAnalysisManager &AM)
 {
-    Intrinsic::ID ID = F.getIntrinsicID();
-    bool Changed = false;
+    /* Iterate over all instructions in the function, looking for memcpy,
+     * memmove, and memset.  When we find one, expand it into a loop. */
 
-    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);
+    for (auto &BB : F) {
+        for (auto &Inst : BB) {
+            if (auto *Memcpy = dyn_cast<MemCpyInst>(&Inst)) {
                 Function *ParentFunc = Memcpy->getParent()->getParent();
                 const TargetTransformInfo &TTI =
                     AM.getResult<TargetIRAnalysis>(*ParentFunc);
                 expandMemCpyAsLoop(Memcpy, TTI);
                 Memcpy->eraseFromParent();
-                Changed = true;
                 break;
             }
-            case Intrinsic::memmove:
-            {
-                auto *Memmove = cast<MemMoveInst>(Inst);
+            else if (auto *Memmove = dyn_cast<MemMoveInst>(&Inst)) {
                 expandMemMoveAsLoop(Memmove);
                 Memmove->eraseFromParent();
-                Changed = true;
                 break;
             }
-            case Intrinsic::memset:
-            {
-                auto *Memset = cast<MemSetInst>(Inst);
+            else if (auto *Memset = dyn_cast<MemSetInst>(&Inst)) {
                 expandMemSetAsLoop(Memset);
                 Memset->eraseFromParent();
-                Changed = true;
                 break;
             }
-            default:
-                break;
         }
     }
 
@@ -297,13 +283,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 +320,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);