examples/nvidia_cuda_q/1_simulation_with_GPUs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4c0fd98e-cee0-47a3-84da-ca65338b6633",
   "metadata": {},
   "source": [
    "# Simulating quantum programs on GPUs\n",
    "\n",
    "In this notebook, you will learn how to simulate quantum circuits using GPUs with NVIDIA CUDA-Q and Braket Hybrid Jobs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d431129d-2036-4dd3-ab6e-b4b5f6c078bd",
   "metadata": {},
   "source": [
    "We start with necessary imports that are used in the examples below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab7965df-47cf-4d92-b745-252bcbf4417f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "from random_circuits import random_circuit_generator_factory\n",
    "\n",
    "from braket.jobs import hybrid_job\n",
    "from braket.jobs.config import InstanceConfig\n",
    "from braket.jobs.environment_variables import get_job_device_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bd2cd58-420c-4bf3-b161-09c27a272f8f",
   "metadata": {},
   "source": [
    "Next, specify the URI of your container image that supports CUDA-Q.\n",
    "\n",
    "If you don't have this URI already, see the notebook \"0_hello_cudaq_jobs.ipynb\", where you can learn about Braket Hybrid Jobs and how to build a container image that supports CUDA-Q. After following the steps in that notebook to upload the container image, you can use the same image URI here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "358b4350-9122-43c6-a52f-3e2ace50c7b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_uri = \"<image-uri>\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a35e065f-ad2c-4a31-9f0b-bae051b04279",
   "metadata": {},
   "source": [
    "## Running hybrid jobs on GPUs\n",
    "To use GPUs for circuit simulation, you can set the target backend to `nvidia`. You also need to select an instance that has NVIDIA GPUs. In the code snippet below, the instance type \"ml.p3.2xlarge\" is used as an example. The instance type \"ml.p3.2xlarge\" has a single NVIDIA V100 GPU. You can check [this page](https://aws.amazon.com/braket/pricing/) and [this page](https://docs.aws.amazon.com/braket/latest/developerguide/braket-jobs-configure-job-instance-for-script.html) to view available instances for Braket Hybrid Jobs. The `include_modules` keyword is set to `random_circuits` in order to use the random circuit generator defined in \"random_circuits.py\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c64ae222-e821-470b-9e7d-624299cf26cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "@hybrid_job(\n",
    "    device=\"local:nvidia/nvidia\",\n",
    "    image_uri=image_uri,\n",
    "    include_modules=\"random_circuits\",\n",
    "    instance_config=InstanceConfig(instanceType=\"ml.p3.2xlarge\"),\n",
    ")\n",
    "def gpu_job(n_qubits, n_gates, n_terms, n_shots):\n",
    "    import cudaq\n",
    "\n",
    "    # Define backend\n",
    "    device = get_job_device_arn()\n",
    "    cudaq.set_target(device.split(\"/\")[-1])\n",
    "    print(\"CUDA-Q backend: \", cudaq.get_target())\n",
    "\n",
    "    # Define circuit and observables\n",
    "    get_random_circuit = random_circuit_generator_factory()\n",
    "    circuit = get_random_circuit(n_qubits, n_gates)\n",
    "    hamiltonian = cudaq.SpinOperator.random(n_qubits, n_terms)\n",
    "\n",
    "    # Time the circuit simulation\n",
    "    t0 = time.time()\n",
    "    result = cudaq.observe(circuit, hamiltonian, shots_count=n_shots)\n",
    "    t1 = time.time()\n",
    "    print(f\"result: {result.expectation()} | duration: {t1-t0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "724501fd-062c-4aca-ad0a-d8374ff6a49e",
   "metadata": {},
   "source": [
    "When the `gpu_job` function is called, it creates a Braket Hybrid Job that runs on AWS, performing circuit simulations with the GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25abc0d9-84bf-40d5-9ec6-5230d9349e68",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_qubits = 20\n",
    "n_gates = 150\n",
    "n_terms = 100\n",
    "n_shots = 1000\n",
    "\n",
    "job = gpu_job(n_qubits, n_gates, n_terms, n_shots)\n",
    "print(\"Job ARN: \", job.arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98cc0783-dfcd-4890-9241-81a785d8d742",
   "metadata": {},
   "source": [
    "## Tensor network backend\n",
    "The `nvidia` backend shown in the example above is a state vector simulator. CUDA-Q also supports tensor network simulation with the `tensornet` backend. You can view [this page](https://nvidia.github.io/cuda-quantum/latest/using/backends/simulators.html#tensor-network-simulators) to learn more about the tensor network simulator. The list of CUDA-Q backends can be viewed in the [CUDA-Q documentation](https://nvidia.github.io/cuda-quantum/latest/using/backends/backends.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aadd5af0-ed35-46d9-8c12-610847cfbed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "@hybrid_job(\n",
    "    device=\"local:nvidia/tensornet\",\n",
    "    image_uri=image_uri,\n",
    "    include_modules=\"random_circuits\",\n",
    "    instance_config=InstanceConfig(instanceType=\"ml.p3.2xlarge\"),\n",
    ")\n",
    "def gpu_tn_job(n_qubits, n_gates, n_terms, n_shots):\n",
    "    import cudaq\n",
    "\n",
    "    # Define backend\n",
    "    device = get_job_device_arn()\n",
    "    cudaq.set_target(device.split(\"/\")[-1])\n",
    "    print(\"CUDA-Q backend: \", cudaq.get_target())\n",
    "\n",
    "    # Define circuit and observables\n",
    "    get_random_circuit = random_circuit_generator_factory()\n",
    "    circuit = get_random_circuit(n_qubits, n_gates)\n",
    "    hamiltonian = cudaq.SpinOperator.random(n_qubits, n_terms)\n",
    "\n",
    "    # Time the circuit simulation\n",
    "    t0 = time.time()\n",
    "    result = cudaq.observe(circuit, hamiltonian, shots_count=n_shots)\n",
    "    t1 = time.time()\n",
    "    print(f\"result: {result.expectation()} | duration: {t1-t0}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "942d183a-2dd5-4a42-800a-8e89e5530267",
   "metadata": {},
   "source": [
    "When circuits are shallow, the tensor network simulator can run circuits with high qubit count. For example, the code snippet below runs a simulation of 50 qubits, and the simulation finishes in a few seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e2b5698-d9ca-4385-bab3-dd2d82d9bab3",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_qubits = 40\n",
    "n_gates = 100\n",
    "n_terms = 10\n",
    "n_shots = 1000\n",
    "\n",
    "tn_job = gpu_tn_job(n_qubits, n_gates, n_terms, n_shots)\n",
    "print(\"Tensornet Job ARN: \", tn_job.arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d696e91",
   "metadata": {},
   "source": [
    "You can download results from both hybrid jobs with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbfdff3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "tn_result = tn_job.result()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "802c3c13-a820-42f8-9a05-4dfea3954fb1",
   "metadata": {},
   "source": [
    "## Summary\n",
    "This notebook shows you how to target CUDA-Q GPU simulators for circuit simulation. In particular, this notebook shows an example of using a CUDA-Q state vector simulator,  the `nvidia` backend. If you have shallow circuits with high qubit count, you can use the CUDA-Q tensor network simulator, the `tensornet` backend, which may execute your circuits faster."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}