homework 3 sml (lesson 9).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b00f9c89-8da3-4c77-bcf2-596e03833501",
   "metadata": {
    "libroCellType": "markdown",
    "libroFormatter": "formatter-string"
   },
   "source": [
    "# 1. Import and install"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b869657c-90ea-4ebd-b7b8-3c7e5db2b448",
   "metadata": {
    "libroCellType": "markdown",
    "libroFormatter": "formatter-string"
   },
   "source": [
    "**<span style=\"color: rgb(224, 0, 0)\">Important notes</span>**: we should add another field if running with docker-compose.  (docker-compose.yaml)\n",
    "\n",
    "```bash\n",
    "services:\n",
    "\n",
    "alice:\n",
    "\n",
    "... (some other settings)\n",
    "\n",
    "cap_add:\n",
    "\n",
    "  - NET_ADMIN\n",
    "\n",
    "\n",
    "bob:\n",
    "\n",
    "... (some other settings)\n",
    "\n",
    "cap_add:\n",
    "\n",
    "  - NET_ADMIN\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7c4094ab-ea3a-48bd-be28-738dd43aa112",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:39.138846Z",
     "shell.execute_reply.started": "2024-06-23T08:48:36.719810Z",
     "to_execute": "2024-06-23T08:48:36.702Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-23 08:48:37,781\tINFO worker.py:1540 -- Connecting to existing Ray cluster at address: 172.23.0.2:6379...\n",
      "2024-06-23 08:48:37,874\tINFO worker.py:1724 -- Connected to Ray cluster.\n",
      "2024-06-23 08:48:37.890 INFO api.py:233 [alice] -- [Anonymous_job] Started rayfed with {'CLUSTER_ADDRESSES': {'alice': 'alice:8000', 'bob': 'bob:8000'}, 'CURRENT_PARTY_NAME': 'alice', 'TLS_CONFIG': {}}\n",
      "2024-06-23 08:48:38.502 INFO barriers.py:284 [alice] -- [Anonymous_job] Succeeded to create receiver proxy actor.\n",
      "\u001b[36m(ReceiverProxyActor pid=2943)\u001b[0m 2024-06-23 08:48:38.499 INFO grpc_proxy.py:359 [alice] -- [Anonymous_job] ReceiverProxy binding port 8000, options: (('grpc.enable_retries', 1), ('grpc.so_reuseport', 0), ('grpc.max_send_message_length', 524288000), ('grpc.max_receive_message_length', 524288000), ('grpc.service_config', '{\"methodConfig\": [{\"name\": [{\"service\": \"GrpcService\"}], \"retryPolicy\": {\"maxAttempts\": 5, \"initialBackoff\": \"5s\", \"maxBackoff\": \"30s\", \"backoffMultiplier\": 2, \"retryableStatusCodes\": [\"UNAVAILABLE\"]}}]}'))...\n",
      "\u001b[36m(ReceiverProxyActor pid=2943)\u001b[0m 2024-06-23 08:48:38.502 INFO grpc_proxy.py:379 [alice] -- [Anonymous_job] Successfully start Grpc service without credentials.\n",
      "2024-06-23 08:48:39.086 INFO barriers.py:333 [alice] -- [Anonymous_job] SenderProxyActor has successfully created.\n",
      "2024-06-23 08:48:39.087 INFO barriers.py:520 [alice] -- [Anonymous_job] Try ping ['bob'] at 0 attemp, up to 3600 attemps.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-23 08:48:37,808\tINFO worker.py:1540 -- Connecting to existing Ray cluster at address: 172.23.0.3:6379...\n",
      "2024-06-23 08:48:37,915\tINFO worker.py:1724 -- Connected to Ray cluster.\n",
      "2024-06-23 08:48:37.936 INFO api.py:233 [bob] -- [Anonymous_job] Started rayfed with {'CLUSTER_ADDRESSES': {'alice': 'alice:8000', 'bob': 'bob:8000'}, 'CURRENT_PARTY_NAME': 'bob', 'TLS_CONFIG': {}}\n",
      "2024-06-23 08:48:38.554 INFO barriers.py:284 [bob] -- [Anonymous_job] Succeeded to create receiver proxy actor.\n",
      "\u001b[36m(ReceiverProxyActor pid=3800)\u001b[0m 2024-06-23 08:48:38.548 INFO grpc_proxy.py:359 [bob] -- [Anonymous_job] ReceiverProxy binding port 8000, options: (('grpc.enable_retries', 1), ('grpc.so_reuseport', 0), ('grpc.max_send_message_length', 524288000), ('grpc.max_receive_message_length', 524288000), ('grpc.service_config', '{\"methodConfig\": [{\"name\": [{\"service\": \"GrpcService\"}], \"retryPolicy\": {\"maxAttempts\": 5, \"initialBackoff\": \"5s\", \"maxBackoff\": \"30s\", \"backoffMultiplier\": 2, \"retryableStatusCodes\": [\"UNAVAILABLE\"]}}]}'))...\n",
      "\u001b[36m(ReceiverProxyActor pid=3800)\u001b[0m 2024-06-23 08:48:38.552 INFO grpc_proxy.py:379 [bob] -- [Anonymous_job] Successfully start Grpc service without credentials.\n",
      "2024-06-23 08:48:39.124 INFO barriers.py:333 [bob] -- [Anonymous_job] SenderProxyActor has successfully created.\n",
      "2024-06-23 08:48:39.125 INFO barriers.py:520 [bob] -- [Anonymous_job] Try ping ['alice'] at 0 attemp, up to 3600 attemps.\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "import spu\n",
    "import os\n",
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "import jax\n",
    "import jax.lax\n",
    "import spu.utils.simulation as spsim\n",
    "import spu.spu_pb2 as spu_pb2\n",
    "from functools import partial\n",
    "\n",
    "network_conf = {\n",
    "    \"parties\": {\n",
    "        \"alice\": {\n",
    "            \"address\": \"alice:8000\",\n",
    "        },\n",
    "        \"bob\": {\n",
    "            \"address\": \"bob:8000\",\n",
    "        },\n",
    "    },\n",
    "}\n",
    "\n",
    "party = os.getenv(\"SELF_PARTY\", \"alice\")\n",
    "sf.shutdown()\n",
    "sf.init(\n",
    "    address=\"127.0.0.1:6379\",\n",
    "    cluster_config={**network_conf, \"self_party\": party},\n",
    "    log_to_driver=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5f3d1117-2917-4668-a9dc-d93c7e3b19a8",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:39.755179Z",
     "shell.execute_reply.started": "2024-06-23T08:48:39.098511Z",
     "to_execute": "2024-06-23T08:48:36.702Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/pty.py:89: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
      "  pid, fd = os.forkpty()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last metadata expiration check: 0:21:00 ago on Sun Jun 23 08:27:39 2024.\n",
      "Package iproute-tc-6.2.0-5.0.1.an8.x86_64 is already installed.\n",
      "Dependencies resolved.\n",
      "Nothing to do.\n",
      "Complete!\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/pty.py:89: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
      "  pid, fd = os.forkpty()\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last metadata expiration check: 0:20:59 ago on Sun Jun 23 08:27:40 2024.\n",
      "Package iproute-tc-6.2.0-5.0.1.an8.x86_64 is already installed.\n",
      "Dependencies resolved.\n",
      "Nothing to do.\n",
      "Complete!\n"
     ]
    }
   ],
   "source": [
    "!yum install -y iproute-tc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e701b9d2-1930-415b-bc46-bfb8cf18ae1f",
   "metadata": {
    "libroCellType": "markdown",
    "libroFormatter": "formatter-string"
   },
   "source": [
    "# Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "80d03546-18f9-4030-9bf3-da9151606800",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:39.759740Z",
     "shell.execute_reply.started": "2024-06-23T08:48:39.712838Z",
     "to_execute": "2024-06-23T08:48:36.702Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [],
   "source": [
    "# we know that dk is wrong when |x| is very small\n",
    "# Let us try it. (we only show part here.)\n",
    "# define some test function and data used in simulation\n",
    "def test_square_and_sum_when_x_small(x):\n",
    "    return jnp.sum(jnp.square(x))\n",
    "\n",
    "x = np.array([1e-5] * 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c53e4bee-1538-4a42-a24f-a38ee0d9fb5d",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:39.992957Z",
     "shell.execute_reply.started": "2024-06-23T08:48:39.717552Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-23 08:48:39.740 INFO xla_bridge.py:863 [alice] -- [Anonymous_job] Unable to initialize backend 'cuda': \n",
      "2024-06-23 08:48:39.741 INFO xla_bridge.py:863 [alice] -- [Anonymous_job] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "2024-06-23 08:48:39.742 INFO xla_bridge.py:863 [alice] -- [Anonymous_job] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory\n",
      "2024-06-23 08:48:39.743 INFO xla_bridge.py:863 [alice] -- [Anonymous_job] Unable to initialize backend 'interpreter': jaxlib.xla_extension.Client: no constructor defined!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-23 08:48:39.891] [info] [cheetah_mul.cc:321] CheetahMul uses 4 modulus for 64 bit input over 64 bit ring\n",
      "[2024-06-23 08:48:39.892] [info] [cheetah_mul.cc:321] CheetahMul uses 4 modulus for 64 bit input over 64 bit ring\n",
      "[2024-06-23 08:48:39.899] [info] [thread_pool.cc:30] Create a fixed thread pool with size 19\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array(-1.5258789e-05, dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-23 08:48:39.787 INFO xla_bridge.py:863 [bob] -- [Anonymous_job] Unable to initialize backend 'cuda': \n",
      "2024-06-23 08:48:39.788 INFO xla_bridge.py:863 [bob] -- [Anonymous_job] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "2024-06-23 08:48:39.789 INFO xla_bridge.py:863 [bob] -- [Anonymous_job] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory\n",
      "2024-06-23 08:48:39.790 INFO xla_bridge.py:863 [bob] -- [Anonymous_job] Unable to initialize backend 'interpreter': jaxlib.xla_extension.Client: no constructor defined!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-23 08:48:39.932] [info] [cheetah_mul.cc:321] CheetahMul uses 4 modulus for 64 bit input over 64 bit ring\n",
      "[2024-06-23 08:48:39.932] [info] [cheetah_mul.cc:321] CheetahMul uses 4 modulus for 64 bit input over 64 bit ring\n",
      "[2024-06-23 08:48:39.940] [info] [thread_pool.cc:30] Create a fixed thread pool with size 19\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array(-1.5258789e-05, dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# First, we run SPU with simulator\n",
    "# Indeed, simulation can be run within single node.\n",
    "# a. run with CHEETAH\n",
    "sim_che = spsim.Simulator.simple(2, spu_pb2.ProtocolKind.CHEETAH, spu_pb2.FieldType.FM64)\n",
    "spsim.sim_jax(sim_che, test_square_and_sum_when_x_small)(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e6da1918-9fd4-40fa-9c8e-74bc0123351a",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:40.010300Z",
     "shell.execute_reply.started": "2024-06-23T08:48:39.953165Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-23 08:48:39.962] [info] [api.cc:163] [Profiling] SPU execution test_square_and_sum_when_x_small completed, input processing took 5.22e-07s, execution took 0.000550024s, output processing took 1.205e-06s, total time 0.000551751s.\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:209] HLO profiling: total time 7.35e-07\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - pphlo.constant, executed 1 times, duration 2.41e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - pphlo.add, executed 5 times, duration 2.03e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - pphlo.free, executed 3 times, duration 1.26e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - pphlo.multiply, executed 1 times, duration 5.9e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - pphlo.reduce, executed 1 times, duration 5.4e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - pphlo.0.0\n",
      "convert, executed 1 times, duration 5.2e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:209] HAL profiling: total time 0.00015334\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - f_mul, executed 1 times, duration 0.00013537s, send bytes 240 recv bytes 400\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - f_add, executed 5 times, duration 1.5145e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - seal, executed 1 times, duration 2.825e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:209] MPC profiling: total time 0.000144999\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - trunc_a, executed 1 times, duration 0.000102831s, send bytes 160 recv bytes 320\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - mul_aa, executed 1 times, duration 2.4066e-05s, send bytes 80 recv bytes 80\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - add_aa, executed 5 times, duration 6.803e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - extract_slice, executed 7 times, duration 5.449e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - reshape, executed 3 times, duration 2.295e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - p2a, executed 1 times, duration 1.49e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - transpose, executed 1 times, duration 1.228e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:212] - broadcast, executed 1 times, duration 8.37e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:39.962] [info] [api.cc:222] Link details: total send bytes 240, recv bytes 400, send actions 3\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-23 08:48:40.007] [info] [api.cc:163] [Profiling] SPU execution test_square_and_sum_when_x_small completed, input processing took 4.61e-07s, execution took 0.000435953s, output processing took 1.218e-06s, total time 0.000437632s.\n",
      "0.0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:209] HLO profiling: total time 9.6e-07\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - pphlo.constant, executed 1 times, duration 4.44e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - pphlo.add, executed 5 times, duration 2.1e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - pphlo.free, executed 3 times, duration 1.33e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - pphlo.multiply, executed 1 times, duration 5.9e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - pphlo.reduce, executed 1 times, duration 5.8e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - pphlo.convert, executed 1 times, duration 5.6e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:209] HAL profiling: total time 0.000119299\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - f_mul, executed 1 times, duration 0.000100762s, send bytes 240 recv bytes 400\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - f_add, executed 5 times, duration 1.4842e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - seal, executed 1 times, duration 3.695e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:209] MPC profiling: total time 0.00011736900000000001\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - trunc_a, executed 1 times, duration 6.1661e-05s, send bytes 160 recv bytes 320\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - mul_aa, executed 1 times, duration 2.988e-05s, send bytes 80 recv bytes 80\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - extract_slice, executed 7 times, duration 1.3421e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - add_aa, executed 5 times, duration 6.459e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - reshape, executed 3 times, duration 2.544e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - transpose, executed 1 times, duration 1.256e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - p2a, executed 1 times, duration 1.167e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:212] - broadcast, executed 1 times, duration 9.81e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:40.007] [info] [api.cc:222] Link details: total send bytes 240, recv bytes 400, send actions 3\n"
     ]
    }
   ],
   "source": [
    "# b. run with ABY3\n",
    "# this time, we alse print some profile stats.\n",
    "config_aby = spu.RuntimeConfig(\n",
    "    protocol=spu_pb2.ProtocolKind.ABY3,\n",
    "    field=spu.FieldType.FM64,\n",
    "    fxp_fraction_bits=18,\n",
    "    enable_hal_profile=True,\n",
    "    enable_pphlo_profile=True,\n",
    ")\n",
    "sim_aby = spsim.Simulator(3, config_aby)\n",
    "print(spsim.sim_jax(sim_aby, test_square_and_sum_when_x_small)(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "627e35a4-1123-4c69-af3f-e8e248efdfb1",
   "metadata": {
    "libroCellType": "markdown",
    "libroFormatter": "formatter-string"
   },
   "source": [
    "# Emulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "819a8fd6-bfe3-426c-b480-6b7d12a2038c",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:40.218775Z",
     "shell.execute_reply.started": "2024-06-23T08:48:39.968026Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Cannot delete qdisc with handle of zero.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Cannot delete qdisc with handle of zero.\n"
     ]
    }
   ],
   "source": [
    "# delete all wan settings beforehand\n",
    "!tc qdisc del dev eth0 root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cbbb14eb-6674-4b30-9028-d826e75a5a63",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:43.566538Z",
     "shell.execute_reply.started": "2024-06-23T08:48:40.177851Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PING bob (172.23.0.3) 56(84) bytes of data.\n",
      "64 bytes from secretnote-bob-1.secretnote_default (172.23.0.3): icmp_seq=1 ttl=64 time=0.027 ms\n",
      "64 bytes from secretnote-bob-1.secretnote_default (172.23.0.3): icmp_seq=2 ttl=64 time=0.053 ms\n",
      "64 bytes from secretnote-bob-1.secretnote_default (172.23.0.3): icmp_seq=3 ttl=64 time=0.031 ms\n",
      "64 bytes from secretnote-bob-1.secretnote_default (172.23.0.3): icmp_seq=4 ttl=64 time=0.059 ms\n",
      "\n",
      "--- bob ping statistics ---\n",
      "4 packets transmitted, 4 received, 0% packet loss, time 3111ms\n",
      "rtt min/avg/max/mdev = 0.027/0.042/0.059/0.015 ms\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PING bob (172.23.0.3) 56(84) bytes of data.\n",
      "64 bytes from 2224ad06a363 (172.23.0.3): icmp_seq=1 ttl=64 time=0.091 ms\n",
      "64 bytes from 2224ad06a363 (172.23.0.3): icmp_seq=2 ttl=64 time=0.022 ms\n",
      "64 bytes from 2224ad06a363 (172.23.0.3): icmp_seq=3 ttl=64 time=0.016 ms\n",
      "64 bytes from 2224ad06a363 (172.23.0.3): icmp_seq=4 ttl=64 time=0.017 ms\n",
      "\n",
      "--- bob ping statistics ---\n",
      "4 packets transmitted, 4 received, 0% packet loss, time 3149ms\n",
      "rtt min/avg/max/mdev = 0.016/0.036/0.091/0.032 ms\n"
     ]
    }
   ],
   "source": [
    "!ping -c 4 bob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c0398832-9984-4568-859e-6b460377f6fa",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:46.856686Z",
     "shell.execute_reply.started": "2024-06-23T08:48:43.488705Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PING alice (172.23.0.2) 56(84) bytes of data.\n",
      "64 bytes from a5ceb21c2bad (172.23.0.2): icmp_seq=1 ttl=64 time=0.102 ms\n",
      "64 bytes from a5ceb21c2bad (172.23.0.2): icmp_seq=2 ttl=64 time=0.012 ms\n",
      "64 bytes from a5ceb21c2bad (172.23.0.2): icmp_seq=3 ttl=64 time=0.010 ms\n",
      "64 bytes from a5ceb21c2bad (172.23.0.2): icmp_seq=4 ttl=64 time=0.012 ms\n",
      "\n",
      "--- alice ping statistics ---\n",
      "4 packets transmitted, 4 received, 0% packet loss, time 3158ms\n",
      "rtt min/avg/max/mdev = 0.010/0.034/0.102/0.039 ms\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PING alice (172.23.0.2) 56(84) bytes of data.\n",
      "64 bytes from secretnote-alice-1.secretnote_default (172.23.0.2): icmp_seq=1 ttl=64 time=0.034 ms\n",
      "64 bytes from secretnote-alice-1.secretnote_default (172.23.0.2): icmp_seq=2 ttl=64 time=0.038 ms\n",
      "64 bytes from secretnote-alice-1.secretnote_default (172.23.0.2): icmp_seq=3 ttl=64 time=0.029 ms\n",
      "64 bytes from secretnote-alice-1.secretnote_default (172.23.0.2): icmp_seq=4 ttl=64 time=0.043 ms\n",
      "\n",
      "--- alice ping statistics ---\n",
      "4 packets transmitted, 4 received, 0% packet loss, time 3091ms\n",
      "rtt min/avg/max/mdev = 0.029/0.036/0.043/0.005 ms\n"
     ]
    }
   ],
   "source": [
    "!ping -c 4 alice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2249de57-2454-42bc-935c-bc10d27f6d23",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:46.865583Z",
     "shell.execute_reply.started": "2024-06-23T08:48:46.848616Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [],
   "source": [
    "# Emulation should be run from source in SPU, so we use Secretflow here to do the efficiency experiments.\n",
    "# You can use the similar trick for emulation directly in SPU.\n",
    "\n",
    "def compute_dk_func(x, eps=1e-6):\n",
    "    return jax.lax.rsqrt(jnp.sum(jnp.square(x)) + eps)\n",
    "\n",
    "x = np.random.rand(1_000_000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "63cfd88a-eedc-4437-933f-6fd4600f7e69",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:46.896797Z",
     "shell.execute_reply.started": "2024-06-23T08:48:46.857560Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [],
   "source": [
    "# SPU settings\n",
    "cluster_def = {\n",
    "    'nodes': [\n",
    "        {'party': 'alice', 'id': 'local:0', 'address': 'alice' + ':12945'},\n",
    "        {'party': 'bob', 'id': 'local:1', 'address': 'bob' + ':12945'},\n",
    "    ],\n",
    "    'runtime_config': {\n",
    "        # SEMI2K support 2/3 PC, ABY3 only support 3PC, CHEETAH only support 2PC.\n",
    "        # pls pay attention to size of nodes above. nodes size need match to PC setting.\n",
    "        'protocol': spu.spu_pb2.SEMI2K,\n",
    "        'field': spu.spu_pb2.FM64\n",
    "    },\n",
    "}\n",
    "\n",
    "alice_device = sf.PYU(\"alice\")\n",
    "bob_device = sf.PYU(\"bob\")\n",
    "spu_device = sf.SPU(cluster_def)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "72e2e206-4981-41a2-9276-5a3979658dbf",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:46.903373Z",
     "shell.execute_reply.started": "2024-06-23T08:48:46.888160Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [],
   "source": [
    "# first, load data to PYU\n",
    "alice_data = alice_device(lambda x: x)(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c462b318-8100-4d57-afc6-79285c485134",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:49.516362Z",
     "shell.execute_reply.started": "2024-06-23T08:48:46.900923Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [],
   "source": [
    "# SPU may need some init, so we run this twice...\n",
    "ret = spu_device(compute_dk_func)(alice_data)\n",
    "sf.reveal(ret);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73f89efc-8eb9-4d26-bf35-08e58762728c",
   "metadata": {
    "libroCellType": "markdown",
    "libroFormatter": "formatter-string"
   },
   "source": [
    "Now, we strict the wan settings for 100Mbit and 10ms delay (20ms RTT)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bdc4e866-db25-48e9-9515-68dcd5be208e",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:49.914330Z",
     "shell.execute_reply.started": "2024-06-23T08:48:48.654306Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Specified qdisc kind is unknown.\n",
      "Error: Failed to find specified qdisc.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Specified qdisc kind is unknown.\n",
      "Error: Failed to find specified qdisc.\n"
     ]
    }
   ],
   "source": [
    "!tc qdisc add dev eth0 root handle 1: tbf rate 100mbit burst 128kb limit 10000\n",
    "!tc qdisc add dev eth0 parent 1:1 handle 10: netem delay 10msec limit 8000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "dfad1d45-9639-40fd-894a-42294f4bdaad",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:53.246344Z",
     "shell.execute_reply.started": "2024-06-23T08:48:49.056517Z",
     "to_execute": "2024-06-23T08:48:36.703Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PING bob (172.23.0.3) 56(84) bytes of data.\n",
      "64 bytes from 2224ad06a363 (172.23.0.3): icmp_seq=1 ttl=64 time=0.025 ms\n",
      "64 bytes from 2224ad06a363 (172.23.0.3): icmp_seq=2 ttl=64 time=0.013 ms\n",
      "64 bytes from 2224ad06a363 (172.23.0.3): icmp_seq=3 ttl=64 time=0.023 ms\n",
      "64 bytes from 2224ad06a363 (172.23.0.3): icmp_seq=4 ttl=64 time=0.028 ms\n",
      "\n",
      "--- bob ping statistics ---\n",
      "4 packets transmitted, 4 received, 0% packet loss, time 3120ms\n",
      "rtt min/avg/max/mdev = 0.013/0.022/0.028/0.006 ms\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PING bob (172.23.0.3) 56(84) bytes of data.\n",
      "64 bytes from secretnote-bob-1.secretnote_default (172.23.0.3): icmp_seq=1 ttl=64 time=0.027 ms\n",
      "64 bytes from secretnote-bob-1.secretnote_default (172.23.0.3): icmp_seq=2 ttl=64 time=0.034 ms\n",
      "64 bytes from secretnote-bob-1.secretnote_default (172.23.0.3): icmp_seq=3 ttl=64 time=0.044 ms\n",
      "64 bytes from secretnote-bob-1.secretnote_default (172.23.0.3): icmp_seq=4 ttl=64 time=0.038 ms\n",
      "\n",
      "--- bob ping statistics ---\n",
      "4 packets transmitted, 4 received, 0% packet loss, time 3132ms\n",
      "rtt min/avg/max/mdev = 0.027/0.035/0.044/0.009 ms\n"
     ]
    }
   ],
   "source": [
    "!ping -c 4 bob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4f13c9bf-1499-4fba-ac06-fab82a007af7",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:56.527107Z",
     "shell.execute_reply.started": "2024-06-23T08:48:52.378485Z",
     "to_execute": "2024-06-23T08:48:36.704Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PING alice (172.23.0.2) 56(84) bytes of data.\n",
      "64 bytes from secretnote-alice-1.secretnote_default (172.23.0.2): icmp_seq=1 ttl=64 time=0.027 ms\n",
      "64 bytes from secretnote-alice-1.secretnote_default (172.23.0.2): icmp_seq=2 ttl=64 time=0.034 ms\n",
      "64 bytes from secretnote-alice-1.secretnote_default (172.23.0.2): icmp_seq=3 ttl=64 time=0.037 ms\n",
      "64 bytes from secretnote-alice-1.secretnote_default (172.23.0.2): icmp_seq=4 ttl=64 time=0.030 ms\n",
      "\n",
      "--- alice ping statistics ---\n",
      "4 packets transmitted, 4 received, 0% packet loss, time 3161ms\n",
      "rtt min/avg/max/mdev = 0.027/0.032/0.037/0.003 ms\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PING alice (172.23.0.2) 56(84) bytes of data.\n",
      "64 bytes from a5ceb21c2bad (172.23.0.2): icmp_seq=1 ttl=64 time=0.009 ms\n",
      "64 bytes from a5ceb21c2bad (172.23.0.2): icmp_seq=2 ttl=64 time=0.017 ms\n",
      "64 bytes from a5ceb21c2bad (172.23.0.2): icmp_seq=3 ttl=64 time=0.016 ms\n",
      "64 bytes from a5ceb21c2bad (172.23.0.2): icmp_seq=4 ttl=64 time=0.017 ms\n",
      "\n",
      "--- alice ping statistics ---\n",
      "4 packets transmitted, 4 received, 0% packet loss, time 3082ms\n",
      "rtt min/avg/max/mdev = 0.009/0.014/0.017/0.005 ms\n"
     ]
    }
   ],
   "source": [
    "!ping -c 4 alice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0593a3d6-4a2a-4f7e-a576-23b8a90ad733",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:56.689686Z",
     "shell.execute_reply.started": "2024-06-23T08:48:55.737767Z",
     "to_execute": "2024-06-23T08:48:36.704Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [],
   "source": [
    "ret = spu_device(compute_dk_func)(alice_data)\n",
    "sf.reveal(ret);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "781e625f-7300-4b6e-8c01-24d99d8248b9",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:56.890931Z",
     "shell.execute_reply.started": "2024-06-23T08:48:56.689627Z",
     "to_execute": "2024-06-23T08:48:36.704Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Cannot delete qdisc with handle of zero.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: Cannot delete qdisc with handle of zero.\n"
     ]
    }
   ],
   "source": [
    "!tc qdisc del dev eth0 root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c363e2c6-9791-4aea-9c01-503854191fc0",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:56.978587Z",
     "shell.execute_reply.started": "2024-06-23T08:48:56.888475Z",
     "to_execute": "2024-06-23T08:48:36.704Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Array([1, 1, 4, 3, 2, 5], dtype=int32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Array([1, 1, 4, 3, 2, 5], dtype=int32)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Homework\n",
    "x = np.array([0.0, 0.2, 6.4, 3.0, 1.6, 12.0])\n",
    "bins = np.array([0.0, 1.0, 2.5, 4.0, 10.0])\n",
    "jnp.digitize(x, bins)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "61a7bdb6-9e16-41ac-a3dd-9a7dacfbc4a0",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:48:57.121027Z",
     "shell.execute_reply.started": "2024-06-23T08:48:56.978551Z",
     "to_execute": "2024-06-23T08:48:36.704Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-23 08:48:57.118] [info] [api.cc:163] [Profiling] SPU execution digitize completed, input processing took 5.5e-07s, execution took 0.087821865s, output processing took 5.51e-07s, total time 0.087822966s.\n",
      "[1 1 4 3 2 5]\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:209] HLO profiling: total time 4.9362000000000006e-05\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.free, executed 439 times, duration 1.9659e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.reshape, executed 121 times, duration 5.458e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.less, executed 69 times, duration 3.89e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.dynamic_slice, executed 72 times, duration 3.144e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.add, executed 60 times, duration 3.125e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.select, executed 67 times, duration 3.103e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.slice, executed 38 times, duration 1.915e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.dynamic_update_slice, executed 36 times, duration 1.906e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.equal, executed 36 times, duration 1.809e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.negate, executed 18 times, duration 1.019e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.constant, executed 11 times, duration 6.96e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.not, executed 13 times, duration 6.77e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.prefer_a, executed 12 times, duration 6.69e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.and, executed 12 times, duration 5.77e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.while, executed 8 times, duration 4.88e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.convert, executed 8 times, duration 4.18e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.sign, executed 6 times, duration 3.4e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.shift_right_logical, executed 6 times, duration 2.94e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.reverse, executed 1 times, duration 6.3e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.subtract, executed 1 times, duration 5.7e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.broadcast, executed 1 times, duration 5.5e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:209] HAL profiling: total time 0.08188369300000001\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_less, executed 134 times, duration 0.028897605s, send bytes 5184 recv bytes 5184\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - _mux, executed 145 times, duration 0.016973054s, send bytes 26544 recv bytes 35136\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mixed_mmul, executed 36 times, duration 0.010881415s, send bytes 2328 recv bytes 2088\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_equal, executed 54 times, duration 0.009920008s, send bytes 6008 recv bytes 5944\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - f_equal, executed 24 times, duration 0.004570654s, send bytes 3072 recv bytes 2688\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - f_less, executed 7 times, duration 0.002355702s, send bytes 1332 recv bytes 1332\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - _sign, executed 6 times, duration 0.002025159s, send bytes 1680 recv bytes 1680\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_negate, executed 18 times, duration 0.001938665s, send bytes 4800 recv bytes 4896\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - _and, executed 12 times, duration 0.001720044s, send bytes 4068 recv bytes 4068\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - _rshift, executed 6 times, duration 0.001646004s, send bytes 4032 recv bytes 4032\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_add, executed 72 times, duration 0.000681961s, send bytes 816 recv bytes 528\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - seal, executed 44 times, duration 0.000164268s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - logical_not, executed 13 times, duration 0.000100036s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_sub, executed 1 times, duration 9.118e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:209] MPC profiling: total time 0.07926754299999998\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - msb_a2b, executed 97 times, duration 0.031371222s, send bytes 7812 recv bytes 7812\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - equal_sp, executed 66 times, duration 0.011929187s, send bytes 7544 recv bytes 7288\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mmul_aa, executed 36 times, duration 0.009430355s, send bytes 288 recv bytes 288\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - b2a, executed 126 times, duration 0.009066552s, send bytes 28680 recv bytes 38664\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mul_a1b, executed 79 times, duration 0.007081277s, send bytes 2736 recv bytes 912\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - a2b, executed 12 times, duration 0.003260473s, send bytes 8064 recv bytes 8064\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - equal_ss, executed 12 times, duration 0.002375865s, send bytes 1536 recv bytes 1344\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mul_aa, executed 66 times, duration 0.002038737s, send bytes 3168 recv bytes 3168\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - add_ap, executed 394 times, duration 0.0003936s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - add_pp, executed 274 times, duration 0.000359732s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - make_p, executed 336 times, duration 0.000338445s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - not_p, executed 182 times, duration 0.000242483s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - reshape, executed 229 times, duration 0.000232151s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - add_aa, executed 153 times, duration 0.000191911s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - extract_slice, executed 146 times, duration 0.000175502s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pad, executed 36 times, duration 0.000149396s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - not_a, executed 129 times, duration 0.000145468s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - update_slice, executed 36 times, duration 0.000119905s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - p2b, executed 36 times, duration 7.2303e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - broadcast, executed 56 times, duration 6.0461e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - and_bb, executed 6 times, duration 5.3183e-05s, send bytes 36 recv bytes 36\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - cast_type_b, executed 36 times, duration 4.6628e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - msb_p, executed 50 times, duration 3.9357e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - xor_bp, executed 13 times, duration 2.8892e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - common_type_b, executed 36 times, duration 2.7992e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - and_bp, executed 6 times, duration 1.0502e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mul_ap, executed 6 times, duration 9.556e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - p2a, executed 8 times, duration 8.924e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - rshift_b, executed 6 times, duration 6.06e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - reverse, executed 1 times, duration 1.424e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:222] Link details: total send bytes 59864, recv bytes 67576, send actions 2102\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-23 08:48:57.118] [info] [api.cc:163] [Profiling] SPU execution digitize completed, input processing took 7.22e-07s, execution took 0.087801274s, output processing took 7.45e-07s, total time 0.087802741s.\n",
      "[1 1 4 3 2 5]\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:209] HLO profiling: total time 5.0784000000000014e-05\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.free, executed 439 times, duration 2.0416e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.reshape, executed 121 times, duration 5.408e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.less, executed 69 times, duration 4.161e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.dynamic_slice, executed 72 times, duration 3.362e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.add, executed 60 times, duration 3.268e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.select, executed 67 times, duration 3.119e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.slice, executed 38 times, duration 2.013e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.dynamic_update_slice, executed 36 times, duration 1.883e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.equal, executed 36 times, duration 1.728e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.negate, executed 18 times, duration 9.87e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.constant, executed 11 times, duration 7.85e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.prefer_a, executed 12 times, duration 6.74e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.not, executed 13 times, duration 6.6e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.and, executed 12 times, duration 6.08e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.while, executed 8 times, duration 4.43e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.convert, executed 8 times, duration 4.19e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.shift_right_logical, executed 6 times, duration 3.42e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.sign, executed 6 times, duration 3.29e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.reverse, executed 1 times, duration 6.6e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.subtract, executed 1 times, duration 5.9e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pphlo.broadcast, executed 1 times, duration 5.4e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:209] HAL profiling: total time 0.081245041\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_less, executed 134 times, duration 0.027869449s, send bytes 5184 recv bytes 5184\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - _mux, executed 145 times, duration 0.016655197s, send bytes 30624 recv bytes 25440\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_equal, executed 54 times, duration 0.010236404s, send bytes 6032 recv bytes 5616\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mixed_mmul, executed 36 times, duration 0.009841381s, send bytes 2408 recv bytes 1768\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - f_equal, executed 24 times, duration 0.004746792s, send bytes 3168 recv bytes 2688\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - f_less, executed 7 times, duration 0.00310348s, send bytes 1332 recv bytes 1332\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_negate, executed 18 times, duration 0.002401493s, send bytes 4656 recv bytes 5040\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - _sign, executed 6 times, duration 0.00217171s, send bytes 1632 recv bytes 1728\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - _and, executed 12 times, duration 0.001895271s, send bytes 4068 recv bytes 4068\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - _rshift, executed 6 times, duration 0.001500865s, send bytes 4032 recv bytes 4032\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_add, executed 72 times, duration 0.000552518s, send bytes 864 recv bytes 480\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - seal, executed 44 times, duration 0.000158446s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - logical_not, executed 13 times, duration 0.000102996s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - i_sub, executed 1 times, duration 9.039e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:209] MPC profiling: total time 0.078636474\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - msb_a2b, executed 97 times, duration 0.03113628s, send bytes 7812 recv bytes 7812\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - equal_sp, executed 66 times, duration 0.012193129s, send bytes 7616 recv bytes 6912\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - b2a, executed 126 times, duration 0.009290111s, send bytes 32696 recv bytes 28792\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mmul_aa, executed 36 times, duration 0.008415687s, send bytes 288 recv bytes 288\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mul_a1b, executed 79 times, duration 0.006876106s, send bytes 2736 recv bytes 912\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - a2b, executed 12 times, duration 0.003101118s, send bytes 8064 recv bytes 8064\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - equal_ss, executed 12 times, duration 0.002579575s, send bytes 1584 recv bytes 1392\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mul_aa, executed 66 times, duration 0.002022213s, send bytes 3168 recv bytes 3168\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - add_ap, executed 394 times, duration 0.000408157s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - make_p, executed 336 times, duration 0.000359658s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - add_pp, executed 274 times, duration 0.000338346s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - not_p, executed 182 times, duration 0.000290664s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - reshape, executed 229 times, duration 0.000240675s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - and_bb, executed 6 times, duration 0.00023897s, send bytes 36 recv bytes 36\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - extract_slice, executed 146 times, duration 0.000188849s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - pad, executed 36 times, duration 0.00017449s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - add_aa, executed 153 times, duration 0.000171558s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - not_a, executed 129 times, duration 0.00015228s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - update_slice, executed 36 times, duration 0.000147775s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - p2b, executed 36 times, duration 7.3991e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - broadcast, executed 56 times, duration 6.254e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - cast_type_b, executed 36 times, duration 4.7685e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - msb_p, executed 50 times, duration 3.7112e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - xor_bp, executed 13 times, duration 2.955e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - common_type_b, executed 36 times, duration 2.5014e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - mul_ap, executed 6 times, duration 1.018e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - and_bp, executed 6 times, duration 9.249e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - p2a, executed 8 times, duration 8.721e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - rshift_b, executed 6 times, duration 5.88e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:212] - reverse, executed 1 times, duration 9.11e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:48:57.118] [info] [api.cc:222] Link details: total send bytes 64000, recv bytes 57376, send actions 2104\n"
     ]
    }
   ],
   "source": [
    "config_aby = spu.RuntimeConfig(\n",
    "    protocol=spu_pb2.ProtocolKind.ABY3,\n",
    "    field=spu.FieldType.FM64,\n",
    "    fxp_fraction_bits=18,\n",
    "    enable_hal_profile=True,\n",
    "    enable_pphlo_profile=True,\n",
    ")\n",
    "sim_aby = spsim.Simulator(3, config_aby)\n",
    "print(spsim.sim_jax(sim_aby, jnp.digitize)(x, bins))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "2c513875-e5cb-4304-ada3-8540947b86b1",
   "metadata": {
    "execution": {
     "parties": "[\"alice\",\"bob\"]",
     "shell.execute_reply.end": "2024-06-23T08:50:09.856076Z",
     "shell.execute_reply.started": "2024-06-23T08:50:09.701889Z",
     "to_execute": "2024-06-23T08:50:09.692Z"
    },
    "libroFormatter": "formatter-string"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h4>alice's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-23 08:50:09.726] [info] [api.cc:163] [Profiling] SPU execution my_digitize_no_endpoint completed, input processing took 1.29e-06s, execution took 0.001513132s, output processing took 1.755e-06s, total time 0.001516177s.\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:209] HLO profiling: total time 1.1920000000000002e-06\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.constant, executed 1 times, duration 3.27e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.free, executed 7 times, duration 3.12e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.add, executed 3 times, duration 1.75e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.convert, executed 2 times, duration 9.9e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.broadcast, executed 2 times, duration 9.4e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.reduce, executed 1 times, duration 6.6e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.not, executed 1 times, duration 6.2e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.less, executed 1 times, duration 5.7e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:209] HAL profiling: total time 0.001011159\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - f_less, executed 1 times, duration 0.000735072s, send bytes 1080 recv bytes 1080\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - i_add, executed 3 times, duration 0.00026092s, send bytes 288 recv bytes 384\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - logical_not, executed 1 times, duration 1.2653e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - seal, executed 1 times, duration 2.514e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:209] MPC profiling: total time 0.0009917259999999998\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - msb_a2b, executed 1 times, duration 0.000709433s, send bytes 1080 recv bytes 1080\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - b2a, executed 3 times, duration 0.000239603s, send bytes 288 recv bytes 384\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - extract_slice, executed 5 times, duration 1.0002e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - add_aa, executed 4 times, duration 7.764e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - add_ap, executed 1 times, duration 6.491e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - broadcast, executed 2 times, duration 4.011e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - not_a, executed 1 times, duration 3.608e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - reshape, executed 3 times, duration 3.13e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - make_p, executed 2 times, duration 3.067e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - xor_bp, executed 1 times, duration 2.088e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - transpose, executed 1 times, duration 1.377e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - p2a, executed 1 times, duration 1.152e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:222] Link details: total send bytes 1368, recv bytes 1464, send actions 12\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:163] [Profiling] SPU execution my_digitize_with_endpoint completed, input processing took 7.41e-07s, execution took 0.001189213s, output processing took 1.524e-06s, total time 0.001191478s.\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:209] HLO profiling: total time 9.64e-07\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.constant, executed 1 times, duration 2.67e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.free, executed 6 times, duration 2.47e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.add, executed 3 times, duration 1.62e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.broadcast, executed 2 times, duration 9.7e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.convert, executed 2 times, duration 8.5e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.greater, executed 1 times, duration 5.7e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.reduce, executed 1 times, duration 4.9e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:209] HAL profiling: total time 0.000901949\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - f_less, executed 1 times, duration 0.000719744s, send bytes 1080 recv bytes 1080\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - i_add, executed 3 times, duration 0.000180285s, send bytes 432 recv bytes 144\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - seal, executed 1 times, duration 1.92e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:209] MPC profiling: total time 0.000900323\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - msb_a2b, executed 1 times, duration 0.000698133s, send bytes 1080 recv bytes 1080\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - b2a, executed 3 times, duration 0.000170065s, send bytes 432 recv bytes 144\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - extract_slice, executed 5 times, duration 7.228e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - add_aa, executed 4 times, duration 7.034e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - not_a, executed 1 times, duration 5.078e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - broadcast, executed 2 times, duration 3.732e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - reshape, executed 3 times, duration 2.952e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - add_ap, executed 1 times, duration 1.996e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - make_p, executed 1 times, duration 1.937e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - transpose, executed 1 times, duration 1.286e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - p2a, executed 1 times, duration 8.82e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:222] Link details: total send bytes 1512, recv bytes 1224, send actions 13\n",
      "res of digitize: [1 1 4 3 2 5], jnp.digitize: [1 1 4 3 2 5]\n",
      "res of digitize: [0 1 4 3 2 5],  jnp.digitize:[0 1 4 3 2 5]\n",
      "------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<h4>bob's Output:</h4>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2024-06-23 08:50:09.726] [info] [api.cc:163] [Profiling] SPU execution my_digitize_no_endpoint completed, input processing took 7.07e-07s, execution took 0.001461329s, output processing took 1.737e-06s, total time 0.001463773s.\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:209] HLO profiling: total time 1.1380000000000002e-06\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.free, executed 7 times, duration 3.41e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.constant, executed 1 times, duration 2.27e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.add, executed 3 times, duration 2.14e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.broadcast, executed 2 times, duration 9.6e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.convert, executed 2 times, duration 9.2e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.less, executed 1 times, duration 5.7e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.not, executed 1 times, duration 5.7e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - pphlo.reduce, executed 1 times, duration 5.4e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:209] HAL profiling: total time 0.001071887\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - f_less, executed 1 times, duration 0.000786636s, send bytes 1080 recv bytes 1080\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - i_add, executed 3 times, duration 0.00026897s, send bytes 336 recv bytes 336\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - logical_not, executed 1 times, duration 1.3659e-05s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - seal, executed 1 times, duration 2.622e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:209] MPC profiling: total time 0.001014497\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - msb_a2b, executed 1 times, duration 0.000757685s, send bytes 1080 recv bytes 1080\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - b2a, executed 3 times, duration 0.000210583s, send bytes 336 recv bytes 336\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - add_aa, executed 4 times, duration 9.866e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - extract_slice, executed 5 times, duration 9.545e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - not_a, executed 1 times, duration 4.549e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - make_p, executed 2 times, duration 4.206e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - reshape, executed 3 times, duration 4.132e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - broadcast, executed 2 times, duration 3.781e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - add_ap, executed 1 times, duration 3.341e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - xor_bp, executed 1 times, duration 3.085e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - transpose, executed 1 times, duration 2.334e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:212] - p2a, executed 1 times, duration 1.39e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.726] [info] [api.cc:222] Link details: total send bytes 1416, recv bytes 1416, send actions 12\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:163] [Profiling] SPU execution my_digitize_with_endpoint completed, input processing took 7.22e-07s, execution took 0.000890652s, output processing took 1.543e-06s, total time 0.000892917s.\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:209] HLO profiling: total time 1.116e-06\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.constant, executed 1 times, duration 3.59e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.free, executed 6 times, duration 2.91e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.add, executed 3 times, duration 1.66e-07s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.broadcast, executed 2 times, duration 9.6e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.convert, executed 2 times, duration 9.2e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.greater, executed 1 times, duration 5.7e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - pphlo.reduce, executed 1 times, duration 5.5e-08s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:209] HAL profiling: total time 0.000561869\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - f_less, executed 1 times, duration 0.000407595s, send bytes 1080 recv bytes 1080\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - i_add, executed 3 times, duration 0.000152036s, send bytes 336 recv bytes 432\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - seal, executed 1 times, duration 2.238e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:209] MPC profiling: total time 0.000555024\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - msb_a2b, executed 1 times, duration 0.000377585s, send bytes 1080 recv bytes 1080\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - b2a, executed 3 times, duration 0.000138579s, send bytes 336 recv bytes 432\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - extract_slice, executed 5 times, duration 8.708e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - add_aa, executed 4 times, duration 8.291e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - not_a, executed 1 times, duration 5.932e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - broadcast, executed 2 times, duration 5.563e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - reshape, executed 3 times, duration 3.725e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - make_p, executed 1 times, duration 1.985e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - add_ap, executed 1 times, duration 1.966e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - transpose, executed 1 times, duration 1.633e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:212] - p2a, executed 1 times, duration 1.057e-06s, send bytes 0 recv bytes 0\n",
      "[2024-06-23 08:50:09.738] [info] [api.cc:222] Link details: total send bytes 1416, recv bytes 1512, send actions 12\n",
      "res of digitize: [1 1 4 3 2 5], jnp.digitize: [1 1 4 3 2 5]\n",
      "res of digitize: [0 1 4 3 2 5],  jnp.digitize:[0 1 4 3 2 5]\n",
      "------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "# MPC-friendly jnp.digitize example\n",
    "# 根据是否有right变量重载函数结果无效，只能用不同的函数名\n",
    "def my_digitize_no_endpoint(x, bins,right=False):\n",
    "    # vectorize\n",
    "    com = x.reshape(*x.shape, -1) >= bins\n",
    "    # count the number of x that exceeds bins\n",
    "    return jnp.sum(com, axis=-1)\n",
    "def my_digitize_with_endpoint(x, bins,right=True):\n",
    "    # vectorize\n",
    "    com = x.reshape(*x.shape, -1) > bins\n",
    "    # count the number of x that exceeds bins\n",
    "    return jnp.sum(com, axis=-1)\n",
    "\n",
    "res_without_right = spsim.sim_jax(sim_aby, my_digitize_no_endpoint)(x, bins,False)\n",
    "res_with_right = spsim.sim_jax(sim_aby, my_digitize_with_endpoint)(x, bins,True)\n",
    "print(f\"res of digitize: {res_without_right}, jnp.digitize: {jnp.digitize(x, bins, False)}\")\n",
    "print(f\"res of digitize: {res_with_right},  jnp.digitize:{jnp.digitize(x, bins, True)}\")\n",
    "print(\"-\"*60)"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}