Skip to content
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

feat[next]: Extend DaCe support for offset providers #1353

Merged
merged 4 commits into from
Oct 30, 2023
Merged

feat[next]: Extend DaCe support for offset providers #1353

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b20ec2c
[dace] Extend support for offset providers
edopao Oct 20, 2023
289cd75
[dace] Cleanup shift visitor
edopao Oct 25, 2023
61c538e
[dace] Minor edit
edopao Oct 26, 2023
77a2010
[dace] Review comments
edopao Oct 30, 2023
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
1 change: 0 additions & 1 deletion src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,6 @@ def visit_StencilClosure(
self, node: itir.StencilClosure, array_table: dict[str, dace.data.Array]
) -> tuple[dace.SDFG, list[str], list[str]]:
assert ItirToSDFG._check_no_lifts(node)
assert ItirToSDFG._check_shift_offsets_are_literals(node)

# Create the closure's nested SDFG and single state.
closure_sdfg = dace.SDFG(name="closure")
Expand Down
114 changes: 47 additions & 67 deletions src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -478,17 +478,7 @@ def visit_FunCall(self, node: itir.FunCall) -> list[ValueExpr] | IteratorExpr:
return self._visit_deref(node)
if isinstance(node.fun, itir.FunCall) and isinstance(node.fun.fun, itir.SymRef):
if node.fun.fun.id == "shift":
offset = node.fun.args[0]
assert isinstance(offset, itir.OffsetLiteral)
offset_name = offset.value
assert isinstance(offset_name, str)
if offset_name not in self.offset_provider:
raise ValueError(f"offset provider for `{offset_name}` is missing")
offset_provider = self.offset_provider[offset_name]
if isinstance(offset_provider, Dimension):
return self._visit_direct_addressing(node)
else:
return self._visit_indirect_addressing(node)
return self._visit_shift(node)
elif node.fun.fun.id == "reduce":
return self._visit_reduce(node)

Expand Down Expand Up @@ -653,39 +643,7 @@ def _make_shift_for_rest(self, rest, iterator):
fun=itir.FunCall(fun=itir.SymRef(id="shift"), args=rest), args=[iterator]
)

def _visit_direct_addressing(self, node: itir.FunCall) -> IteratorExpr:
assert isinstance(node.fun, itir.FunCall)
shift = node.fun
assert isinstance(shift, itir.FunCall)

tail, rest = self._split_shift_args(shift.args)
if rest:
iterator = self.visit(self._make_shift_for_rest(rest, node.args[0]))
else:
iterator = self.visit(node.args[0])

assert isinstance(tail[0], itir.OffsetLiteral)
offset = tail[0].value
assert isinstance(offset, str)
shifted_dim = self.offset_provider[offset].value

assert isinstance(tail[1], itir.OffsetLiteral)
shift_amount = tail[1].value
assert isinstance(shift_amount, int)

args = [ValueExpr(iterator.indices[shifted_dim], dace.int64)]
internals = [f"{arg.value.data}_v" for arg in args]
expr = f"{internals[0]} + {shift_amount}"
shifted_value = self.add_expr_tasklet(
list(zip(args, internals)), expr, dace.dtypes.int64, "dir_addr"
)[0].value

shifted_index = {dim: value for dim, value in iterator.indices.items()}
shifted_index[shifted_dim] = shifted_value

return IteratorExpr(iterator.field, shifted_index, iterator.dtype, iterator.dimensions)

def _visit_indirect_addressing(self, node: itir.FunCall) -> IteratorExpr:
def _visit_shift(self, node: itir.FunCall) -> IteratorExpr:
shift = node.fun
assert isinstance(shift, itir.FunCall)
tail, rest = self._split_shift_args(shift.args)
Expand All @@ -695,40 +653,48 @@ def _visit_indirect_addressing(self, node: itir.FunCall) -> IteratorExpr:
iterator = self.visit(node.args[0])

assert isinstance(tail[0], itir.OffsetLiteral)
offset = tail[0].value
assert isinstance(offset, str)
offset_dim = tail[0].value
assert isinstance(offset_dim, str)
offset_node = self.visit(tail[1])[0]

assert isinstance(tail[1], itir.OffsetLiteral)
element = tail[1].value
assert isinstance(element, int)

if isinstance(self.offset_provider[offset], NeighborTableOffsetProvider):
table = self.offset_provider[offset]
shifted_dim = table.origin_axis.value
target_dim = table.neighbor_axis.value

conn = self.context.state.add_access(connectivity_identifier(offset))
if isinstance(self.offset_provider[offset_dim], NeighborTableOffsetProvider):
offset_provider = self.offset_provider[offset_dim]
connectivity = self.context.state.add_access(connectivity_identifier(offset_dim))

shifted_dim = offset_provider.origin_axis.value
target_dim = offset_provider.neighbor_axis.value
args = [
ValueExpr(conn, table.table.dtype),
ValueExpr(connectivity, offset_provider.table.dtype),
ValueExpr(iterator.indices[shifted_dim], dace.int64),
offset_node,
]

internals = [f"{arg.value.data}_v" for arg in args]
expr = f"{internals[0]}[{internals[1]}, {element}]"
else:
offset_provider = self.offset_provider[offset]
assert isinstance(offset_provider, StridedNeighborOffsetProvider)
expr = f"{internals[0]}[{internals[1]}, {internals[2]}]"
elif isinstance(self.offset_provider[offset_dim], StridedNeighborOffsetProvider):
offset_provider = self.offset_provider[offset_dim]

shifted_dim = offset_provider.origin_axis.value
target_dim = offset_provider.neighbor_axis.value
offset_value = iterator.indices[shifted_dim]
args = [ValueExpr(offset_value, dace.int64)]
internals = [f"{offset_value.data}_v"]
expr = f"{internals[0]} * {offset_provider.max_neighbors} + {element}"
args = [
ValueExpr(iterator.indices[shifted_dim], dace.int64),
offset_node,
]
internals = [f"{arg.value.data}_v" for arg in args]
expr = f"{internals[0]} * {offset_provider.max_neighbors} + {internals[1]}"
else:
assert isinstance(self.offset_provider[offset_dim], Dimension)

shifted_dim = self.offset_provider[offset_dim].value
target_dim = shifted_dim
args = [
ValueExpr(iterator.indices[shifted_dim], dace.int64),
offset_node,
]
internals = [f"{arg.value.data}_v" for arg in args]
expr = f"{internals[0]} + {internals[1]}"

shifted_value = self.add_expr_tasklet(
list(zip(args, internals)), expr, dace.dtypes.int64, "ind_addr"
list(zip(args, internals)), expr, dace.dtypes.int64, "shift"
)[0].value

shifted_index = {dim: value for dim, value in iterator.indices.items()}
Expand All @@ -737,6 +703,20 @@ def _visit_indirect_addressing(self, node: itir.FunCall) -> IteratorExpr:

return IteratorExpr(iterator.field, shifted_index, iterator.dtype, iterator.dimensions)

def visit_OffsetLiteral(self, node: itir.OffsetLiteral) -> list[ValueExpr]:
offset = node.value
assert isinstance(offset, int)
offset_var = unique_var_name()
self.context.body.add_scalar(offset_var, dace.dtypes.int64, transient=True)
offset_node = self.context.state.add_access(offset_var)
tasklet_node = self.context.state.add_tasklet(
"get_offset", {}, {"__out"}, f"__out = {offset}"
)
self.context.state.add_edge(
tasklet_node, "__out", offset_node, None, dace.Memlet.simple(offset_var, "0")
)
return [ValueExpr(offset_node, self.context.body.arrays[offset_var].dtype)]

def _visit_reduce(self, node: itir.FunCall):
result_name = unique_var_name()
result_access = self.context.state.add_access(result_name)
Expand Down