{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"# Using Moose with PyMoose Computations\n",
"\n",
"This example demonstrates how to convert a PyMoose computation (i.e. a function decorated with `pm.computation`) to its equivalent in textual form. The resulting text file can be used directly with [Moose command line tools](https://github.com/tf-encrypted/moose/tree/main/moose/src/bin) to analyze, compile, and evaluate computations.\n",
"\n",
"\n",
" Table of Contents
\n",
"\n",
" - Generating a Moose computation in textual form
\n",
" - Using Elk to compile and analyze the textual form
\n",
" \n",
" - Evaluating the textual form against a Moose runtime
\n",
"
\n",
" \n"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import pathlib\n",
"\n",
"import numpy as np\n",
"\n",
"import pymoose as pm\n",
"from pymoose.computation import utils"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Generating a Moose computation in textual form\n",
"\n",
"We want to work with some secure computation, e.g. the function below that securely computes a dot-product between two public constants."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"FIXED = pm.fixed(24, 40)\n",
"\n",
"player0 = pm.host_placement(\"player0\")\n",
"player1 = pm.host_placement(\"player1\")\n",
"player2 = pm.host_placement(\"player2\")\n",
"repl = pm.replicated_placement(\"replicated\", [player0, player1, player2])\n",
"\n",
"@pm.computation\n",
"def my_computation():\n",
" with player0:\n",
" x = pm.constant(np.array([1., 2., 3.]).reshape((1, 3)))\n",
" x = pm.cast(x, dtype=FIXED)\n",
" with player1:\n",
" w = pm.constant(np.array([4., 5., 6.]).reshape((3, 1)))\n",
" w = pm.cast(w, dtype=FIXED)\n",
" \n",
" with repl:\n",
" y_hat = pm.dot(x, w)\n",
"\n",
" with player2:\n",
" result = pm.cast(y_hat, dtype=pm.float64)\n",
" return result"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"(Note that using constants this way is not generally secure, since constants are embedded in the computation graph in plaintext. This is just a simple, pedagogical example of a replicated computation.)\n",
"\n",
"Next we'll implement a function that writes this high-level, abstract computation into a text file at some location `filepath`."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"def comp_to_moose(computation_func, filepath):\n",
" traced_comp: pm.edsl.base.AbstractComputation = pm.trace(computation_func)\n",
" comp_bin: bytes = utils.serialize_computation(traced_comp)\n",
" rust_comp: pm.MooseComputation = pm.elk_compiler.compile_computation(comp_bin, passes=[])\n",
" textual_comp: str = rust_comp.to_textual()\n",
" with open(filepath, \"w\") as f:\n",
" f.write(textual_comp)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Use of the `compile_computation` function here might seem a bit strange at first glance. Since we are passing an empty list of compiler passes to run (`passes=[]`), this will be a no-op for the Elk compiler. But by using `compile_computation` we are implicitly marshaling the Python computation into its canonical Moose computation in Rust, since this function is actually a Rust binding. We get a `MooseComputation` object back, which is just a reference to the Rust-managed Moose computation.\n",
"\n",
"This object has a few methods, including `to_textual()`, which parses the computation into its canonical string. We call this string the \"textual form\" or \"textual representation\" of a Moose computation. When computations in textual form are written to a file, we use the `.moose` extension by convention. Working with computations in textual form allows us to use them with the lower-level Moose and Elk machinery that the Moose command line tools offer, since those are built with the textual representation in mind."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"After writing the computation to disk, we'll exclusively use Moose command line tools to work with it."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"\n",
"this_dir = pathlib.Path.cwd()\n",
"comp_to_moose(my_computation, this_dir / \"dotprod.moose\")\n",
"print(this_dir)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Using Elk to compile and analyze the textual form\n",
"Since we passed an empty list of compiler passes instead of using the default passes, this computation remains un-compiled, and its textual form reflects that:"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"constant_0 = Constant{value = HostFloat64Tensor([[1.0, 2.0, 3.0]])}: () -> Tensor () @Host(player0)\n",
"cast_0 = Cast: (Tensor) -> Tensor (constant_0) @Host(player0)\n",
"constant_1 = Constant{value = HostFloat64Tensor([[4.0], [5.0], [6.0]])}: () -> Tensor () @Host(player1)\n",
"cast_1 = Cast: (Tensor) -> Tensor (constant_1) @Host(player1)\n",
"dot_0 = Dot: (Tensor, Tensor) -> Tensor (cast_0, cast_1) @Replicated(player0, player1, player2)\n",
"cast_2 = Cast: (Tensor) -> Tensor (dot_0) @Host(player2)\n",
"output_0 = Output{tag = \"output_0\"}: (Tensor) -> Tensor (cast_2) @Host(player2)"
]
}
],
"source": [
"!cat dotprod.moose"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"For example, all our non-constant values in the computation have generic type `Tensor` -- this is a higher-level type that is erased during the `Lowering` pass of Elk's default compilation. Another example is that our dot product is placed on the `@Replicated(player0, player1, player2)` placement. Replicated placements are virtual placements, meaning their operations will always have to be compiled into a series of operations against a concrete placement like `HostPlacement` before the computation can be evaluated against a Moose runtime.\n",
"\n",
"Next, let's run a lowering pass on our computation, to see what a compiled computation might look like:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"op_0 = Constant{value = HostFloat64Tensor([[1.0, 2.0, 3.0]])}: () -> HostFloat64Tensor () @Host(player0)\n",
"op_1 = RingFixedpointEncode{scaling_base = 2, scaling_exp = 40}: (HostFloat64Tensor) -> HostRing128Tensor (op_0) @Host(player0)\n",
"op_2 = Constant{value = HostFloat64Tensor([[4.0], [5.0], [6.0]])}: () -> HostFloat64Tensor () @Host(player1)\n",
"op_3 = RingFixedpointEncode{scaling_base = 2, scaling_exp = 40}: (HostFloat64Tensor) -> HostRing128Tensor (op_2) @Host(player1)\n",
"op_4 = PrfKeyGen: () -> HostPrfKey () @Host(player0)\n",
"op_5 = PrfKeyGen: () -> HostPrfKey () @Host(player1)\n",
"op_6 = PrfKeyGen: () -> HostPrfKey () @Host(player2)\n",
"op_7 = Shape: (HostRing128Tensor) -> HostShape (op_1) @Host(player0)\n",
"op_8 = DeriveSeed{sync_key = 909e3e87e8db5c18176e40eebc0ec34b}: (HostPrfKey) -> HostSeed (op_4) @Host(player0)\n",
"op_9 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_7, op_8) @Host(player0)\n",
"op_10 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_1, op_9) @Host(player0)\n",
"op_11 = DeriveSeed{sync_key = 909e3e87e8db5c18176e40eebc0ec34b}: (HostPrfKey) -> HostSeed (op_4) @Host(player2)\n",
"op_12 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (op_7) @Host(player2)\n",
"op_13 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_7, op_11) @Host(player2)\n",
"op_14 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (op_7) @Host(player1)\n",
"op_15 = Shape: (HostRing128Tensor) -> HostShape (op_3) @Host(player1)\n",
"op_16 = DeriveSeed{sync_key = 9b73e4da141ba80305cd7c47a1e270e0}: (HostPrfKey) -> HostSeed (op_5) @Host(player1)\n",
"op_17 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_15, op_16) @Host(player1)\n",
"op_18 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_3, op_17) @Host(player1)\n",
"op_19 = DeriveSeed{sync_key = 9b73e4da141ba80305cd7c47a1e270e0}: (HostPrfKey) -> HostSeed (op_5) @Host(player0)\n",
"op_20 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (op_15) @Host(player0)\n",
"op_21 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_15, op_19) @Host(player0)\n",
"op_22 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (op_15) @Host(player2)\n",
"op_23 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_9, op_20) @Host(player0)\n",
"op_24 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_9, op_21) @Host(player0)\n",
"op_25 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_23, op_24) @Host(player0)\n",
"op_26 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_10, op_20) @Host(player0)\n",
"op_27 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_25, op_26) @Host(player0)\n",
"op_28 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_10, op_17) @Host(player1)\n",
"op_29 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_10, op_18) @Host(player1)\n",
"op_30 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_28, op_29) @Host(player1)\n",
"op_31 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_14, op_17) @Host(player1)\n",
"op_32 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_30, op_31) @Host(player1)\n",
"op_33 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_12, op_18) @Host(player2)\n",
"op_34 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_12, op_22) @Host(player2)\n",
"op_35 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_33, op_34) @Host(player2)\n",
"op_36 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_13, op_18) @Host(player2)\n",
"op_37 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_35, op_36) @Host(player2)\n",
"op_38 = Shape: (HostRing128Tensor) -> HostShape (op_27) @Host(player0)\n",
"op_39 = Shape: (HostRing128Tensor) -> HostShape (op_32) @Host(player1)\n",
"op_40 = Shape: (HostRing128Tensor) -> HostShape (op_37) @Host(player2)\n",
"op_41 = DeriveSeed{sync_key = 7bf21f004f43a654120314d0c1f79d5a}: (HostPrfKey) -> HostSeed (op_4) @Host(player0)\n",
"op_42 = DeriveSeed{sync_key = c7cb29c6001aae64ade8e0c1dad8b1d7}: (HostPrfKey) -> HostSeed (op_5) @Host(player0)\n",
"op_43 = DeriveSeed{sync_key = c7cb29c6001aae64ade8e0c1dad8b1d7}: (HostPrfKey) -> HostSeed (op_5) @Host(player1)\n",
"op_44 = DeriveSeed{sync_key = a96a9f06dab83dbc9267a4e3c609cec7}: (HostPrfKey) -> HostSeed (op_6) @Host(player1)\n",
"op_45 = DeriveSeed{sync_key = a96a9f06dab83dbc9267a4e3c609cec7}: (HostPrfKey) -> HostSeed (op_6) @Host(player2)\n",
"op_46 = DeriveSeed{sync_key = 7bf21f004f43a654120314d0c1f79d5a}: (HostPrfKey) -> HostSeed (op_4) @Host(player2)\n",
"op_47 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_38, op_41) @Host(player0)\n",
"op_48 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_38, op_42) @Host(player0)\n",
"op_49 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_47, op_48) @Host(player0)\n",
"op_50 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_39, op_43) @Host(player1)\n",
"op_51 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_39, op_44) @Host(player1)\n",
"op_52 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_50, op_51) @Host(player1)\n",
"op_53 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_40, op_45) @Host(player2)\n",
"op_54 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_40, op_46) @Host(player2)\n",
"op_55 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_53, op_54) @Host(player2)\n",
"op_56 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_27, op_49) @Host(player0)\n",
"op_57 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_32, op_52) @Host(player1)\n",
"op_58 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_37, op_55) @Host(player2)\n",
"op_59 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_56, op_57) @Host(player0)\n",
"op_60 = Shape: (HostRing128Tensor) -> HostShape (op_59) @Host(player0)\n",
"op_61 = Sample{}: (HostShape) -> HostRing128Tensor (op_60) @Host(player2)\n",
"op_62 = Shr{amount = 127}: (HostRing128Tensor) -> HostRing128Tensor (op_61) @Host(player2)\n",
"op_63 = Shl{amount = 1}: (HostRing128Tensor) -> HostRing128Tensor (op_61) @Host(player2)\n",
"op_64 = Shr{amount = 41}: (HostRing128Tensor) -> HostRing128Tensor (op_63) @Host(player2)\n",
"op_65 = PrfKeyGen: () -> HostPrfKey () @Host(player2)\n",
"op_66 = DeriveSeed{sync_key = 667169f78edb927b50eb369de6deb6d9}: (HostPrfKey) -> HostSeed (op_65) @Host(player2)\n",
"op_67 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_60, op_66) @Host(player2)\n",
"op_68 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_61, op_67) @Host(player2)\n",
"op_69 = DeriveSeed{sync_key = d6dde441ad3684b0a9d08288772f0912}: (HostPrfKey) -> HostSeed (op_65) @Host(player2)\n",
"op_70 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_60, op_69) @Host(player2)\n",
"op_71 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_64, op_70) @Host(player2)\n",
"op_72 = DeriveSeed{sync_key = aff94776cab60aceccd166e3bb23cfa6}: (HostPrfKey) -> HostSeed (op_65) @Host(player2)\n",
"op_73 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_60, op_72) @Host(player2)\n",
"op_74 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_62, op_73) @Host(player2)\n",
"op_75 = Fill{value = Ring128(1)}: (HostShape) -> HostRing128Tensor (op_60) @Host(player0)\n",
"op_76 = Shl{amount = 126}: (HostRing128Tensor) -> HostRing128Tensor (op_75) @Host(player0)\n",
"op_77 = Shl{amount = 86}: (HostRing128Tensor) -> HostRing128Tensor (op_75) @Host(player0)\n",
"op_78 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_59, op_76) @Host(player0)\n",
"op_79 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_78, op_67) @Host(player0)\n",
"op_80 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_58, op_68) @Host(player1)\n",
"op_81 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_79, op_80) @Host(player0)\n",
"op_82 = Shl{amount = 1}: (HostRing128Tensor) -> HostRing128Tensor (op_81) @Host(player0)\n",
"op_83 = Shr{amount = 41}: (HostRing128Tensor) -> HostRing128Tensor (op_82) @Host(player0)\n",
"op_84 = Shr{amount = 127}: (HostRing128Tensor) -> HostRing128Tensor (op_81) @Host(player0)\n",
"op_85 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_73, op_84) @Host(player0)\n",
"op_86 = Mul: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_73, op_84) @Host(player0)\n",
"op_87 = Mul: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_74, op_84) @Host(player1)\n",
"op_88 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_85, op_86) @Host(player0)\n",
"op_89 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_74, op_87) @Host(player1)\n",
"op_90 = Mul: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_73, op_84) @Host(player0)\n",
"op_91 = Mul: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_74, op_84) @Host(player1)\n",
"op_92 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_88, op_90) @Host(player0)\n",
"op_93 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_89, op_91) @Host(player1)\n",
"op_94 = Shl{amount = 87}: (HostRing128Tensor) -> HostRing128Tensor (op_92) @Host(player0)\n",
"op_95 = Shl{amount = 87}: (HostRing128Tensor) -> HostRing128Tensor (op_93) @Host(player1)\n",
"op_96 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_83, op_70) @Host(player0)\n",
"op_97 = Neg: (HostRing128Tensor) -> HostRing128Tensor (op_71) @Host(player1)\n",
"op_98 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_96, op_94) @Host(player0)\n",
"op_99 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_97, op_95) @Host(player1)\n",
"op_100 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_98, op_77) @Host(player0)\n",
"op_101 = PrfKeyGen: () -> HostPrfKey () @Host(player2)\n",
"op_102 = DeriveSeed{sync_key = 09e29329702bc80417def9391fc402fc}: (HostPrfKey) -> HostSeed (op_101) @Host(player2)\n",
"op_103 = DeriveSeed{sync_key = 58a2af91873777cddb83499358faafed}: (HostPrfKey) -> HostSeed (op_101) @Host(player2)\n",
"op_104 = Shape: (HostRing128Tensor) -> HostShape (op_100) @Host(player0)\n",
"op_105 = Shape: (HostRing128Tensor) -> HostShape (op_99) @Host(player1)\n",
"op_106 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_104, op_102) @Host(player0)\n",
"op_107 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_105, op_103) @Host(player1)\n",
"op_108 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_104, op_102) @Host(player2)\n",
"op_109 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_104, op_103) @Host(player2)\n",
"op_110 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_100, op_106) @Host(player0)\n",
"op_111 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_99, op_107) @Host(player1)\n",
"op_112 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_110, op_111) @Host(player0)\n",
"op_113 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_108, op_112) @Host(player2)\n",
"op_114 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_113, op_109) @Host(player2)\n",
"op_115 = RingFixedpointDecode{scaling_base = 2, scaling_exp = 40}: (HostRing128Tensor) -> HostFloat64Tensor (op_114) @Host(player2)\n",
"op_116 = Output{tag = \"output_0\"}: (HostFloat64Tensor) -> HostFloat64Tensor (op_115) @Host(player2)\n"
]
}
],
"source": [
"!elk compile dotprod.moose --passes lowering"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Using Elk to get an evaluation-ready computation\n",
"\n",
"Examining this compiled computation, we can see that all types are concrete, e.g. they all have the `Host` prefix, denoting that they are values owned by particular host devices. All operations are also pinned against HostPlacements; there are no more virtual placements.\n",
"\n",
"However, this computation is not yet ready for evaluation as we can see by examining the following two lines in the compiled output:\n",
"```console\n",
"op_7 = Shape: (HostRing128Tensor) -> HostShape (op_1) @Host(player0)\n",
"...\n",
"op_12 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (op_7) @Host(player2)\n",
"```\n",
"Taking this apart step-by-step, we can see that `op_7` evaluates the shape of a tensor located on `player0`, then `op_12` uses that shape to generate a tensor on `player2`.\n",
"\n",
"Since this is the first time the output of `op_7` is used outside of `player0`, we should ask ourselves: how will this shape data make its way from `player0` to `player2` at runtime? The answer is that we have to run another compiler pass to insert the networking ops that will be required at runtime, which we can do now:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"op_0 = Constant{value = HostFloat64Tensor([[1.0, 2.0, 3.0]])}: () -> HostFloat64Tensor () @Host(player0)\n",
"op_1 = RingFixedpointEncode{scaling_base = 2, scaling_exp = 40}: (HostFloat64Tensor) -> HostRing128Tensor (op_0) @Host(player0)\n",
"op_2 = Constant{value = HostFloat64Tensor([[4.0], [5.0], [6.0]])}: () -> HostFloat64Tensor () @Host(player1)\n",
"op_3 = RingFixedpointEncode{scaling_base = 2, scaling_exp = 40}: (HostFloat64Tensor) -> HostRing128Tensor (op_2) @Host(player1)\n",
"op_4 = PrfKeyGen: () -> HostPrfKey () @Host(player0)\n",
"op_5 = PrfKeyGen: () -> HostPrfKey () @Host(player1)\n",
"op_6 = PrfKeyGen: () -> HostPrfKey () @Host(player2)\n",
"op_7 = Shape: (HostRing128Tensor) -> HostShape (op_1) @Host(player0)\n",
"op_8 = DeriveSeed{sync_key = 732bb872ba911aa46838af1291715b45}: (HostPrfKey) -> HostSeed (op_4) @Host(player0)\n",
"op_9 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_7, op_8) @Host(player0)\n",
"op_10 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_1, op_9) @Host(player0)\n",
"op_11 = DeriveSeed{sync_key = 732bb872ba911aa46838af1291715b45}: (HostPrfKey) -> HostSeed (receive_0) @Host(player2)\n",
"op_12 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (receive_4) @Host(player2)\n",
"op_13 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (receive_4, op_11) @Host(player2)\n",
"op_14 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (receive_3) @Host(player1)\n",
"op_15 = Shape: (HostRing128Tensor) -> HostShape (op_3) @Host(player1)\n",
"op_16 = DeriveSeed{sync_key = e84a0e505706ab471e7731d57128a9de}: (HostPrfKey) -> HostSeed (op_5) @Host(player1)\n",
"op_17 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_15, op_16) @Host(player1)\n",
"op_18 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_3, op_17) @Host(player1)\n",
"op_19 = DeriveSeed{sync_key = e84a0e505706ab471e7731d57128a9de}: (HostPrfKey) -> HostSeed (receive_1) @Host(player0)\n",
"op_20 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (receive_7) @Host(player0)\n",
"op_21 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (receive_7, op_19) @Host(player0)\n",
"op_22 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (receive_6) @Host(player2)\n",
"op_23 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_9, op_20) @Host(player0)\n",
"op_24 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_9, op_21) @Host(player0)\n",
"op_25 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_23, op_24) @Host(player0)\n",
"op_26 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_10, op_20) @Host(player0)\n",
"op_27 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_25, op_26) @Host(player0)\n",
"op_28 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_5, op_17) @Host(player1)\n",
"op_29 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_5, op_18) @Host(player1)\n",
"op_30 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_28, op_29) @Host(player1)\n",
"op_31 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_14, op_17) @Host(player1)\n",
"op_32 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_30, op_31) @Host(player1)\n",
"op_33 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_12, receive_8) @Host(player2)\n",
"op_34 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_12, op_22) @Host(player2)\n",
"op_35 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_33, op_34) @Host(player2)\n",
"op_36 = Dot: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_13, receive_8) @Host(player2)\n",
"op_37 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_35, op_36) @Host(player2)\n",
"op_38 = Shape: (HostRing128Tensor) -> HostShape (op_27) @Host(player0)\n",
"op_39 = Shape: (HostRing128Tensor) -> HostShape (op_32) @Host(player1)\n",
"op_40 = Shape: (HostRing128Tensor) -> HostShape (op_37) @Host(player2)\n",
"op_41 = DeriveSeed{sync_key = ddea463ece88b0ad989145dea88b6eee}: (HostPrfKey) -> HostSeed (op_4) @Host(player0)\n",
"op_42 = DeriveSeed{sync_key = a8168949e08f07336a54475ed061d6ac}: (HostPrfKey) -> HostSeed (receive_1) @Host(player0)\n",
"op_43 = DeriveSeed{sync_key = a8168949e08f07336a54475ed061d6ac}: (HostPrfKey) -> HostSeed (op_5) @Host(player1)\n",
"op_44 = DeriveSeed{sync_key = f8b8f465f1e9ac823d13fde16b3cb151}: (HostPrfKey) -> HostSeed (receive_2) @Host(player1)\n",
"op_45 = DeriveSeed{sync_key = f8b8f465f1e9ac823d13fde16b3cb151}: (HostPrfKey) -> HostSeed (op_6) @Host(player2)\n",
"op_46 = DeriveSeed{sync_key = ddea463ece88b0ad989145dea88b6eee}: (HostPrfKey) -> HostSeed (receive_0) @Host(player2)\n",
"op_47 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_38, op_41) @Host(player0)\n",
"op_48 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_38, op_42) @Host(player0)\n",
"op_49 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_47, op_48) @Host(player0)\n",
"op_50 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_39, op_43) @Host(player1)\n",
"op_51 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_39, op_44) @Host(player1)\n",
"op_52 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_50, op_51) @Host(player1)\n",
"op_53 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_40, op_45) @Host(player2)\n",
"op_54 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_40, op_46) @Host(player2)\n",
"op_55 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_53, op_54) @Host(player2)\n",
"op_56 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_27, op_49) @Host(player0)\n",
"op_57 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_32, op_52) @Host(player1)\n",
"op_58 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_37, op_55) @Host(player2)\n",
"op_59 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_56, receive_9) @Host(player0)\n",
"op_60 = Shape: (HostRing128Tensor) -> HostShape (op_59) @Host(player0)\n",
"op_61 = Sample{}: (HostShape) -> HostRing128Tensor (receive_11) @Host(player2)\n",
"op_62 = Shr{amount = 127}: (HostRing128Tensor) -> HostRing128Tensor (op_61) @Host(player2)\n",
"op_63 = Shl{amount = 1}: (HostRing128Tensor) -> HostRing128Tensor (op_61) @Host(player2)\n",
"op_64 = Shr{amount = 41}: (HostRing128Tensor) -> HostRing128Tensor (op_63) @Host(player2)\n",
"op_65 = PrfKeyGen: () -> HostPrfKey () @Host(player2)\n",
"op_66 = DeriveSeed{sync_key = a253b0e2321f99c37ffe778852687910}: (HostPrfKey) -> HostSeed (op_65) @Host(player2)\n",
"op_67 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (receive_11, op_66) @Host(player2)\n",
"op_68 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_61, op_67) @Host(player2)\n",
"op_69 = DeriveSeed{sync_key = 51512c9c05600558dc8552d2990eb5fb}: (HostPrfKey) -> HostSeed (op_65) @Host(player2)\n",
"op_70 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (receive_11, op_69) @Host(player2)\n",
"op_71 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_64, op_70) @Host(player2)\n",
"op_72 = DeriveSeed{sync_key = cdc71aee6053c89dede727260609ffe6}: (HostPrfKey) -> HostSeed (op_65) @Host(player2)\n",
"op_73 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (receive_11, op_72) @Host(player2)\n",
"op_74 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_62, op_73) @Host(player2)\n",
"op_75 = Fill{value = Ring128(1)}: (HostShape) -> HostRing128Tensor (op_60) @Host(player0)\n",
"op_76 = Shl{amount = 126}: (HostRing128Tensor) -> HostRing128Tensor (op_75) @Host(player0)\n",
"op_77 = Shl{amount = 86}: (HostRing128Tensor) -> HostRing128Tensor (op_75) @Host(player0)\n",
"op_78 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_59, op_76) @Host(player0)\n",
"op_79 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_78, receive_12) @Host(player0)\n",
"op_80 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_10, receive_13) @Host(player1)\n",
"op_81 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_79, receive_18) @Host(player0)\n",
"op_82 = Shl{amount = 1}: (HostRing128Tensor) -> HostRing128Tensor (op_81) @Host(player0)\n",
"op_83 = Shr{amount = 41}: (HostRing128Tensor) -> HostRing128Tensor (op_82) @Host(player0)\n",
"op_84 = Shr{amount = 127}: (HostRing128Tensor) -> HostRing128Tensor (op_81) @Host(player0)\n",
"op_85 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_16, op_84) @Host(player0)\n",
"op_86 = Mul: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_16, op_84) @Host(player0)\n",
"op_87 = Mul: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_17, receive_19) @Host(player1)\n",
"op_88 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_85, op_86) @Host(player0)\n",
"op_89 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_17, op_87) @Host(player1)\n",
"op_90 = Mul: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_16, op_84) @Host(player0)\n",
"op_91 = Mul: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (receive_17, receive_19) @Host(player1)\n",
"op_92 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_88, op_90) @Host(player0)\n",
"op_93 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_89, op_91) @Host(player1)\n",
"op_94 = Shl{amount = 87}: (HostRing128Tensor) -> HostRing128Tensor (op_92) @Host(player0)\n",
"op_95 = Shl{amount = 87}: (HostRing128Tensor) -> HostRing128Tensor (op_93) @Host(player1)\n",
"op_96 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_83, receive_14) @Host(player0)\n",
"op_97 = Neg: (HostRing128Tensor) -> HostRing128Tensor (receive_15) @Host(player1)\n",
"op_98 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_96, op_94) @Host(player0)\n",
"op_99 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_97, op_95) @Host(player1)\n",
"op_100 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_98, op_77) @Host(player0)\n",
"op_101 = PrfKeyGen: () -> HostPrfKey () @Host(player2)\n",
"op_102 = DeriveSeed{sync_key = 1c6bbc89ad75b5e4d55a101e530dd370}: (HostPrfKey) -> HostSeed (op_101) @Host(player2)\n",
"op_103 = DeriveSeed{sync_key = ffccecbd671030e991711fc13bdd7d7b}: (HostPrfKey) -> HostSeed (op_101) @Host(player2)\n",
"op_104 = Shape: (HostRing128Tensor) -> HostShape (op_100) @Host(player0)\n",
"op_105 = Shape: (HostRing128Tensor) -> HostShape (op_99) @Host(player1)\n",
"op_106 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_104, receive_20) @Host(player0)\n",
"op_107 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (op_105, receive_21) @Host(player1)\n",
"op_108 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (receive_22, op_102) @Host(player2)\n",
"op_109 = SampleSeeded{}: (HostShape, HostSeed) -> HostRing128Tensor (receive_22, op_103) @Host(player2)\n",
"op_110 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_100, op_106) @Host(player0)\n",
"op_111 = Sub: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_99, op_107) @Host(player1)\n",
"op_112 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_110, receive_23) @Host(player0)\n",
"op_113 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_108, receive_24) @Host(player2)\n",
"op_114 = Add: (HostRing128Tensor, HostRing128Tensor) -> HostRing128Tensor (op_113, op_109) @Host(player2)\n",
"op_115 = RingFixedpointDecode{scaling_base = 2, scaling_exp = 40}: (HostRing128Tensor) -> HostFloat64Tensor (op_114) @Host(player2)\n",
"op_116 = Output{tag = \"output_0\"}: (HostFloat64Tensor) -> HostFloat64Tensor (op_115) @Host(player2)\n",
"send_0 = Send{rendezvous_key = 00000000000000000000000000000000, receiver = \"player2\"}: (HostPrfKey) -> HostUnit (op_4) @Host(player0)\n",
"receive_0 = Receive{rendezvous_key = 00000000000000000000000000000000, sender = \"player0\"}: () -> HostPrfKey () @Host(player2)\n",
"send_1 = Send{rendezvous_key = 01000000000000000000000000000000, receiver = \"player0\"}: (HostPrfKey) -> HostUnit (op_5) @Host(player1)\n",
"receive_1 = Receive{rendezvous_key = 01000000000000000000000000000000, sender = \"player1\"}: () -> HostPrfKey () @Host(player0)\n",
"send_2 = Send{rendezvous_key = 02000000000000000000000000000000, receiver = \"player1\"}: (HostPrfKey) -> HostUnit (op_6) @Host(player2)\n",
"receive_2 = Receive{rendezvous_key = 02000000000000000000000000000000, sender = \"player2\"}: () -> HostPrfKey () @Host(player1)\n",
"send_3 = Send{rendezvous_key = 03000000000000000000000000000000, receiver = \"player1\"}: (HostShape) -> HostUnit (op_7) @Host(player0)\n",
"receive_3 = Receive{rendezvous_key = 03000000000000000000000000000000, sender = \"player0\"}: () -> HostShape () @Host(player1)\n",
"send_4 = Send{rendezvous_key = 04000000000000000000000000000000, receiver = \"player2\"}: (HostShape) -> HostUnit (op_7) @Host(player0)\n",
"receive_4 = Receive{rendezvous_key = 04000000000000000000000000000000, sender = \"player0\"}: () -> HostShape () @Host(player2)\n",
"send_5 = Send{rendezvous_key = 05000000000000000000000000000000, receiver = \"player1\"}: (HostRing128Tensor) -> HostUnit (op_10) @Host(player0)\n",
"receive_5 = Receive{rendezvous_key = 05000000000000000000000000000000, sender = \"player0\"}: () -> HostRing128Tensor () @Host(player1)\n",
"send_6 = Send{rendezvous_key = 06000000000000000000000000000000, receiver = \"player2\"}: (HostShape) -> HostUnit (op_15) @Host(player1)\n",
"receive_6 = Receive{rendezvous_key = 06000000000000000000000000000000, sender = \"player1\"}: () -> HostShape () @Host(player2)\n",
"send_7 = Send{rendezvous_key = 07000000000000000000000000000000, receiver = \"player0\"}: (HostShape) -> HostUnit (op_15) @Host(player1)\n",
"receive_7 = Receive{rendezvous_key = 07000000000000000000000000000000, sender = \"player1\"}: () -> HostShape () @Host(player0)\n",
"send_8 = Send{rendezvous_key = 08000000000000000000000000000000, receiver = \"player2\"}: (HostRing128Tensor) -> HostUnit (op_18) @Host(player1)\n",
"receive_8 = Receive{rendezvous_key = 08000000000000000000000000000000, sender = \"player1\"}: () -> HostRing128Tensor () @Host(player2)\n",
"send_9 = Send{rendezvous_key = 09000000000000000000000000000000, receiver = \"player0\"}: (HostRing128Tensor) -> HostUnit (op_57) @Host(player1)\n",
"receive_9 = Receive{rendezvous_key = 09000000000000000000000000000000, sender = \"player1\"}: () -> HostRing128Tensor () @Host(player0)\n",
"send_10 = Send{rendezvous_key = 0a000000000000000000000000000000, receiver = \"player1\"}: (HostRing128Tensor) -> HostUnit (op_58) @Host(player2)\n",
"receive_10 = Receive{rendezvous_key = 0a000000000000000000000000000000, sender = \"player2\"}: () -> HostRing128Tensor () @Host(player1)\n",
"send_11 = Send{rendezvous_key = 0b000000000000000000000000000000, receiver = \"player2\"}: (HostShape) -> HostUnit (op_60) @Host(player0)\n",
"receive_11 = Receive{rendezvous_key = 0b000000000000000000000000000000, sender = \"player0\"}: () -> HostShape () @Host(player2)\n",
"send_12 = Send{rendezvous_key = 0c000000000000000000000000000000, receiver = \"player0\"}: (HostRing128Tensor) -> HostUnit (op_67) @Host(player2)\n",
"receive_12 = Receive{rendezvous_key = 0c000000000000000000000000000000, sender = \"player2\"}: () -> HostRing128Tensor () @Host(player0)\n",
"send_13 = Send{rendezvous_key = 0d000000000000000000000000000000, receiver = \"player1\"}: (HostRing128Tensor) -> HostUnit (op_68) @Host(player2)\n",
"receive_13 = Receive{rendezvous_key = 0d000000000000000000000000000000, sender = \"player2\"}: () -> HostRing128Tensor () @Host(player1)\n",
"send_14 = Send{rendezvous_key = 0e000000000000000000000000000000, receiver = \"player0\"}: (HostRing128Tensor) -> HostUnit (op_70) @Host(player2)\n",
"receive_14 = Receive{rendezvous_key = 0e000000000000000000000000000000, sender = \"player2\"}: () -> HostRing128Tensor () @Host(player0)\n",
"send_15 = Send{rendezvous_key = 0f000000000000000000000000000000, receiver = \"player1\"}: (HostRing128Tensor) -> HostUnit (op_71) @Host(player2)\n",
"receive_15 = Receive{rendezvous_key = 0f000000000000000000000000000000, sender = \"player2\"}: () -> HostRing128Tensor () @Host(player1)\n",
"send_16 = Send{rendezvous_key = 10000000000000000000000000000000, receiver = \"player0\"}: (HostRing128Tensor) -> HostUnit (op_73) @Host(player2)\n",
"receive_16 = Receive{rendezvous_key = 10000000000000000000000000000000, sender = \"player2\"}: () -> HostRing128Tensor () @Host(player0)\n",
"send_17 = Send{rendezvous_key = 11000000000000000000000000000000, receiver = \"player1\"}: (HostRing128Tensor) -> HostUnit (op_74) @Host(player2)\n",
"receive_17 = Receive{rendezvous_key = 11000000000000000000000000000000, sender = \"player2\"}: () -> HostRing128Tensor () @Host(player1)\n",
"send_18 = Send{rendezvous_key = 12000000000000000000000000000000, receiver = \"player0\"}: (HostRing128Tensor) -> HostUnit (op_80) @Host(player1)\n",
"receive_18 = Receive{rendezvous_key = 12000000000000000000000000000000, sender = \"player1\"}: () -> HostRing128Tensor () @Host(player0)\n",
"send_19 = Send{rendezvous_key = 13000000000000000000000000000000, receiver = \"player1\"}: (HostRing128Tensor) -> HostUnit (op_84) @Host(player0)\n",
"receive_19 = Receive{rendezvous_key = 13000000000000000000000000000000, sender = \"player0\"}: () -> HostRing128Tensor () @Host(player1)\n",
"send_20 = Send{rendezvous_key = 14000000000000000000000000000000, receiver = \"player0\"}: (HostSeed) -> HostUnit (op_102) @Host(player2)\n",
"receive_20 = Receive{rendezvous_key = 14000000000000000000000000000000, sender = \"player2\"}: () -> HostSeed () @Host(player0)\n",
"send_21 = Send{rendezvous_key = 15000000000000000000000000000000, receiver = \"player1\"}: (HostSeed) -> HostUnit (op_103) @Host(player2)\n",
"receive_21 = Receive{rendezvous_key = 15000000000000000000000000000000, sender = \"player2\"}: () -> HostSeed () @Host(player1)\n",
"send_22 = Send{rendezvous_key = 16000000000000000000000000000000, receiver = \"player2\"}: (HostShape) -> HostUnit (op_104) @Host(player0)\n",
"receive_22 = Receive{rendezvous_key = 16000000000000000000000000000000, sender = \"player0\"}: () -> HostShape () @Host(player2)\n",
"send_23 = Send{rendezvous_key = 17000000000000000000000000000000, receiver = \"player0\"}: (HostRing128Tensor) -> HostUnit (op_111) @Host(player1)\n",
"receive_23 = Receive{rendezvous_key = 17000000000000000000000000000000, sender = \"player1\"}: () -> HostRing128Tensor () @Host(player0)\n",
"send_24 = Send{rendezvous_key = 18000000000000000000000000000000, receiver = \"player2\"}: (HostRing128Tensor) -> HostUnit (op_112) @Host(player0)\n",
"receive_24 = Receive{rendezvous_key = 18000000000000000000000000000000, sender = \"player0\"}: () -> HostRing128Tensor () @Host(player2)\n"
]
}
],
"source": [
"!elk compile dotprod.moose --passes lowering,networking"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Examining the output, we can see that our operations `op_7` and `op_12` have been updated to allow for sending the shape data from `player0` to `player2`:\n",
"```console\n",
"op_7 = Shape: (HostRing128Tensor) -> HostShape (op_1) @Host(player0)\n",
"...\n",
"op_12 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (receive_4) @Host(player2)\n",
"...\n",
"send_4 = Send{rendezvous_key = 04000000000000000000000000000000, receiver = \"player2\"}: (HostShape) -> HostUnit (op_7) @Host(player0)\n",
"receive_4 = Receive{rendezvous_key = 04000000000000000000000000000000, sender = \"player0\"}: () -> HostShape () @Host(player2)\n",
"```"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"It's easier to see what's happening by reordering these ops to respect input-output relations:\n",
"```console\n",
"op_7 = Shape: (HostRing128Tensor) -> HostShape (op_1) @Host(player0)\n",
"send_4 = Send{rendezvous_key = 04000000000000000000000000000000, receiver = \"player2\"}: (HostShape) -> HostUnit (op_7) @Host(player0)\n",
"receive_4 = Receive{rendezvous_key = 04000000000000000000000000000000, sender = \"player0\"}: () -> HostShape () @Host(player2)\n",
"op_12 = Fill{value = Ring128(0)}: (HostShape) -> HostRing128Tensor (receive_4) @Host(player2)\n",
"```\n",
"The result of adding the networking pass to compilation is that our line for `op_12` has been modified to take its input from a `receive_4 @ Host(player2)` operation. This operation shares a rendezvous key with `send_4 @ Host(player0)`, which takes `op_7` as an input. Thus the computation now instructs that the output of `op_7` should be sent from `player0` to `player2`, and used as the input to `op_12`."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Common compilation passes\n",
"The following is a sane default list of compiler passes to compile arbitrary PyMoose computations into their runtime-ready versions:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"!elk compile dotprod.moose -p typing,lowering,prune,networking,toposort -o dotprod-networked.moose"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"- The `typing` pass looks for Values of `UnknownType` in a computation and attempts to fill in type information using a simple one-hop inference rule (e.g. use the next operation's signature to fill in the output type). It's only necessary if you were lazy with how you specified output vtypes of a few kinds of operations in your PyMoose computation.\n",
"- The `prune` pass looks for all subgraphs that aren't connected to an Output operation, and removes those subgraphs from the graph.\n",
"- The `toposort` pass reorders the resulting textual form of the computation so that the directed graph it represents is in topological order, i.e. any operation that takes a set of inputs comes after those inputs. Note this means that you cannot generally treat the textual form of a computation as unique."
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"### Other compilation passes\n",
"\n",
"#### WellFormed\n",
"If you only want to a textual computation for correctness without actually returning its compiled form, the `wellformed` pass can be faster than performing the full `lowering`.\n",
"\n",
"#### Print\n",
"\n",
"The `print` pass allows you to convert your `.moose` file into `.dot` files for GraphViz rendering:"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"digraph {\n",
" 0 [ label = \"constant_0 = Constant\\l@Host(player0)\" shape = rectangle color = \"#336699\"]\n",
" 1 [ label = \"cast_0 = Cast\\l@Host(player0)\" shape = rectangle color = \"#336699\"]\n",
" 2 [ label = \"constant_1 = Constant\\l@Host(player1)\" shape = rectangle color = \"#ff0000\"]\n",
" 3 [ label = \"cast_1 = Cast\\l@Host(player1)\" shape = rectangle color = \"#ff0000\"]\n",
" 4 [ label = \"dot_0 = Dot\\l@Replicated(player0, player1, player2)\" shape = rectangle color = \"#ff6600\"]\n",
" 5 [ label = \"cast_2 = Cast\\l@Host(player2)\" shape = rectangle color = \"#92cd00\"]\n",
" 6 [ label = \"output_0 = Output\\l@Host(player2)\" shape = house color = \"#92cd00\"]\n",
" 0 -> 1 [ ]\n",
" 2 -> 3 [ ]\n",
" 1 -> 4 [ ]\n",
" 3 -> 4 [ ]\n",
" 4 -> 5 [ ]\n",
" 5 -> 6 [ ]\n",
"}\n",
"\n"
]
}
],
"source": [
"!elk compile dotprod.moose -p print -o dotprod.moose"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Using [GraphvizOnline](https://dreampuf.github.io/GraphvizOnline), we can render this output as png:\n",
"\n",
"![](\"./_static/dotprod.png\")
"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"The pass can be mixed and matched with other passes like `lowering`, `prune`, and `networking`, although in general compiled/networked computation graphs are often too large to be rendered this way in a useful manner. Here is the compiled version of the computation above without any networking ops:"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"%%capture\n",
"!elk compile dotprod.moose -p lowering,prune,toposort,print -o dotprod-compiled.moose"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Rendered with GraphViz Online:\n",
"\n",
"![](\"./_static/dotprod-compiled.png\")
"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"And finally, the same computation with networking ops added:"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"%%capture\n",
"!elk compile dotprod.moose -p lowering,prune,networking,toposort,print -o dotprod-networked.moose"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Rendered with GraphvizOnline:\n",
"\n",
"![](\"./_static/dotprod-networked.png\")
"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Evaluating the textual form against a Moose runtime\n",
"We can use the `dasher` binary from the Moose command line tools to locally evaluate the networked computation from above in a simulated setting:"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {"mystnb": {"merge_streams": true}},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Roles found: [\"player0\", \"player1\", \"player2\"]\n",
"Output 'op_116' ready:\n",
"Ok(HostFloat64Tensor(HostTensor([[32.0]], shape=[1, 1], strides=[1, 1], layout=CFcf (0xf), dynamic ndim=2, HostPlacement { owner: Role(\"player2\") })))\n",
"\n"
]
}
],
"source": [
"!dasher dotprod-networked.moose"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Although we often refer to as \"the Moose runtime\", Moose is actually a framework that can be used to build your own secure computation runtime. The three components of a runtime implementation include storage, networking, and choreography.\n",
"- Storage is the mechanism by which Save and Load operations are implemented.\n",
"- Networking is the mechanism by which Send and Receive operations are implemented. This determines how values get communicated from one host to another at runtime.\n",
"- Choreography is the mechanism by which computations are actually launched and executed against a set of Moose workers/executors. This covers several important aspects of execution, including:\n",
" - How each worker gets a copy of the computation\n",
" - How each worker provides input values as arguments to the computation\n",
" - How each worker reports (or doesn't report) its computation outputs\n",
"\n",
"The Moose command line tools ship with several canonical implementations of these packaged into binaries that we call \"reindeer\".\n",
"1. [`dasher`](https://github.com/tf-encrypted/moose/blob/main/moose/src/bin/dasher/main.rs) is a reindeer used for simulation, somewhat similar to PyMoose's `LocalMooseRuntime`.\n",
"2. [`rudolph`](https://github.com/tf-encrypted/moose/tree/main/moose/src/bin/rudolph) is a reindeer that uses filesystem choreography, in-memory storage, and gRPC networking.\n",
"3. [`comet`](https://github.com/tf-encrypted/moose/tree/main/moose/src/bin/comet) is similar to `rudolph` but uses a gRPC client `cometctl` for choreography. This is similar to PyMoose's `GrpcMooseRuntime`.\n",
"4. [`vixen`](https://github.com/tf-encrypted/moose/tree/main/moose/src/bin/vixen/main.rs) is a naive, single-worker implementation that doesn't implement any choreography at all. It should not be used, and is only there for legacy/educational reasons."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "moose",
"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.4"
},
"orig_nbformat": 4,
"vscode": {
"interpreter": {
"hash": "377d2daf6546a2dd25b0fbd54faa498522b36d8a6837c782d0bfb30722f17a46"
}
}
},
"nbformat": 4,
"nbformat_minor": 2
}