Repository: exo-lang/exo

Issue:
Dependent procs should be static
Procs that are included in compilation only because they are called by a "root" proc should be marked `static` in C to promote inlining and reduce code size.

Discussion:

Patch:
diff --git a/src/exo/LoopIR_compiler.py b/src/exo/LoopIR_compiler.py
index 512f67d..e783cba 100644
--- a/src/exo/LoopIR_compiler.py
+++ b/src/exo/LoopIR_compiler.py
@@ -263,12 +263,12 @@ def compile_to_strings(lib_name, proc_list):
     # Header contents
     ctxt_def = _compile_context_struct(find_all_configs(proc_list), ctxt_name)
     struct_defns = set()
-    fwd_decls = []
+    public_fwd_decls = []
 
     # Body contents
     memory_code = _compile_memories(find_all_mems(proc_list))
     builtin_code = _compile_builtins(find_all_builtins(proc_list))
-    priv_decls = []
+    private_fwd_decls = []
     proc_bodies = []
 
     # Compile proc bodies
@@ -289,18 +289,20 @@ def compile_to_strings(lib_name, proc_list):
                 "*/",
             ])
         else:
-            orig_p = id(p)
+            is_public_decl = id(p) in orig_procs
+
             p = PrecisionAnalysis(p).result()
             p = WindowAnalysis(p).result()
             p = MemoryAnalysis(p).result()
-            comp = Compiler(p, ctxt_name)
+            comp = Compiler(p, ctxt_name, is_public_decl=is_public_decl)
             d, b = comp.comp_top()
             struct_defns |= comp.struct_defns()
-            # only dump .h-file forward declarations for requested procedures
-            if orig_p in orig_procs:
-                fwd_decls.append(d)
+
+            if is_public_decl:
+                public_fwd_decls.append(d)
             else:
-                priv_decls.append(d)
+                private_fwd_decls.append(d)
+
             proc_bodies.append(b)
 
     # Structs are just blobs of code... still sort them for output stability
@@ -330,7 +332,7 @@ def compile_to_strings(lib_name, proc_list):
 
 {from_lines(ctxt_def)}
 {from_lines(struct_defns)}
-{from_lines(fwd_decls)}
+{from_lines(public_fwd_decls)}
 '''
 
     body_contents = f'''
@@ -345,7 +347,7 @@ static int8_t _clamp_32to8(int32_t x) {{
 
 {from_lines(memory_code)}
 {from_lines(builtin_code)}
-{from_lines(priv_decls)}
+{from_lines(private_fwd_decls)}
 {from_lines(proc_bodies)}
 '''
 
@@ -395,7 +397,7 @@ def _compile_context_struct(configs, ctxt_name):
 # Loop IR Compiler
 
 class Compiler:
-    def __init__(self, proc, ctxt_name, **kwargs):
+    def __init__(self, proc, ctxt_name, *, is_public_decl):
         assert isinstance(proc, LoopIR.proc)
 
         self.proc = proc
@@ -448,15 +450,17 @@ class Compiler:
 
         self.comp_stmts(self.proc.body)
 
+        static_kwd = '' if is_public_decl else 'static '
+
         # Generate headers here?
         comment = (f"// {name}(\n" +
                    ',\n'.join(['//     ' + s for s in typ_comments]) +
                    '\n'
                    "// )\n")
         proc_decl = (comment +
-                     f"void {name}( {', '.join(arg_strs)} );\n")
+                     f"{static_kwd}void {name}( {', '.join(arg_strs)} );\n")
         proc_def = (comment +
-                    f"void {name}( {', '.join(arg_strs)} ) {{\n" +
+                    f"{static_kwd}void {name}( {', '.join(arg_strs)} ) {{\n" +
                     "\n".join(self._lines) +
                     "\n"
                     "}\n")


Test patch:
diff --git a/tests/golden/test_apps/test_neon_sgemm.txt b/tests/golden/test_apps/test_neon_sgemm.txt
index d8f7e68..0003297 100644
--- a/tests/golden/test_apps/test_neon_sgemm.txt
+++ b/tests/golden/test_apps/test_neon_sgemm.txt
@@ -80,7 +80,7 @@ static int8_t _clamp_32to8(int32_t x) {
 //     B : [f32][K,16]  @DRAM,
 //     C : [f32][4,16]  @DRAM
 // )
-void neon_microkernel( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void neon_microkernel( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 
 /* relying on the following instruction..."
@@ -93,7 +93,7 @@ neon_broadcast_4xf32(dst,src)
 //     B : [f32][K,16]  @DRAM,
 //     C : [f32][4,16]  @DRAM
 // )
-void neon_microkernel( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void neon_microkernel( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
 EXO_ASSUME(B.strides[1] == 1);
diff --git a/tests/golden/test_apps/test_x86_sgemm.txt b/tests/golden/test_apps/test_x86_sgemm.txt
index 9a95a86..8414b87 100644
--- a/tests/golden/test_apps/test_x86_sgemm.txt
+++ b/tests/golden/test_apps/test_x86_sgemm.txt
@@ -84,7 +84,7 @@ static int8_t _clamp_32to8(int32_t x) {
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][M,64]  @DRAM
 // )
-void bottom_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t M, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void bottom_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t M, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // right_panel_kernel_scheduled(
 //     N : size,
@@ -93,7 +93,7 @@ void bottom_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t M, int
 //     B : [f32][K,N]  @DRAM,
 //     C : [f32][6,N]  @DRAM
 // )
-void right_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t N, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void right_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t N, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // sgemm_above_kernel(
 //     M : size,
@@ -103,7 +103,7 @@ void right_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t N, int_
 //     B : [f32][K,N]  @DRAM,
 //     C : [f32][M,N]  @DRAM
 // )
-void sgemm_above_kernel( test_case_Context *ctxt, int_fast32_t M, int_fast32_t N, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void sgemm_above_kernel( test_case_Context *ctxt, int_fast32_t M, int_fast32_t N, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // sgemm_kernel_avx512_1x4(
 //     K : size,
@@ -111,7 +111,7 @@ void sgemm_above_kernel( test_case_Context *ctxt, int_fast32_t M, int_fast32_t N
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][1,64]  @DRAM
 // )
-void sgemm_kernel_avx512_1x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void sgemm_kernel_avx512_1x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // sgemm_kernel_avx512_2x4(
 //     K : size,
@@ -119,7 +119,7 @@ void sgemm_kernel_avx512_1x4( test_case_Context *ctxt, int_fast32_t K, struct ex
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][2,64]  @DRAM
 // )
-void sgemm_kernel_avx512_2x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void sgemm_kernel_avx512_2x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // sgemm_kernel_avx512_3x4(
 //     K : size,
@@ -127,7 +127,7 @@ void sgemm_kernel_avx512_2x4( test_case_Context *ctxt, int_fast32_t K, struct ex
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][3,64]  @DRAM
 // )
-void sgemm_kernel_avx512_3x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void sgemm_kernel_avx512_3x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // sgemm_kernel_avx512_4x4(
 //     K : size,
@@ -135,7 +135,7 @@ void sgemm_kernel_avx512_3x4( test_case_Context *ctxt, int_fast32_t K, struct ex
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][4,64]  @DRAM
 // )
-void sgemm_kernel_avx512_4x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void sgemm_kernel_avx512_4x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // sgemm_kernel_avx512_5x4(
 //     K : size,
@@ -143,7 +143,7 @@ void sgemm_kernel_avx512_4x4( test_case_Context *ctxt, int_fast32_t K, struct ex
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][5,64]  @DRAM
 // )
-void sgemm_kernel_avx512_5x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void sgemm_kernel_avx512_5x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // sgemm_kernel_avx512_6x4(
 //     K : size,
@@ -151,7 +151,7 @@ void sgemm_kernel_avx512_5x4( test_case_Context *ctxt, int_fast32_t K, struct ex
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][6,64]  @DRAM
 // )
-void sgemm_kernel_avx512_6x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
+static void sgemm_kernel_avx512_6x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C );
 
 // bottom_panel_kernel_scheduled(
 //     M : size,
@@ -160,7 +160,7 @@ void sgemm_kernel_avx512_6x4( test_case_Context *ctxt, int_fast32_t K, struct ex
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][M,64]  @DRAM
 // )
-void bottom_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t M, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void bottom_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t M, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(M >= 1);
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
@@ -243,7 +243,7 @@ _mm512_storeu_ps(&{dst_data}, {src_data});
 //     B : [f32][K,N]  @DRAM,
 //     C : [f32][6,N]  @DRAM
 // )
-void right_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t N, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void right_panel_kernel_scheduled( test_case_Context *ctxt, int_fast32_t N, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(N >= 1);
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
@@ -409,7 +409,7 @@ if (((N) / (16)) == 0) {
 //     B : [f32][K,N]  @DRAM,
 //     C : [f32][M,N]  @DRAM
 // )
-void sgemm_above_kernel( test_case_Context *ctxt, int_fast32_t M, int_fast32_t N, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void sgemm_above_kernel( test_case_Context *ctxt, int_fast32_t M, int_fast32_t N, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(M >= 1);
 EXO_ASSUME(N >= 1);
 EXO_ASSUME(K >= 1);
@@ -575,7 +575,7 @@ if (K % 512 > 0) {
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][1,64]  @DRAM
 // )
-void sgemm_kernel_avx512_1x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void sgemm_kernel_avx512_1x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
 EXO_ASSUME(B.strides[1] == 1);
@@ -610,7 +610,7 @@ for (int i = 0; i < 1; i++) {
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][2,64]  @DRAM
 // )
-void sgemm_kernel_avx512_2x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void sgemm_kernel_avx512_2x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
 EXO_ASSUME(B.strides[1] == 1);
@@ -645,7 +645,7 @@ for (int i = 0; i < 2; i++) {
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][3,64]  @DRAM
 // )
-void sgemm_kernel_avx512_3x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void sgemm_kernel_avx512_3x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
 EXO_ASSUME(B.strides[1] == 1);
@@ -680,7 +680,7 @@ for (int i = 0; i < 3; i++) {
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][4,64]  @DRAM
 // )
-void sgemm_kernel_avx512_4x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void sgemm_kernel_avx512_4x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
 EXO_ASSUME(B.strides[1] == 1);
@@ -715,7 +715,7 @@ for (int i = 0; i < 4; i++) {
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][5,64]  @DRAM
 // )
-void sgemm_kernel_avx512_5x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void sgemm_kernel_avx512_5x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
 EXO_ASSUME(B.strides[1] == 1);
@@ -750,7 +750,7 @@ for (int i = 0; i < 5; i++) {
 //     B : [f32][K,64]  @DRAM,
 //     C : [f32][6,64]  @DRAM
 // )
-void sgemm_kernel_avx512_6x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
+static void sgemm_kernel_avx512_6x4( test_case_Context *ctxt, int_fast32_t K, struct exo_win_2f32 A, struct exo_win_2f32 B, struct exo_win_2f32 C ) {
 EXO_ASSUME(K >= 1);
 EXO_ASSUME(A.strides[1] == 1);
 EXO_ASSUME(B.strides[1] == 1);
diff --git a/tests/golden/test_precision/test_good_prec2.txt b/tests/golden/test_precision/test_good_prec2.txt
index 99983e8..bafdb65 100644
--- a/tests/golden/test_precision/test_good_prec2.txt
+++ b/tests/golden/test_precision/test_good_prec2.txt
@@ -52,7 +52,7 @@ static int8_t _clamp_32to8(int32_t x) {
 //     y : f32[m]  @DRAM,
 //     r : f32  @DRAM
 // )
-void dot( c_code_str_Context *ctxt, int_fast32_t m, float* x, float* y, float* r );
+static void dot( c_code_str_Context *ctxt, int_fast32_t m, float* x, float* y, float* r );
 
 // dot(
 //     m : size,
@@ -60,7 +60,7 @@ void dot( c_code_str_Context *ctxt, int_fast32_t m, float* x, float* y, float* r
 //     y : f32[m]  @DRAM,
 //     r : f32  @DRAM
 // )
-void dot( c_code_str_Context *ctxt, int_fast32_t m, float* x, float* y, float* r ) {
+static void dot( c_code_str_Context *ctxt, int_fast32_t m, float* x, float* y, float* r ) {
 *r = 0.0;
 for (int i = 0; i < m; i++) {
   *r += x[(i) * (1)] * y[(i) * (1)];
diff --git a/tests/golden/test_window/test_normalize.txt b/tests/golden/test_window/test_normalize.txt
index a7a829d..61bb3a1 100644
--- a/tests/golden/test_window/test_normalize.txt
+++ b/tests/golden/test_window/test_normalize.txt
@@ -56,7 +56,7 @@ static int8_t _clamp_32to8(int32_t x) {
 //     y : [f32][m]  @DRAM,
 //     r : f32  @DRAM
 // )
-void dot( c_code_str_Context *ctxt, int_fast32_t m, struct exo_win_1f32 x, struct exo_win_1f32 y, float* r );
+static void dot( c_code_str_Context *ctxt, int_fast32_t m, struct exo_win_1f32 x, struct exo_win_1f32 y, float* r );
 
 // dot(
 //     m : size,
@@ -64,7 +64,7 @@ void dot( c_code_str_Context *ctxt, int_fast32_t m, struct exo_win_1f32 x, struc
 //     y : [f32][m]  @DRAM,
 //     r : f32  @DRAM
 // )
-void dot( c_code_str_Context *ctxt, int_fast32_t m, struct exo_win_1f32 x, struct exo_win_1f32 y, float* r ) {
+static void dot( c_code_str_Context *ctxt, int_fast32_t m, struct exo_win_1f32 x, struct exo_win_1f32 y, float* r ) {
 *r = 0.0;
 for (int i = 0; i < m; i++) {
   *r += x.data[(i) * (x.strides[0])] * y.data[(i) * (y.strides[0])];