Skip to content

Update mlir-aie to v1.4.0 and migrate the operator library to its APIs #145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
atassis wants to merge 1 commit into
base: devel
Choose a base branch
from
Open

Update mlir-aie to v1.4.0 and migrate the operator library to its APIs #145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 5 additions & 7 deletions iron/common/compilation/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -518,7 +518,6 @@ def compile(self, graph):
str(self.aiecc_path),
"-v",
f"-j{os.environ.get('AIECC_JOBS', '1')}",
"--no-compile-host",
]
if self.use_chess:
compile_cmd += [
Expand All @@ -534,7 +533,7 @@ def compile(self, graph):
]
compile_cmd += [
"--expand-load-pdis",
"--generate-full-elf",
"--get-full-elf",
"--full-elf-name",
os.path.abspath(artifact.filename),
*artifact.extra_flags,
Expand Down Expand Up @@ -573,7 +572,6 @@ def compile(self, graph):
str(self.aiecc_path),
"-v",
f"-j{os.environ.get('AIECC_JOBS', '1')}",
"--no-compile-host",
]
if self.use_chess:
compile_cmd += [
Expand All @@ -597,7 +595,7 @@ def compile(self, graph):
0
] # TODO: this does not handle the case of multiple xclbins with different kernel names or flags from the same MLIR
compile_cmd += first_xclbin.extra_flags + [
"--aie-generate-xclbin",
"--get-xclbin",
"--xclbin-name=" + os.path.abspath(first_xclbin.filename),
"--xclbin-kernel-name=" + first_xclbin.kernel_name,
]
Expand All @@ -610,10 +608,10 @@ def compile(self, graph):
first_insts_bin = mlir_sources_to_insts[mlir_source][
0
] # TODO: this does not handle the case of multiple insts.bins with different flags from the same MLIR
if not do_compile_xclbin:
compile_cmd += ["--no-compile"]
# Outputs are selected by --get-<name>; asking only for the insts is what
# "--no-compile" used to mean, so there is nothing to opt out of here.
compile_cmd += first_insts_bin.extra_flags + [
"--aie-generate-npu-insts",
"--get-npu-insts",
"--npu-insts-name=" + os.path.abspath(first_insts_bin.filename),
]
compile_cmd += [os.path.abspath(mlir_source.filename)]
Expand Down
10 changes: 4 additions & 6 deletions iron/operators/_trace.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,13 +39,11 @@ def _default_coretile_events():
]


def maybe_enable_trace(rt, trace_size, workers, coretile_events=None):
"""Configure per-op hardware trace on ``rt`` if tracing is requested.

Call inside the ``rt.sequence(...)`` block, before ``rt.start(...)``.
def maybe_enable_trace(prog, trace_size, workers, coretile_events=None):
"""Configure per-op hardware trace if tracing is requested.

Args:
rt: the ``Runtime`` being built.
prog: the ``Program`` being built.
trace_size: the design's ``trace_size`` argument (may be None/0).
workers: the design's workers; the first ``IRON_TRACE_NTILES`` are traced.
coretile_events: override the default core-tile event set.
Expand All @@ -61,7 +59,7 @@ def maybe_enable_trace(rt, trace_size, workers, coretile_events=None):
# meaningless (a negative slice index would silently drop the LAST worker).
ntiles = max(0, int(os.environ.get("IRON_TRACE_NTILES", "1")))

rt.enable_trace(
prog.enable_trace(
ts,
workers=list(workers)[:ntiles],
coretile_events=(
Expand Down
43 changes: 25 additions & 18 deletions iron/operators/axpy/design.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from ml_dtypes import bfloat16
import numpy as np

from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker
from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
from aie.helpers.taplib.tap import TensorAccessPattern
from aie.iron.controlflow import range_
from iron.operators._trace import maybe_enable_trace
Expand Down Expand Up @@ -84,38 +84,45 @@ def core_body(of_in1, of_in2, of_out, axpy):
]

# Runtime operations to move data to/from the AIE-array
rt = Runtime()
with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C):
maybe_enable_trace(rt, trace_size, my_workers)
rt.start(*my_workers)

def sequence(A, B, C, in1_prods, in2_prods, out_conses):
# Initialize a group for parallel drain tasks, with fill resources free'd when drains complete.
tg = rt.task_group()
tg = TaskGroup()

# Fill the input objectFIFOs with data
for i in range(num_columns):
rt.fill(
of_in1s[i].prod(),
in1_prods[i].fill(
A,
taps[i],
task_group=tg,
group=tg,
)
rt.fill(
of_in2s[i].prod(),
in2_prods[i].fill(
B,
taps[i],
task_group=tg,
group=tg,
)
# Drain the output objectFIFOs with data
for i in range(num_columns):
rt.drain(
of_outs[i].cons(),
out_conses[i].drain(
C,
taps[i],
wait=True, # wait for the transfer to complete and data to be available
task_group=tg,
group=tg,
)
rt.finish_task_group(tg)
tg.finish()

rt = Runtime(
sequence,
[
tensor_ty,
tensor_ty,
tensor_ty,
[of_in1s[i].prod() for i in range(num_columns)],
[of_in2s[i].prod() for i in range(num_columns)],
[of_outs[i].cons() for i in range(num_columns)],
],
)

# Place program components (assign them resources on the device) and generate an MLIR module
return Program(dev, rt).resolve_program()
prog = Program(dev, rt, workers=my_workers)
maybe_enable_trace(prog, trace_size, my_workers)
return prog.resolve_program()
43 changes: 25 additions & 18 deletions iron/operators/binary_elementwise_design.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from ml_dtypes import bfloat16
import numpy as np

from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker
from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
from aie.helpers.taplib.tap import TensorAccessPattern
from aie.iron.controlflow import range_
from iron.operators._trace import maybe_enable_trace
Expand Down Expand Up @@ -83,37 +83,44 @@ def core_body(of_in1, of_in2, of_out, eltwise_fn):
]

# Runtime operations to move data to/from the AIE-array
rt = Runtime()
with rt.sequence(tensor_ty, tensor_ty, tensor_ty) as (A, B, C):
maybe_enable_trace(rt, trace_size, my_workers)
rt.start(*my_workers)

tg = rt.task_group()
def sequence(A, B, C, in1_prods, in2_prods, out_conses):
tg = TaskGroup()

# Fill the input objectFIFOs with data
for i in range(num_columns):
rt.fill(
of_in1s[i].prod(),
in1_prods[i].fill(
A,
taps[i],
task_group=tg,
group=tg,
)
rt.fill(
of_in2s[i].prod(),
in2_prods[i].fill(
B,
taps[i],
task_group=tg,
group=tg,
)
# Drain the output objectFIFOs with data
for i in range(num_columns):
rt.drain(
of_outs[i].cons(),
out_conses[i].drain(
C,
taps[i],
wait=True,
task_group=tg,
group=tg,
)
rt.finish_task_group(tg)
tg.finish()

rt = Runtime(
sequence,
[
tensor_ty,
tensor_ty,
tensor_ty,
[of_in1s[i].prod() for i in range(num_columns)],
[of_in2s[i].prod() for i in range(num_columns)],
[of_outs[i].cons() for i in range(num_columns)],
],
)

# Place program components and generate an MLIR module
return Program(dev, rt).resolve_program()
prog = Program(dev, rt, workers=my_workers)
maybe_enable_trace(prog, trace_size, my_workers)
return prog.resolve_program()
36 changes: 21 additions & 15 deletions iron/operators/channeled_unary_design.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from ml_dtypes import bfloat16
import numpy as np

from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker
from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
from aie.helpers.taplib.tap import TensorAccessPattern
from aie.iron.controlflow import range_
from iron.operators._trace import maybe_enable_trace
Expand Down Expand Up @@ -97,33 +97,39 @@ def core_fn(of_in, of_out, kernel_line):
]

# Runtime operations to move data to/from the AIE-array
rt = Runtime()
with rt.sequence(transfer_type, transfer_type) as (a_in, b_out):
maybe_enable_trace(rt, trace_size, my_workers)
rt.start(*my_workers)

tg = rt.task_group()
def sequence(a_in, b_out, in_prods, out_conses):
tg = TaskGroup()

# Fill the input objectFIFOs with data
for i in range(num_columns):
for j in range(num_channels):
rt.fill(
of_ins[i * num_channels + j].prod(),
in_prods[i * num_channels + j].fill(
a_in,
taps[i * num_channels + j],
task_group=tg,
group=tg,
)
# Drain the output objectFIFOs with data
for i in range(num_columns):
for j in range(num_channels):
rt.drain(
of_outs[i * num_channels + j].cons(),
out_conses[i * num_channels + j].drain(
b_out,
taps[i * num_channels + j],
wait=True,
task_group=tg,
group=tg,
)
rt.finish_task_group(tg)
tg.finish()

rt = Runtime(
sequence,
[
transfer_type,
transfer_type,
[of.prod() for of in of_ins],
[of.cons() for of in of_outs],
],
)

# Place components and generate an MLIR module
return Program(dev, rt).resolve_program()
prog = Program(dev, rt, workers=my_workers)
maybe_enable_trace(prog, trace_size, my_workers)
return prog.resolve_program()
38 changes: 22 additions & 16 deletions iron/operators/dequant/design.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from ml_dtypes import bfloat16
import numpy as np

from aie.iron import Kernel, ObjectFifo, Program, Runtime, Worker
from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker
from aie.helpers.taplib.tap import TensorAccessPattern
from aie.iron.controlflow import range_

Expand Down Expand Up @@ -118,35 +118,41 @@ def core_body(of_in1, of_out, dequant_kernel):
]

# Runtime operations to move data to/from the AIE-array
rt = Runtime()
with rt.sequence(in_tensor_ty, out_tensor_ty) as (A, C):
if enable_trace:
rt.enable_trace(trace_size)
rt.start(*my_workers)
def sequence(A, C, of_in1s_prods, of_outs_conss):

# Initialize a group for parallel drain tasks, with fill resources free'd when drains complete.
tg = rt.task_group()
tg = TaskGroup()

# Fill the input objectFIFOs with data
for i in range(num_columns):
for j in range(num_channels):
rt.fill(
of_in1s[i * num_channels + j].prod(),
of_in1s_prods[i * num_channels + j].fill(
A,
taps_in[i * num_channels + j],
task_group=tg,
group=tg,
)
# Drain the output objectFIFOs with data
for i in range(num_columns):
for j in range(num_channels):
rt.drain(
of_outs[i * num_channels + j].cons(),
of_outs_conss[i * num_channels + j].drain(
C,
taps_out[i * num_channels + j],
wait=True, # wait for the transfer to complete and data to be available
task_group=tg,
group=tg,
)
rt.finish_task_group(tg)

tg.finish()

rt = Runtime(
sequence,
[
in_tensor_ty,
out_tensor_ty,
[of.prod() for of in of_in1s],
[of.cons() for of in of_outs],
],
)
# Place program components (assign them resources on the device) and generate an MLIR module
return Program(dev, rt).resolve_program()
prog = Program(dev, rt, workers=my_workers)
if enable_trace:
prog.enable_trace(trace_size)
return prog.resolve_program()
Loading
Loading