parameters: ()
return_statement: return compiler, runner
docstring: Benchmark for mlir sparse matrix multiplication. Because its an
      MLIR benchmark we need to return both a `compiler` function and a `runner`
      function.
docstring_summary: Benchmark for mlir sparse matrix multiplication. Because its an
      MLIR benchmark we need to return both a `compiler` function and a `runner`
      function.
docstring_tokens: Benchmark for mlir sparse matrix multiplication . Because its an MLIR benchmark we need to return both a compiler function and a runner function .
function: def benchmark_sparse_mlir_multiplication():
      """Benchmark for mlir sparse matrix multiplication. Because its an
      MLIR benchmark we need to return both a `compiler` function and a `runner`
      function.
      """
      with ir.Context(), ir.Location.unknown():
          module = ir.Module.create()
          f64 = ir.F64Type.get()
          param1_type = ir.RankedTensorType.get([1000, 1500], f64)
          param2_type = ir.RankedTensorType.get([1500, 2000], f64)
          result_type = ir.RankedTensorType.get([1000, 2000], f64)
          with ir.InsertionPoint(module.body):
              @builtin.FuncOp.from_py_func(param1_type, param2_type, result_type)
              def sparse_kernel(x, y, z):
                  return matmul_dsl(x, y, outs=[z])
  
      def compiler():
          with ir.Context(), ir.Location.unknown():
              kernel_func = get_kernel_func_from_module(module)
              timer_func = emit_timer_func()
              wrapped_func = emit_benchmark_wrapped_main_func(
                  kernel_func,
                  timer_func
              )
              main_module_with_benchmark = ir.Module.parse(
                  str(timer_func) + str(wrapped_func) + str(kernel_func)
              )
              setup_passes(main_module_with_benchmark)
              c_runner_utils = os.getenv("MLIR_C_RUNNER_UTILS", "")
              assert os.path.exists(c_runner_utils),\
                  f"{c_runner_utils} does not exist." \
                  f" Please pass a valid value for" \
                  f" MLIR_C_RUNNER_UTILS environment variable."
              runner_utils = os.getenv("MLIR_RUNNER_UTILS", "")
              assert os.path.exists(runner_utils),\
                  f"{runner_utils} does not exist." \
                  f" Please pass a valid value for MLIR_RUNNER_UTILS" \
                  f" environment variable."
  
              engine = ExecutionEngine(
                  main_module_with_benchmark,
                  3,
                  shared_libs=[c_runner_utils, runner_utils]
              )
              return engine.invoke
  
      def runner(engine_invoke):
          compiled_program_args = []
          for argument_type in [
              result_type, param1_type, param2_type, result_type
          ]:
              argument_type_str = str(argument_type)
              dimensions_str = re.sub("<|>|tensor", "", argument_type_str)
              dimensions = [int(dim) for dim in dimensions_str.split("x")[:-1]]
              if argument_type == result_type:
                  argument = np.zeros(dimensions, np.float64)
              else:
                  argument = create_sparse_np_tensor(dimensions, 1000)
              compiled_program_args.append(
                  ctypes.pointer(
                      ctypes.pointer(rt.get_ranked_memref_descriptor(argument))
                  )
              )
          np_timers_ns = np.array([0], dtype=np.int64)
          compiled_program_args.append(
              ctypes.pointer(
                  ctypes.pointer(rt.get_ranked_memref_descriptor(np_timers_ns))
              )
          )
          engine_invoke("main", *compiled_program_args)
          return int(np_timers_ns[0])
  
      return compiler, runner
function_tokens: def benchmark_sparse_mlir_multiplication ( ) : with ir . Context ( ) , ir . Location . unknown ( ) : module = ir . Module . create ( ) f64 = ir . F64Type . get ( ) param1_type = ir . RankedTensorType . get ( [ 1000 , 1500 ] , f64 ) param2_type = ir . RankedTensorType . get ( [ 1500 , 2000 ] , f64 ) result_type = ir . RankedTensorType . get ( [ 1000 , 2000 ] , f64 ) with ir . InsertionPoint ( module . body ) : @ builtin . FuncOp . from_py_func ( param1_type , param2_type , result_type ) def sparse_kernel ( x , y , z ) : return matmul_dsl ( x , y , outs = [ z ] ) def compiler ( ) : with ir . Context ( ) , ir . Location . unknown ( ) : kernel_func = get_kernel_func_from_module ( module ) timer_func = emit_timer_func ( ) wrapped_func = emit_benchmark_wrapped_main_func ( kernel_func , timer_func ) main_module_with_benchmark = ir . Module . parse ( str ( timer_func ) + str ( wrapped_func ) + str ( kernel_func ) ) setup_passes ( main_module_with_benchmark ) c_runner_utils = os . getenv ( "MLIR_C_RUNNER_UTILS" , "" ) assert os . path . exists ( c_runner_utils ) , f"{c_runner_utils} does not exist." f" Please pass a valid value for" f" MLIR_C_RUNNER_UTILS environment variable." runner_utils = os . getenv ( "MLIR_RUNNER_UTILS" , "" ) assert os . path . exists ( runner_utils ) , f"{runner_utils} does not exist." f" Please pass a valid value for MLIR_RUNNER_UTILS" f" environment variable." engine = ExecutionEngine ( main_module_with_benchmark , 3 , shared_libs = [ c_runner_utils , runner_utils ] ) return engine . invoke def runner ( engine_invoke ) : compiled_program_args = [ ] for argument_type in [ result_type , param1_type , param2_type , result_type ] : argument_type_str = str ( argument_type ) dimensions_str = re . sub ( "<|>|tensor" , "" , argument_type_str ) dimensions = [ int ( dim ) for dim in dimensions_str . split ( "x" ) [ : - 1 ] ] if argument_type == result_type : argument = np . zeros ( dimensions , np . float64 ) else : argument = create_sparse_np_tensor ( dimensions , 1000 ) compiled_program_args . append ( ctypes . pointer ( ctypes . pointer ( rt . get_ranked_memref_descriptor ( argument ) ) ) ) np_timers_ns = np . array ( [ 0 ] , dtype = np . int64 ) compiled_program_args . append ( ctypes . pointer ( ctypes . pointer ( rt . get_ranked_memref_descriptor ( np_timers_ns ) ) ) ) engine_invoke ( "main" , * compiled_program_args ) return int ( np_timers_ns [ 0 ] ) return compiler , runner