wip: sink alloca insts

LuisaGroup · Dec 19, 2024 · d2ef6db · d2ef6db
1 parent 17c0769
commit d2ef6db
Showing 1 changed file with 25 additions and 3 deletions.
diff --git a/src/xir/passes/sink_alloca.cpp b/src/xir/passes/sink_alloca.cpp
@@ -1,23 +1,45 @@
 #include <luisa/core/logging.h>
+#include <luisa/core/stl/unordered_map.h>
 #include <luisa/xir/function.h>
+#include <luisa/xir/instructions/alloca.h>
 #include <luisa/xir/passes/sink_alloca.h>
 
 namespace luisa::compute::xir {
 
 class AllocaInstSinker {
 
 private:
-    static void _collect_alloca_insts_in_block(luisa::vector<AllocaInst *> collected, BasicBlock *bb) noexcept {
+    static void _collect_alloca_insts_in_block(luisa::vector<AllocaInst *> &collected,
+                                               luisa::unordered_set<BasicBlock *> &visited,
+                                               BasicBlock *bb) noexcept {
+        if (visited.emplace(bb).second) {
+            for (auto &&inst : bb->instructions()) {
+                if (inst.derived_instruction_tag() == DerivedInstructionTag::ALLOCA) {
+                    collected.emplace_back(static_cast<AllocaInst *>(&inst));
+                } else {// find basic blocks
+                    for (auto &&o : inst.operand_uses()) {
+                        if (auto v = o->value(); v->derived_value_tag() == DerivedValueTag::BASIC_BLOCK) {
+                            _collect_alloca_insts_in_block(collected, visited, static_cast<BasicBlock *>(v));
+                        }
+                    }
+                }
+            }
+        }
     }
 
 private:
     [[nodiscard]] static bool _try_sink_alloca_inst(AllocaInst *inst) noexcept {
+        return false;
     }
 
 public:
     [[nodiscard]] static SinkAllocaInfo run(FunctionDefinition *f) noexcept {
-        luisa::vector<AllocaInst *> collected;
-        _collect_alloca_insts_in_block(collected, f->body_block());
+        auto collected = [f] {
+            luisa::vector<AllocaInst *> collected;
+            luisa::unordered_set<BasicBlock *> visited;
+            _collect_alloca_insts_in_block(collected, visited, f->body_block());
+            return collected;
+        }();
         collected.erase(std::remove_if(collected.begin(), collected.end(), [](AllocaInst *inst) noexcept {
                             return !_try_sink_alloca_inst(inst);
                         }),