rank_k_reduce_6x16 (exo/examples /avx2_matmul) example Error

[UmeshDeshmukh]

I am following scheduling tutorial for rank_k_reduce_6x16 kernel and after using replace_all schedule with mm256_loadu_ps, I started getting compilation error. To rule out the possibility of typo I checked on https://godbolt.exo-lang.dev/ by removing all the code below line no. 50
The code on godbolt:
`"""
from future import annotations

import os
import sys

from exo import proc
from exo.platforms.x86 import *
from exo.stdlib.scheduling import *

@proc
def rank_k_reduce_6x16(
K: size, A: f32[6, K] @ DRAM, B: f32[K, 16] @ DRAM, C: f32[6, 16] @ DRAM
):
for i in seq(0, 6):
for j in seq(0, 16):
for k in seq(0, K):
C[i, j] += A[i, k] * B[k, j]

print("=============Original Matmul==============")
print(rank_k_reduce_6x16)

avx = rename(rank_k_reduce_6x16, "rank_k_reduce_6x16_scheduled")
avx = reorder_loops(avx, "j k")
avx = reorder_loops(avx, "i k")
avx = divide_loop(avx, "for j in _: ", 8, ["jo", "ji"], perfect=True)
avx = stage_mem(avx, "for k in :", "C[0:6, 0:16]", "C_reg")
avx = simplify(avx)
avx = divide_dim(avx, "C_reg:", 1, 8)
avx = repeat(divide_loop)(avx, "for i1 in _: ", 8, ["i2", "i3"], perfect=True)
avx = simplify(avx)
avx = set_memory(avx, "C_reg:", AVX2)
avx = replace_all(avx, mm256_loadu_ps)
avx = replace_all(avx, mm256_storeu_ps)
avx = simplify(avx)
print("=============Optimized Matmul==============")
print(avx)"""
`

Output from exocc on godbolt:
""" =============Original Matmul============== def rank_k_reduce_6x16(K: size, A: f32[6, K] @ DRAM, B: f32[K, 16] @ DRAM, C: f32[6, 16] @ DRAM): for i in seq(0, 6): for j in seq(0, 16): for k in seq(0, K): C[i, j] += A[i, k] * B[k, j] =============Optimized Matmul============== def rank_k_reduce_6x16_scheduled(K: size, A: f32[6, K] @ DRAM, B: f32[K, 16] @ DRAM, C: f32[6, 16] @ DRAM): C_reg: f32[6, 2, 8] @ AVX2 for i0 in seq(0, 6): for i2 in seq(0, 2): mm256_loadu_ps(C_reg[i0, i2, 0:8], C[i0, 8 * i2:8 + 8 * i2]) for k in seq(0, K): for i in seq(0, 6): for jo in seq(0, 2): for ji in seq(0, 8): C_reg[i, jo, ji] += A[i, k] * B[k, ji + 8 * jo] for i0 in seq(0, 6): for i2 in seq(0, 2): mm256_storeu_ps(C[i0, 8 * i2:8 + 8 * i2], C_reg[i0, i2, 0:8]) Traceback (most recent call last): File "/home/ubuntu/exo/exo/bin/exocc", line 8, in <module> sys.exit(main()) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/main.py", line 61, in main exo.compile_procs(library, outdir, f"{stem}.c", f"{stem}.h") File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/API.py", line 145, in compile_procs c_data, h_data = compile_procs_to_strings(proc_list, h_file) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/API.py", line 153, in compile_procs_to_strings return run_compile([p._loopir_proc for p in proc_list], h_file_name) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 326, in run_compile fwd_decls, body = compile_to_strings(lib_name, proc_list) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 415, in compile_to_strings comp = Compiler(p, ctxt_name, is_public_decl=is_public_decl) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 599, in __init__ self.comp_stmts(self.proc.body) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 658, in comp_stmts self.comp_s(b) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 906, in comp_s self.comp_stmts(s.body) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 658, in comp_stmts self.comp_s(b) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 906, in comp_s self.comp_stmts(s.body) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 658, in comp_stmts self.comp_s(b) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 906, in comp_s self.comp_stmts(s.body) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 658, in comp_stmts self.comp_s(b) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 906, in comp_s self.comp_stmts(s.body) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 658, in comp_stmts self.comp_s(b) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/backend/LoopIR_compiler.py", line 850, in comp_s self.add_line(mem.reduce(s, lhs, rhs)) File "/home/ubuntu/exo/exo/lib/python3.10/site-packages/exo/core/memory.py", line 129, in reduce raise MemGenError( exo.core.memory.MemGenError: <source>:24:16: cannot reduce to buffer 'C_reg' in memory 'AVX2' Compiler returned: 1"""