{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Transformations\n",
"\n",
"Apart from modifying graphs manually using the SDFG API, you can also define more complex behavior that matches certain patterns in SDFGs and makes changes. This behavior is expressed in classes called `Transformation`s. Transformations are a powerful tool to optimize applications in DaCe. You can go from naive code to state-of-the-art performance using only transformations.\n",
"\n",
"There are two general types of transformations: **pattern-matching transformations** (extending the `Transformation` class) and **subgraph transformations** (extending the `SubgraphTransformation` class). The former is based on one or more specific subgraph expressions, and the latter can be applied to any subgraph that passes the conditions. Internally, there are two methods to a Transformation: `can_be_applied`, which returns True if it can be applied on a given subgraph; and `apply`, which modifies the graph using the SDFG API. Some transformations run automatically on each graph (as part of the simplification pass), and others have to be called manually. \n",
"\n",
"You can find a list of the built-in standard transformations [here](https://spcldace.readthedocs.io/en/latest/source/dace.transformation.html). While these transformations cover many use-cases, they cannot cover everything (for example, domain-specific behavior). Thus, Transformations are easily extensible and [below](#Creating-New-Transformations) we show how to register a new one. You can register your transformations to be pattern-based, subgraph-based, or not, upon defining a new class.\n",
"\n",
"This tutorial will deal with the different ways transformations can be applied on code:\n",
"\n",
"* [Apply anywhere or everywhere](#Applying-Transformations)\n",
"* [Enumerate matches of a transformation](#Enumerating-Matches)\n",
"* [Apply at a specific location](#Apply-to-Specific-Location)\n",
"* [Interactive optimization](#Interactive-Optimization) ([command-line](#Command-line-interface), [Visual Studio Code](#Visual-Studio-Code))\n",
"* [Create your own transformation](#Creating-New-Transformations)\n",
"\n",
"We will use the following example throughout this tutorial:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"import dace\n",
"@dace.function\n",
"def dbladd(A: dace.float64[1000, 1000], B: dace.float64[1000, 1000]):\n",
" dbl = B\n",
" return A + dbl * B"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Applying Transformations\n",
"\n",
"The easiest way to apply a transformation is to import the Transformation subclass and call `apply_transformations` or `apply_transformations_repeated` on the SDFG. The methods accept a transformation or a list of transformations, their parameters, and other parameters (for more info, see its [documentation](https://spcldace.readthedocs.io/en/latest/source/dace.sdfg.html#dace.sdfg.sdfg.SDFG.apply_transformations)).\n",
"\n",
"To demonstrate transformations, we will first show the raw SDFG of `dbladd`, without simplification:"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (dbladd)"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sdfg = dbladd.to_sdfg(simplify=False)\n",
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"There are four states in this SDFG, and an unnecessary copy to `dbl`. In order to fuse the states, we can apply the `StateFusion` transformation:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (dbladd)"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from dace.transformation.interstate import StateFusion\n",
"sdfg.apply_transformations(StateFusion)\n",
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This fused the first two states. Since we want to fuse the entire graph, we can use the method that repeatedly applies the transformation until no more states can be fused without breaking the correctness of the graph:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (dbladd)"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sdfg.apply_transformations_repeated(StateFusion)\n",
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that the dataflow aspects of the graph are clearer, it is easy to see that the `dbl` array is redundant. One of the simplification pass transformations in the DaCe standard library deals with redundant array copies when one array is transient and unused anywhere else. To invoke this transformation, we can either import it directly, or simply try to apply all simplification transformations. Note that this happens automatically when a Python function is defined without our special `simplify=False` argument:"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (dbladd)"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sdfg.simplify()\n",
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Enumerating Matches\n",
"\n",
"As graphs grow larger, there will be more than one match to a transformation. For pattern matching transformations, we can enumerate the matching subgraphs in an SDFG for a given transformation using the `dace.transformation.optimizer.Optimizer` class.\n",
"\n",
"In the example below, we try to find which maps can be tiled (to increase the locality of the computation) within the current SDFG:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Match: MapTiling in [MapEntry (_Mult__map[__i0=0:1000, __i1=0:1000])]\n",
"Match: MapTiling in [MapEntry (_Add__map[__i0=0:1000, __i1=0:1000])]\n"
]
}
],
"source": [
"from dace.transformation.dataflow import MapTiling\n",
"from dace.transformation.optimizer import Optimizer\n",
"\n",
"for xform in Optimizer(sdfg).get_pattern_matches(patterns=[MapTiling]):\n",
" print('Match:', xform.print_match(sdfg))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The transformation can then be applied by calling `xform.apply(graph, sdfg)`, where `graph` may be the SDFG state in which the pattern is found, or the SDFG (for multi-state transformations)."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Custom Subgraph Enumeration\n",
"\n",
"If you want to match a certain subgraph (in the SDFG or within a state) without creating a transformation, you can use the `enumerate_matches` API in `dace.transformation.pattern_matching`. It uses the same mechanism as transformations but can be used for general pattern matching:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Match found in state BinOp_5 . Nodes: [MapExit (_Mult__map[__i0=0:1000, __i1=0:1000]), AccessNode (__tmp0), MapEntry (_Add__map[__i0=0:1000, __i1=0:1000])]\n"
]
}
],
"source": [
"from dace.transformation.passes.pattern_matching import enumerate_matches\n",
"from dace.sdfg import utils as sdutil # For creating path graphs\n",
"\n",
"# Construct subgraph pattern (MapExit -> AccessNode -> MapEntry)\n",
"pattern = sdutil.node_path_graph(dace.nodes.MapExit, dace.nodes.AccessNode,\n",
" dace.nodes.MapEntry)\n",
"\n",
"# Loop over matches\n",
"for subgraph in enumerate_matches(sdfg, pattern):\n",
" print(\"Match found in state\", subgraph.graph.label, \". Nodes:\", subgraph.nodes())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Apply to Specific Location\n",
"\n",
"We can also invoke transformations at specific locations using the `apply_to` static function of each transformation. In each transformation class, a list of statically constructed nodes define the transformation structure. For example, in `MapFusion`, there are three such nodes, called `first_map_exit`, `array`, and `second_map_entry`. If you know which maps to apply the transformation to, simply use these three names as keyword arguments to the `apply_to` function. Below is an example that finds the multiplication map exit, addition map entry, and the array between them:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"_Mult__map[__i0=0:1000, __i1=0:1000] -> __tmp0 -> _Add__map[__i0=0:1000, __i1=0:1000]\n"
]
}
],
"source": [
"# Since there is only one state (thanks to StateFusion), we can use the first one in the SDFG\n",
"state = sdfg.node(0)\n",
"\n",
"# The multiplication map is called \"_Mult__map\" (see above graph), we can query it\n",
"mult_exit = next(n for n in state.nodes() if isinstance(n, dace.nodes.MapExit) and n.label == '_Mult__map')\n",
"# Same goes for the addition entry\n",
"add_entry = next(n for n in state.nodes() if isinstance(n, dace.nodes.MapEntry) and n.label == '_Add__map')\n",
"# Since all redundant arrays have been removed by simplification, we can get the only transient\n",
"# array that remains in the graph\n",
"transient = next(aname for aname, desc in sdfg.arrays.items() if desc.transient)\n",
"access_node = next(n for n in state.nodes() if isinstance(n, dace.nodes.AccessNode) and n.data == transient)\n",
"\n",
"# We will apply the transformation on these three nodes\n",
"print(mult_exit, '->', access_node, '->', add_entry)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {
"scrolled": true
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (dbladd)"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from dace.transformation.dataflow import MapFusion\n",
"\n",
"MapFusion.apply_to(sdfg,\n",
" first_map_exit=mult_exit,\n",
" array=access_node,\n",
" second_map_entry=add_entry)\n",
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"and now the two maps are fused.\n",
"\n",
"### Subgraph Transformations\n",
"\n",
"The same can be applied with subgraph transformations, but a subgraph is required instead of a pattern. In the following example we reload the SDFG (in the default mode, namely with SDFG simplification), and then apply `SubgraphFusion` on the entire state (namely all of its nodes, or `state.nodes()`). The result will be similar to fusing the two maps as above, but can generalize to fuse any number of maps, parallel regions, and other cases."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (dbladd)"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from dace.transformation.subgraph import SubgraphFusion\n",
"\n",
"sdfg = dbladd.to_sdfg()\n",
"\n",
"# Single-state SDFG\n",
"state = sdfg.node(0)\n",
"\n",
"SubgraphFusion.apply_to(sdfg, state.nodes())\n",
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Interactive Optimization\n",
"\n",
"Sometimes it is useful to apply transformations in sequence interactively. To open the prompt from Jupyter or Python, call `SDFGOptimizer`:\n"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"0. Transformation ElementWiseArrayOperation in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"1. Transformation ElementWiseArrayOperation2D in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"2. Transformation FPGATransformSDFG in []\n",
"3. Transformation FPGATransformState in [SDFGState (BinOp_5)]\n",
"4. Transformation GPUTransformLocalStorage in outer_fused[__i0=0:1000, __i1=0:1000]\n",
"5. Transformation GPUTransformMap in outer_fused[__i0=0:1000, __i1=0:1000]\n",
"6. Transformation GPUTransformSDFG in []\n",
"7. Transformation MapDimShuffle in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"8. Transformation MapExpansion in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"9. Transformation MapFission in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"10. Transformation MapTiling in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"11. Transformation MapTilingWithOverlap in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"12. Transformation MapUnroll in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"13. Transformation NestSDFG in []\n",
"14. Transformation ReductionNOperation in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])]\n",
"15. Transformation StripMining in outer_fused: ['__i0', '__i1']\n",
"16. Transformation TaskletFusion in [Tasklet (_Mult_), AccessNode (__tmp0), Tasklet (_Add_)]\n"
]
},
{
"name": "stdin",
"output_type": "stream",
"text": [
"Select the pattern to apply (0 - 16 or name$id): MapExpansion$0\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"You selected (MapExpansion$0) pattern MapExpansion in [MapEntry (outer_fused[__i0=0:1000, __i1=0:1000])] with parameters {}\n",
"0. Transformation ElementWiseArrayOperation in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"1. Transformation FPGATransformSDFG in []\n",
"2. Transformation FPGATransformState in [SDFGState (BinOp_5)]\n",
"3. Transformation GPUGridStridedTiling in [MapEntry (outer_fused[__i0=0:1000]), MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"4. Transformation GPUTransformLocalStorage in outer_fused[__i0=0:1000]\n",
"5. Transformation GPUTransformMap in outer_fused[__i0=0:1000]\n",
"6. Transformation GPUTransformMap in outer_fused___i1[__i1=0:1000]\n",
"7. Transformation GPUTransformSDFG in []\n",
"8. Transformation InLocalStorage in outer_fused[__i0=0:1000] -> outer_fused___i1[__i1=0:1000]\n",
"9. Transformation MPITransformMap in [MapEntry (outer_fused[__i0=0:1000])]\n",
"10. Transformation MPITransformMap in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"11. Transformation MapDimShuffle in [MapEntry (outer_fused[__i0=0:1000])]\n",
"12. Transformation MapDimShuffle in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"13. Transformation MapFission in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"14. Transformation MapInterchange in [MapEntry (outer_fused[__i0=0:1000]), MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"15. Transformation MapTiling in [MapEntry (outer_fused[__i0=0:1000])]\n",
"16. Transformation MapTiling in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"17. Transformation MapTilingWithOverlap in [MapEntry (outer_fused[__i0=0:1000])]\n",
"18. Transformation MapTilingWithOverlap in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"19. Transformation MapToForLoop in [MapEntry (outer_fused[__i0=0:1000])]\n",
"20. Transformation MapToForLoop in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"21. Transformation MapUnroll in [MapEntry (outer_fused[__i0=0:1000])]\n",
"22. Transformation NestSDFG in []\n",
"23. Transformation OutLocalStorage in outer_fused___i1[__i1=0:1000] -> outer_fused[__i0=0:1000]\n",
"24. Transformation ReductionNOperation in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"25. Transformation StripMining in outer_fused: ['__i0']\n",
"26. Transformation StripMining in outer_fused___i1: ['__i1']\n",
"27. Transformation TaskletFusion in [Tasklet (_Mult_), AccessNode (__tmp0), Tasklet (_Add_)]\n"
]
},
{
"name": "stdin",
"output_type": "stream",
"text": [
"Select the pattern to apply (0 - 27 or name$id): MapTiling$0(tile_sizes=(128,))\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"You selected (MapTiling$0) pattern MapTiling in [MapEntry (outer_fused[__i0=0:1000])] with parameters {'tile_sizes': (128,)}\n",
"0. Transformation ElementWiseArrayOperation in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"1. Transformation FPGATransformSDFG in []\n",
"2. Transformation FPGATransformState in [SDFGState (BinOp_5)]\n",
"3. Transformation GPUGridStridedTiling in [MapEntry (outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1]), MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"4. Transformation GPUGridStridedTiling in [MapEntry (outer_fused[tile___i0=0:1000:128]), MapEntry (outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1])]\n",
"5. Transformation GPUTransformLocalStorage in outer_fused[tile___i0=0:1000:128]\n",
"6. Transformation GPUTransformMap in outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1]\n",
"7. Transformation GPUTransformMap in outer_fused___i1[__i1=0:1000]\n",
"8. Transformation GPUTransformMap in outer_fused[tile___i0=0:1000:128]\n",
"9. Transformation GPUTransformSDFG in []\n",
"10. Transformation InLocalStorage in outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1] -> outer_fused___i1[__i1=0:1000]\n",
"11. Transformation InLocalStorage in outer_fused[tile___i0=0:1000:128] -> outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1]\n",
"12. Transformation MPITransformMap in [MapEntry (outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1])]\n",
"13. Transformation MPITransformMap in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"14. Transformation MPITransformMap in [MapEntry (outer_fused[tile___i0=0:1000:128])]\n",
"15. Transformation MapDimShuffle in [MapEntry (outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1])]\n",
"16. Transformation MapDimShuffle in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"17. Transformation MapDimShuffle in [MapEntry (outer_fused[tile___i0=0:1000:128])]\n",
"18. Transformation MapFission in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"19. Transformation MapInterchange in [MapEntry (outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1]), MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"20. Transformation MapTiling in [MapEntry (outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1])]\n",
"21. Transformation MapTiling in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"22. Transformation MapTiling in [MapEntry (outer_fused[tile___i0=0:1000:128])]\n",
"23. Transformation MapTilingWithOverlap in [MapEntry (outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1])]\n",
"24. Transformation MapTilingWithOverlap in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"25. Transformation MapTilingWithOverlap in [MapEntry (outer_fused[tile___i0=0:1000:128])]\n",
"26. Transformation MapToForLoop in [MapEntry (outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1])]\n",
"27. Transformation MapToForLoop in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"28. Transformation MapToForLoop in [MapEntry (outer_fused[tile___i0=0:1000:128])]\n",
"29. Transformation MapUnroll in [MapEntry (outer_fused[tile___i0=0:1000:128])]\n",
"30. Transformation NestSDFG in []\n",
"31. Transformation OutLocalStorage in outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1] -> outer_fused[tile___i0=0:1000:128]\n",
"32. Transformation OutLocalStorage in outer_fused___i1[__i1=0:1000] -> outer_fused[__i0=tile___i0:Min(999, tile___i0 + 127) + 1]\n",
"33. Transformation ReductionNOperation in [MapEntry (outer_fused___i1[__i1=0:1000])]\n",
"34. Transformation StripMining in outer_fused: ['__i0']\n",
"35. Transformation StripMining in outer_fused___i1: ['__i1']\n",
"36. Transformation StripMining in outer_fused: ['tile___i0']\n",
"37. Transformation TaskletFusion in [Tasklet (_Mult_), AccessNode (__tmp0), Tasklet (_Add_)]\n"
]
},
{
"name": "stdin",
"output_type": "stream",
"text": [
"Select the pattern to apply (0 - 37 or name$id): \n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"You did not select a valid option. Quitting optimization ...\n"
]
}
],
"source": [
"from dace.transformation.optimizer import SDFGOptimizer\n",
"sdfg = SDFGOptimizer(sdfg).optimize()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The prompt can be used with numbers (e.g., `7` for the 8th transformation) or names (`MapExpansion$0` is the first occurrence of `MapExpansion`). If parameters are given, as in the above example, they are called as it was a Python dictionary: `MapTiling$0(tile_sizes=(128,))`.\n",
"\n",
"If the \"Enter\" key is pressed, the SDFG is no longer transformed and the function returns with the resulting graph.\n",
"\n",
"Internally, the default behavior calls the SDFG console optimizer (`dace.transformation.optimizer.SDFGOptimizer`). This can be modified in the configuration value `optimizer.interface`."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Command-line interface\n",
"\n",
"The console interactive optimizer can also be called from the command line directly, through the `sdfgcc` tool:\n",
"\n",
"```sh\n",
"sdfgcc --optimize path/to/file.sdfg\n",
"```\n",
"\n",
"You could also trigger the command-line interface every time a `@dace` function is called. It can be done by configuring a call hook, for example by setting the environment variable `DACE_call_hooks` to `dace.cli_optimize_on_call`.\n",
"\n",
"Example:\n",
"\n",
"```sh\n",
"DACE_call_hooks=dace.cli_optimize_on_call python myfile.py\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Visual Studio Code\n",
"\n",
"An extension that integrates DaCe into VSCode can be used to interactively edit and transform SDFGs.\n",
"\n",
"To install it, go to the [VSCode Marketplace](https://marketplace.visualstudio.com/items?itemName=phschaad.sdfv) and download the DaCe plugin. Alternatively, open an .sdfg file in VSCode, or search for the extension directly.\n",
"\n",
"Upon opening an SDFG file, the viewer will prompt for transformations in the \"SDFG Optimization\" pane. As you change the view (panning, zooming, collapsing nodes), the transformations under \"Viewport\" will change. \n",
"\n",
"Selecting nodes (through single click, ctrl-click, or the box select mode at the top pane) will add transformations and matching subgraph transformations to the \"Selection\" pane. History appears at the bottom and saves as part of the SDFG file, which you can then use to revert and apply new transformations. See the example below:\n",
"\n",
"![vscode plugin](https://raw.githubusercontent.com/spcl/dace-vscode/master/images/sdfg_optimization.gif \"vscode plugin\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Creating New Transformations\n",
"\n",
"Extending the standard transformations is easy. New transformations can be used for domain-specific optimizations, or simply wrapping expertise gathered over time. In pattern-based transformations, there are three main parts to an implementation: the **pattern(s)**, the **match** function, and the **replacement** (apply) method. Subgraph transformations behave exactly the same, but without the pattern part. \n",
"\n",
"A pattern-based transformation extends the `PatternTransformation` class, whereas subgraph transformation extends `SubgraphTransformation` from the same module. There are several helper functions in `dace.transformation.helpers` and `dace.transformation.subgraph.helpers` that may make your life easier while writing transformations, please check them out in the [documentation](https://spcldace.readthedocs.io/en/latest/source/dace.transformation.html).\n",
"\n",
"In this section, we will make a new pattern-based transformation that takes care of the SDFG we made in the last section:"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (dbladd)"
]
},
"execution_count": 12,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As you can see, there is an unnecessary `__tmp1` transient array between the two tasklets, as a result of `SubgraphFusion`. We can make a \"Redundant array between tasklets\" transformation that checks for such cases, and removes that array if it is not used anywhere else, directly connecting the two tasklets instead.\n",
"\n",
"Using the three parts, we can define a pattern-based transformation:\n",
"1. **Pattern**: Our pattern graph should be `Tasklet -> AccessNode -> Tasklet`, which is the minimal subgraph in which this occurs. Since the transformation matches within a state, we will register it as a single-state transformation by extending `SingleStateTransformation`.\n",
"2. **Match**: We should only match such a subgaph if the array (a) is only a scalar or size 1, (b) is transient (i.e., not defined outside this SDFG), and (c) never used again apart from between those two tasklets. Only this way can we guarantee the correctness of the transformed program.\n",
"3. **Replacement**: If the transformation is matched, we will remove the array and reconnect the tasklets together.\n",
"\n",
"For the pattern part, we must construct node objects as static fields within the class (that must start with an underscore in order to not be recognized as properties). Some of them require arguments, but anything can be set there. Then, we should define an implementation of the `expressions` class method to state how those node objects should be connected.\n",
"\n",
"The match and replacement parts are implemented by the `can_be_applied` class method, returning a boolean for a specific subgraph candidate, and the `apply` instance method, which modifies the given SDFG using the SDFG API.\n",
"\n",
"We can now start implementing the transformation:"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"from dace.transformation import transformation as xf\n",
"from dace.sdfg import utils as sdutil\n",
"from dace import registry\n",
"\n",
"\n",
"class RedundantArrayTasklet(xf.SingleStateTransformation):\n",
" \"\"\" Removes a redundant transient array if and only if it is between two tasklets. \"\"\"\n",
"\n",
" # Pattern: define pattern nodes\n",
" tasklet_a = xf.PatternNode(dace.nodes.Tasklet)\n",
" array = xf.PatternNode(dace.nodes.AccessNode)\n",
" tasklet_b = xf.PatternNode(dace.nodes.Tasklet)\n",
"\n",
" # This method returns a list of graphs that represent the pattern\n",
" @classmethod\n",
" def expressions(cls):\n",
" # We have only one expression, and it is a path graph between the three nodes\n",
" return [sdutil.node_path_graph(cls.tasklet_a, cls.array, cls.tasklet_b)]\n",
"\n",
" # Match function\n",
" def can_be_applied(self, graph, expr_index, sdfg, permissive=False):\n",
" # Getting the actual node in the graph\n",
" array_node = self.array\n",
"\n",
" # Get data descriptor from SDFG registry and access node\n",
" desc = sdfg.arrays[array_node.data]\n",
"\n",
" # Match (a): Check array size\n",
" if desc.total_size != 1:\n",
" return False\n",
"\n",
" # Match (b): Check if transient\n",
" if not desc.transient:\n",
" return False\n",
"\n",
" # Match (c): Check if used again in any state in this SDFG\n",
" if len([node for state in sdfg.nodes()\n",
" for node in state.nodes()\n",
" if isinstance(node, dace.nodes.AccessNode)\n",
" and node.data == array_node.data]) > 1:\n",
" return False\n",
"\n",
" # Match (c): Check if any other tasklet uses this transient\n",
" if graph.in_degree(array_node) + graph.out_degree(array_node) != 2:\n",
" return False\n",
"\n",
" # Match found!\n",
" return True\n",
"\n",
" # Replacement (note that this is a method of a transformation instance)\n",
" def apply(self, graph, sdfg):\n",
" # Query the pattern match for the actual node\n",
" array_node = self.array\n",
"\n",
" # Get incoming and outgoing edges of the access node\n",
" # (there are only one of each according to `can_be_applied`)\n",
" in_edge = state.in_edges(array_node)[0]\n",
" out_edge = state.out_edges(array_node)[0]\n",
"\n",
" # Construct a new edge between the two nodes, using the input/output\n",
" # nodes and connectors\n",
" state.add_edge(in_edge.src, in_edge.src_conn,\n",
" out_edge.dst, out_edge.dst_conn,\n",
" dace.Memlet())\n",
"\n",
" # Finally, remove the redundant array node from the graph\n",
" state.remove_node(array_node)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now that the transformation is implemented, we can check if it works:"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"1"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sdfg.apply_transformations(RedundantArrayTasklet)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As the return value states the number of times a transformation was applied, the transformation was applied once. Let's look at the graph:"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (dbladd)"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The node is now removed."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Composing Transformations\n",
"\n",
"Transformations can call other transformations to create powerful compositions and make verification easier. Simply use one of the APIs to invoke a transformation (such as `apply_to`) within the code of another transformation to do so.\n",
"\n",
"### Optimization and Customization\n",
"\n",
"There are more methods you can implement to optimize transformations:\n",
" * `annotates_memlets`: If the method returns True, memlets will not be propagated after transformation is applied (for performance and/or overriding default propagation behavior).\n",
" * `match_to_str`: Returns a string-representation of the match to customize printout in the command-line interface.\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.1"
}
},
"nbformat": 4,
"nbformat_minor": 4
}