set_memory unable to handle

[ashriram]

I am trying the following example below. Unclear why I am unable to allocate AVX2 register.

exo.memory.MemGenError: /private/tmp/exo/examples/example2.py:37:8: cannot write to buffer 'xyy' in memory 'AVX2'

$ exocc -o out --stem example example2.py
    raise MemGenError(
exo.memory.MemGenError: /private/tmp/exo/examples/example2.py:37:8: cannot write to buffer 'xyy' in memory 'AVX2'

from __future__ import annotations

import os
import sys

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

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

old_split = repeat(divide_loop)
old_unroll = repeat(unroll_loop)


def old_fission_after(proc, stmt_pattern, n_lifts=1):
    def find_stmts(p):
        return [c.after() for c in p.find_all(stmt_pattern)]

    return loop_hack(autofission, find_stmts)(proc, n_lifts)



# Hide output when running through exocc.
if __name__ != "__main__" and hasattr(os, "devnull"):
    sys.stdout = open(os.devnull, "w")



@proc
def simple_math_avx2_sched(
    n: size, x: R[n] @ DRAM, y: R[n] @ DRAM
):  # pragma: no cover
    for i in seq(0, n):
        x[i] = x[i] * y[i] * y[i]
p = rename(simple_math_avx2_sched , "math_avx2_scheduled")
p = old_split(p, "i", 8, ["io", "ii"], tail="cut_and_guard")
p = stage_assn(p, "x[_] = _ #0", "xyy")
p = autolift_alloc(p, "xyy: _", keep_dims=True)
p = set_memory(p, "xyy", AVX2)
p = old_fission_after(p, "xyy[_] = _")

[yamaguchi1024]

Hi @ashriram,

Thanks for trying out Exo!

The code after old_fission_after (last line from your code) is this:

def math_avx2_scheduled(n: size, x: R[n] @ DRAM, y: R[n] @ DRAM):
    for io in seq(0, n / 8):
        xyy: R[8] @ AVX2
        for ii in seq(0, 8):
            xyy[ii] = x[8 * io + ii] * y[8 * io + ii] * y[8 * io + ii]
        for ii in seq(0, 8):
            x[8 * io + ii] = xyy[ii]
    if n % 8 > 0:
        for ii in seq(0, n % 8):
            x[ii + n / 8 *
              8] = x[ii + n / 8 * 8] * y[ii + n / 8 * 8] * y[ii + n / 8 * 8]

You can see that the buffer xyy: R[8] @AVX2 is written and read in for ii loops. You cannot read and write directly from AVX2 registers, so Exo's C codegen is complaining that you need to use AVX instructions if you want to read or write from AVX2 registers.

Therefore, if you continue scheduling this (and use replace to load and store from AVX2 registers), this error will go away. You can ignore this error until you're done with scheduling.

Let us know if anything is unclear!

[ashriram]

Aah; got it.
Thank you