fixed lift_indirect_targets

Hugobros3 · Hugobros3 · commit 853af249ab99 · 2024-04-17T17:19:36.000+02:00
diff --git a/src/shady/passes/lift_indirect_targets.c b/src/shady/passes/lift_indirect_targets.c
@@ -15,6 +15,7 @@
 #include "../analysis/free_variables.h"
 #include "../analysis/uses.h"
 #include "../analysis/leak.h"
+#include "../analysis/verify.h"
 
 #include <assert.h>
 #include <string.h>
@@ -31,6 +32,8 @@ typedef struct Context_ {
     struct Dict* lifted;
     bool disable_lowering;
     const CompilerConfig* config;
+
+    bool* todo;
 } Context;
 
 static const Node* process_node(Context* ctx, const Node* node);
@@ -66,6 +69,20 @@ static const Node* add_spill_instrs(Context* ctx, BodyBuilder* builder, struct L
     return sp;
 }
 
+static void add_to_recover_context(struct List* recover_context, struct Dict* set, const Node* except) {
+    Nodes params = get_abstraction_params(except);
+    size_t i = 0;
+    const Node* item;
+    while (dict_iter(set, &i, &item, NULL)) {
+        for (size_t j = 0; j < params.count; j++) {
+            if (item == params.nodes[j])
+                goto skip;
+        }
+        append_list(const Node*, recover_context, item );
+        skip:;
+    }
+}
+
 static LiftedCont* lambda_lift(Context* ctx, const Node* cont, String given_name) {
     assert(is_basic_block(cont) || is_case(cont));
     LiftedCont** found = find_value_dict(const Node*, LiftedCont*, ctx->lifted, cont);
@@ -82,20 +99,19 @@ static LiftedCont* lambda_lift(Context* ctx, const Node* cont, String given_name
     CFNode* cf_node = scope_lookup(ctx->scope, cont);
     CFNodeVariables* node_vars = *find_value_dict(CFNode*, CFNodeVariables*, ctx->scope_vars, cf_node);
     struct List* recover_context = new_list(const Node*);
-    size_t recover_context_size = entries_count_dict(node_vars->free_set);
-
-    {
-        debugv_print("lambda_lift: free (to-be-spilled) variables at '%s' (count=%d): ", name, entries_count_dict(node_vars->free_set));
-        size_t i = 0;
-        const Node* item;
-        while (dict_iter(node_vars->free_set, &i, &item, NULL)) {
-            append_list(const Node*, recover_context, item );
-            debugv_print(get_value_name_safe(item));
-            if (i + 1 < recover_context_size)
-                debugv_print(", ");
-        }
-        debugv_print("\n");
+
+    // add_to_recover_context(recover_context, node_vars->free_set, cont);
+    add_to_recover_context(recover_context, node_vars->bound_set, cont);
+    size_t recover_context_size = entries_count_list(recover_context);
+
+    debugv_print("lambda_lift: free (to-be-spilled) variables at '%s' (count=%d): ", name, recover_context_size);
+    for (size_t i = 0; i < recover_context_size; i++) {
+        const Node* item = read_list(const Node*, recover_context)[i];
+        debugv_print(get_value_name_safe(item));
+        if (i + 1 < recover_context_size)
+            debugv_print(", ");
     }
+    debugv_print("\n");
 
     // Create and register new parameters for the lifted continuation
     Nodes new_params = recreate_variables(&ctx->rewriter, oparams);
@@ -106,7 +122,13 @@ static LiftedCont* lambda_lift(Context* ctx, const Node* cont, String given_name
     insert_dict(const Node*, LiftedCont*, ctx->lifted, cont, lifted_cont);
 
     Context lifting_ctx = *ctx;
-    lifting_ctx.rewriter = create_rewriter(ctx->rewriter.src_module, ctx->rewriter.dst_module, (RewriteNodeFn) process_node);
+    // struct Dict* old_map = lifting_ctx.rewriter.map;
+    // lifting_ctx.rewriter.map = clone_dict(lifting_ctx.rewriter.map);
+
+    // lifting_ctx.rewriter = create_rewriter(ctx->rewriter.src_module, ctx->rewriter.dst_module, (RewriteNodeFn) process_node);
+    // lifting_ctx.rewriter.decls_map = NULL;
+    lifting_ctx.rewriter.map = new_dict(const Node*, Node*, (HashFn) hash_node, (CmpFn) compare_node);
+    lifting_ctx.rewriter.parent = &ctx->rewriter;
     register_processed_list(&lifting_ctx.rewriter, oparams, new_params);
 
     const Node* payload = var(a, qualified_type_helper(uint32_type(a), false), "sp");
@@ -140,6 +162,7 @@ static LiftedCont* lambda_lift(Context* ctx, const Node* cont, String given_name
     const Node* substituted = rewrite_node(&lifting_ctx.rewriter, obody);
     //destroy_dict(lifting_ctx.rewriter.processed);
     destroy_rewriter(&lifting_ctx.rewriter);
+    // lifting_ctx.rewriter.map = old_map;
 
     assert(is_terminator(substituted));
     new_fn->payload.fun.body = finish_body(bb, substituted);
@@ -151,27 +174,18 @@ static const Node* process_node(Context* ctx, const Node* node) {
     const Node* found = search_processed(&ctx->rewriter, node);
     if (found) return found;
 
-    // TODO: share this code
-    if (is_declaration(node)) {
-        String name = get_declaration_name(node);
-        Nodes decls = get_module_declarations(ctx->rewriter.dst_module);
-        for (size_t i = 0; i < decls.count; i++) {
-            if (strcmp(get_declaration_name(decls.nodes[i]), name) == 0)
-                return decls.nodes[i];
-        }
-    }
-
     IrArena* a = ctx->rewriter.dst_arena;
 
-    if (ctx->disable_lowering)
-         return recreate_node_identity(&ctx->rewriter, node);
-
-    switch (node->tag) {
+    switch (is_declaration(node)) {
         case Function_TAG: {
+            while (ctx->rewriter.parent)
+                ctx = (Context*) ctx->rewriter.parent;
+
             Context fn_ctx = *ctx;
             fn_ctx.scope = new_scope(node);
             fn_ctx.scope_uses = create_uses_map(node, (NcDeclaration | NcType));
             fn_ctx.scope_vars = compute_scope_variables_map(fn_ctx.scope);
+            fn_ctx.disable_lowering = lookup_annotation(node, "Internal");
             ctx = &fn_ctx;
 
             Node* new = recreate_decl_header_identity(&ctx->rewriter, node);
@@ -182,51 +196,88 @@ static const Node* process_node(Context* ctx, const Node* node) {
             destroy_scope(ctx->scope);
             return new;
         }
+        default:
+            break;
+    }
+
+    if (ctx->disable_lowering)
+         return recreate_node_identity(&ctx->rewriter, node);
+
+    switch (node->tag) {
         case Let_TAG: {
             const Node* oinstruction = get_let_instruction(node);
             if (oinstruction->tag == Control_TAG) {
                 const Node* oinside = oinstruction->payload.control.inside;
                 assert(is_case(oinside));
                 if (!is_control_static(ctx->scope_uses, oinstruction) || ctx->config->hacks.force_join_point_lifting) {
+                    *ctx->todo = true;
+
                     const Node* otail = get_let_tail(node);
                     BodyBuilder* bb = begin_body(a);
                     LiftedCont* lifted_tail = lambda_lift(ctx, otail, unique_name(a, format_string_arena(a->arena, "post_control_%s", get_abstraction_name(ctx->scope->entry->node))));
                     const Node* sp = add_spill_instrs(ctx, bb, lifted_tail->save_values);
                     const Node* tail_ptr = fn_addr_helper(a, lifted_tail->lifted_fn);
 
                     const Node* jp = gen_primop_e(bb, create_joint_point_op, rewrite_nodes(&ctx->rewriter, oinstruction->payload.control.yield_types), mk_nodes(a, tail_ptr, sp));
+                    // dumbass hack
+                    jp = gen_primop_e(bb, subgroup_assume_uniform_op, empty(a), singleton(jp));
 
                     return finish_body(bb, let(a, quote_helper(a, singleton(jp)), rewrite_node(&ctx->rewriter, oinside)));
                 }
             }
-
-            return recreate_node_identity(&ctx->rewriter, node);
+            break;
         }
-        default: return recreate_node_identity(&ctx->rewriter, node);
+        default: break;
     }
+    return recreate_node_identity(&ctx->rewriter, node);
 }
 
 Module* lift_indirect_targets(const CompilerConfig* config, Module* src) {
     ArenaConfig aconfig = get_arena_config(get_module_arena(src));
+    IrArena* a = NULL;
+    Module* dst;
+
+    int round = 0;
+    while (true) {
+        debugv_print("lift_indirect_target: round %d\n", round++);
+        IrArena* oa = a;
+        a = new_ir_arena(aconfig);
+        dst = new_module(a, get_module_name(src));
+        bool todo = false;
+        Context ctx = {
+            .rewriter = create_rewriter(src, dst, (RewriteNodeFn) process_node),
+            .lifted = new_dict(const Node*, LiftedCont*, (HashFn) hash_node, (CmpFn) compare_node),
+            .config = config,
+
+            .todo = &todo
+        };
+
+        rewrite_module(&ctx.rewriter);
+
+        size_t iter = 0;
+        LiftedCont* lifted_cont;
+        while (dict_iter(ctx.lifted, &iter, NULL, &lifted_cont)) {
+            destroy_list(lifted_cont->save_values);
+            free(lifted_cont);
+        }
+        destroy_dict(ctx.lifted);
+        destroy_rewriter(&ctx.rewriter);
+        log_module(DEBUGVV, config, dst);
+        verify_module(config, dst);
+        src = dst;
+        if (oa)
+            destroy_ir_arena(oa);
+        if (!todo) {
+            break;
+        }
+    }
+
     // this will be safe now since we won't lift any more code after this pass
     aconfig.optimisations.weaken_non_leaking_allocas = true;
-    IrArena* a = new_ir_arena(aconfig);
-    Module* dst = new_module(a, get_module_name(src));
-    Context ctx = {
-        .rewriter = create_rewriter(src, dst, (RewriteNodeFn) process_node),
-        .lifted = new_dict(const Node*, LiftedCont*, (HashFn) hash_node, (CmpFn) compare_node),
-        .config = config,
-    };
-
-    rewrite_module(&ctx.rewriter);
-
-    size_t iter = 0;
-    LiftedCont* lifted_cont;
-    while (dict_iter(ctx.lifted, &iter, NULL, &lifted_cont)) {
-        destroy_list(lifted_cont->save_values);
-        free(lifted_cont);
-    }
-    destroy_dict(ctx.lifted);
-    destroy_rewriter(&ctx.rewriter);
+    IrArena* a2 = new_ir_arena(aconfig);
+    dst = new_module(a2, get_module_name(src));
+    Rewriter r = create_importer(src, dst);
+    rewrite_module(&r);
+    destroy_ir_arena(a);
     return dst;
 }
diff --git a/src/shady/rewrite.c b/src/shady/rewrite.c
@@ -33,7 +33,8 @@ Rewriter create_rewriter(Module* src, Module* dst, RewriteNodeFn fn) {
 void destroy_rewriter(Rewriter* r) {
     assert(r->map);
     destroy_dict(r->map);
-    destroy_dict(r->decls_map);
+    if (!r->parent)
+        destroy_dict(r->decls_map);
 }
 
 Rewriter create_importer(Module* src, Module* dst) {
@@ -137,11 +138,27 @@ static Nodes rewrite_ops_helper(Rewriter* rewriter, NodeClass class, String op_n
     return rewrite_nodes_with_fn(rewriter, old_nodes, rewriter->rewrite_fn);
 }
 
+static const Node* search_processed_(const Rewriter* ctx, const Node* old, bool deep) {
+    if (is_declaration(old)) {
+        const Node** found = find_value_dict(const Node*, const Node*, ctx->decls_map, old);
+        return found ? *found : NULL;
+    }
+
+    while (ctx) {
+        assert(ctx->map && "this rewriter has no processed cache");
+        const Node** found = find_value_dict(const Node*, const Node*, ctx->map, old);
+        if (found)
+            return *found;
+        if (deep)
+            ctx = ctx->parent;
+        else
+            ctx = NULL;
+    }
+    return NULL;
+}
+
 const Node* search_processed(const Rewriter* ctx, const Node* old) {
-    struct Dict* map = is_declaration(old) ? ctx->decls_map : ctx->map;
-    assert(map && "this rewriter has no processed cache");
-    const Node** found = find_value_dict(const Node*, const Node*, map, old);
-    return found ? *found : NULL;
+    return search_processed_(ctx, old, false);
 }
 
 const Node* find_processed(const Rewriter* ctx, const Node* old) {
@@ -154,7 +171,7 @@ void register_processed(Rewriter* ctx, const Node* old, const Node* new) {
     assert(old->arena == ctx->src_arena);
     assert(new->arena == ctx->dst_arena);
 #ifndef NDEBUG
-    const Node* found = search_processed(ctx, old);
+    const Node* found = search_processed_(ctx, old, false);
     if (found) {
         error_print("Trying to replace ");
         log_node(ERROR, old);
@@ -190,6 +207,7 @@ bool compare_node(Node**, Node**);
 #include "rewrite_generated.c"
 
 void rewrite_module(Rewriter* rewriter) {
+    assert(rewriter->dst_module != rewriter->src_module);
     Nodes old_decls = get_module_declarations(rewriter->src_module);
     for (size_t i = 0; i < old_decls.count; i++) {
         if (old_decls.nodes[i]->tag == NominalType_TAG) continue;
diff --git a/src/shady/rewrite.h b/src/shady/rewrite.h
@@ -37,6 +37,9 @@ struct Rewriter_ {
         bool fold_quote;
         bool process_variables;
     } config;
+
+    Rewriter* parent;
+
     struct Dict* map;
     struct Dict* decls_map;
 };
