Tutorials/Accelerated Python: Fixes and improvements to memory spaces, asynchrony, and kernel authoring notebooks

tutorials/accelerated-python/notebooks/fundamentals/06__asynchrony__power_iteration.ipynb

@@ -31,11 +31,9 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {
         "id": "rO4kOPuP_0JG"
       },
-      "outputs": [],
       "source": [
         "import os\n",
         "\n",
@@ -46,7 +44,9 @@
         "  !pip install \"nvtx\" \"nsightful[notebook] @ git+https://github.com/brycelelbach/nsightful.git\" > /dev/null 2>&1\n",
         "\n",
         "print(\"Environment setup complete.\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -77,18 +77,17 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {
         "id": "sEicxhLO_9G9"
       },
-      "outputs": [],
       "source": [
         "%%writefile power_iteration__baseline.py\n",
         "\n",
         "import numpy as np\n",
         "import cupy as cp\n",
         "import cupyx as cpx\n",
         "import nvtx\n",
+        "import time\n",
         "from dataclasses import dataclass\n",
         "\n",
         "@dataclass\n",
@@ -140,17 +139,18 @@
         "A_device = generate_device()\n",
         "\n",
         "# Warmup to ensure modules are loaded and code is JIT compiled before timing.\n",
-        "estimate_device(A_device, cfg=PowerIterationConfig(progress=False))\n",
+        "estimate_device(A_device, cfg=PowerIterationConfig(max_steps=1, check_frequency=1, progress=False))\n",
+        "cp.cuda.get_current_stream().synchronize()\n",
         "\n",
-        "start = cp.cuda.get_current_stream().record()\n",
+        "start = time.perf_counter()\n",
         "lam_est_device = estimate_device(A_device).item()\n",
-        "stop = cp.cuda.get_current_stream().record()\n",
-        "\n",
-        "duration = cp.cuda.get_elapsed_time(start, stop) / 1e3\n",
+        "stop = time.perf_counter()\n",
         "\n",
         "print()\n",
-        "print(f\"GPU Execution Time: {duration:.3f} s\")"
-      ]
+        "print(f\"{(stop - start) * 1000:.3f} ms\")"
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -165,14 +165,14 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {
         "id": "1HU5p1IhAkTA"
       },
-      "outputs": [],
       "source": [
         "!nsys profile --cuda-event-trace=false --force-overwrite true -o power_iteration__baseline python power_iteration__baseline.py"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -189,17 +189,17 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {
         "id": "s6VVOnGQR3Ph"
       },
-      "outputs": [],
       "source": [
         "import nsightful\n",
         "\n",
         "!nsys export --type sqlite --quiet true --force-overwrite true power_iteration__baseline.nsys-rep\n",
         "nsightful.display_nsys_sqlite_file_in_notebook(\"power_iteration__baseline.sqlite\", title=\"Power Iteration - Baseline\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -216,7 +216,7 @@
         "The first is to limit when we start and stop profiling in the program. In Python, we can do this with `cupyx.profiler.profile()`, which give us a Python context manager. Any CUDA code used during scope will be included in the profile.\n",
         "\n",
         "```\n",
-        "not_in_the profile()\n",
+        "not_in_the_profile()\n",
         "with cpx.profiler.profile():\n",
         "  in_the_profile()\n",
         "not_in_the_profile()\n",
@@ -227,7 +227,7 @@
         "We can also annotate specific regions of our code, which will show up in the profiler. We can even add categories, domains, and colors to these regions, and they can be nested. To add these annotations, we use `nvtx.annnotate()`, another Python context manager, this time from a library called NVTX.\n",
         "\n",
         "```\n",
-        "with nvtx.annotate(\"Loop\")\n",
+        "with nvtx.annotate(\"Loop\"):\n",
         "  for i in range(20):\n",
         "     with nvtx.annotate(f\"Step {i}\"):\n",
         "       pass\n",
@@ -269,16 +269,15 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "%%writefile power_iteration__async.py\n",
         "\n",
         "import numpy as np\n",
         "import cupy as cp\n",
         "import cupyx as cpx\n",
         "import nvtx\n",
+        "import time\n",
         "from dataclasses import dataclass\n",
         "\n",
         "@dataclass\n",
@@ -305,17 +304,18 @@
         "A_device = generate_device()\n",
         "\n",
         "# Warmup to ensure modules are loaded and code is JIT compiled before timing.\n",
-        "estimate_device(A_device, cfg=PowerIterationConfig(progress=False))\n",
+        "estimate_device(A_device, cfg=PowerIterationConfig(max_steps=1, check_frequency=1, progress=False))\n",
+        "cp.cuda.get_current_stream().synchronize()\n",
         "\n",
-        "start = cp.cuda.get_current_stream().record()\n",
+        "start = time.perf_counter()\n",
         "lam_est_device = estimate_device(A_device).item()\n",
-        "stop = cp.cuda.get_current_stream().record()\n",
-        "\n",
-        "duration = cp.cuda.get_elapsed_time(start, stop) / 1e3\n",
+        "stop = time.perf_counter()\n",
         "\n",
         "print()\n",
-        "print(f\"GPU Execution Time: {duration:.3f} s\")"
-      ]
+        "print(f\"{(stop - start) * 1000:.3f} ms\")"
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -326,14 +326,14 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {
         "id": "pszz-k8cDfqy"
       },
-      "outputs": [],
       "source": [
         "!python power_iteration__async.py"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -348,24 +348,22 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {
         "id": "uSPFNIb9KcPb"
       },
-      "outputs": [],
       "source": [
         "power_iteration_baseline_output   = !python power_iteration__baseline.py\n",
-        "power_iteration_baseline_duration = float(power_iteration_baseline_output[-1].split()[-2])\n",
+        "power_iteration_baseline_duration = float(power_iteration_baseline_output[-1].split()[0])\n",
         "power_iteration_async_output      = !python power_iteration__async.py\n",
-        "power_iteration_async_duration    = float(power_iteration_async_output[-1].split()[-2])\n",
+        "power_iteration_async_duration    = float(power_iteration_async_output[-1].split()[0])\n",
         "speedup = power_iteration_baseline_duration / power_iteration_async_duration\n",
         "\n",
-        "print(f\"GPU Execution Time\")\n",
-        "print()\n",
-        "print(f\"power_iteration_baseline: {power_iteration_baseline_duration:.3f} s\")\n",
-        "print(f\"power_iteration_async:    {power_iteration_async_duration:.3f} s\")\n",
+        "print(f\"power_iteration_baseline: {power_iteration_baseline_output[-1]}\")\n",
+        "print(f\"power_iteration_async:    {power_iteration_async_output[-1]}\")\n",
         "print(f\"power_iteration_async speedup over power_iteration_baseline: {speedup:.2f}\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -378,14 +376,14 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {
         "id": "BtQR4CHikWFK"
       },
-      "outputs": [],
       "source": [
         "!nsys profile --cuda-event-trace=false --capture-range=cudaProfilerApi --capture-range-end=stop --force-overwrite true -o power_iteration__async python power_iteration__async.py"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -398,15 +396,15 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {
         "id": "mWXBvi-hFGhU"
       },
-      "outputs": [],
       "source": [
         "!nsys export --type sqlite --quiet true --force-overwrite true power_iteration__async.nsys-rep\n",
         "nsightful.display_nsys_sqlite_file_in_notebook(\"power_iteration__async.sqlite\", title=\"Power Iteration - Async Event\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     }
   ],
   "metadata": {
@@ -423,4 +421,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
+}