fix handling of SUBGROUP_SIZE

Author: Hugobros3

src/shady/compile.c

@@ -113,8 +113,7 @@ CompilationResult run_compiler_passes(CompilerConfig* config, Module** pmod) {
     RUN_PASS(lift_indirect_targets)
     RUN_PASS(opt_mem2reg) // run because we can now weaken non-leaking allocas
 
-    if (config->specialization.execution_model != EmNone)
-        RUN_PASS(specialize_execution_model)
+    RUN_PASS(specialize_execution_model)
 
     RUN_PASS(opt_stack)
 

src/shady/passes/specialize_entry_point.c

@@ -50,9 +50,7 @@ static const Node* process(Context* ctx, const Node* node) {
         }
         case Constant_TAG: {
             Node* ncnst = (Node*) recreate_node_identity(&ctx->rewriter, node);
-            if (strcmp(get_declaration_name(ncnst), "SUBGROUP_SIZE") == 0) {
-                ncnst->payload.constant.instruction = quote_helper(a, singleton(uint32_literal(a, ctx->config->specialization.subgroup_size)));
-            } else if (strcmp(get_declaration_name(ncnst), "SUBGROUPS_PER_WG") == 0) {
+            if (strcmp(get_declaration_name(ncnst), "SUBGROUPS_PER_WG") == 0) {
                 // SUBGROUPS_PER_WG = (NUMBER OF INVOCATIONS IN SUBGROUP / SUBGROUP SIZE)
                 // Note: this computations assumes only full subgroups are launched, if subgroups can launch partially filled then this relationship does not hold.
                 uint32_t wg_size[3];
@@ -87,9 +85,6 @@ static const Node* find_entry_point(Module* m, const CompilerConfig* config) {
 }
 
 static void specialize_arena_config(const CompilerConfig* config, Module* src, ArenaConfig* target) {
-    size_t subgroup_size = config->specialization.subgroup_size;
-    assert(subgroup_size);
-
     const Node* old_entry_point_decl = find_entry_point(src, config);
     if (old_entry_point_decl->tag != Function_TAG)
         error("%s is not a function", config->specialization.entry_point);

src/shady/passes/specialize_execution_model.c

@@ -22,6 +22,13 @@ static const Node* process(Context* ctx, const Node* node) {
 
     IrArena* a = ctx->rewriter.dst_arena;
     switch (node->tag) {
+        case Constant_TAG: {
+            Node* ncnst = (Node*) recreate_node_identity(&ctx->rewriter, node);
+            if (strcmp(get_declaration_name(ncnst), "SUBGROUP_SIZE") == 0) {
+                ncnst->payload.constant.instruction = quote_helper(a, singleton(uint32_literal(a, ctx->config->specialization.subgroup_size)));
+            }
+            return ncnst;
+        }
         default: break;
     }
     return recreate_node_identity(&ctx->rewriter, node);
@@ -44,6 +51,9 @@ Module* specialize_execution_model(const CompilerConfig* config, Module* src) {
     IrArena* a = new_ir_arena(aconfig);
     Module* dst = new_module(a, get_module_name(src));
 
+    size_t subgroup_size = config->specialization.subgroup_size;
+    assert(subgroup_size);
+
     Context ctx = {
         .rewriter = create_rewriter(src, dst, (RewriteNodeFn) process),
         .config = config,

src/shady/transform/internal_constants.h

@@ -4,7 +4,7 @@
 #include "shady/ir.h"
 
 #define INTERNAL_CONSTANTS(X) \
-X(SUBGROUP_SIZE, uint32_type(arena), uint32_literal(arena, config->specialization.subgroup_size)) \
+X(SUBGROUP_SIZE, uint32_type(arena), uint32_literal(arena, 64)) \
 X(SUBGROUPS_PER_WG, uint32_type(arena), uint32_literal(arena, 1)) \
 
 void generate_dummy_constants(const CompilerConfig* config, Module*);