How to use the dace.sdfg function in dace

#!/usr/bin/env python
from __future__ import print_function

import gc
import argparse
import dace
import numpy as np
import dace.frontend.common as np_frontend

import os
from timeit import default_timer as timer

SDFG = dace.sdfg.SDFG

M = dace.symbol('M')
N = dace.symbol('N')
K = dace.symbol('K')
L = dace.symbol('L')
O = dace.symbol('O')

alpha = dace.ndarray([L, O], dtype=dace.float64)
A = dace.ndarray([M, N, K], dtype=dace.float64)
B = dace.ndarray([M, N, K], dtype=dace.float64)

if __name__ == "__main__":
    print("==== Program start ====")

    parser = argparse.ArgumentParser()
    parser.add_argument("M", type=int, nargs="?", default=128)

parent_nested_axpy.add_memlet_path(in_data_C,
                                           nested_axpy_2_node,
                                           memlet=dace.Memlet.simple(in_data_C.data, '0:N'),
                                           dst_conn="mem_C")
    parent_nested_axpy.add_memlet_path(stream_rd,
                                           nested_axpy_2_node,
                                           memlet=dace.Memlet.simple(stream_rd.data, '0:N'),
                                           dst_conn="stream_in")
    parent_nested_axpy.add_memlet_path(nested_axpy_2_node,
                                 stream_wr,
                                 memlet=dace.Memlet.simple(stream_wr.data, '0:N'),
                                 src_conn="stream_out")

    parent_sdfg.add_edge(copy_in_state, parent_nested_axpy, dace.sdfg.edges.InterstateEdge())
    parent_sdfg.validate()

    return parent_sdfg

#!/usr/bin/env python
from __future__ import print_function

import gc
import argparse
import dace
import numpy as np
import dace.frontend.common as np_frontend

import os
from timeit import default_timer as timer

SDFG = dace.sdfg.SDFG

M = dace.symbol('M')
N = dace.symbol('N')
K = dace.symbol('K')

A = dace.ndarray([3, 7, 9, M, N], dtype=dace.float64)
B = dace.ndarray([2, 5, 8, 4, N, K], dtype=dace.float64)
C = dace.ndarray([3, 3, M, K], dtype=dace.float64)

if __name__ == "__main__":
    print("==== Program start ====")

    parser = argparse.ArgumentParser()
    parser.add_argument("M", type=int, nargs="?", default=128)
    parser.add_argument("N", type=int, nargs="?", default=128)
    parser.add_argument("K", type=int, nargs="?", default=128)

# Skip map in input edges
        for edge in nstate.out_edges(map_entry):
            src_node = nstate.memlet_path(edge)[0].src
            nstate.add_edge(src_node, None, edge.dst, edge.dst_conn, edge.data)
            nstate.remove_edge(edge)

        # Skip map in output edges
        for edge in nstate.in_edges(map_exit):
            dst_node = nstate.memlet_path(edge)[-1].dst
            nstate.add_edge(edge.src, edge.src_conn, dst_node, None, edge.data)
            nstate.remove_edge(edge)

        # Remove nodes from dynamic map range
        nstate.remove_nodes_from(
            [e.src for e in dace.sdfg.dynamic_map_inputs(nstate, map_entry)])
        # Remove scope nodes
        nstate.remove_nodes_from([map_entry, map_exit])

        return node, nstate

def make_sdfg(specialize):

    if specialize:
        name = "gemv_transposed_{}x{}".format(N.get(), M.get())
    else:
        name = "gemv_transposed_{}xM".format(N.get())

    sdfg = dace.SDFG(name)

    init_state = make_init_state(sdfg)
    fpga_state = make_outer_compute_state(sdfg)
    finalize_state = make_finalize_state(sdfg)

    sdfg.add_edge(init_state, fpga_state, dace.sdfg.InterstateEdge())
    sdfg.add_edge(fpga_state, finalize_state, dace.sdfg.InterstateEdge())

    return sdfg

def _get_transpose_input(node, state, sdfg):
    """Returns the transpose input edge, array, and shape."""
    for edge in state.in_edges(node):
        if edge.dst_conn == "_inp":
            subset = dc(edge.data.subset)
            subset.squeeze()
            size = subset.size()
            outer_array = sdfg.data(
                dace.sdfg.find_input_arraynode(state, edge).data)
            return edge, outer_array, (size[0], size[1])
    raise ValueError("Transpose input connector \"_inp\" not found.")

y_begin = sdfg.add_state("y_begin")
    y_entry = sdfg.add_state("y_entry")
    y_end = sdfg.add_state("y_end")

    x_begin = sdfg.add_state("x_begin")
    x_entry = sdfg.add_state("x_entry")
    x_end = sdfg.add_state("x_end")

    pre_shift = sdfg.add_state("pre_shift")
    loop_body = sdfg.add_state("compute_body")
    post_shift = sdfg.add_state("post_shift")

    sdfg.add_edge(time_begin, time_entry,
                  dace.sdfg.InterstateEdge(assignments={"t": 0}))
    sdfg.add_edge(y_begin, y_entry,
                  dace.sdfg.InterstateEdge(assignments={"y": 0}))
    sdfg.add_edge(x_begin, x_entry,
                  dace.sdfg.InterstateEdge(assignments={"x": 0}))

    sdfg.add_edge(
        time_entry, y_begin,
        dace.sdfg.InterstateEdge(
            condition=dace.properties.CodeProperty.from_string(
                "t < T", language=dace.dtypes.Language.Python)))
    sdfg.add_edge(
        y_entry, x_begin,
        dace.sdfg.InterstateEdge(
            condition=dace.properties.CodeProperty.from_string(
                "y < H", language=dace.dtypes.Language.Python)))
    sdfg.add_edge(
        x_entry, pre_shift,
        dace.sdfg.InterstateEdge(

''' DaCe SDFG linear algebra operation library. '''

import copy
import dace
import dace.sdfg as sd
import dace.subsets as sbs
from dace import symbolic
import typing
import numpy as np
from dace.libraries.blas.blas_helpers import to_blastype as _to_blastype

State = dace.sdfg.SDFGState
Shape = typing.List[typing.Union[int, symbolic.symbol]]
Index = typing.List[typing.Union[int, str, symbolic.symbol]]
Node = dace.sdfg.nodes.Node
DNode = dace.sdfg.nodes.AccessNode

# TODO: Most of the external operations here emit Z (complex double) ops, fix


def _to_cudatype(dtype):
    """ Returns a CUDA typename that corresponds to the input type.
        Used in CUBLAS calls. """

    if dtype == np.float16:
        return '__half'
    elif dtype == np.float32:
        return 'float'

time_end = sdfg.add_state("time_end")

    y_begin = sdfg.add_state("y_begin")
    y_entry = sdfg.add_state("y_entry")
    y_end = sdfg.add_state("y_end")

    x_begin = sdfg.add_state("x_begin")
    x_entry = sdfg.add_state("x_entry")
    x_end = sdfg.add_state("x_end")

    loop_body = sdfg.add_state("read_memory")

    sdfg.add_edge(time_begin, time_entry,
                  dace.sdfg.InterstateEdge(assignments={"t": 0}))
    sdfg.add_edge(y_begin, y_entry,
                  dace.sdfg.InterstateEdge(assignments={"y": 0}))
    sdfg.add_edge(x_begin, x_entry,
                  dace.sdfg.InterstateEdge(assignments={"x": 0}))

    sdfg.add_edge(
        time_entry, y_begin,
        dace.sdfg.InterstateEdge(
            condition=dace.properties.CodeProperty.from_string(
                "t < T", language=dace.dtypes.Language.Python)))
    sdfg.add_edge(
        y_entry, x_begin,
        dace.sdfg.InterstateEdge(
            condition=dace.properties.CodeProperty.from_string(
                "y < H", language=dace.dtypes.Language.Python)))
    sdfg.add_edge(
        x_entry, loop_body,
        dace.sdfg.InterstateEdge(

def expressions():
        graph = dace.sdfg.graph.OrderedDiGraph()
        graph.add_node(MatrixProductTranspose._transpose_a)
        graph.add_node(MatrixProductTranspose._at)
        graph.add_node(MatrixProductTranspose._transpose_b)
        graph.add_node(MatrixProductTranspose._bt)
        graph.add_node(MatrixProductTranspose._a_times_b)
        graph.add_edge(MatrixProductTranspose._transpose_a,
                       MatrixProductTranspose._at, None)
        graph.add_edge(MatrixProductTranspose._at,
                       MatrixProductTranspose._a_times_b, None)
        graph.add_edge(MatrixProductTranspose._transpose_b,
                       MatrixProductTranspose._bt, None)
        graph.add_edge(MatrixProductTranspose._bt,
                       MatrixProductTranspose._a_times_b, None)
        return [graph]