minor improvements to peephole store-load forwarding

LuisaGroup · Jan 10, 2025 · 44dc415 · 44dc415
1 parent dd93d02
commit 44dc415
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 36 deletions.
diff --git a/include/luisa/xir/passes/peephole_store_forward.h b/include/luisa/xir/passes/peephole_store_forward.h
@@ -23,7 +23,6 @@ struct PeepholeStoreForwardInfo {
     luisa::unordered_map<LoadInst *, StoreInst *> forwarded_instructions;
 };
 
-[[nodiscard]] LC_XIR_API PeepholeStoreForwardInfo peephole_store_forward_pass_run_on_basic_block(BasicBlock *block) noexcept;
 [[nodiscard]] LC_XIR_API PeepholeStoreForwardInfo peephole_store_forward_pass_run_on_function(Function *function) noexcept;
 [[nodiscard]] LC_XIR_API PeepholeStoreForwardInfo peephole_store_forward_pass_run_on_module(Module *module) noexcept;
 

diff --git a/src/backends/fallback/fallback_shader.cpp b/src/backends/fallback/fallback_shader.cpp
@@ -173,12 +173,13 @@ FallbackShader::FallbackShader(FallbackDevice *device, const ShaderOption &optio
 
     // run some simple optimization passes on XIR to reduce the size of LLVM IR
     Clock opt_clk;
+    auto dce1_info = xir::dce_pass_run_on_module(xir_module);
     auto store_forward_info = xir::peephole_store_forward_pass_run_on_module(xir_module);
-    auto dce_info = xir::dce_pass_run_on_module(xir_module);
+    auto dce2_info = xir::dce_pass_run_on_module(xir_module);
     LUISA_INFO("Forwarded {} store instruction(s), "
                "removed {} dead instructions in {} ms.",
                store_forward_info.forwarded_instructions.size(),
-               dce_info.removed_instructions.size(),
+               dce1_info.removed_instructions.size() + dce2_info.removed_instructions.size(),
                opt_clk.toc());
 
     // dump for debugging

diff --git a/src/xir/passes/peephole_store_forward.cpp b/src/xir/passes/peephole_store_forward.cpp
@@ -27,8 +27,9 @@ namespace detail {
     return nullptr;
 }
 
-// TODO: we only handle scalars for now
-static void run_peephole_store_forward_on_basic_block(BasicBlock *block, PeepholeStoreForwardInfo &info) noexcept {
+// TODO: we only handle local alloca's in straight-line code for now
+static void run_peephole_store_forward_on_basic_block(luisa::unordered_set<BasicBlock *> &visited,
+                                                      BasicBlock *block, PeepholeStoreForwardInfo &info) noexcept {
 
     luisa::unordered_map<AllocaInst *, luisa::vector<Value *>> variable_pointers;// maps variables to pointers
     luisa::unordered_map<Value *, StoreInst *> latest_stores;                    // maps pointers to the latest store instruction
@@ -45,34 +46,51 @@ static void run_peephole_store_forward_on_basic_block(BasicBlock *block, Peephol
         return nullptr;
     };
 
-    for (auto &&inst : block->instructions()) {
-        switch (inst.derived_instruction_tag()) {
-            case DerivedInstructionTag::LOAD: {
-                auto load = static_cast<LoadInst *>(&inst);
-                if (auto iter = latest_stores.find(load->variable()); iter != latest_stores.end()) {
-                    removable_loads.emplace(load, iter->second);
+    // we visit the block and all of its single straight-line successors
+    while (visited.emplace(block).second) {
+        // process the instructions in the block
+        for (auto &&inst : block->instructions()) {
+            switch (inst.derived_instruction_tag()) {
+                case DerivedInstructionTag::LOAD: {
+                    auto load = static_cast<LoadInst *>(&inst);
+                    if (auto iter = latest_stores.find(load->variable()); iter != latest_stores.end()) {
+                        removable_loads.emplace(load, iter->second);
+                    }
+                    break;
                 }
-                break;
-            }
-            case DerivedInstructionTag::STORE: {
-                auto store = static_cast<StoreInst *>(&inst);
-                // if this is a store to (part of) a local alloca, we might be able to forward it
-                if (auto pointer = store->variable(); invalidate_interfering_stores(pointer)) {
-                    latest_stores[pointer] = store;
+                case DerivedInstructionTag::STORE: {
+                    auto store = static_cast<StoreInst *>(&inst);
+                    // if this is a store to (part of) a local alloca, we might be able to forward it
+                    if (auto pointer = store->variable(); invalidate_interfering_stores(pointer)) {
+                        latest_stores[pointer] = store;
+                    }
+                    break;
                 }
-                break;
-            }
-            case DerivedInstructionTag::GEP: {
-                // users of GEPs will handle the forwarding, so we don't need to do anything here
-                break;
-            }
-            default: {// for other instructions, we invalidate possibly interfering stores
-                for (auto op_use : inst.operand_uses()) {
-                    invalidate_interfering_stores(op_use->value());
+                case DerivedInstructionTag::GEP: {
+                    // users of GEPs will handle the forwarding, so we don't need to do anything here
+                    break;
+                }
+                default: {// for other instructions, we invalidate possibly interfering stores
+                    for (auto op_use : inst.operand_uses()) {
+                        invalidate_interfering_stores(op_use->value());
+                    }
+                    break;
                 }
-                break;
             }
         }
+        // move to the next block if it is the only successor and only has a single predecessor
+        BasicBlock *next = nullptr;
+        auto successor_count = 0u;
+        block->traverse_successors(true, [&](BasicBlock *succ) noexcept {
+            successor_count++;
+            next = succ;
+        });
+        if (successor_count != 1) { break; }
+        // check if the next block has a single predecessor
+        auto pred_count = 0u;
+        next->traverse_predecessors(false, [&](BasicBlock *) noexcept { pred_count++; });
+        if (pred_count != 1) { break; }
+        block = next;
     }
     for (auto &&[load, store] : removable_loads) {
         load->replace_all_uses_with(store->value());
@@ -83,20 +101,15 @@ static void run_peephole_store_forward_on_basic_block(BasicBlock *block, Peephol
 
 void run_peephole_store_forward_on_function(Function *function, PeepholeStoreForwardInfo &info) noexcept {
     if (auto definition = function->definition()) {
-        definition->traverse_basic_blocks([&](BasicBlock *block) noexcept {
-            run_peephole_store_forward_on_basic_block(block, info);
+        luisa::unordered_set<BasicBlock *> visited;
+        definition->traverse_basic_blocks(BasicBlockTraversalOrder::REVERSE_POST_ORDER, [&](BasicBlock *block) noexcept {
+            run_peephole_store_forward_on_basic_block(visited, block, info);
         });
     }
 }
 
 }// namespace detail
 
-PeepholeStoreForwardInfo peephole_store_forward_pass_run_on_basic_block(BasicBlock *block) noexcept {
-    PeepholeStoreForwardInfo info;
-    detail::run_peephole_store_forward_on_basic_block(block, info);
-    return info;
-}
-
 PeepholeStoreForwardInfo peephole_store_forward_pass_run_on_function(Function *function) noexcept {
     PeepholeStoreForwardInfo info;
     detail::run_peephole_store_forward_on_function(function, info);