diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt --- a/mlir/CMakeLists.txt +++ b/mlir/CMakeLists.txt @@ -188,3 +188,7 @@ endif() add_subdirectory(cmake/modules) + +if (MLIR_ENABLE_CXX_BENCHMARKS) + add_subdirectory(benchmark) +endif() diff --git a/mlir/benchmark/CMakeLists.txt b/mlir/benchmark/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/benchmark/CMakeLists.txt @@ -0,0 +1,32 @@ +include(ExternalProject) + +set(LLVM_LINK_COMPONENTS + Support) + +set(GOOGLE_BENCHMARK_TARGET_FLAGS ${BENCHMARK_DIALECT_FLAG}) +if (MLIR_BENCHMARK_GCC_TOOLCHAIN) + set(GOOGLE_BENCHMARK_TARGET_FLAGS + --gcc-toolchain=${MLIR_BENCHMARK_GCC_TOOLCHAIN}) +endif() +string(REPLACE ";" " " GOOGLE_BENCHMARK_TARGET_FLAGS "${GOOGLE_BENCHMARK_TARGET_FLAGS}") + +ExternalProject_Add(google-benchmark + EXCLUDE_FROM_ALL ON + PREFIX google-benchmark + SOURCE_DIR ${MLIR_SOURCE_DIR}/../llvm/utils/benchmark + INSTALL_DIR ${CMAKE_CURRENT_BINARY_DIR}/google-benchmark + CMAKE_CACHE_ARGS + -DBUILD_SHARED_LIBS:BOOL=OFF + -DCMAKE_POSITION_INDEPENDENT_CODE:BOOL=ON + -DCMAKE_C_COMPILER:STRING=${CMAKE_C_COMPILER} + -DCMAKE_CXX_COMPILER:STRING=${CMAKE_CXX_COMPILER} + -DCMAKE_CXX_FLAGS:STRING=${GOOGLE_BENCHMARK_TARGET_FLAGS} + -DCMAKE_CXX_STANDARD:STRING=14 + -DCMAKE_BUILD_TYPE:STRING=RELEASE + -DCMAKE_INSTALL_PREFIX:PATH= + -DBENCHMARK_ENABLE_TESTING:BOOL=OFF) + +set(GOOGLE_BENCHMARK_MLIR_INSTALL ${CMAKE_CURRENT_BINARY_DIR}/google-benchmark) +set(GOOGLE_BENCHMARK_LINK_FLAGS -L${GOOGLE_BENCHMARK_MLIR_INSTALL}/lib/) + +add_benchmark(MLIRStarter MLIRStarter.cpp) diff --git a/mlir/benchmark/Dialect/SparseTensor/python/common.py b/mlir/benchmark/Dialect/SparseTensor/python/common.py new file mode 100644 --- /dev/null +++ b/mlir/benchmark/Dialect/SparseTensor/python/common.py @@ -0,0 +1,15 @@ +import timeit + + +def run_benchmark(execution_time_upper_limit, ntimes=10, verbose=False): + def wrapper(f): + print(f"BENCHMARK: {f.__name__}") + time_taken_on_runs = timeit.repeat(f, repeat=ntimes, number=1) + average_time_taken = sum(time_taken_on_runs) / len(time_taken_on_runs) + if verbose: + print(f"Time taken on all runs: {time_taken_on_runs}") + if average_time_taken < execution_time_upper_limit: + print(f"Average time taken less than {execution_time_upper_limit} seconds") + else: + print(f"Average time taken exceeds {execution_time_upper_limit} seconds") + return wrapper \ No newline at end of file diff --git a/mlir/benchmark/Dialect/SparseTensor/python/lit.local.cfg b/mlir/benchmark/Dialect/SparseTensor/python/lit.local.cfg new file mode 100644 --- /dev/null +++ b/mlir/benchmark/Dialect/SparseTensor/python/lit.local.cfg @@ -0,0 +1,5 @@ +# Disable ASAN's leak detection for python OpsDSL tests. +config.environment['ASAN_OPTIONS'] = 'detect_leaks=0' +# Only run when python bindings are enabled. +if not config.enable_bindings_python: + config.unsupported = True diff --git a/mlir/benchmark/Dialect/SparseTensor/python/test_SpMM.py b/mlir/benchmark/Dialect/SparseTensor/python/test_SpMM.py new file mode 100644 --- /dev/null +++ b/mlir/benchmark/Dialect/SparseTensor/python/test_SpMM.py @@ -0,0 +1,170 @@ +# RUN: SUPPORT_LIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s + +import ctypes +import numpy as np +import os + +import mlir.all_passes_registration + +from mlir import ir +from mlir import runtime as rt +from mlir import execution_engine +from mlir import passmanager + +from mlir.dialects import sparse_tensor as st +from mlir.dialects import builtin +from mlir.dialects.linalg.opdsl import lang as dsl +from common import run_benchmark + + +@dsl.linalg_structured_op +def matmul_dsl( + A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K), + B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N), + C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True)): + C[dsl.D.m, dsl.D.n] += A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n] + + +def build_SpMM(attr: st.EncodingAttr): + """Build SpMM kernel. + + This method generates a linalg op with for matrix multiplication using + just the Python API. Effectively, a generic linalg op is constructed + that computes C(i,j) += A(i,k) * B(k,j) for annotated matrix A. + """ + module = ir.Module.create() + f64 = ir.F64Type.get() + a = ir.RankedTensorType.get([3, 4], f64, attr) + b = ir.RankedTensorType.get([4, 2], f64) + c = ir.RankedTensorType.get([3, 2], f64) + arguments = [a, b, c] + with ir.InsertionPoint(module.body): + + @builtin.FuncOp.from_py_func(*arguments) + def spMxM(*args): + return matmul_dsl(args[0], args[1], outs=[args[2]]) + + return module + + +def boilerplate(attr: st.EncodingAttr): + """Returns boilerplate main method. + + This method sets up a boilerplate main method that takes three tensors + (a, b, c), converts the first tensor a into s sparse tensor, and then + calls the sparse kernel for matrix multiplication. For convenience, + this part is purely done as string input. + """ + return f""" +func @main(%ad: tensor<3x4xf64>, %b: tensor<4x2xf64>, %c: tensor<3x2xf64>) -> tensor<3x2xf64> + attributes {{ llvm.emit_c_interface }} {{ + %a = sparse_tensor.convert %ad : tensor<3x4xf64> to tensor<3x4xf64, {attr}> + %0 = call @spMxM(%a, %b, %c) : (tensor<3x4xf64, {attr}>, + tensor<4x2xf64>, + tensor<3x2xf64>) -> tensor<3x2xf64> + return %0 : tensor<3x2xf64> +}} +""" + + +def build_compile_and_run_SpMM(attr: st.EncodingAttr, support_lib: str, + compiler): + # Build. + module = build_SpMM(attr) + func = str(module.operation.regions[0].blocks[0].operations[0].operation) + module = ir.Module.parse(func + boilerplate(attr)) + + # Compile. + compiler(module) + engine = execution_engine.ExecutionEngine( + module, opt_level=0, shared_libs=[support_lib]) + + # Set up numpy input and buffer for output. + a = np.array( + [[1.1, 0.0, 0.0, 1.4], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 3.3, 0.0]], + np.float64) + b = np.array([[1.0, 2.0], [4.0, 3.0], [5.0, 6.0], [8.0, 7.0]], np.float64) + c = np.zeros((3, 2), np.float64) + out = np.zeros((3, 2), np.float64) + + mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a))) + mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b))) + mem_c = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(c))) + mem_out = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(out))) + + # Invoke the kernel and get numpy output. + # Built-in bufferization uses in-out buffers. + # TODO: replace with inplace comprehensive bufferization. + engine.invoke('main', mem_out, mem_a, mem_b, mem_c) + + # Sanity check on computed result. + expected = np.matmul(a, b); + c = rt.ranked_memref_to_numpy(mem_out[0]) + if np.allclose(c, expected): + pass + else: + quit(f'FAILURE') + + +class SparseCompiler: + """Sparse compiler passes.""" + + def __init__(self, options: str): + pipeline = ( + f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),' + f'sparsification{{{options}}},' + f'sparse-tensor-conversion,' + f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),' + f'convert-scf-to-std,' + f'func-bufferize,' + f'tensor-constant-bufferize,' + f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),' + f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},' + f'lower-affine,' + f'convert-memref-to-llvm,' + f'convert-std-to-llvm,' + f'reconcile-unrealized-casts') + self.pipeline = pipeline + + def __call__(self, module: ir.Module): + passmanager.PassManager.parse(self.pipeline).run(module) + + +# CHECK-LABEL: TEST: testSpMM +# CHECK: Passed 8 tests +@run_benchmark(1) +def testSpMM(): + # Obtain path to runtime support library. + support_lib = os.getenv('SUPPORT_LIB') + assert os.path.exists(support_lib), f'{support_lib} does not exist' + + with ir.Context() as ctx, ir.Location.unknown(): + count = 0 + # Loop over various ways to compile and annotate the SpMM kernel with + # a *single* sparse tensor. Note that we deliberate do not exhaustively + # search the full state space to reduce runtime of the test. It is + # straightforward to adapt the code below to explore more combinations. + par = 0 + vec = 0 + vl = 1 + e = False + opt = (f'parallelization-strategy={par} ' + f'vectorization-strategy={vec} ' + f'vl={vl} enable-simd-index32={e}') + levels = [[st.DimLevelType.dense, st.DimLevelType.dense], + [st.DimLevelType.dense, st.DimLevelType.compressed], + [st.DimLevelType.compressed, st.DimLevelType.dense], + [st.DimLevelType.compressed, st.DimLevelType.compressed]] + orderings = [ + ir.AffineMap.get_permutation([0, 1]), + ir.AffineMap.get_permutation([1, 0]) + ] + bitwidths = [0] + for level in levels: + for ordering in orderings: + for pwidth in bitwidths: + for iwidth in bitwidths: + attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth) + compiler = SparseCompiler(options=opt) + build_compile_and_run_SpMM(attr, support_lib, compiler) + count = count + 1 diff --git a/mlir/benchmark/MLIRStarter.cpp b/mlir/benchmark/MLIRStarter.cpp new file mode 100644 --- /dev/null +++ b/mlir/benchmark/MLIRStarter.cpp @@ -0,0 +1,11 @@ +#include "benchmark/benchmark.h" + +static void BM_MLIRStarter(benchmark::State &state) { + std::string x = "hello"; + for (auto _ : state) { + std::string copy(x); + } +} +BENCHMARK(BM_MLIRStarter); + +BENCHMARK_MAIN();