metagraph-mlir Documentation¶
metagraph-mlir is a plugin for `Metagraph`_ that enables algorithms to be just-in-time (JIT) compiled with MLIR. To learn more about compiler plugins in Metagraph, see the Compiler Plugins section of the Metagraph Plugin Author Guide.
metagraph-mlir is licensed under the Apache 2.0 license and the source code can be found on Github.
Installation¶
metagraph-mlir is currently only distributed via conda. To install:
conda install -c metagraph -c conda-forge metagraph-mlir
Note that metagraph-mlir requires a development version of MLIR from the unreleased LLVM 12.0. This version is automatically installed from the metagraph channel by conda.
Implementing Algorithms with metagraph-mlir¶
The metagraph-mlir compiler requires that the concrete algorithm (learn more
about Metagraph algorithms here) function return a
metagraph_mlir.compiler.MLIRFunc
object with the implementation of the
algorithm in MLIR. For example, suppose an abstract algorithm has already
been defined to add two vectors:
@abstract_algorithm("example.add")
def example_add(a: Vector, b: Vector) -> Vector:
pass
A concrete implementation of this algorithm in MLIR be written this way:
@concrete_algorithm("example.add", compiler="mlir")
def compiled_add(
a: NumpyVectorType, b: NumpyVectorType
) -> NumpyVectorType:
return MLIRFunc(
name="example_add",
arg_types=["tensor<?xf32>", "tensor<?xf32>"],
ret_type="tensor<?xf32>",
mlir=b"""\
#trait_testing_add = {
indexing_maps = [
affine_map<(i) -> (i)>, // A
affine_map<(i) -> (i)>, // B
affine_map<(i) -> (i)> // X (out)
],
iterator_types = ["parallel"],
doc = "X(i) = A(i) OP B(i)"
}
func private @example_add(%arga: tensor<?xf32>, %argb: tensor<?xf32>) -> tensor<?xf32> {
%0 = linalg.generic #trait_testing_add
ins(%arga, %argb: tensor<?xf32>, tensor<?xf32>)
outs(%arga: tensor<?xf32>) {
^bb(%a: f32, %b: f32, %s: f32):
%0 = addf %a, %b : f32
linalg.yield %0 : f32
} -> tensor<?xf32>
return %0 : tensor<?xf32>
}
""",
)
This MLIRFunc object returned by compile_add
in this example contains a
number of special attributes describing the function, including:
name
: Name of the main entry point function.
arg_types
: A list of strings containing the MLIR types for each function argument.
ret_type
: A string containing the MLIR type of return value
mlir
: A bytes object (not string) containing the text of the MLIR code.
To minimize the size of the JIT compiled modules (and to simplify inspection
of the MLIR after inlining into the generated wrapper function), all MLIR
functions in the body should be defined as private
.
The MLIR functions are JIT compiled using the JIT engine in mlir-graphblas. See that documentation for the examples of how to write functions with scalars, dense tensors, and sparse tensors.
Translating between ScipyGraph and MLIRGraphBLASGraph¶
metagraph-mlir currently supports translations between graphs of type ScipyGraph
(provided as a core type in Metagraph) and MLIRGraphBLASGraph
(provided by metagraph-mlir).
MLIRGraphBLASGraph
is an adjacency matrix representation of a graph implemented via MLIR’s current sparse tensor support. This tutorial provides an overview of how MLIR currently supports sparse tensors. It also shows examples of how to generate an instance of an MLIR sparse tensor.
MLIRGraphBLASGraph
alows us to wrap an MLIR sparse tensor as a graph, e.g.
import numpy as np
import mlir_graphblas
from mlir_graphblas.sparse_utils import MLIRSparseTensor
import metagraph as mg
# The sparse adjacency matrix below looks like this (where the underscores represent zeros):
#
# [[ 1.2, ___, ___, ___ ],
# [ ___, ___, ___, 3.4 ],
# [ ___, ___, 5.6, ___ ],
# [ ___, ___, ___, ___ ]]
#
indices = np.array([
[0, 0],
[1, 3],
[2, 2],
], dtype=np.uint64)
values = np.array([1.2, 3.4, 5.6], dtype=np.float32)
sizes = np.array([4, 4], dtype=np.uint64)
sparsity = np.array([False, True], dtype=np.bool8)
sparse_tensor = mlir_graphblas.sparse_utils.MLIRSparseTensor(indices, values, sizes, sparsity)
has_weighted_edges = True
graph = mg.wrappers.Graph.MLIRGraphBLASGraph(sparse_tensor, has_weighted_edges, aprops={
"node_type": "set",
"node_dtype": None,
"edge_type": "map",
"edge_dtype": "float",
"is_directed": True,
})
We can translate this graph into a ScipyGraph
like so:
scipy_graph = mg.translate(g, mg.wrappers.Graph.ScipyGraph)
This will allow us verify the following:
assert np.isclose(
scipy_graph.value.toarray(),
np.array([
[1.2, 0. , 0. , 0. ],
[0. , 0. , 0. , 3.4],
[0. , 0. , 5.6, 0. ],
[0. , 0. , 0. , 0. ]
])).all()
We can also translate back easily:
graph_round_trip = mg.translate(scipy_graph, mg.wrappers.Graph.MLIRGraphBLASGraph)
There are some limitations:
Currently, we can only translate from
MLIRGraphBLASGraph
toScipyGraph
if theMLIRGraphBLASGraph
instance’s underlying sparse tensor is dense in the first dimension and sparse in the second dimension, i.e. if it is in CSR format.
MLIRGraphBLASGraph
currently only supports 32-bit and 64-bit floating point edge weights. Thus, when translating toMLIRGraphBLASGraph
, exceptions will be raised if the source type instance has boolean, string, or integer edge weights.