def test_obim_python(threads_1):
# pylint: disable=unused-argument
order = []
def metric(out, i):
return out[i]
def f(out, i, ctx):
order.append(i)
orig = out[i]
out[i] = 10 - i
if orig == 0:
ctx.push(i)
out = np.zeros(10, dtype=int)
for_each(
range(10),
partial(f, out),
worklist=OrderedByIntegerMetric(partial(metric, out)),
)
assert np.allclose(out, np.array([10, 9, 8, 7, 6, 5, 4, 3, 2, 1]))
assert order == [
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0
]
def kcore_async(graph: Graph, k_core_num, property_name):
num_nodes = graph.num_nodes()
initial_worklist = InsertBag[np.uint64]()
current_degree = NUMAArray[np.uint64](num_nodes, AllocationPolicy.INTERLEAVED)
timer = StatTimer("Kcore: Property Graph Numba: " + property_name)
timer.start()
# Initialize
do_all(
range(num_nodes), compute_degree_count_operator(graph, current_degree.as_numpy()), steal=True,
)
# Setup initial worklist
do_all(
range(num_nodes),
setup_initial_worklist_operator(initial_worklist, current_degree.as_numpy(), k_core_num),
steal=True,
)
# Compute k-core
for_each(
initial_worklist,
compute_async_kcore_operator(graph, current_degree.as_numpy(), k_core_num),
steal=True,
disable_conflict_detection=True,
)
timer.stop()
# Add the ranks as a new property to the property graph
graph.add_node_property(pyarrow.table({property_name: current_degree}))
def test_for_each_wrong_closure():
@do_all_operator()
def f(out, i):
out[i] = i + 1
out = np.zeros(10, dtype=int)
with pytest.raises(TypeError):
for_each(range(10), f(out))
def test_for_each_no_push(modes):
@for_each_operator()
def f(out, i, ctx):
# pylint: disable=unused-argument
out[i] += i + 1
out = np.zeros(10, dtype=int)
for_each(range(10), f(out), **modes)
assert np.allclose(out, np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]))
def test_for_each(modes):
@for_each_operator()
def f(out, i, ctx):
out[i] += i + 1
if i == 8:
ctx.push(10)
out = np.zeros(11, dtype=int)
for_each(range(10), f(out), **modes)
assert np.allclose(out, np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]))
def test_for_each_conflict_detection_unsupported(lock):
total = 0
def f(i, ctx):
# pylint: disable=unused-argument
nonlocal total
with lock:
total += i
with pytest.raises(ValueError):
for_each(range(10), f, disable_conflict_detection=False)
def test_for_each_python_with_push(modes, lock):
total = 0
def f(i, ctx):
nonlocal total
with lock:
total += i
if i == 8:
ctx.push(100)
for_each(range(10), f, **modes)
assert total == 145
def test_per_socket_chunk_fifo(threads_1):
# pylint: disable=unused-argument
order = []
def f(out, i, ctx):
order.append(i)
orig = out[i]
out[i] = 10 - i
if orig == 0:
ctx.push(9 - i)
out = np.zeros(10, dtype=int)
for_each(range(10), partial(f, out), worklist=PerSocketChunkFIFO())
assert np.allclose(out, np.array([10, 9, 8, 7, 6, 5, 4, 3, 2, 1]))
assert order == [
0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0
]
def test_for_each_opaque(modes):
from katana.local import InsertBag
@for_each_operator()
def f(out, s, ctx):
out[s.y] += s.x
if s.y < 10:
ctx.push((s.x + 1, s.y + 1))
dt = np.dtype([
("x", np.float32),
("y", np.int8),
], align=True)
data = InsertBag[dt]()
data.push((1.1, 0))
out = np.zeros(10, dtype=float)
for_each(data, f(out), **modes)
assert np.allclose(
out, np.array([1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1, 9.1, 10.1]))
def sssp(graph: Graph, source, length_property, shift, property_name):
dists = create_distance_array(graph, source, length_property)
# Define the struct type here so it can depend on the type of the weight property
UpdateRequest = np.dtype([("src", np.uint32), ("dist", dists.dtype)])
init_bag = InsertBag[UpdateRequest]()
init_bag.push((source, 0))
t = StatTimer("Total SSSP")
t.start()
for_each(
init_bag,
sssp_operator(graph, dists, graph.get_edge_property(length_property)),
worklist=OrderedByIntegerMetric(obim_indexer(shift)),
disable_conflict_detection=True,
loop_name="SSSP",
)
t.stop()
print("Elapsed time: ", t.get(), "milliseconds.")
graph.add_node_property(pyarrow.table({property_name: dists}))