def test_random_walk(): # fmt: off indices = np.array( [[0, 1], [0, 2], [1, 0], [1, 3], [2, 0], [2, 4], [3, 2]], dtype=np.uint64, ) values = np.array([100, 200, 300, 400, 175, 222, 333], dtype=np.float64) # fmt: on sizes = np.array([5, 5], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) graph = MLIRSparseTensor(indices, values, sizes, sparsity) assert graph.verify() paths = mlalgo.random_walk(graph, 20, [0, 1, 2, 3, 4, 4, 3, 2, 1, 0]) # Perform validation for each row to verify correctness valid_steps = {} for start, end in indices: valid_steps.setdefault(start, set()).add(end) pointers = paths.pointers[-1] nodes = paths.values for irow, (istart, iend) in enumerate(zip(pointers[:-1], pointers[1:])): assert istart != iend, f"Initial node missing for row {irow}" for jstart, jend in zip(nodes[istart:int(iend - 1)], nodes[int(istart + 1):iend]): assert jstart in valid_steps, f"Row [{irow}] {jstart} is a terminator" assert (jend in valid_steps[jstart] ), f"Row [{irow}] {jstart}->{jend} is not a valid step"
def test_pagerank(special_passes): # fmt: off indices = np.array( [[0, 1], [0, 2], [1, 3], [2, 3], [2, 4], [3, 4], [4, 0]], dtype=np.uint64, ) # fmt: on values = np.array([1.1, 9.8, 4.2, 7.1, 0.2, 6.9, 2.2], dtype=np.float64) sizes = np.array([5, 5], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) m = MLIRSparseTensor(indices, values, sizes, sparsity) assert m.verify() expected = np.array( [0.2541917746, 0.1380315018, 0.1380315018, 0.2059901768, 0.2637550447]) # Test success pr, niters = mlalgo.pagerank(m, tol=1e-7) assert np.abs(pr.values - expected).sum() < 1e-5, pr.values # Test maxiter reached, failed to converge pr, niters = mlalgo.pagerank(m, tol=1e-7, maxiter=6, compile_with_passes=special_passes) assert niters == 6 assert (np.abs(pr.values - expected).sum() > 1e-5), "Unexpectedly converged in 6 iterations"
def test_bfs(special_passes): # 0 - 1 5 - 6 # | X | | / # 3 - 4 -- 2 - 7 # fmt: off indices = np.array( [[0, 1], [0, 3], [0, 4], [1, 0], [1, 3], [1, 4], [2, 4], [2, 5], [2, 6], [2, 7], [3, 0], [3, 1], [3, 4], [4, 0], [4, 1], [4, 2], [4, 3], [5, 2], [5, 6], [6, 2], [6, 5], [7, 2]], dtype=np.uint64, ) values = np.array( [ 100, 200, 300, 100, 400, 500, 99, 50, 55, 75, 200, 400, 600, 300, 500, 99, 600, 50, 60, 55, 60, 75 ], dtype=np.float64, ) # fmt: on sizes = np.array([8, 8], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) a = MLIRSparseTensor(indices, values, sizes, sparsity) assert a.verify() parents, levels = mlalgo.bfs(0, a, compile_with_passes=special_passes) expected_parents = np.array([0, 0, 4, 0, 0, 2, 2, 2]) expected_levels = np.array([0, 1, 2, 1, 1, 3, 3, 3]) assert np.all(parents.toarray() == expected_parents) assert np.all(levels.toarray() == expected_levels)
def test_triangle_count(special_passes): # 0 - 1 5 - 6 # | X | | / # 3 - 4 -- 2 - 7 # fmt: off indices = np.array( [[0, 1], [0, 3], [0, 4], [1, 0], [1, 3], [1, 4], [2, 4], [2, 5], [2, 6], [2, 7], [3, 0], [3, 1], [3, 4], [4, 0], [4, 1], [4, 2], [4, 3], [5, 2], [5, 6], [6, 2], [6, 5], [7, 2]], dtype=np.uint64, ) values = np.array( [ 100, 200, 300, 100, 400, 500, 99, 50, 55, 75, 200, 400, 600, 300, 500, 99, 600, 50, 60, 55, 60, 75 ], dtype=np.float64, ) # fmt: on sizes = np.array([8, 8], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) a = MLIRSparseTensor(indices, values, sizes, sparsity) assert a.verify() num_triangles = mlalgo.triangle_count(a, compile_with_passes=special_passes) assert num_triangles == 5, num_triangles
def test_haversine_distance(): # Sanity check # https://www.igismap.com/haversine-formula-calculate-geographic-distance-earth/ # Nebraska # v1 = Vector.from_values([0], [41.507483]) # w1 = Vector.from_values([0], [-99.436554]) # Kansas # v2 = Vector.from_values([0], [38.504048]) # w2 = Vector.from_values([0], [-98.315949]) # Build a function to call the haversine_distance utility from mlir_graphblas.mlir_builder import MLIRFunctionBuilder irb = MLIRFunctionBuilder( "haversine_distance", input_types=[ "tensor<?xf64, #CV64>", "tensor<?xf64, #CV64>", "tensor<?xf64, #CV64>", "tensor<?xf64, #CV64>", ], return_types=["tensor<?xf64, #CV64>"], aliases=mlalgo._build_common_aliases(), ) v1, w1, v2, w2 = irb.inputs result = algo_utils.haversine_distance(irb, v1, w1, v2, w2) irb.return_vars(result) compiled_func = irb.compile() # haversine_distance(v1, w1, v2, w2)[0].new().isclose(347.3, abs_tol=0.1) v1 = MLIRSparseTensor( np.array([[0]], dtype=np.uint64), np.array([41.507483], dtype=np.float64), np.array([1], dtype=np.uint64), np.array([True], dtype=np.bool8), ) w1 = MLIRSparseTensor( np.array([[0]], dtype=np.uint64), np.array([-99.436554], dtype=np.float64), np.array([1], dtype=np.uint64), np.array([True], dtype=np.bool8), ) v2 = MLIRSparseTensor( np.array([[0]], dtype=np.uint64), np.array([38.504048], dtype=np.float64), np.array([1], dtype=np.uint64), np.array([True], dtype=np.bool8), ) w2 = MLIRSparseTensor( np.array([[0]], dtype=np.uint64), np.array([-98.315949], dtype=np.float64), np.array([1], dtype=np.uint64), np.array([True], dtype=np.bool8), ) dist = compiled_func(v1, w1, v2, w2) assert math.isclose(dist.values[0], 347.3, abs_tol=0.1)
def test_vertex_nomination(special_passes): # fmt: off indices = np.array( [[0, 1], [0, 3], [1, 4], [1, 6], [2, 5], [3, 0], [3, 2], [4, 5], [5, 1], [6, 2], [6, 3], [6, 4]], dtype=np.uint64, ) # fmt: on values = np.array([2, 3, 8, 4, 1, 3, 3, 7, 1, 5, 7, 3], dtype=np.float64) sizes = np.array([7, 7], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) m = MLIRSparseTensor(indices, values, sizes, sparsity) assert m.verify() indices = np.array([[6]], dtype=np.uint64) values = np.array([0], dtype=np.float64) sizes = np.array([7], dtype=np.uint64) sparsity = np.array([True], dtype=np.bool8) v = MLIRSparseTensor(indices, values, sizes, sparsity) assert v.verify() # Compute Vertex Nomination # correct answer for node #6 is node #4 w = mlalgo.vertex_nomination(m, v, compile_with_passes=special_passes) assert w == 4 # correct answer for nodes #0,1,5 is node #3 indices = np.array([[0], [1], [5]], dtype=np.uint64) values = np.array([0, 0, 0], dtype=np.float64) v2 = MLIRSparseTensor(indices, values, sizes, sparsity) assert v2.verify() w2 = mlalgo.vertex_nomination(m, v2, compile_with_passes=special_passes) assert w2 == 3
def test_jit_engine_sparse_tensor(engine, mlir_type): mlir_template = r""" #trait_sum_reduction = {{ indexing_maps = [ affine_map<(i,j,k) -> (i,j,k)>, // A affine_map<(i,j,k) -> ()> // x (scalar out) ], iterator_types = ["reduction", "reduction", "reduction"], doc = "x += SUM_ijk A(i,j,k)" }} #sparseTensor = #sparse_tensor.encoding<{{ dimLevelType = [ "compressed", "compressed", "compressed" ], dimOrdering = affine_map<(i,j,k) -> (i,j,k)>, pointerBitWidth = 64, indexBitWidth = 64 }}> func @{func_name}(%argA: tensor<10x20x30x{mlir_type}, #sparseTensor>) -> {mlir_type} {{ %out_tensor = arith.constant dense<0.0> : tensor<{mlir_type}> %reduction = linalg.generic #trait_sum_reduction ins(%argA: tensor<10x20x30x{mlir_type}, #sparseTensor>) outs(%out_tensor: tensor<{mlir_type}>) {{ ^bb(%a: {mlir_type}, %x: {mlir_type}): %0 = arith.addf %x, %a : {mlir_type} linalg.yield %0 : {mlir_type} }} -> tensor<{mlir_type}> %answer = tensor.extract %reduction[] : tensor<{mlir_type}> return %answer : {mlir_type} }} """ np_type = MLIR_TYPE_TO_NP_TYPE[mlir_type] func_name = f"func_{mlir_type}" mlir_text = mlir_template.format(func_name=func_name, mlir_type=mlir_type) indices = np.array([[0, 0, 0], [1, 1, 1]], dtype=np.uint64) values = np.array([1.2, 3.4], dtype=np_type) sizes = np.array([10, 20, 30], dtype=np.uint64) sparsity = np.array([True, True, True], dtype=np.bool8) a = MLIRSparseTensor(indices, values, sizes, sparsity) assert a.verify() assert engine.add(mlir_text, GRAPHBLAS_PASSES) == [func_name] result = engine[func_name](a) expected_result = 4.6 assert ( abs(result - expected_result) < 1e-6 ), f"""
def test_graph_search(): # fmt: off indices = np.array( [[0, 1], [0, 2], [1, 0], [1, 3], [2, 0], [2, 4], [3, 2], [4, 4]], dtype=np.uint64, ) values = np.array([100, 200, 300, 400, 175, 222, 333, 200], dtype=np.float64) # fmt: on sizes = np.array([5, 5], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) graph = MLIRSparseTensor(indices, values, sizes, sparsity) assert graph.verify() # Random Uniform (no seed, so truly random) count = mlalgo.graph_search(graph, 3, [2, 4], "random") # Check for one of the possible solutions: # [0, 1, 4] or [0, 1, 3, 4] or [0, 2, 4] or [0, 2, 4] or [4] # [2, 1, 3] [1, 1, 1, 3] [2, 1, 3] [1, 1, 4] [6] for idx, vals in [ ([0, 1, 4], [2, 1, 3]), ([0, 1, 3, 4], [1, 1, 1, 3]), ([0, 2, 4], [2, 1, 3]), ([0, 2, 4], [1, 1, 4]), ([4], [6]), ]: if len(count.indices[0]) == len(idx): if (count.indices[0] == idx).all() and (count.values == vals).all(): break else: assert False, f"Invalid solution: idx={count.indices[0]}, vals={count.values}" # Random weighted count = mlalgo.graph_search(graph, 5, [0, 2], "random_weighted", rand_seed=14) assert (count.indices[0] == [0, 2, 4]).all() assert (count.values == [2, 3, 5]).all() # argmin count = mlalgo.graph_search(graph, 3, [0, 3], "argmin") assert (count.indices[0] == [0, 1, 2]).all() assert (count.values == [2, 3, 1]).all() # argmax count = mlalgo.graph_search(graph, 3, [0, 1], "argmax") assert (count.indices[0] == [2, 3, 4]).all() assert (count.values == [2, 1, 3]).all()
def test_ir_builder_triangle_count(): # 0 - 1 5 - 6 # | X | | / # 3 - 4 -- 2 - 7 # fmt: off indices = np.array( [ [0, 1], [0, 3], [0, 4], [1, 0], [1, 3], [1, 4], [2, 4], [2, 5], [2, 6], [2, 7], [3, 0], [3, 1], [3, 4], [4, 0], [4, 1], [4, 2], [4, 3], [5, 2], [5, 6], [6, 2], [6, 5], [7, 2], ], dtype=np.uint64, ) values = np.array( [ 100, 200, 300, 100, 400, 500, 99, 50, 55, 75, 200, 400, 600, 300, 500, 99, 600, 50, 60, 55, 60, 75 ], dtype=np.float64, ) # fmt: on sizes = np.array([8, 8], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) input_tensor = MLIRSparseTensor(indices, values, sizes, sparsity) assert input_tensor.verify() assert 5 == triangle_count(input_tensor) return
def test_matrix_to_from_coo(engine: MlirJitEngine, aliases: AliasMap): irb = MLIRFunctionBuilder( "matrix_to_coo", input_types=["tensor<?x?xf64, #CSR64>"], return_types=["tensor<?x?xindex>", "tensor<?xf64>"], aliases=aliases, ) (tensor, ) = irb.inputs indices, values = irb.graphblas.to_coo(tensor) irb.return_vars(indices, values) matrix_to_coo = irb.compile(engine=engine, passes=GRAPHBLAS_PASSES) # Test results input_indices = np.array( [[0, 1], [0, 2], [1, 0], [1, 3], [2, 0], [2, 4], [6, 2]], dtype=np.uint64) input_values = np.array([-100, 200, 300, 400, 175, 222, 333.333], dtype=np.float64) sizes = np.array([7, 5], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) graph = MLIRSparseTensor(input_indices, input_values, sizes, sparsity) graph.verify() new_indices, new_values = matrix_to_coo(graph) np.testing.assert_equal(input_indices, new_indices) np.testing.assert_allclose(input_values, new_values) irb = MLIRFunctionBuilder( "matrix_from_coo", input_types=["tensor<?x?xindex>", "tensor<?xf64>", "index", "index"], return_types=["tensor<?x?xf64, #CSR64>"], aliases=aliases, ) (indices, values, nrows, ncols) = irb.inputs tensor = irb.graphblas.from_coo(indices, values, (nrows, ncols)) irb.return_vars(tensor) matrix_from_coo = irb.compile(engine=engine, passes=GRAPHBLAS_PASSES) # Test results g2 = matrix_from_coo(input_indices, input_values, 7, 5) g2.verify() g2_indices, g2_values = matrix_to_coo(g2) np.testing.assert_equal(input_indices, g2_indices) np.testing.assert_allclose(input_values, g2_values)
def sparse_from_dense(arr): nrows, ncols = arr.shape indices = np.array([[i // ncols, i % ncols] for i in range(nrows * ncols)], dtype=np.uint64) values = arr.flatten() sizes = np.array(arr.shape, dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) return MLIRSparseTensor(indices, values, sizes, sparsity)
def test_vector_to_from_coo(engine: MlirJitEngine, aliases: AliasMap): irb = MLIRFunctionBuilder( "vector_to_coo", input_types=["tensor<?xf64, #CV64>"], return_types=["tensor<?x?xindex>", "tensor<?xf64>"], aliases=aliases, ) (tensor, ) = irb.inputs indices, values = irb.graphblas.to_coo(tensor) irb.return_vars(indices, values) vector_to_coo = irb.compile(engine=engine, passes=GRAPHBLAS_PASSES) # Test results input_indices = np.array([[1], [6], [8]], dtype=np.uint64) input_values = np.array([19.234, 1.1, 2.2], dtype=np.float64) sizes = np.array([20], dtype=np.uint64) sparsity = np.array([True], dtype=np.bool8) vec = MLIRSparseTensor(input_indices, input_values, sizes, sparsity) vec.verify() new_indices, new_values = vector_to_coo(vec) np.testing.assert_equal(input_indices, new_indices) np.testing.assert_allclose(input_values, new_values) irb = MLIRFunctionBuilder( "vector_from_coo", input_types=["tensor<?x?xindex>", "tensor<?xf64>", "index"], return_types=["tensor<?xf64, #CV64>"], aliases=aliases, ) (indices, values, size) = irb.inputs tensor = irb.graphblas.from_coo(indices, values, (size, )) irb.return_vars(tensor) vector_from_coo = irb.compile(engine=engine, passes=GRAPHBLAS_PASSES) # Test results v2 = vector_from_coo(input_indices, input_values, 20) v2.verify() v2_indices, v2_values = vector_to_coo(v2) np.testing.assert_equal(input_indices, v2_indices) np.testing.assert_allclose(input_values, v2_values)
def test_ties(): # fmt: off indices = np.array( [[0, 1], [0, 2], [1, 0], [1, 3], [2, 0], [2, 4], [3, 2], [4, 4]], dtype=np.uint64, ) values = np.array([100, 200, 300, 400, 175, 222, 333, 200], dtype=np.float64) # fmt: on sizes = np.array([5, 5], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) graph = MLIRSparseTensor(indices, values, sizes, sparsity) assert graph.verify() subgraph = mlalgo.totally_induced_edge_sampling(graph, 0.25, rand_seed=2021) assert np.all(subgraph.pointers[1] == [0, 1, 1, 3, 4, 5]) assert np.all(subgraph.indices[1] == [2, 0, 4, 2, 4])
def test_ir_builder_convert_layout_wrapper(engine: MlirJitEngine, aliases: AliasMap): ir_builder = MLIRFunctionBuilder( "convert_layout_wrapper", input_types=["tensor<?x?xf64, #CSR64>"], return_types=("tensor<?x?xf64, #CSC64>", ), aliases=aliases, ) (input_var, ) = ir_builder.inputs convert_layout_result = ir_builder.graphblas.convert_layout( input_var, "tensor<?x?xf64, #CSC64>") ir_builder.return_vars(convert_layout_result) assert ir_builder.get_mlir() # Test Compiled Function convert_layout_wrapper_callable = ir_builder.compile( engine=engine, passes=GRAPHBLAS_PASSES) indices = np.array( [ [1, 2], [4, 3], ], dtype=np.uint64, ) values = np.array([1.2, 4.3], dtype=np.float64) sizes = np.array([8, 8], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) input_tensor = MLIRSparseTensor(indices, values, sizes, sparsity) assert input_tensor.verify() dense_input_tensor = np.zeros([8, 8], dtype=np.float64) dense_input_tensor[1, 2] = 1.2 dense_input_tensor[4, 3] = 4.3 assert np.isclose(dense_input_tensor, input_tensor.toarray()).all() output_tensor = convert_layout_wrapper_callable(input_tensor) assert output_tensor.verify() assert np.isclose(dense_input_tensor, output_tensor.toarray()).all()
def test_mssp(special_passes): # This must be in sorted-for-CSR format. Random order breaks the constructor in strange ways. # fmt: off indices = np.array( [[0, 1], [0, 3], [1, 4], [1, 6], [2, 5], [3, 0], [3, 2], [4, 5], [5, 1], [6, 2], [6, 3], [6, 4]], dtype=np.uint64, ) # fmt: on values = np.array([2, 3, 8, 4, 1, 3, 3, 7, 1, 5, 7, 3], dtype=np.float64) sizes = np.array([7, 7], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) m = MLIRSparseTensor(indices, values, sizes, sparsity) assert m.verify() indices = np.array([[0, 1], [1, 3]], dtype=np.uint64) values = np.array([0, 0], dtype=np.float64) sizes = np.array([2, 7], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) v = MLIRSparseTensor(indices, values, sizes, sparsity) assert v.verify() # Compute MSSP # correct answer from node #1 -- [14, 0, 9, 11, 7, 10, 4] # correct answer from node #3 -- [3, 5, 3, 0, 12, 4, 9] w = mlalgo.mssp(m, v, compile_with_passes=special_passes) assert (w.indices[1] == [0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6]).all() assert (w.values == [14, 0, 9, 11, 7, 10, 4, 3, 5, 3, 0, 12, 4, 9]).all()
def sparsify_array(input_array: np.ndarray, sparsity_values: Sequence[bool], missing=0) -> MLIRSparseTensor: """Converts a numpy array into a MLIRSparseTensor.""" indices = np.array( list(zip(*(input_array != missing).nonzero())), dtype=np.uint64, ) values = np.array( [input_array[coordinate] for coordinate in map(tuple, indices)], dtype=input_array.dtype, ) sizes = np.array(input_array.shape, dtype=np.uint64) sparsity = np.array(sparsity_values, dtype=np.bool8) sparse_tensor = MLIRSparseTensor(indices, values, sizes, sparsity) return sparse_tensor
def mlir_from_scipy(x: ScipyGraph, **props) -> MLIRGraphBLASGraph: aprops = ScipyGraph.Type.compute_abstract_properties( x, { "node_type", "edge_type", "node_dtype", "edge_dtype", "is_directed" }) if aprops["edge_type"] != "map": raise NotImplementedError( f"MLIRGraphBLASGraph requires edge_type=='map', not '{aprops['edge_type']}'" ) if aprops["edge_type"] == "map" and x.value.dtype not in ( np.float32, np.float64, ): raise NotImplementedError( f"{MLIRGraphBLASGraph} instances with {aprops['edge_dtype']} edge weights not supported." ) matrix = x.value.tocoo() indices = np.column_stack([matrix.row, matrix.col]).astype(np.uint64) values = matrix.data sizes = np.array(matrix.shape, dtype=np.uint64) sparsity = np.array( [False, True], dtype=np.bool8 ) # TODO Make this sparse in all dimensions; [False, True] is equivalent to CSR mlir_sparse_tensor = MLIRSparseTensor(indices, values, sizes, sparsity) is_weighted = aprops["edge_type"] == "map" node_list = np.copy(x.node_list) if aprops["node_type"] == "map": node_vals = np.copy(x.node_vals) else: node_vals = None return MLIRGraphBLASGraph(mlir_sparse_tensor, is_weighted, node_list, node_vals, aprops=aprops)
def test_geolocation(): # fmt: off indices = np.array([ [0, 1], [0, 2], [1, 0], [1, 3], [2, 0], [2, 4], [2, 5], [3, 1], [3, 4], [3, 6], [3, 7], [4, 2], [4, 3], [5, 2], [6, 3], [7, 3], ], dtype=np.uint64) values = np.array([ 100, 200, 100, 300, 200, 400, 500, 300, 600, 700, 800, 400, 600, 500, 700, 800 ], dtype=np.float64) sizes = np.array([8, 8], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) graph = MLIRSparseTensor(indices, values, sizes, sparsity) indices = np.array([[1], [2], [4], [6], [7]], dtype=np.uint64) values = np.array([ 37.7449063493, 37.8668048274, 9.4276164485, 33.9774659389, 39.2598884729 ], dtype=np.float64) sizes = np.array([8], dtype=np.uint64) sparsity = np.array([True], dtype=np.bool8) lat = MLIRSparseTensor(indices, values, sizes, sparsity) indices = np.array([[1], [2], [4], [6], [7]], dtype=np.uint64) values = np.array([ -122.009432884, -122.257973253, -110.640705659, -114.886512278, -106.804662071 ], dtype=np.float64) lon = MLIRSparseTensor(indices, values, sizes, sparsity) expected_lat = np.array([ 37.80592086, 37.74490635, 37.86680483, 33.97769335, 9.42761645, 37.86680483, 33.97746594, 39.25988847 ]) expected_lon = np.array([ -122.13360051, -122.00943288, -122.25797325, -114.88638873, -110.64070566, -122.25797325, -114.88651228, -106.80466207 ]) # fmt: on lat, lon = mlalgo.geolocation(graph, lat, lon) np.testing.assert_allclose(lat.toarray(), expected_lat) np.testing.assert_allclose(lon.toarray(), expected_lon)
def test_verify(nrows, ncols): for indices in itertools.chain.from_iterable( itertools.combinations(list(range(nrows * ncols)), n) for n in range(nrows * ncols + 1)): rows = [x // ncols for x in indices] cols = [x % ncols for x in indices] vals = np.arange(len(indices)) M = gb.Matrix.from_values(rows, cols, vals, nrows=nrows, ncols=ncols) M.wait() # CSR d = M.ss.export("csr", sort=True) r, c, v = M.to_values() mt = MLIRSparseTensor( np.stack([r, c]).T.copy(), v, np.array(M.shape, dtype=np.uint64), np.array([False, True], dtype=np.bool8), ) verify(mt) assert mt.verify() np.testing.assert_array_equal(mt.get_pointers(0), []) np.testing.assert_array_equal(mt.get_pointers(1), d["indptr"]) np.testing.assert_array_equal(mt.get_indices(0), []) np.testing.assert_array_equal(mt.get_indices(1), d["col_indices"]) np.testing.assert_array_equal(mt.values, d["values"]) assert mt.shape == M.shape assert mt.sizes == M.shape # HyperCSR d = M.ss.export("hypercsr", sort=True) r, c, v = M.to_values() mt = MLIRSparseTensor( np.stack([r, c]).T.copy(), v, np.array(M.shape, dtype=np.uint64), np.array([True, True], dtype=np.bool8), ) verify(mt) assert mt.verify() np.testing.assert_array_equal(mt.get_pointers(0), [0, len(d["rows"])]) np.testing.assert_array_equal(mt.get_pointers(1), d["indptr"]) np.testing.assert_array_equal(mt.get_indices(0), d["rows"]) np.testing.assert_array_equal(mt.get_indices(1), d["col_indices"]) np.testing.assert_array_equal(mt.values, d["values"]) assert mt.shape == M.shape assert mt.sizes == M.shape # CSC d = M.ss.export("csc", sort=True) M2 = M.T.new() M2.wait() c, r, v = M2.to_values() mt = MLIRSparseTensor( np.stack([c, r]).T.copy(), v, np.array(M.shape[::-1], dtype=np.uint64), np.array([False, True], dtype=np.bool8), np.array([1, 0], dtype=np.uint64), ) verify(mt) assert mt.verify() np.testing.assert_array_equal(mt.get_pointers(0), []) np.testing.assert_array_equal(mt.get_pointers(1), d["indptr"]) np.testing.assert_array_equal(mt.get_indices(0), []) np.testing.assert_array_equal(mt.get_indices(1), d["row_indices"]) np.testing.assert_array_equal(mt.values, d["values"]) assert mt.shape == M.shape assert mt.sizes[::-1] == M.shape # HyperCSC d = M.ss.export("hypercsc", sort=True) M2 = M.T.new() M2.wait() c, r, v = M2.to_values() mt = MLIRSparseTensor( np.stack([c, r]).T.copy(), v, np.array(M.shape[::-1], dtype=np.uint64), np.array([True, True], dtype=np.bool8), np.array([1, 0], dtype=np.uint64), ) verify(mt) assert mt.verify() np.testing.assert_array_equal(mt.get_pointers(0), [0, len(d["cols"])]) np.testing.assert_array_equal(mt.get_pointers(1), d["indptr"]) np.testing.assert_array_equal(mt.get_indices(0), d["cols"]) np.testing.assert_array_equal(mt.get_indices(1), d["row_indices"]) np.testing.assert_array_equal(mt.values, d["values"]) assert mt.shape == M.shape assert mt.sizes[::-1] == M.shape if vals.size == nrows * ncols: # FullR d = M.ss.export("fullr", sort=True) r, c, v = M.to_values() mt = MLIRSparseTensor( np.stack([r, c]).T.copy(), v, np.array(M.shape, dtype=np.uint64), np.array([False, False], dtype=np.bool8), ) verify(mt) assert mt.verify() np.testing.assert_array_equal(mt.values, d["values"].ravel()) assert mt.shape == M.shape assert mt.sizes == M.shape # FullC d = M.ss.export("fullc", sort=True) M2 = M.T.new() M2.wait() c, r, v = M2.to_values() mt = MLIRSparseTensor( np.stack([c, r]).T.copy(), v, np.array(M.shape[::-1], dtype=np.uint64), np.array([False, False], dtype=np.bool8), np.array([1, 0], dtype=np.uint64), ) verify(mt) assert mt.verify() np.testing.assert_array_equal(mt.values, d["values"].T.ravel()) assert mt.shape == M.shape assert mt.sizes[::-1] == M.shape
def test_jit_engine_zero_values(engine): mlir_text = """ module { func private @_mlir_ciface_sparseValuesF64(!llvm.ptr<i8>) -> memref<?xf64> func private @_mlir_ciface_sparseIndices64(!llvm.ptr<i8>, index) -> memref<?xindex> func private @_mlir_ciface_sparsePointers64(!llvm.ptr<i8>, index) -> memref<?xindex> func private @sparseDimSize(!llvm.ptr<i8>, index) -> index func @transpose(%output: !llvm.ptr<i8>, %input: !llvm.ptr<i8>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %n_row = call @sparseDimSize(%input, %c0) : (!llvm.ptr<i8>, index) -> index %n_col = call @sparseDimSize(%input, %c1) : (!llvm.ptr<i8>, index) -> index %Ap = call @_mlir_ciface_sparsePointers64(%input, %c1) : (!llvm.ptr<i8>, index) -> memref<?xindex> %Aj = call @_mlir_ciface_sparseIndices64(%input, %c1) : (!llvm.ptr<i8>, index) -> memref<?xindex> %Ax = call @_mlir_ciface_sparseValuesF64(%input) : (!llvm.ptr<i8>) -> memref<?xf64> %Bp = call @_mlir_ciface_sparsePointers64(%output, %c1) : (!llvm.ptr<i8>, index) -> memref<?xindex> %Bi = call @_mlir_ciface_sparseIndices64(%output, %c1) : (!llvm.ptr<i8>, index) -> memref<?xindex> %Bx = call @_mlir_ciface_sparseValuesF64(%output) : (!llvm.ptr<i8>) -> memref<?xf64> %nnz = memref.load %Ap[%n_row] : memref<?xindex> // compute number of non-zero entries per column of A scf.for %arg2 = %c0 to %n_col step %c1 { memref.store %c0, %Bp[%arg2] : memref<?xindex> } scf.for %n = %c0 to %nnz step %c1 { %colA = memref.load %Aj[%n] : memref<?xindex> %colB = memref.load %Bp[%colA] : memref<?xindex> %colB1 = arith.addi %colB, %c1 : index memref.store %colB1, %Bp[%colA] : memref<?xindex> } // cumsum the nnz per column to get Bp memref.store %c0, %Bp[%n_col] : memref<?xindex> scf.for %col = %c0 to %n_col step %c1 { %temp = memref.load %Bp[%col] : memref<?xindex> %cumsum = memref.load %Bp[%n_col] : memref<?xindex> memref.store %cumsum, %Bp[%col] : memref<?xindex> %cumsum2 = arith.addi %cumsum, %temp : index memref.store %cumsum2, %Bp[%n_col] : memref<?xindex> } scf.for %row = %c0 to %n_row step %c1 { %j_start = memref.load %Ap[%row] : memref<?xindex> %row_plus1 = arith.addi %row, %c1 : index %j_end = memref.load %Ap[%row_plus1] : memref<?xindex> scf.for %jj = %j_start to %j_end step %c1 { %col = memref.load %Aj[%jj] : memref<?xindex> %dest = memref.load %Bp[%col] : memref<?xindex> memref.store %row, %Bi[%dest] : memref<?xindex> %axjj = memref.load %Ax[%jj] : memref<?xf64> memref.store %axjj, %Bx[%dest] : memref<?xf64> // Bp[col]++ %bp_inc = memref.load %Bp[%col] : memref<?xindex> %bp_inc1 = arith.addi %bp_inc, %c1 : index memref.store %bp_inc1, %Bp[%col]: memref<?xindex> } } %last_last = memref.load %Bp[%n_col] : memref<?xindex> memref.store %c0, %Bp[%n_col] : memref<?xindex> scf.for %col = %c0 to %n_col step %c1 { %temp = memref.load %Bp[%col] : memref<?xindex> %last = memref.load %Bp[%n_col] : memref<?xindex> memref.store %last, %Bp[%col] : memref<?xindex> memref.store %temp, %Bp[%n_col] : memref<?xindex> } memref.store %last_last, %Bp[%n_col] : memref<?xindex> return } } """ indices = np.array( [ [0, 0], [1, 0], ], dtype=np.uint64, ) values = np.array([8, 9], dtype=np.float64) sizes = np.array([2, 2], dtype=np.uint64) sparsity = np.array([False, True], dtype=np.bool8) input_tensor = MLIRSparseTensor(indices, values, sizes, sparsity) output_tensor = MLIRSparseTensor(indices, values, sizes, sparsity) assert input_tensor.verify() assert output_tensor.verify() assert engine.add(mlir_text, GRAPHBLAS_PASSES) == ["transpose"] assert engine.transpose(output_tensor, input_tensor) is None assert input_tensor.verify() assert np.isclose(input_tensor.toarray(), np.array([[8, 0], [9, 0]])).all() assert output_tensor.verify() assert np.isclose(output_tensor.toarray(), np.array([[8, 9], [0, 0]])).all() return
def test_jit_engine_sequence_of_sparse_tensors_input(engine): mlir_text = """ #trait_sum_reduction = { indexing_maps = [ affine_map<(i,j) -> (i,j)>, affine_map<(i,j) -> ()> ], iterator_types = ["reduction", "reduction"], doc = "Sparse Tensor Sum" } #sparseTensor = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ], pointerBitWidth = 64, indexBitWidth = 64 }> func private @ptr8_to_matrix_csr_f64_p64i64(!llvm.ptr<i8>) -> tensor<2x3xf64, #sparseTensor> func @sparse_tensors_summation(%sequence: !llvm.ptr<!llvm.ptr<i8>>, %sequence_length: index) -> f64 { // Take an array of sparse 2x3 matrices %output_storage = arith.constant dense<0.0> : tensor<f64> %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %ci0 = arith.constant 0 : i64 %ci1 = arith.constant 1 : i64 %cf0 = arith.constant 0.0 : f64 %sum_memref = memref.alloc() : memref<f64> memref.store %cf0, %sum_memref[] : memref<f64> scf.for %i = %c0 to %sequence_length step %c1 iter_args(%iter=%ci0) -> (i64) { // llvm.getelementptr just does pointer arithmetic %sparse_tensor_ptr_ptr = llvm.getelementptr %sequence[%iter] : (!llvm.ptr<!llvm.ptr<i8>>, i64) -> !llvm.ptr<!llvm.ptr<i8>> // dereference %sparse_tensor_ptr_ptr to get an !llvm.ptr<i8> %sparse_tensor_ptr = llvm.load %sparse_tensor_ptr_ptr : !llvm.ptr<!llvm.ptr<i8>> %sparse_tensor = call @ptr8_to_matrix_csr_f64_p64i64(%sparse_tensor_ptr) : (!llvm.ptr<i8>) -> tensor<2x3xf64, #sparseTensor> %reduction = linalg.generic #trait_sum_reduction ins(%sparse_tensor: tensor<2x3xf64, #sparseTensor>) outs(%output_storage: tensor<f64>) { ^bb(%a: f64, %x: f64): %0 = arith.addf %x, %a : f64 linalg.yield %0 : f64 } -> tensor<f64> %reduction_value = tensor.extract %reduction[] : tensor<f64> %current_sum = memref.load %sum_memref[] : memref<f64> %updated_sum = arith.addf %reduction_value, %current_sum : f64 memref.store %updated_sum, %sum_memref[] : memref<f64> %plus_one = arith.addi %iter, %ci1 : i64 scf.yield %plus_one : i64 } %sum = memref.load %sum_memref[] : memref<f64> return %sum : f64 } """ assert engine.add(mlir_text, GRAPHBLAS_PASSES) == ["sparse_tensors_summation"] num_sparse_tensors = 10 # generate values sqrt_values = np.sqrt(np.arange(1, num_sparse_tensors + 1, 0.5)) value_iter = iter(sqrt_values) next_values = lambda: np.array( [next(value_iter), next(value_iter)], dtype=np.float64 ) # generate coordinates coordinate_iter = itertools.count() next_indices = lambda: np.array( sorted( [ (next(coordinate_iter) % 2, next(coordinate_iter) % 3), (next(coordinate_iter) % 2, next(coordinate_iter) % 3), ] ), dtype=np.uint64, ) sparse_tensors = [] for _ in range(num_sparse_tensors): indices = next_indices() values = next_values() sizes = np.array([2, 3], dtype=np.uint64) sparsity = np.array([True, True], dtype=np.bool8) sparse_tensor = MLIRSparseTensor(indices, values, sizes, sparsity) assert sparse_tensor.verify() sparse_tensors.append(sparse_tensor) expected_sum = sqrt_values.sum() assert np.isclose( expected_sum, engine.sparse_tensors_summation(sparse_tensors, num_sparse_tensors), ) return