{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Using the Python-Numpy Frontend in DaCe\n",
"\n",
"In this tutorial, we will see how one can write code using the python-numpy frontend. The frontend supports a subset of the python language and array/matrix operations inspired by the numpy module."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Let's start with a first example that will showcase the basic elements of the DaCe program. First, we import the `dace` and `numpy` modules:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"import dace\n",
"import numpy as np"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then, we declare the program parameters, which can be either symbols or constants:"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"M, N, K = 24, 24, 24"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We proceed by writing the DaCe program as a regular python method, that is annotated with the `dace.program` annotation. The parameters of the python method must have type annotations. For example, below we define a `gemm` method, for implementing the generalized matrix-matrix multiplication operation. The first 3 parameters are the 32-bit floating-point matrices `A`, `B` and `C`. The last 2 parameters are the 32-bit floating-point scalar values `alpha` and `beta`. The implementation of the method is written the same way as in python, using the numpy module."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"@dace.program\n",
"def gemm(A: dace.float32[M, K], B: dace.float32[K, N], C: dace.float32[M, N],\n",
" alpha: dace.float32, beta: dace.float32):\n",
" C[:] = alpha * A @ B + beta * C"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The `[:]` slice expression is representing the whole range of the array/matrix. Note that in DaCe you are not allowed to redefine data in the same DaCe program. Therefore, if you define an array `C`, you may not assign to it a different value. For example, if we changed the implementation of the `gemm` method to `C = alpha * A @ B + beta * C`, we would get an error."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The DaCe program may be parsed to an SDFG and/or compiled using the the same methods from the SDFG API:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"sdfg = gemm.to_sdfg()"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (gemm)"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"sdfg"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Supported Python/Numpy operators"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The frontend supports the unary operators `{+, -, not, ~}`. Note that the `not` operator works the same way as the `numpy.logical_not` method."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"@dace.program\n",
"def uadd(A: dace.int64[5, 5], B: dace.int64[5, 5]):\n",
" B[:] = +A\n",
"\n",
"@dace.program\n",
"def usub(A: dace.int64[5, 5], B: dace.int64[5, 5]):\n",
" B[:] = -A\n",
"\n",
"@dace.program\n",
"def logicalnot(A: dace.bool[5, 5], B: dace.bool[5, 5]):\n",
" B[:] = not A\n",
"\n",
"@dace.program\n",
"def invert(A: dace.int64[5, 5], B: dace.int64[5, 5]):\n",
" B[:] = ~A"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The frontend support the binary operators `{+, -, *, /, //, %, **, @, <<, >>, |, ^, &, and, or, ==, !=, <, <=, >, >=}`. Note that the boolean operators `{and, or}` work the same way as the methods `numpy.logical_and` and `numpy.logical_or`. Apart from the matrix-multiplication operator `@`, all the other operators are point-wise. Note that the return type of the operators is the one returned by [`numpy.result_type`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.result_type.html). Exception to that are the boolean operators, which return boolean values."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Python augmented assignments are supported for all operators. Some examples follow:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"@dace.program\n",
"def augfloordiv(A: dace.int64[5, 5], B: dace.int64[5, 5]):\n",
" B //= A\n",
"\n",
"@dace.program\n",
"def augmod(A: dace.int64[5, 5], B: dace.int64[5, 5]):\n",
" B %= A\n",
"\n",
"@dace.program\n",
"def augpow(A: dace.int64[5, 5], B: dace.int64[5, 5]):\n",
" B **= A\n",
"\n",
"@dace.program\n",
"def auglshift(A: dace.int64[5, 5], B: dace.int64[5, 5]):\n",
" B <<= A"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Operations between arrays/matrices and scalars or arrays of size 1 behave the same way as with numpy:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"@dace.program\n",
"def addscalar(A: dace.int64[5, 5], B: dace.int64, C: dace.int64[5, 5]):\n",
" C[:] = A + B\n",
"\n",
"@dace.program\n",
"def floordivnumber(A: dace.int64[5, 5], C: dace.int64[5, 5]):\n",
" C[:] = A // 5"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Defining Data, Maps and Sequential-Loops\n",
"\n",
"Transient arrays can be defined with `dace.define_local` or just `numpy.ndarray`. Furthermore, transient scalars can be defined with `dace.define_local_scalar`."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"@dace.program\n",
"def transient(A: dace.float32[M, N, K]):\n",
" s = np.ndarray(shape=(M, N, K), dtype=np.int32)\n",
" t = dace.define_local(A.shape, A.dtype)\n",
" s[:] = A\n",
" t[:] = A"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Python for-loops are automatically converted to control-flow:"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (forloop)"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"N, BS = (dace.symbol(name) for name in ['N', 'BS'])\n",
"\n",
"@dace.program\n",
"def forloop(HD: dace.complex128[N, BS, BS], HE: dace.complex128[N, BS, BS],\n",
" HF: dace.complex128[N, BS, BS],\n",
" sigmaRSD: dace.complex128[N, BS, BS],\n",
" sigmaRSE: dace.complex128[N, BS, BS],\n",
" sigmaRSF: dace.complex128[N, BS, BS]):\n",
"\n",
" for n in range(N):\n",
" if n < N - 1:\n",
" HE[n] -= sigmaRSE[n]\n",
" else:\n",
" HE[n] = -sigmaRSE[n]\n",
" if n > 0:\n",
" HF[n] -= sigmaRSF[n]\n",
" else:\n",
" HF[n] = -sigmaRSF[n]\n",
" HD[n] = HD[n] - sigmaRSD[n]\n",
" \n",
"forloop.to_sdfg()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Maps (parallel for-loops) can be created with `dace.map`:"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"\n",
""
],
"text/plain": [
"SDFG (maptest)"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"Nkz, NE, Nqz, Nw, N3D, NA, NB, Norb = (\n",
" dace.symbol(name)\n",
" for name in ['Nkz', 'NE', 'Nqz', 'Nw',\n",
" 'N3D', 'NA', 'NB', 'Norb'])\n",
"\n",
"@dace.program\n",
"def maptest(neigh_idx: dace.int32[NA, NB],\n",
" dH: dace.complex128[NA, NB, N3D, Norb, Norb],\n",
" G: dace.complex128[Nkz, NE, NA, Norb, Norb],\n",
" D: dace.complex128[Nqz, Nw, NA, NB, N3D, N3D],\n",
" Sigma: dace.complex128[Nkz, NE, NA, Norb, Norb]):\n",
"\n",
" for k, E, q, w, i, j, a, b in dace.map[0:Nkz, 0:NE,\n",
" 0:Nqz, 0:Nw,\n",
" 0:N3D, 0:N3D,\n",
" 0:NA, 0:NB]:\n",
" dHG = G[k-q, E-w, neigh_idx[a, b]] @ dH[a, b, i]\n",
" dHD = dH[a, b, j] * D[q, w, a, b, i, j]\n",
" Sigma[k, E, a] += dHG @ dHD\n",
" \n",
"maptest.to_sdfg()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Combining explicit dataflow with numpy\n",
"\n",
"The python-numpy syntax can be used in combination with the explicit dataflow syntax:"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"N = dace.symbol('N')\n",
"\n",
"@dace.program\n",
"def slicetest(A: dace.float64[N, N - 1], B: dace.float64[N - 1, N],\n",
" C: dace.float64[N - 1, N - 1]):\n",
" tmp = A[1:N] * B[:, 0:N - 1]\n",
" for i, j in dace.map[0:4, 0:4]:\n",
" with dace.tasklet:\n",
" t << tmp[i, j]\n",
" c >> C[i, j]\n",
" c = t"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"@dace.program\n",
"def saoptest(A: dace.float64[5, 5], alpha: dace.float64,\n",
" B: dace.float64[5, 5]):\n",
" tmp = alpha * A * 5\n",
" for i, j in dace.map[0:5, 0:5]:\n",
" with dace.tasklet:\n",
" t << tmp[i, j]\n",
" c >> B[i, j]\n",
" c = t"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Other operations and methods\n",
"\n",
"Reductions can be defined with `dace.reduce`:"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"W = dace.symbol('W')\n",
"H = dace.symbol('H')\n",
"\n",
"@dace.program(dace.float32[H, W], dace.float32[H, W], dace.float32[1])\n",
"def mapreduce_test(A, B, sum):\n",
" tmp = dace.define_local([H, W], dace.float32)\n",
"\n",
" @dace.map(_[0:H, 0:W])\n",
" def compute_tile(i, j):\n",
" a << A[i, j]\n",
" b >> B[i, j]\n",
" t >> tmp[i, j]\n",
"\n",
" b = a * 5\n",
" t = a * 5\n",
"\n",
" sum[:] = dace.reduce(lambda a, b: a + b, tmp)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The frontend also supports basic math methods such as `{exp, sin, cos, sqrt, log, conj, real, imag}`. See the documentation for details."
]
}
],
"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
}