{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#using Pkg\n",
    "#Pkg.add(url=\"https://github.com/QEDjl-project/QEDprocesses.jl/\")\n",
    "\n",
    "using MetagraphOptimization\n",
    "using CUDA\n",
    "using UUIDs\n",
    "using BenchmarkTools\n",
    "\n",
    "println(\"Threads: $(Threads.nthreads())\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preparation of graph\n",
    "machine = Machine([MetagraphOptimization.NumaNode(0, 1, MetagraphOptimization.default_strategy(MetagraphOptimization.NumaNode), -1.0, UUIDs.uuid1())], [-1.0;;])\n",
    "model = QEDModel()\n",
    "process = parse_process(\"ke->kke\", model)\n",
    "graph = gen_graph(process)\n",
    "n_inputs = 10_000\n",
    "inputs = [gen_process_input(process) for _ in 1:n_inputs]\n",
    "cu_inputs = Vector()#CuArray(inputs)\n",
    "optimizer = ReductionOptimizer()\n",
    "\n",
    "get_compute_function(graph, process, machine) # run once for compilation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function bench(func, inputs, cu_inputs)\n",
    "    compile_time = @elapsed func(inputs[1])\n",
    "\n",
    "    single_thread = @elapsed func.(inputs)\n",
    "    multi_threaded = @elapsed Threads.@threads for i in eachindex(inputs)\n",
    "        func(inputs[i]) \n",
    "    end\n",
    "    \n",
    "    gpu_compile = 0 #@elapsed CUDA.@sync func.(cu_inputs[1:2])\n",
    "    gpu = 0 #@elapsed CUDA.@sync func.(cu_inputs)\n",
    "    return (cpu_compile_time = compile_time, gpu_compile_time = gpu_compile, cpu_single_thread_time = single_thread, cpu_multi_thread_time = multi_threaded, gpu_time = gpu)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bench and produce data\n",
    "using DataFrames\n",
    "\n",
    "STEPSIZE = 10\n",
    "n = 0\n",
    "\n",
    "df = DataFrame(operations=Int[], graph_nodes=Int[], graph_edges=Int[], graph_ce=Float64[], graph_dt=Float64[], graph_ci=Float64[], gen_func_t=Float64[], cpu_compile_t=Float64[], cpu_st_t=Float64[], cpu_mt_t=Float64[], gpu_compile_t=Float64[], gpu_t=Float64[])\n",
    "\n",
    "while true\n",
    "    func_gen_time = @elapsed func = get_compute_function(graph, process, machine)\n",
    "    res = bench(func, inputs, cu_inputs)\n",
    "\n",
    "    graph_properties = get_properties(graph)\n",
    "    push!(df, (\n",
    "        n,\n",
    "        graph_properties.noNodes,\n",
    "        graph_properties.noEdges,\n",
    "        graph_properties.computeEffort,\n",
    "        graph_properties.data,\n",
    "        graph_properties.computeIntensity,\n",
    "        func_gen_time,\n",
    "        res.cpu_compile_time,\n",
    "        res.cpu_single_thread_time,\n",
    "        res.cpu_multi_thread_time,\n",
    "        res.gpu_compile_time,\n",
    "        res.gpu_time\n",
    "    ))\n",
    "\n",
    "    if fixpoint_reached(optimizer, graph)\n",
    "        break\n",
    "    end\n",
    "\n",
    "    optimize!(optimizer, graph, STEPSIZE)\n",
    "    n += STEPSIZE\n",
    "end\n",
    ";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot data\n",
    "using Plots\n",
    "using StatsPlots\n",
    "\n",
    "img = @df df scatter(\n",
    "    :operations, \n",
    "    [:gen_func_t, :cpu_st_t, :cpu_mt_t], \n",
    "    label=[\"Function generation (s)\" \"Single threaded execution (s)\" \"$(Threads.nthreads())-threaded execution (s)\"], \n",
    "    title=\"$process using $optimizer ($(n_inputs) inputs)\",\n",
    "    linewidth=2,\n",
    "    xlabel=\"optimizer steps\",\n",
    "    ylabel=\"time (s)\",\n",
    "    yscale=:log10,\n",
    "    minorgrid=true,\n",
    "    size=(800, 600),\n",
    "    fmt=:pdf\n",
    ")\n",
    "\n",
    "savefig(img, \"../images/$(String(process))_exec_$(n_inputs)_inputs.pdf\")\n",
    "\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = @df df scatter(\n",
    "    :operations,\n",
    "    [:graph_nodes, :graph_edges],\n",
    "    label=[\"Graph Nodes (#)\" \"Graph Edges (#)\"],\n",
    "    title=\"$process using $optimizer\",\n",
    "    linewidth=2,\n",
    "    xlabel=\"optimizer steps\",\n",
    "    ylims=(0.0, 1.05 * maximum(df.graph_edges)),\n",
    "    fmt=:pdf,\n",
    "    size=(800, 600)\n",
    ")\n",
    "\n",
    "savefig(img, \"../images/$(String(process))_graph_properties.pdf\")\n",
    "\n",
    "img"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.4",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}