def engine():
jit_engine = MlirJitEngine()
jit_engine.add(
"""
#CSR64 = #sparse_tensor.encoding<{
dimLevelType = [ "dense", "compressed" ],
dimOrdering = affine_map<(i,j) -> (i,j)>,
pointerBitWidth = 64,
indexBitWidth = 64
}>
#CSC64 = #sparse_tensor.encoding<{
dimLevelType = [ "dense", "compressed" ],
dimOrdering = affine_map<(i,j) -> (j,i)>,
pointerBitWidth = 64,
indexBitWidth = 64
}>
func @csr_to_csc(%matrix: tensor<?x?xf64, #CSR64>) -> tensor<?x?xf64, #CSC64> {
%converted = graphblas.convert_layout %matrix : tensor<?x?xf64, #CSR64> to tensor<?x?xf64, #CSC64>
return %converted : tensor<?x?xf64, #CSC64>
}
""",
GRAPHBLAS_PASSES,
)
return jit_engine
def test_profile_multiple_returns(engine: MlirJitEngine, ):
result_dir = tempfile.TemporaryDirectory()
num_iterations = 12345678
function_name = f"multivalue_slow_mul_{int(uuid.uuid4())}"
mlir_text = f"""
func @{function_name}(%input_val: i64) -> (i64, i64) {{
%ci0 = constant 0 : i64
%c0 = constant 0 : index
%c1 = constant 1 : index
%num_iterations = constant {num_iterations} : index
%answer = scf.for %i = %c0 to %num_iterations step %c1 iter_args(%current_sum=%ci0) -> (i64) {{
%new_sum = addi %current_sum, %input_val : i64
scf.yield %new_sum : i64
}}
return %input_val, %answer : i64, i64
}}
"""
try:
engine.add(
mlir_text,
GRAPHBLAS_PASSES,
profile=True,
profile_result_directory=result_dir.name,
)
except NotImplementedError as e:
if not sys.platform.startswith("linux"):
pytest.skip(f"Profiling not supported on {sys.platform}.")
except RuntimeError as e:
(err_string, ) = e.args
if err_string == ("Profiling not permitted since the contents of "
"/proc/sys/kernel/perf_event_paranoid must be -1."):
pytest.skip(f"Profiling not permitted.")
else:
raise e
compiled_func = getattr(engine, function_name)
input_val = 12
with captured_stdout() as stdout_string_container:
redundant_input_val, result = compiled_func(input_val)
assert redundant_input_val == input_val
assert result == input_val * num_iterations
(stdout_string, ) = stdout_string_container
assert "Source code & Disassembly of" in stdout_string
assert "Sorted summary for file" in stdout_string
assert "Percent" in stdout_string
assert (len([
file_name for file_name in os.listdir(result_dir.name)
if file_name.startswith("perf-") and file_name.endswith(".data")
]) == 1)
return
def engine():
return MlirJitEngine()
def test_ir_diag(
matrix_type_template: str,
vector_type_template: str,
mlir_type: str,
engine: MlirJitEngine,
aliases: AliasMap,
):
matrix_type = matrix_type_template.format(scalar_type=mlir_type)
vector_type = vector_type_template.format(scalar_type=mlir_type)
np_type = MLIR_TYPE_TO_NP_TYPE[mlir_type]
# Build Functions
ir_builder = MLIRFunctionBuilder(
f"diag_func_{mlir_type}",
input_types=[vector_type, matrix_type],
return_types=[
matrix_type,
vector_type,
],
aliases=aliases,
)
(input_vector, input_matrix) = ir_builder.inputs
output_matrix = ir_builder.graphblas.diag(input_vector, matrix_type)
output_vector = ir_builder.graphblas.diag(input_matrix, vector_type)
ir_builder.return_vars(output_matrix, output_vector)
diag_func = ir_builder.compile(engine=engine, passes=GRAPHBLAS_PASSES)
# Test Results
dense_input_vector = np.array(
[0, 0, 0, 1, 0, -2, 0, 0],
dtype=np_type,
)
input_vector = sparsify_array(dense_input_vector, [True])
assert input_vector.verify()
dense_input_matrix = np.array(
[
[0, 7, 7, 0, 7],
[0, 1, 7, 0, 0],
[0, 1, 0, 7, 0],
[0, 7, 0, 2, 0],
[7, 7, 0, 0, 0],
],
dtype=np_type,
)
input_matrix = sparsify_array(dense_input_matrix, [False, True])
assert input_matrix.verify()
matrix_type_is_csc = [1, 0] == SparseTensorType.parse(
matrix_type, aliases).encoding.ordering
if matrix_type_is_csc:
input_matrix = engine.csr_to_csc(input_matrix)
output_matrix, output_vector = diag_func(input_vector, input_matrix)
assert output_matrix.verify()
assert output_vector.verify()
output_matrix = output_matrix.toarray()
output_vector = output_vector.toarray()
expected_output_matrix = np.diagflat(dense_input_vector)
expected_output_vector = np.diag(dense_input_matrix)
assert np.all(expected_output_matrix == output_matrix)
assert np.all(expected_output_vector == output_vector)
def test_ir_reduce_to_vector(
input_type_template: str,
reduce_rows_output_type_template: str,
reduce_columns_output_type_template: str,
mlir_type: str,
engine: MlirJitEngine,
aliases: AliasMap,
):
input_type = input_type_template.format(scalar_type=mlir_type)
reduce_rows_output_type = reduce_rows_output_type_template.format(
scalar_type=mlir_type)
reduce_columns_output_type = reduce_columns_output_type_template.format(
scalar_type=mlir_type)
np_type = MLIR_TYPE_TO_NP_TYPE[mlir_type]
# Build Functions
ir_builder = MLIRFunctionBuilder(
f"reduce_func_{mlir_type}",
input_types=[input_type],
return_types=[
reduce_rows_output_type,
"tensor<?xi64, #CV64>",
reduce_rows_output_type,
"tensor<?xi64, #CV64>",
"tensor<?xi64, #CV64>",
"tensor<?xi64, #CV64>",
],
aliases=aliases,
)
(matrix, ) = ir_builder.inputs
reduced_rows = ir_builder.graphblas.reduce_to_vector(matrix, "plus", 1)
reduced_columns = ir_builder.graphblas.reduce_to_vector(matrix, "count", 0)
zero_scalar = ir_builder.arith.constant(0, mlir_type)
reduced_rows_clamped = ir_builder.graphblas.apply(reduced_rows,
"min",
right=zero_scalar)
reduced_rows_clamped = ir_builder.graphblas.apply(reduced_rows_clamped,
"identity")
reduced_columns_abs = ir_builder.graphblas.apply(reduced_columns, "abs")
reduced_columns_abs = ir_builder.graphblas.apply(reduced_columns_abs,
"identity")
reduced_columns_negative_abs = ir_builder.graphblas.apply(
reduced_columns, "ainv")
reduced_columns_negative_abs = ir_builder.graphblas.apply(
reduced_columns_negative_abs, "identity")
reduced_rows_argmin = ir_builder.graphblas.reduce_to_vector(
matrix, "argmin", 1)
reduced_columns_argmax = ir_builder.graphblas.reduce_to_vector(
matrix, "argmax", 0)
ir_builder.return_vars(
reduced_rows,
reduced_columns,
reduced_rows_clamped,
reduced_columns_negative_abs,
reduced_rows_argmin,
reduced_columns_argmax,
)
reduce_func = ir_builder.compile(engine=engine, passes=GRAPHBLAS_PASSES)
# Test Results
dense_input_tensor = np.array(
[
[1, 0, 0, 0],
[-2, 0, 3, -4],
[0, 0, 0, 0],
[0, 0, 5, -6],
[0, -7, 0, 8],
],
dtype=np_type,
)
input_tensor = sparsify_array(dense_input_tensor, [False, True])
input_type_is_csc = [1, 0
] == SparseTensorType.parse(input_type,
aliases).encoding.ordering
if input_type_is_csc:
input_tensor = engine.csr_to_csc(input_tensor)
(
reduced_rows,
reduced_columns,
reduced_rows_clamped,
reduced_columns_negative_abs,
reduced_rows_argmin,
reduced_columns_argmax,
) = reduce_func(input_tensor)
assert reduced_rows.verify()
assert reduced_columns.verify()
assert reduced_rows_clamped.verify()
assert reduced_columns_negative_abs.verify()
assert reduced_rows_argmin.verify()
assert reduced_columns_argmax.verify()
reduced_rows = reduced_rows.toarray()
reduced_columns = reduced_columns.toarray()
reduced_rows_clamped = reduced_rows_clamped.toarray()
reduced_columns_negative_abs = reduced_columns_negative_abs.toarray()
reduced_rows_argmin = reduced_rows_argmin.toarray()
reduced_columns_argmax = reduced_columns_argmax.toarray()
expected_reduced_rows = dense_input_tensor.sum(axis=1)
expected_reduced_columns = (dense_input_tensor.astype(bool).sum(
axis=0).astype(np_type))
expected_reduced_rows_clamped = np.copy(expected_reduced_rows)
expected_reduced_rows_clamped[expected_reduced_rows_clamped > 0] = 0
expected_reduced_columns_negative_abs = -np.abs(expected_reduced_columns)
M = dense_input_tensor.copy()
M[dense_input_tensor == 0] = dense_input_tensor.max() + 1
expected_reduced_rows_argmin = np.argmin(M, axis=1)
M = dense_input_tensor.copy()
M[dense_input_tensor == 0] = dense_input_tensor.min() - 1
expected_reduced_columns_argmax = np.argmax(M, axis=0)
assert np.all(reduced_rows == expected_reduced_rows)
assert np.all(reduced_columns == expected_reduced_columns)
assert np.all(reduced_rows_clamped == expected_reduced_rows_clamped)
assert np.all(
reduced_columns_negative_abs == expected_reduced_columns_negative_abs)
assert np.all(reduced_rows_argmin == expected_reduced_rows_argmin)
assert np.all(reduced_columns_argmax == expected_reduced_columns_argmax)
def engine():
jit_engine = MlirJitEngine()
return jit_engine