, _module: 'Any | None' = None, _hash: 'Optional[int]' = None, _bind_postprocess: 'None | Callable' = None, _phantom: 'bool' = False) -> None\n",
- " | Initialize self. See help(type(self)) for accurate signature.\n",
- " | \n",
- " | __repr__(self) -> 'str'\n",
- " | Return repr(self).\n",
- " | \n",
- " | __setattr__(self, name, value)\n",
- " | Implement setattr(self, name, value).\n",
- " | \n",
- " | __setstate__ = _dataclass_setstate(self, state)\n",
- " | \n",
- " | bind(self, *args, output, subsymbols=(), _call_ctx=None, **kwargs) -> 'BoundSymbol'\n",
- " | \n",
- " | name_with_module(self)\n",
- " | \n",
- " | normalize(self, *args, **kwargs)\n",
- " | \n",
- " | ----------------------------------------------------------------------\n",
- " | Readonly properties defined here:\n",
- " | \n",
- " | module\n",
- " | \n",
- " | ----------------------------------------------------------------------\n",
- " | Data descriptors defined here:\n",
- " | \n",
- " | __weakref__\n",
- " | list of weak references to the object (if defined)\n",
- " | \n",
- " | id\n",
- " | \n",
- " | is_fusion\n",
- " | \n",
- " | is_prim\n",
- " | \n",
- " | meta\n",
- " | \n",
- " | name\n",
- " | \n",
- " | python_impl\n",
- " | \n",
- " | python_printer\n",
- " | \n",
- " | ----------------------------------------------------------------------\n",
- " | Data and other attributes defined here:\n",
- " | \n",
- " | __annotations__ = {'_bind_postprocess': 'None | Callable', '_hash': 'O...\n",
- " | \n",
- " | __dataclass_fields__ = {'_bind_postprocess': Field(name='_bind_postpro...\n",
- " | \n",
- " | __dataclass_params__ = _DataclassParams(init=True,repr=True,eq=True,or...\n",
- " | \n",
- " | __match_args__ = ('name', 'meta', 'python_impl', 'id', 'is_prim', 'is_...\n",
- "\n"
- ]
+ "data": {
+ "text/plain": [
+ "[sincos_executor, sdpa]"
+ ]
+ },
+ "execution_count": 3,
+ "metadata": {},
+ "output_type": "execute_result"
}
],
"source": [
- "from thunder.core.symbol import Symbol\n",
- "\n",
- "help(Symbol)"
+ "sincos_executor = thunder.extend.OperatorExecutor(\"sincos_executor\", version='0.1')\n",
+ "thunder.add_default_executor(sincos_executor)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "4f147274",
+ "metadata": {},
+ "source": [
+ "We define meta and implementation: "
]
},
{
"cell_type": "code",
"execution_count": 4,
- "id": "ba10b306",
+ "id": "d5a72aff",
"metadata": {},
"outputs": [],
"source": [
"@log\n",
- "def sincos_meta(input):\n",
- " return (TensorProxy(like=input), TensorProxy(like=input))\n",
+ "def sincos_meta(inp):\n",
+ " return (TensorProxy(like=inp), TensorProxy(like=inp))\n",
"\n",
- "class CustomOps(Enum):\n",
- " sincos = 0\n",
- "\n",
- "sincos = Symbol(\n",
- " id=CustomOps.sincos,\n",
- " name=\"sincos\",\n",
- " meta=sincos_meta,\n",
- " is_prim=True,\n",
- ")"
+ "@log\n",
+ "def sincos_impl(inp):\n",
+ " return torch.sin(inp), torch.cos(inp)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "a06c6260",
+ "metadata": {},
+ "source": [
+ "And register it as `sincos`:\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 5,
+ "id": "03516b03",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "[Symbol name=sincos]"
+ ]
+ },
+ "execution_count": 5,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "sincos = sincos_executor.register_operator('sincos', meta=sincos_meta, fn=sincos_impl)\n",
+ "sincos"
]
},
{
@@ -224,7 +188,7 @@
},
{
"cell_type": "code",
- "execution_count": 5,
+ "execution_count": 6,
"id": "8c5da6f2",
"metadata": {},
"outputs": [],
@@ -236,13 +200,13 @@
},
{
"cell_type": "code",
- "execution_count": 6,
+ "execution_count": 7,
"id": "aef98360",
"metadata": {},
"outputs": [],
"source": [
- "a = torch.randn(1, device=\"cuda\")\n",
- "b = torch.randn(1, device=\"cuda\")"
+ "a = torch.randn(1)\n",
+ "b = torch.randn(1)"
]
},
{
@@ -255,7 +219,7 @@
},
{
"cell_type": "code",
- "execution_count": 7,
+ "execution_count": 8,
"id": "87f9f6e7",
"metadata": {},
"outputs": [
@@ -263,7 +227,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "Couldn't find an eager implementation for sincos\n"
+ "Attempting to execute outside of a tracing context, which is not supported\n"
]
}
],
@@ -284,7 +248,7 @@
},
{
"cell_type": "code",
- "execution_count": 8,
+ "execution_count": 9,
"id": "f938dff7-bac6-4807-b79d-a16cb5c6d90c",
"metadata": {},
"outputs": [
@@ -292,23 +256,25 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "call sincos_meta(TensorProxy(name=a, shape=(1,), dtype=float32, device=cuda:0))\n",
- "|<- sincos_meta = (TensorProxy(name=t0, shape=(1,), dtype=float32, device=cuda:0), TensorProxy(name=t1, shape=(1,), dtype=float32, device=cuda:0))\n",
+ "call sincos_meta(TensorProxy(name=a, shape=(1,), dtype=float32, device=cpu))\n",
+ "|<- sincos_meta = (TensorProxy(name=t0, shape=(1,), dtype=float32, device=cpu), TensorProxy(name=t1, shape=(1,), dtype=float32, device=cpu))\n",
"\n",
- "# import __main__ as __main__\n",
- "# import thunder as thunder\n",
- "# import thunder.torch as ltorch\n",
+ "# Constructed by Dead Code Elimination (took 0 milliseconds)\n",
+ "import thunder\n",
+ "import thunder.torch as ltorch\n",
"import torch\n",
+ "from thunder.executors.torchex import no_autocast\n",
"\n",
"@torch.no_grad()\n",
+ "@no_autocast()\n",
"def fun(a, b):\n",
- " # a: \"cuda:0 f32[1]\" \n",
- " # b: \"cuda:0 f32[1]\" \n",
- " (t0, t1) = __main__.sincos(a)\n",
- " t2 = ltorch.add(t0, t1, alpha=None) # t2: \"cuda:0 f32[1]\"\n",
- " # t2 = prims.add(t0, t1) # t2: \"cuda:0 f32[1]\"\n",
- " t3 = ltorch.add(t2, b, alpha=None) # t3: \"cuda:0 f32[1]\"\n",
- " # t3 = prims.add(t2, b) # t3: \"cuda:0 f32[1]\"\n",
+ " # a: \"cpu f32[1]\" \n",
+ " # b: \"cpu f32[1]\" \n",
+ " (t0, t1) = sincos(a)\n",
+ " t2 = ltorch.add(t0, t1, alpha=None) # t2: \"cpu f32[1]\"\n",
+ " # t2 = prims.add(t0, t1) # t2: \"cpu f32[1]\"\n",
+ " t3 = ltorch.add(t2, b, alpha=None) # t3: \"cpu f32[1]\"\n",
+ " # t3 = prims.add(t2, b) # t3: \"cpu f32[1]\"\n",
" return t3\n"
]
}
@@ -321,7 +287,7 @@
},
{
"cell_type": "code",
- "execution_count": 9,
+ "execution_count": 10,
"id": "6eb4818b",
"metadata": {},
"outputs": [
@@ -329,17 +295,17 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "Bound symbol with id=PrimIDs.UNPACK_TRIVIAL is represented in the trace as |# a: \"cuda:0 f32[1]\" |\n",
- "Bound symbol with id=PrimIDs.UNPACK_TRIVIAL is represented in the trace as |# b: \"cuda:0 f32[1]\" |\n",
- "Bound symbol with id=CustomOps.sincos is represented in the trace as |(t0, t1) = __main__.sincos(a)|\n",
- "Bound symbol with id=torch.add is represented in the trace as |t2 = ltorch.add(t0, t1, alpha=None) # t2: \"cuda:0 f32[1]\"\n",
- " # t2 = prims.add(t0, t1) # t2: \"cuda:0 f32[1]\"|\n",
+ "Bound symbol with id=PrimIDs.UNPACK_TRIVIAL is represented in the trace as |# a: \"cpu f32[1]\" |\n",
+ "Bound symbol with id=PrimIDs.UNPACK_TRIVIAL is represented in the trace as |# b: \"cpu f32[1]\" |\n",
+ "Bound symbol with id=sincos is represented in the trace as |(t0, t1) = sincos(a)|\n",
+ "Bound symbol with id=torch.add is represented in the trace as |t2 = ltorch.add(t0, t1, alpha=None) # t2: \"cpu f32[1]\"\n",
+ " # t2 = prims.add(t0, t1) # t2: \"cpu f32[1]\"|\n",
" It has the following subsymbols:\n",
- " id=PrimIDs.ADD |t2 = prims.add(t0, t1) # t2: \"cuda:0 f32[1]\"|\n",
- "Bound symbol with id=torch.add is represented in the trace as |t3 = ltorch.add(t2, b, alpha=None) # t3: \"cuda:0 f32[1]\"\n",
- " # t3 = prims.add(t2, b) # t3: \"cuda:0 f32[1]\"|\n",
+ " id=PrimIDs.ADD |t2 = prims.add(t0, t1) # t2: \"cpu f32[1]\"|\n",
+ "Bound symbol with id=torch.add is represented in the trace as |t3 = ltorch.add(t2, b, alpha=None) # t3: \"cpu f32[1]\"\n",
+ " # t3 = prims.add(t2, b) # t3: \"cpu f32[1]\"|\n",
" It has the following subsymbols:\n",
- " id=PrimIDs.ADD |t3 = prims.add(t2, b) # t3: \"cuda:0 f32[1]\"|\n",
+ " id=PrimIDs.ADD |t3 = prims.add(t2, b) # t3: \"cpu f32[1]\"|\n",
"Bound symbol with id=PrimIDs.RETURN is represented in the trace as |return t3|\n"
]
}
@@ -364,7 +330,7 @@
},
{
"cell_type": "code",
- "execution_count": 10,
+ "execution_count": 11,
"id": "41566de2-a60f-4c87-a3d6-58e6a89dc38b",
"metadata": {},
"outputs": [],
@@ -374,151 +340,292 @@
},
{
"cell_type": "code",
- "execution_count": 11,
- "id": "bbbb90c2",
+ "execution_count": 12,
+ "id": "24af4b99",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
- "call sincos_meta(TensorProxy(name=a, shape=(1,), dtype=float32, device=cuda:0))\n",
- "|<- sincos_meta = (TensorProxy(name=t0, shape=(1,), dtype=float32, device=cuda:0), TensorProxy(name=t1, shape=(1,), dtype=float32, device=cuda:0))\n",
+ "call sincos_meta(TensorProxy(name=t_0, shape=(1,), dtype=float32, device=cpu))\n",
+ "|<- sincos_meta = (TensorProxy(name=t0, shape=(1,), dtype=float32, device=cpu), TensorProxy(name=t1, shape=(1,), dtype=float32, device=cpu))\n",
"\n",
- "Could not find executor for bound symbol (t0, t1) = __main__.sincos(a)\n"
+ "call sincos_impl(Tensor(shape=torch.Size([1]), stride=(1,), dtype=torch.float32, device=cpu) with values tensor([0.1413]))\n",
+ "|<- sincos_impl = (Tensor(shape=torch.Size([1]), stride=(1,), dtype=torch.float32, device=cpu) with values tensor([0.1408]), Tensor(shape=torch.Size([1]), stride=(1,), dtype=torch.float32, device=cpu) with values tensor([0.9900]))\n",
+ "\n"
]
+ },
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/home/tv/firma/grid/thunder/lightning-thunder/thunder/core/jit_ext.py:478: UserWarning: We are using a (non-const) value of unknown type NoneType, which may or may not be safe. This is currently considered a sharp edge even with interpretation=INTERPRETATION_OPTIONS.TRANSLATE_PYTHON. For cases in which we are overly strict, please file an issue. Thank you!\n",
+ " warnings.warn(s)\n",
+ "/home/tv/firma/grid/thunder/lightning-thunder/thunder/core/jit_ext.py:478: UserWarning: We are using a (non-const) value of type bool, which is not identified as an input. This is currently considered a sharp edge even with interpretation=INTERPRETATION_OPTIONS.TRANSLATE_PYTHON. For cases in which we are overly strict, please file an issue. Thank you!\n",
+ " warnings.warn(s)\n",
+ "/home/tv/firma/grid/thunder/lightning-thunder/thunder/core/jit_ext.py:478: UserWarning: We are using a (non-const) value of unknown type SequenceIter, which may or may not be safe. This is currently considered a sharp edge even with interpretation=INTERPRETATION_OPTIONS.TRANSLATE_PYTHON. For cases in which we are overly strict, please file an issue. Thank you!\n",
+ " warnings.warn(s)\n",
+ "/home/tv/firma/grid/thunder/lightning-thunder/thunder/core/jit_ext.py:478: UserWarning: We are using a (non-const) value of type int, which is not identified as an input. This is currently considered a sharp edge even with interpretation=INTERPRETATION_OPTIONS.TRANSLATE_PYTHON. For cases in which we are overly strict, please file an issue. Thank you!\n",
+ " warnings.warn(s)\n",
+ "/home/tv/firma/grid/thunder/lightning-thunder/thunder/core/jit_ext.py:478: UserWarning: We are using a (non-const) value of unknown type NotImplementedType, which may or may not be safe. This is currently considered a sharp edge even with interpretation=INTERPRETATION_OPTIONS.TRANSLATE_PYTHON. For cases in which we are overly strict, please file an issue. Thank you!\n",
+ " warnings.warn(s)\n",
+ "/home/tv/firma/grid/thunder/lightning-thunder/thunder/core/jit_ext.py:478: UserWarning: We are using a (non-const) value of unknown type StopIteration, which may or may not be safe. This is currently considered a sharp edge even with interpretation=INTERPRETATION_OPTIONS.TRANSLATE_PYTHON. For cases in which we are overly strict, please file an issue. Thank you!\n",
+ " warnings.warn(s)\n"
+ ]
+ },
+ {
+ "data": {
+ "text/plain": [
+ "tensor([0.7666])"
+ ]
+ },
+ "execution_count": 12,
+ "metadata": {},
+ "output_type": "execute_result"
}
],
"source": [
- "try:\n",
- " cfun(a, b)\n",
- "except RuntimeError as e:\n",
- " print(e)"
+ "cfun(a, b)"
]
},
{
"cell_type": "markdown",
- "id": "3b1fd6e3",
+ "id": "d7cec09d",
"metadata": {},
"source": [
- "There's no registered executor for `sincos` so we need to register an executor for our new primitive. Let's do that."
+ "Let's check how our function is represented in the execution trace now (change to `thunder.last_traces(cfun)[0]` to see the trace before transformations)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "id": "a7ff30ef",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "# Constructed by Delete Last Used (took 0 milliseconds)\n",
+ "import torch\n",
+ "from thunder.executors.torchex import no_autocast\n",
+ "\n",
+ "@torch.no_grad()\n",
+ "@no_autocast()\n",
+ "def computation(a, b):\n",
+ " # a: \"cpu f32[1]\" \n",
+ " # b: \"cpu f32[1]\" \n",
+ " (res, cos) = sincos(a)\n",
+ " del a\n",
+ " result = torch.add(res, cos) # result: \"cpu f32[1]\"\n",
+ " # result = ltorch.add(res, cos, alpha=None) # result: \"cpu f32[1]\"\n",
+ " # result = prims.add(res, cos) # result: \"cpu f32[1]\"\n",
+ " del res, cos\n",
+ " t3 = torch.add(result, b) # t3: \"cpu f32[1]\"\n",
+ " # t3 = ltorch.add(result, b, alpha=None) # t3: \"cpu f32[1]\"\n",
+ " # t3 = prims.add(result, b) # t3: \"cpu f32[1]\"\n",
+ " del result, b\n",
+ " return t3"
+ ]
+ },
+ "execution_count": 13,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "thunder.last_traces(cfun)[-1]"
]
},
{
"cell_type": "markdown",
- "id": "026680b3-7b46-4f4b-b16b-641fa9bdcdf4",
+ "id": "35b71375",
"metadata": {},
"source": [
- "Check out the \"adding-operator-executor.ipynb\" notebook to see how to implement an executor for a Symbol."
+ "For a peek under the hood, we can also first create a new symbol (without reference to an executor) and then register an executor for that.\n"
]
},
{
"cell_type": "code",
- "execution_count": 12,
- "id": "2460f808-eacb-4a0f-8f62-6a17e3dce6e8",
+ "execution_count": 14,
+ "id": "f28094bb",
"metadata": {},
"outputs": [],
"source": [
- "from thunder.executors import add_operator_executor\n",
- "\n",
- "@log\n",
- "def checker_sincos(a):\n",
- " # We allow the sincos function to be called with any tensor\n",
- " return True\n",
- "\n",
+ "from thunder.core.symbol import Symbol\n",
"@log\n",
- "def executor_sincos(a):\n",
- " return torch.sin(a), torch.cos(a)\n",
+ "def sincos_meta(input):\n",
+ " return (TensorProxy(like=input), TensorProxy(like=input))\n",
"\n",
- "op_map = {\n",
- " CustomOps.sincos: (\"sincos\", checker_sincos, executor_sincos)\n",
- "}\n",
+ "# this gives a nice, unique, printable id\n",
+ "class CustomOps(Enum):\n",
+ " sincos2 = 0\n",
"\n",
- "add_operator_executor(\"sincos_executor\", op_map, add_to_default_executors=True)"
+ "sincos2 = Symbol(\n",
+ " id=CustomOps.sincos2,\n",
+ " name=\"sincos2\",\n",
+ " meta=sincos_meta,\n",
+ " is_prim=True,\n",
+ ")"
]
},
{
"cell_type": "code",
- "execution_count": 13,
- "id": "d864fa05",
+ "execution_count": 15,
+ "id": "7fbab758",
"metadata": {},
"outputs": [],
"source": [
- "# Let's try again\n",
- "cfun = thunder.compile(fun, disable_preprocessing=True)"
+ "def fun2(a, b):\n",
+ " sin, cos = sincos2(a)\n",
+ " return sin + cos + b\n",
+ "\n",
+ "cfun2 = thunder.jit(fun2)"
]
},
{
"cell_type": "code",
- "execution_count": 14,
- "id": "24af4b99",
+ "execution_count": 16,
+ "id": "950d74ad",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
- "call sincos_meta(TensorProxy(name=a, shape=(1,), dtype=float32, device=cuda:0))\n",
- "|<- sincos_meta = (TensorProxy(name=t0, shape=(1,), dtype=float32, device=cuda:0), TensorProxy(name=t1, shape=(1,), dtype=float32, device=cuda:0))\n",
+ "call sincos_meta(TensorProxy(name=t_0, shape=(1,), dtype=float32, device=cpu))\n",
+ "|<- sincos_meta = (TensorProxy(name=t0, shape=(1,), dtype=float32, device=cpu), TensorProxy(name=t1, shape=(1,), dtype=float32, device=cpu))\n",
"\n",
- "call checker_sincos(TensorProxy(name=a, shape=(1,), dtype=float32, device=cuda:0))\n",
- "|<- checker_sincos = True\n",
+ "Failed to find an executor for bound symbol bsym=(res, cos) = __main__.sincos2(a)\n"
+ ]
+ }
+ ],
+ "source": [
+ "try:\n",
+ " cfun2(a, b)\n",
+ "except RuntimeError as e:\n",
+ " print(e)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "aadcf2a9",
+ "metadata": {},
+ "source": [
+ "There's no registered executor for `sincos` so we need to register an executor for our new primitive. Let's do that."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "995febba",
+ "metadata": {},
+ "source": [
+ "Check out the \"adding-operator-executor.ipynb\" notebook to see how to implement an executor for a Symbol."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 17,
+ "id": "956a4a6b",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "@log\n",
+ "def checker_sincos2(a):\n",
+ " # We allow the sincos function to be called with any tensor\n",
+ " return True\n",
+ "\n",
+ "@log\n",
+ "def executor_sincos2(a):\n",
+ " # we need to have something here works with TensorProxies during the transformations,\n",
+ " # so we need to functions from thunder.torch or thunder.clang or other Symbols \n",
+ " return thunder.torch.sin(a), thunder.torch.cos(a)\n",
+ "\n",
+ "sincos_executor.register_implementation(sincos2, checker=checker_sincos2, execution_transform=executor_sincos2)\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 18,
+ "id": "1c77c508",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "call sincos_meta(TensorProxy(name=t_0, shape=(1,), dtype=float32, device=cpu))\n",
+ "|<- sincos_meta = (TensorProxy(name=t0, shape=(1,), dtype=float32, device=cpu), TensorProxy(name=t1, shape=(1,), dtype=float32, device=cpu))\n",
"\n",
- "call checker_sincos(TensorProxy(name=a, shape=(1,), dtype=float32, device=cuda:0))\n",
- "|<- checker_sincos = True\n",
+ "call checker_sincos2(TensorProxy(name=a, shape=(1,), dtype=float32, device=cpu))\n",
+ "|<- checker_sincos2 = True\n",
"\n",
- "call executor_sincos(Tensor(shape=torch.Size([1]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([-0.6296], device='cuda:0'))\n",
- "|<- executor_sincos = (Tensor(shape=torch.Size([1]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([-0.5889], device='cuda:0'), Tensor(shape=torch.Size([1]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([0.8082], device='cuda:0'))\n",
+ "call executor_sincos2(TensorProxy(name=a, shape=(1,), dtype=float32, device=cpu))\n",
+ "|<- executor_sincos2 = (TensorProxy(name=t4, shape=(1,), dtype=float32, device=cpu), TensorProxy(name=t5, shape=(1,), dtype=float32, device=cpu))\n",
"\n"
]
},
{
"data": {
"text/plain": [
- "tensor([0.1889], device='cuda:0')"
+ "tensor([0.7666])"
]
},
- "execution_count": 14,
+ "execution_count": 18,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
- "cfun(a, b)"
+ "# Let's try again\n",
+ "cfun2 = thunder.jit(fun2)\n",
+ "cfun2(a, b)"
]
},
{
"cell_type": "code",
- "execution_count": 15,
- "id": "a7ff30ef",
+ "execution_count": 19,
+ "id": "f9797cf2",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
- "# Constructed by Delete Last Used\n",
+ "# Constructed by Delete Last Used (took 0 milliseconds)\n",
"import torch\n",
+ "from thunder.executors.torchex import no_autocast\n",
+ "\n",
"@torch.no_grad()\n",
- "def fun(a, b):\n",
- " # a: \"cuda:0 f32[1]\" \n",
- " # b: \"cuda:0 f32[1]\" \n",
- " (t0, t1) = sincos(a)\n",
- " del [a]\n",
- " (t3,) = nvFusion0(b, t0, t1)\n",
- " # t2 = prims.add(t0, t1) # t2: \"cuda:0 f32[1]\"\n",
- " # t3 = prims.add(t2, b) # t3: \"cuda:0 f32[1]\"\n",
- " del [b, t0, t1]\n",
+ "@no_autocast()\n",
+ "def computation(a, b):\n",
+ " # a: \"cpu f32[1]\" \n",
+ " # b: \"cpu f32[1]\" \n",
+ " res = torch.sin(a) # res: \"cpu f32[1]\"\n",
+ " # res = ltorch.sin(a) # res: \"cpu f32[1]\"\n",
+ " # res = prims.sin(a) # res: \"cpu f32[1]\"\n",
+ " cos = torch.cos(a) # cos: \"cpu f32[1]\"\n",
+ " # cos = ltorch.cos(a) # cos: \"cpu f32[1]\"\n",
+ " # cos = prims.cos(a) # cos: \"cpu f32[1]\"\n",
+ " del a\n",
+ " result = torch.add(res, cos) # result: \"cpu f32[1]\"\n",
+ " # result = ltorch.add(res, cos, alpha=None) # result: \"cpu f32[1]\"\n",
+ " # result = prims.add(res, cos) # result: \"cpu f32[1]\"\n",
+ " del res, cos\n",
+ " t3 = torch.add(result, b) # t3: \"cpu f32[1]\"\n",
+ " # t3 = ltorch.add(result, b, alpha=None) # t3: \"cpu f32[1]\"\n",
+ " # t3 = prims.add(result, b) # t3: \"cpu f32[1]\"\n",
+ " del result, b\n",
" return t3"
]
},
- "execution_count": 15,
+ "execution_count": 19,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# Let's check how our function is represented in the execution trace now\n",
- "thunder.last_traces(cfun)[-1]"
+ "thunder.last_traces(cfun2)[-1]"
]
},
{
@@ -550,7 +657,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
- "version": "3.10.13"
+ "version": "3.11.8"
}
},
"nbformat": 4,
diff --git a/notebooks/adding_custom_operator_backward.ipynb b/notebooks/adding_custom_operator_backward.ipynb
index 2ed02db162..dc44b8c3d7 100644
--- a/notebooks/adding_custom_operator_backward.ipynb
+++ b/notebooks/adding_custom_operator_backward.ipynb
@@ -370,11 +370,11 @@
"\n",
"@torch.no_grad()\n",
"@no_autocast()\n",
- "def computation(a, target):\n",
- " # a: \"cuda:0 f32[2048, 50257]\" \n",
- " # target: \"cuda:0 i64[2048]\" \n",
- " (res, _) = apex_xentropy_forward(a, target, None, None, -100, None, 'none', 0.0)\n",
- " del a, target\n",
+ "def computation(logits, labels):\n",
+ " # logits: \"cuda:0 f32[2048, 50257]\" \n",
+ " # labels: \"cuda:0 i64[2048]\" \n",
+ " (res, _) = apex_xentropy_forward(logits, labels, None, None, -100, None, 'none', 0.0)\n",
+ " del logits, labels\n",
" return res"
]
},
@@ -626,188 +626,31 @@
},
{
"cell_type": "markdown",
- "id": "39fd6fce",
+ "id": "b4ec7c57",
"metadata": {},
"source": [
- "With this, we can use the `grad` transform to get the gradient:"
+ "With these registrations, we can compile a function and it will be automatically transformed into forward and backward and wrapped in a PyTorch autograd.Function calling the backward trace computed by Thunder.\n"
]
},
{
"cell_type": "code",
"execution_count": 12,
- "id": "d9f6dfde",
+ "id": "8c5da6f2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
- "call apex_cross_entropy_grad(TensorProxy(name=logits, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=labels, shape=(2048,), dtype=int64, device=cuda:0), None, None, -100, None, none, 0.0)\n",
- " call apex_xentropy_forward_meta(TensorProxy(name=logits, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=labels, shape=(2048,), dtype=int64, device=cuda:0), None, None, -100, None, none, 0.0)\n",
+ "call apex_cross_entropy_grad(TensorProxy(name=a, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=target, shape=(2048,), dtype=int64, device=cuda:0), None, None, [IntegerProxy name=ignore_index, value=-1], None, none, [FloatProxy name=label_smoothing, value=0.0])\n",
+ " call apex_xentropy_forward_meta(TensorProxy(name=a, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=target, shape=(2048,), dtype=int64, device=cuda:0), None, None, [IntegerProxy name=ignore_index, value=-1], None, none, [FloatProxy name=label_smoothing, value=0.0])\n",
" |<- apex_xentropy_forward_meta = (TensorProxy(name=t1, shape=(2048,), dtype=float32, device=cuda:0), TensorProxy(name=t0, shape=(2048,), dtype=int64, device=cuda:0))\n",
"\n",
- " call apex_xentropy_backward_meta(TensorProxy(name=t2, shape=(2048,), dtype=float32, device=cuda:0), TensorProxy(name=logits, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=labels, shape=(2048,), dtype=int64, device=cuda:0), TensorProxy(name=t0, shape=(2048,), dtype=int64, device=cuda:0), 0.0)\n",
+ " call apex_xentropy_backward_meta(TensorProxy(name=t2, shape=(2048,), dtype=float32, device=cuda:0), TensorProxy(name=a, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=target, shape=(2048,), dtype=int64, device=cuda:0), TensorProxy(name=t0, shape=(2048,), dtype=int64, device=cuda:0), [FloatProxy name=label_smoothing, value=0.0])\n",
" |<- apex_xentropy_backward_meta = TensorProxy(name=t3, shape=(2048, 50257), dtype=float32, device=cuda:0)\n",
"\n",
"|<- apex_cross_entropy_grad = TensorProxy(name=t1, shape=(2048,), dtype=float32, device=cuda:0)\n",
"\n",
- "call apex_xentropy_forward_impl(Tensor(shape=torch.Size([2048, 50257]), stride=(50257, 1), dtype=torch.float32, device=cuda:0) with values tensor([[ 0.7825, -1.1014, -0.9563, ..., 0.2801, 0.5359, -1.4094],\n",
- " [ 1.1592, 0.8128, 0.5846, ..., 1.0255, 0.4217, 0.2548],\n",
- " [ 0.8622, 0.5320, -1.5205, ..., -1.4938, -1.0423, -0.9527],\n",
- " ...,\n",
- " [-0.8978, 2.1914, 0.1603, ..., 0.0704, -0.7642, 1.4002],\n",
- " [ 0.1750, 0.6244, 1.1711, ..., 0.3491, -0.5760, -1.4034],\n",
- " [ 1.3689, -1.5422, 0.8149, ..., 0.9625, 1.0281, 1.4206]],\n",
- " device='cuda:0', requires_grad=True), Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.int64, device=cuda:0) with values tensor([43812, 33387, 31729, ..., 27740, 2907, 8268], device='cuda:0'), None, None, -100, None, none, 0.0)\n",
- "|<- apex_xentropy_forward_impl = (Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([11.8060, 10.7141, 11.4505, ..., 11.2361, 10.6558, 11.2219],\n",
- " device='cuda:0'), Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([11.3327, 11.3294, 11.3304, ..., 11.3231, 11.3170, 11.3209],\n",
- " device='cuda:0'))\n",
- "\n",
- "call apex_xentropy_backward_impl(Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([1., 1., 1., ..., 1., 1., 1.], device='cuda:0'), Tensor(shape=torch.Size([2048, 50257]), stride=(50257, 1), dtype=torch.float32, device=cuda:0) with values tensor([[ 0.7825, -1.1014, -0.9563, ..., 0.2801, 0.5359, -1.4094],\n",
- " [ 1.1592, 0.8128, 0.5846, ..., 1.0255, 0.4217, 0.2548],\n",
- " [ 0.8622, 0.5320, -1.5205, ..., -1.4938, -1.0423, -0.9527],\n",
- " ...,\n",
- " [-0.8978, 2.1914, 0.1603, ..., 0.0704, -0.7642, 1.4002],\n",
- " [ 0.1750, 0.6244, 1.1711, ..., 0.3491, -0.5760, -1.4034],\n",
- " [ 1.3689, -1.5422, 0.8149, ..., 0.9625, 1.0281, 1.4206]],\n",
- " device='cuda:0', requires_grad=True), Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.int64, device=cuda:0) with values tensor([43812, 33387, 31729, ..., 27740, 2907, 8268], device='cuda:0'), Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([11.3327, 11.3294, 11.3304, ..., 11.3231, 11.3170, 11.3209],\n",
- " device='cuda:0'), 0.0)\n",
- "|<- apex_xentropy_backward_impl = Tensor(shape=torch.Size([2048, 50257]), stride=(50257, 1), dtype=torch.float32, device=cuda:0) with values tensor([[2.6187e-05, 3.9804e-06, 4.6022e-06, ..., 1.5845e-05, 2.0466e-05,\n",
- " 2.9255e-06],\n",
- " [3.8294e-05, 2.7081e-05, 2.1557e-05, ..., 3.3501e-05, 1.8315e-05,\n",
- " 1.5500e-05],\n",
- " [2.8425e-05, 2.0432e-05, 2.6236e-06, ..., 2.6946e-06, 4.2325e-06,\n",
- " 4.6290e-06],\n",
- " ...,\n",
- " [4.9265e-06, 1.0818e-04, 1.4192e-05, ..., 1.2971e-05, 5.6303e-06,\n",
- " 4.9039e-05],\n",
- " [1.4491e-05, 2.2712e-05, 3.9235e-05, ..., 1.7247e-05, 6.8383e-06,\n",
- " 2.9895e-06],\n",
- " [4.7630e-05, 2.5919e-06, 2.7372e-05, ..., 3.1723e-05, 3.3876e-05,\n",
- " 5.0159e-05]], device='cuda:0')\n",
- "\n",
- "Difference: 1.3969838619232178e-09\n",
- "# Constructed by Delete Last Used (took 0 milliseconds)\n",
- "import torch\n",
- "from thunder.executors.torchex import no_autocast\n",
- "\n",
- "@torch.no_grad()\n",
- "@no_autocast()\n",
- "def computation(logits, labels):\n",
- " # logits: \"cuda:0 f32[2048, 50257]\" \n",
- " # labels: \"cuda:0 i64[2048]\" \n",
- " (_, t0) = apex_xentropy_forward(logits, labels, None, None, -100, None, 'none', 0.0)\n",
- " t4 = torch.full((2048,), 1.0, device=torch.device(\"cuda:0\"), dtype=torch.float32) # t4: \"cuda:0 f32[2048]\"\n",
- " # t4 = ltorch.full((2048,), 1.0, device=torch.device(\"cuda:0\"), dtype=torch.float32) # t4: \"cuda:0 f32[2048]\"\n",
- " # t4 = prims.full((2048,), 1.0, device=devices.Device(\"cuda:0\"), dtype=dtypes.float32) # t4: \"cuda:0 f32[2048]\"\n",
- " t3 = apex_xentropy_backward(t4, logits, labels, t0, 0.0) # t3: \"cuda:0 f32[2048, 50257]\"\n",
- " del t4, logits, labels, t0\n",
- " return [t3]\n"
- ]
- }
- ],
- "source": [
- "logits = torch.randn([2048, 50257], device=\"cuda\", requires_grad=True)\n",
- "labels = torch.randint(0, 50257, [2048], device=\"cuda\")\n",
- "\n",
- "grad_jfn = thunder.core.transforms.grad(jfn)\n",
- "actual_grad, = grad_jfn(logits, labels)\n",
- "\n",
- "expected_grad, = torch.autograd.grad(loss_fn(logits, labels).sum(), logits)\n",
- "\n",
- "\n",
- "print(\"Difference:\", (actual_grad - expected_grad).abs().max().item())\n",
- "print(thunder.last_traces(grad_jfn)[-1])\n",
- " \n",
- "\n"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "f9f85e3b",
- "metadata": {},
- "source": [
- "But life isn't completely simple. When we noticed that we thought about how to do backward for a long time, this is our previous approach, that is (in March 2024) needed for getting PyTorch Autograd integration.\n",
- "This works by having a _forward rule_ for generating a tuple of result and values saved for backward and a _backward rule_ that takes the saved values and output grad to compute the input grads, much like PyTorch autograd itself, but with the pluggable executor architecture of Thunder.\n",
- "\n",
- "We are working at allowing you to skip this part!"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 13,
- "id": "b0379bc4",
- "metadata": {},
- "outputs": [],
- "source": [
- "from thunder.core.transforms import register_augmented_forward_with_checker, register_backward\n",
- "\n",
- "def apex_xentropy_forward_rule(\n",
- " a,\n",
- " target,\n",
- " weight=None,\n",
- " size_average=None,\n",
- " ignore_index=-100,\n",
- " reduce=None,\n",
- " reduction=\"mean\",\n",
- " label_smoothing=0.0,\n",
- "):\n",
- " loss, max_log_sum_exp = apex_xentropy_forward(\n",
- " a,\n",
- " target,\n",
- " weight,\n",
- " size_average,\n",
- " ignore_index,\n",
- " reduce,\n",
- " reduction,\n",
- " label_smoothing,\n",
- " )\n",
- " primal = loss\n",
- " saved_for_backward = (a, target, max_log_sum_exp, reduction, label_smoothing)\n",
- " return primal, saved_for_backward\n",
- "\n",
- "register_augmented_forward_with_checker(\n",
- " apex_xentropy_ex,\n",
- " \"torch.nn.functional.cross_entropy\",\n",
- " apex_xentropy_checker,\n",
- " apex_xentropy_forward_rule,\n",
- ")\n",
- "\n",
- "@register_backward((apex_xentropy_ex, thunder.torch.cross_entropy.id))\n",
- "def apex_cross_entropy_backward_rule(\n",
- " logits, labels, max_log_sum_exp, reduction, smoothing, grad\n",
- "):\n",
- " if reduction != \"none\":\n",
- " raise ValueError(f\"Invalid reduction: {reduction}\")\n",
- "\n",
- " grad_logits = apex_xentropy_backward(\n",
- " grad,\n",
- " logits,\n",
- " labels,\n",
- " max_log_sum_exp,\n",
- " smoothing,\n",
- " )\n",
- " return grad_logits, *([None] * 7)"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "b4ec7c57",
- "metadata": {},
- "source": [
- "With these registrations, we can compile a function and it will be automatically transformed into forward and backward and wrapped in a PyTorch autograd.Function calling the backward trace computed by Thunder.\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 14,
- "id": "8c5da6f2",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
"call apex_xentropy_forward_meta(TensorProxy(name=logits, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=labels, shape=(2048,), dtype=int64, device=cuda:0), None, None, -1, None, none, 0.0)\n",
"|<- apex_xentropy_forward_meta = (TensorProxy(name=t1, shape=(2048,), dtype=float32, device=cuda:0), TensorProxy(name=t0, shape=(2048,), dtype=int64, device=cuda:0))\n",
"\n",
@@ -861,7 +704,7 @@
" -4.9872e-05, -6.3328e-05]], device='cuda:0')\n",
"\n",
"Max error in loss: 9.5367431640625e-07\n",
- "Max error in logits grad: 1.3969838619232178e-09\n"
+ "Max error in logits grad: 2.384185791015625e-07\n"
]
},
{
@@ -966,13 +809,12 @@
" return (t3, None)]"
]
},
- "execution_count": 14,
+ "execution_count": 12,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
- "from thunder.core.transforms import value_and_grad\n",
"from thunder import torch as ltorch\n",
"\n",
"torch.manual_seed(0)\n",
@@ -988,10 +830,10 @@
"actual_loss = cfn(logits, labels)\n",
"go = torch.randn_like(actual_loss)\n",
"\n",
- "actual_grads, = torch.autograd.grad(actual_loss, logits, go)\n",
+ "actual_grad, = torch.autograd.grad(actual_loss, logits, go)\n",
"\n",
"expected_loss = loss_fn(logits, labels)\n",
- "expected_grads, = torch.autograd.grad(expected_loss, logits, go)\n",
+ "expected_grad, = torch.autograd.grad(expected_loss, logits, go)\n",
"\n",
"print(\"Max error in loss:\", (actual_loss - expected_loss).abs().max().item())\n",
"print(\"Max error in logits grad:\", (actual_grad - expected_grad).abs().max().item())\n",
@@ -999,6 +841,102 @@
"thunder.last_traces(cfn)[-1]"
]
},
+ {
+ "cell_type": "markdown",
+ "id": "54d6a5ea",
+ "metadata": {},
+ "source": [
+ "Alternatively, we can also use the `grad` transform to get the gradient:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "id": "c88118eb",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "call apex_cross_entropy_grad(TensorProxy(name=logits, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=labels, shape=(2048,), dtype=int64, device=cuda:0), None, None, -100, None, none, 0.0)\n",
+ " call apex_xentropy_forward_meta(TensorProxy(name=logits, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=labels, shape=(2048,), dtype=int64, device=cuda:0), None, None, -100, None, none, 0.0)\n",
+ " |<- apex_xentropy_forward_meta = (TensorProxy(name=t1, shape=(2048,), dtype=float32, device=cuda:0), TensorProxy(name=t0, shape=(2048,), dtype=int64, device=cuda:0))\n",
+ "\n",
+ " call apex_xentropy_backward_meta(TensorProxy(name=t2, shape=(2048,), dtype=float32, device=cuda:0), TensorProxy(name=logits, shape=(2048, 50257), dtype=float32, device=cuda:0), TensorProxy(name=labels, shape=(2048,), dtype=int64, device=cuda:0), TensorProxy(name=t0, shape=(2048,), dtype=int64, device=cuda:0), 0.0)\n",
+ " |<- apex_xentropy_backward_meta = TensorProxy(name=t3, shape=(2048, 50257), dtype=float32, device=cuda:0)\n",
+ "\n",
+ "|<- apex_cross_entropy_grad = TensorProxy(name=t1, shape=(2048,), dtype=float32, device=cuda:0)\n",
+ "\n",
+ "call apex_xentropy_forward_impl(Tensor(shape=torch.Size([2048, 50257]), stride=(50257, 1), dtype=torch.float32, device=cuda:0) with values tensor([[ 0.5390, 0.1760, -1.0790, ..., 0.1695, -0.8082, -0.6984],\n",
+ " [ 2.1555, 1.3938, 0.3928, ..., 0.8937, -0.4949, 1.1610],\n",
+ " [ 0.6784, 1.1188, 0.7508, ..., -0.0941, 0.8380, 0.1878],\n",
+ " ...,\n",
+ " [-1.5834, -0.1573, -1.3511, ..., 0.6167, -0.1083, 0.4116],\n",
+ " [-0.5476, 0.5831, 0.0791, ..., -0.4986, -0.5270, 0.0954],\n",
+ " [ 0.2825, -1.0378, -0.5506, ..., 0.0149, 1.3521, -1.0823]],\n",
+ " device='cuda:0', requires_grad=True), Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.int64, device=cuda:0) with values tensor([44917, 35770, 41569, ..., 9798, 33992, 36123], device='cuda:0'), None, None, -100, None, none, 0.0)\n",
+ "|<- apex_xentropy_forward_impl = (Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([10.0233, 11.9095, 11.2898, ..., 10.9289, 10.7487, 10.7455],\n",
+ " device='cuda:0'), Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([11.3241, 11.3207, 11.3283, ..., 11.3224, 11.3186, 11.3205],\n",
+ " device='cuda:0'))\n",
+ "\n",
+ "call apex_xentropy_backward_impl(Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([1., 1., 1., ..., 1., 1., 1.], device='cuda:0'), Tensor(shape=torch.Size([2048, 50257]), stride=(50257, 1), dtype=torch.float32, device=cuda:0) with values tensor([[ 0.5390, 0.1760, -1.0790, ..., 0.1695, -0.8082, -0.6984],\n",
+ " [ 2.1555, 1.3938, 0.3928, ..., 0.8937, -0.4949, 1.1610],\n",
+ " [ 0.6784, 1.1188, 0.7508, ..., -0.0941, 0.8380, 0.1878],\n",
+ " ...,\n",
+ " [-1.5834, -0.1573, -1.3511, ..., 0.6167, -0.1083, 0.4116],\n",
+ " [-0.5476, 0.5831, 0.0791, ..., -0.4986, -0.5270, 0.0954],\n",
+ " [ 0.2825, -1.0378, -0.5506, ..., 0.0149, 1.3521, -1.0823]],\n",
+ " device='cuda:0', requires_grad=True), Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.int64, device=cuda:0) with values tensor([44917, 35770, 41569, ..., 9798, 33992, 36123], device='cuda:0'), Tensor(shape=torch.Size([2048]), stride=(1,), dtype=torch.float32, device=cuda:0) with values tensor([11.3241, 11.3207, 11.3283, ..., 11.3224, 11.3186, 11.3205],\n",
+ " device='cuda:0'), 0.0)\n",
+ "|<- apex_xentropy_backward_impl = Tensor(shape=torch.Size([2048, 50257]), stride=(50257, 1), dtype=torch.float32, device=cuda:0) with values tensor([[2.0706e-05, 1.4403e-05, 4.1058e-06, ..., 1.4309e-05, 5.3827e-06,\n",
+ " 6.0079e-06],\n",
+ " [1.0461e-04, 4.8840e-05, 1.7949e-05, ..., 2.9621e-05, 7.3879e-06,\n",
+ " 3.8697e-05],\n",
+ " [2.3705e-05, 3.6822e-05, 2.5485e-05, ..., 1.0948e-05, 2.7806e-05,\n",
+ " 1.4513e-05],\n",
+ " ...,\n",
+ " [2.4836e-06, 1.0338e-05, 3.1331e-06, ..., 2.2417e-05, 1.0857e-05,\n",
+ " 1.8259e-05],\n",
+ " [7.0235e-06, 2.1758e-05, 1.3145e-05, ..., 7.3762e-06, 7.1699e-06,\n",
+ " 1.3360e-05],\n",
+ " [1.6078e-05, 4.2941e-06, 6.9897e-06, ..., 1.2304e-05, 4.6857e-05,\n",
+ " 4.1070e-06]], device='cuda:0')\n",
+ "\n",
+ "Difference: 1.3969838619232178e-09\n",
+ "# Constructed by Delete Last Used (took 0 milliseconds)\n",
+ "import torch\n",
+ "from thunder.executors.torchex import no_autocast\n",
+ "\n",
+ "@torch.no_grad()\n",
+ "@no_autocast()\n",
+ "def computation(logits, labels):\n",
+ " # logits: \"cuda:0 f32[2048, 50257]\" \n",
+ " # labels: \"cuda:0 i64[2048]\" \n",
+ " (_, t0) = apex_xentropy_forward(logits, labels, None, None, -100, None, 'none', 0.0)\n",
+ " t4 = torch.full((2048,), 1.0, device=torch.device(\"cuda:0\"), dtype=torch.float32) # t4: \"cuda:0 f32[2048]\"\n",
+ " # t4 = ltorch.full((2048,), 1.0, device=torch.device(\"cuda:0\"), dtype=torch.float32) # t4: \"cuda:0 f32[2048]\"\n",
+ " # t4 = prims.full((2048,), 1.0, device=devices.Device(\"cuda:0\"), dtype=dtypes.float32) # t4: \"cuda:0 f32[2048]\"\n",
+ " t3 = apex_xentropy_backward(t4, logits, labels, t0, 0.0) # t3: \"cuda:0 f32[2048, 50257]\"\n",
+ " del t4, logits, labels, t0\n",
+ " return [t3]\n"
+ ]
+ }
+ ],
+ "source": [
+ "logits = torch.randn([2048, 50257], device=\"cuda\", requires_grad=True)\n",
+ "labels = torch.randint(0, 50257, [2048], device=\"cuda\")\n",
+ "\n",
+ "grad_jfn = thunder.core.transforms.grad(jfn)\n",
+ "actual_grad, = grad_jfn(logits, labels)\n",
+ "\n",
+ "expected_grad, = torch.autograd.grad(loss_fn(logits, labels).sum(), logits)\n",
+ "\n",
+ "\n",
+ "print(\"Difference:\", (actual_grad - expected_grad).abs().max().item())\n",
+ "print(thunder.last_traces(grad_jfn)[-1])\n"
+ ]
+ },
{
"cell_type": "markdown",
"id": "e234a47b",
@@ -1008,8 +946,7 @@
"\n",
"- We defined a custom executor with custom operations (Symbols in Thunder language), each with a *Meta-* (data propagation) *function* and an implementation.\n",
"- We defined and registered rules to map existing operations to our new operations. This allows us to use optimizations on our model without changing the model's code! \n",
- "- We defined a gradient rule and saw how we the `grad` transform uses it.\n",
- "- We saw another (older) way to implement forward and backward rules that is currently needed to get automatic integration with PyTorch's autograd.\n",
+ "- We defined a gradient rule and saw how our automatic PyTorch Autograd integration or the explicit `grad` transform uses it.\n",
"\n",
"Now go and implement your favourite optimized operators. We would love to hear about your use-cases!\n"
]
diff --git a/notebooks/adding_operator_executor.ipynb b/notebooks/adding_operator_executor.ipynb
deleted file mode 100644
index eac5685d3e..0000000000
--- a/notebooks/adding_operator_executor.ipynb
+++ /dev/null
@@ -1,688 +0,0 @@
-{
- "cells": [
- {
- "cell_type": "markdown",
- "id": "b6f1f42d-f146-4c9c-8ed8-74f2bcf153f0",
- "metadata": {},
- "source": [
- "# Adding an operator executor\n",
- "\n",
- "We are going to write a simple executor for `prims.add` function that calls NumPy's addition function. Our executor will be restricted to only work with inputs with certain properties. We will use the `add_operator_executor` function to create our executor."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 1,
- "id": "576d267d-9cef-4414-a722-b2cef0665cce",
- "metadata": {},
- "outputs": [],
- "source": [
- "import thunder\n",
- "import torch\n",
- "import numpy as np"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 2,
- "id": "02e16bf5",
- "metadata": {},
- "outputs": [],
- "source": [
- "#@title Helper functions (execute this cell)\n",
- "import functools\n",
- "\n",
- "_indentation = 0\n",
- "def _log(msg=None):\n",
- " \"\"\"Print a message at current indentation.\"\"\"\n",
- " if msg is not None:\n",
- " print(\" \" * _indentation + msg)\n",
- "\n",
- "def _log_indent(msg=None):\n",
- " \"\"\"Print a message and then indent the rest.\"\"\"\n",
- " global _indentation\n",
- " _log(msg)\n",
- " _indentation = 2 + _indentation\n",
- "\n",
- "def _log_unindent(msg=None):\n",
- " \"\"\"Unindent then print a message.\"\"\"\n",
- " global _indentation\n",
- " _indentation = _indentation - 2\n",
- " _log(msg)\n",
- " \n",
- "def log(func):\n",
- " \"\"\"A decorator for functions to log arguments and results.\"\"\"\n",
- " name = func.__name__\n",
- " def pp(v):\n",
- " \"\"\"Print certain values more succinctly\"\"\"\n",
- " vtype = str(type(v))\n",
- " if isinstance(v, tuple):\n",
- " return \"({})\".format(pp_values(v))\n",
- " elif isinstance(v, thunder.core.proxies.TensorProxy):\n",
- " return f\"TensorProxy(name={v.name}, shape={v.shape}, dtype={v.dtype}, device={v.device})\"\n",
- " elif isinstance(v, torch.Tensor):\n",
- " return f\"Tensor(shape={v.shape}, stride={v.stride()}, dtype={v.dtype}, device={v.device}) with values {v}\"\n",
- " else:\n",
- " return str(v)\n",
- " def pp_values(args):\n",
- " return \", \".join([pp(arg) for arg in args])\n",
- "\n",
- " @functools.wraps(func)\n",
- " def func_wrapper(*args):\n",
- " _log_indent(\"call {}({})\".format(name, pp_values(args)))\n",
- " res = func(*args)\n",
- " _log_unindent(\"|<- {} = {}\\n\".format(name, pp(res)))\n",
- " return res\n",
- "\n",
- " return func_wrapper"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 3,
- "id": "666fa494-21f5-4ed7-829e-f8648fddb13a",
- "metadata": {},
- "outputs": [],
- "source": [
- "# This is our test function\n",
- "def fun(a, b):\n",
- " return a + b * a"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 4,
- "id": "edbf395f-9549-4ae3-957c-ba34fc956b3f",
- "metadata": {},
- "outputs": [],
- "source": [
- "# This is our test input\n",
- "a = torch.randn(2, 2, device=\"cuda\")\n",
- "b = torch.randn(2, 1, device=\"cuda\")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 5,
- "id": "f938dff7-bac6-4807-b79d-a16cb5c6d90c",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "# import thunder as thunder\n",
- "# import thunder.torch as ltorch\n",
- "import torch\n",
- "\n",
- "@torch.no_grad()\n",
- "def fun(a, b):\n",
- " # a: \"cuda:0 f32[2, 2]\" \n",
- " # b: \"cuda:0 f32[2, 1]\" \n",
- " t1 = ltorch.mul(b, a) # t1: \"cuda:0 f32[2, 2]\"\n",
- " # t0 = prims.broadcast_in_dim(b, [2, 2], (0, 1)) # t0: \"cuda:0 f32[2, 2]\"\n",
- " # t1 = prims.mul(t0, a) # t1: \"cuda:0 f32[2, 2]\"\n",
- " t2 = ltorch.add(a, t1, alpha=None) # t2: \"cuda:0 f32[2, 2]\"\n",
- " # t2 = prims.add(a, t1) # t2: \"cuda:0 f32[2, 2]\"\n",
- " return t2\n"
- ]
- }
- ],
- "source": [
- "# Let's see first how this function is represented as a trace\n",
- "trace = thunder.trace()(fun, a, b)\n",
- "print(trace)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 6,
- "id": "6eb4818b",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Bound symbol with id=PrimIDs.UNPACK_TRIVIAL is represented in the trace as |# a: \"cuda:0 f32[2, 2]\" |\n",
- "Bound symbol with id=PrimIDs.UNPACK_TRIVIAL is represented in the trace as |# b: \"cuda:0 f32[2, 1]\" |\n",
- "Bound symbol with id=torch.mul is represented in the trace as |t1 = ltorch.mul(b, a) # t1: \"cuda:0 f32[2, 2]\"\n",
- " # t0 = prims.broadcast_in_dim(b, [2, 2], (0, 1)) # t0: \"cuda:0 f32[2, 2]\"\n",
- " # t1 = prims.mul(t0, a) # t1: \"cuda:0 f32[2, 2]\"|\n",
- " It has the following subsymbols:\n",
- " id=PrimIDs.BROADCAST_IN_DIM |t0 = prims.broadcast_in_dim(b, [2, 2], (0, 1)) # t0: \"cuda:0 f32[2, 2]\"|\n",
- " id=PrimIDs.MUL |t1 = prims.mul(t0, a) # t1: \"cuda:0 f32[2, 2]\"|\n",
- "Bound symbol with id=torch.add is represented in the trace as |t2 = ltorch.add(a, t1, alpha=None) # t2: \"cuda:0 f32[2, 2]\"\n",
- " # t2 = prims.add(a, t1) # t2: \"cuda:0 f32[2, 2]\"|\n",
- " It has the following subsymbols:\n",
- " id=PrimIDs.ADD |t2 = prims.add(a, t1) # t2: \"cuda:0 f32[2, 2]\"|\n",
- "Bound symbol with id=PrimIDs.RETURN is represented in the trace as |return t2|\n"
- ]
- }
- ],
- "source": [
- "# We can loop over the recorded operations that we call BoundSymbols\n",
- "for bound_symbol in trace.bound_symbols:\n",
- " print(f\"Bound symbol with id={bound_symbol.sym.id} is represented in the trace as |{bound_symbol}|\")\n",
- " if bound_symbol.subsymbols:\n",
- " print(\" It has the following subsymbols:\")\n",
- " for subsymbol in bound_symbol.subsymbols:\n",
- " print(f\" id={subsymbol.sym.id} |{subsymbol}|\")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 7,
- "id": "41566de2-a60f-4c87-a3d6-58e6a89dc38b",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Help on function add_operator_executor in module thunder.executors:\n",
- "\n",
- "add_operator_executor(name, op_map, *, add_to_default_executors: bool = True) -> None\n",
- "\n"
- ]
- }
- ],
- "source": [
- "from thunder.executors import add_operator_executor\n",
- "\n",
- "help(add_operator_executor)"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "026680b3-7b46-4f4b-b16b-641fa9bdcdf4",
- "metadata": {},
- "source": [
- "The key argument here is `op_map`.\n",
- "\n",
- "`op_map` is a dictionary with the id of the operator we're providing executor for as a key and `(name, checker_fn, implementation_fn)` tuple as a value.\n",
- "\n",
- "* `name` is the name of our execution function that would be appearing in the execution trace.\n",
- "* `checker_fn` accepts the same set of arguments as the operator itself but returns `True` or `False` to signal to the executor orchestrator whether this particular set of inputs is supported or not.\n",
- "* `implementation_fn` accepts real PyTorch tensors and expected to return PyTorch tensors."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 8,
- "id": "e02aaf0d",
- "metadata": {},
- "outputs": [],
- "source": [
- "# Let's define the addition function that can work only with NumPy's ndarrays\n",
- "\n",
- "@log\n",
- "def add_numpy(a, b):\n",
- " assert isinstance(a, np.ndarray), \"a must be a NumPy ndarray\"\n",
- " assert isinstance(b, np.ndarray), \"b must be a NumPy ndarray\"\n",
- " return np.add(a, b)\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 9,
- "id": "7ddbe12e",
- "metadata": {},
- "outputs": [],
- "source": [
- "# We also need conversion functions from PyTorch to NumPy and back\n",
- "@log\n",
- "def torch_to_numpy(tensors):\n",
- " return tuple(t.detach().cpu().numpy() for t in tensors)\n",
- "\n",
- "@log\n",
- "def numpy_to_torch(arrays, device):\n",
- " return tuple(torch.from_numpy(arr).to(device) for arr in arrays)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 10,
- "id": "2460f808-eacb-4a0f-8f62-6a17e3dce6e8",
- "metadata": {},
- "outputs": [],
- "source": [
- "@log\n",
- "def checker_add_numpy(a, b):\n",
- " # Suppose we only support float32 dtype, 2D, and (2, N) shape\n",
- " first_condition = a.dtype == b.dtype == thunder.dtypes.float32\n",
- " second_condition = a.ndim == b.ndim == 2\n",
- " third_condition = a.shape[0] == b.shape[0] == 2\n",
- " return first_condition and second_condition and third_condition"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 11,
- "id": "3a61b05f",
- "metadata": {},
- "outputs": [],
- "source": [
- "@log\n",
- "def executor_add_numpy(a, b):\n",
- " np_a, np_b = torch_to_numpy((a, b))\n",
- " np_res = add_numpy(np_a, np_b)\n",
- " res, = numpy_to_torch((np_res,), a.device)\n",
- " return res"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "c502944e",
- "metadata": {},
- "source": [
- "Now we have all the pieces to create our executor."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 12,
- "id": "e7d3eb2f",
- "metadata": {},
- "outputs": [],
- "source": [
- "op_map = {\n",
- " thunder.prims.PrimIDs.ADD: (\"add_numpy\", checker_add_numpy, executor_add_numpy)\n",
- "}"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 13,
- "id": "11f71c82",
- "metadata": {},
- "outputs": [],
- "source": [
- "# Let's send our operator map to `add_operator_executor` to register our executor under the name \"custom_add_executor\"\n",
- "\n",
- "add_operator_executor(\"custom_add_executor\", op_map, add_to_default_executors=False)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 14,
- "id": "d864fa05",
- "metadata": {},
- "outputs": [],
- "source": [
- "# Let's test our executor\n",
- "\n",
- "cfun = thunder.compile(fun, executors_list=[\"custom_add_executor\"])"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 15,
- "id": "24af4b99",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Could not find executor for bound symbol t1 = ltorch.mul(b, a) # t1: \"cuda:0 f32[2, 2]\"\n",
- " # t0 = prims.broadcast_in_dim(b, [2, 2], (0, 1)) # t0: \"cuda:0 f32[2, 2]\"\n",
- " # t1 = prims.mul(t0, a) # t1: \"cuda:0 f32[2, 2]\"\n"
- ]
- }
- ],
- "source": [
- "try:\n",
- " cfun(a, b)\n",
- "except RuntimeError as e:\n",
- " print(e)"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "d74d0c97",
- "metadata": {},
- "source": [
- "The above function errors out because we haven't provided an executor for `ltorch.mul` yet. Let's do that."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 16,
- "id": "f1ff48d0",
- "metadata": {},
- "outputs": [],
- "source": [
- "cfun = thunder.compile(fun, executors_list=[\"custom_add_executor\", thunder.executors.TORCH])"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 17,
- "id": "b1527d5e",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "call checker_add_numpy(TensorProxy(name=a, shape=(2, 2), dtype=float32, device=cuda:0), TensorProxy(name=t1, shape=(2, 2), dtype=float32, device=cuda:0))\n",
- "|<- checker_add_numpy = True\n",
- "\n",
- "call checker_add_numpy(TensorProxy(name=a, shape=(2, 2), dtype=float32, device=cuda:0), TensorProxy(name=t1, shape=(2, 2), dtype=float32, device=cuda:0))\n",
- "|<- checker_add_numpy = True\n",
- "\n",
- "call executor_add_numpy(Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-0.6906, -0.9761],\n",
- " [ 0.9819, -0.1328]], device='cuda:0'), Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-1.3271, -1.8759],\n",
- " [-0.2897, 0.0392]], device='cuda:0'))\n",
- " call torch_to_numpy((Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-0.6906, -0.9761],\n",
- " [ 0.9819, -0.1328]], device='cuda:0'), Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-1.3271, -1.8759],\n",
- " [-0.2897, 0.0392]], device='cuda:0')))\n",
- " |<- torch_to_numpy = ([[-0.6905969 -0.97613984]\n",
- " [ 0.98193294 -0.13276565]], [[-1.3271405 -1.8758768 ]\n",
- " [-0.28966585 0.03916528]])\n",
- "\n",
- " call add_numpy([[-0.6905969 -0.97613984]\n",
- " [ 0.98193294 -0.13276565]], [[-1.3271405 -1.8758768 ]\n",
- " [-0.28966585 0.03916528]])\n",
- " |<- add_numpy = [[-2.0177374 -2.8520167 ]\n",
- " [ 0.69226706 -0.09360038]]\n",
- "\n",
- " call numpy_to_torch(([[-2.0177374 -2.8520167 ]\n",
- " [ 0.69226706 -0.09360038]]), cuda:0)\n",
- " |<- numpy_to_torch = (Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-2.0177, -2.8520],\n",
- " [ 0.6923, -0.0936]], device='cuda:0'))\n",
- "\n",
- "|<- executor_add_numpy = Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-2.0177, -2.8520],\n",
- " [ 0.6923, -0.0936]], device='cuda:0')\n",
- "\n"
- ]
- },
- {
- "data": {
- "text/plain": [
- "tensor([[-2.0177, -2.8520],\n",
- " [ 0.6923, -0.0936]], device='cuda:0')"
- ]
- },
- "execution_count": 17,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "cfun(a, b)"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "c55b5ed6",
- "metadata": {},
- "source": [
- "Our logging decorator shows us that the `checker_add_numpy` function got called twice with `TensorProxy` as arguments and both times the function returned `True`. This means that our executor is going to be used for this particular execution trace.\n",
- "\n",
- "Then we see that the `executor_add_numpy` function is called with regular PyTorch tensors as arguments and it returns a regular PyTorch tensor."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 18,
- "id": "a7ff30ef",
- "metadata": {},
- "outputs": [
- {
- "data": {
- "text/plain": [
- "# Constructed by Delete Last Used\n",
- "# import torch as torch\n",
- "import torch\n",
- "\n",
- "@torch.no_grad()\n",
- "def fun(a, b):\n",
- " # a: \"cuda:0 f32[2, 2]\" \n",
- " # b: \"cuda:0 f32[2, 1]\" \n",
- " t1 = torch.mul(b, a) # t1: \"cuda:0 f32[2, 2]\"\n",
- " del [b]\n",
- " t2 = add_numpy(a, t1) # t2: \"cuda:0 f32[2, 2]\"\n",
- " del [a, t1]\n",
- " return t2"
- ]
- },
- "execution_count": 18,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "# Let's check how our function is represented in the execution trace now\n",
- "thunder.last_traces(cfun)[-1]"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 19,
- "id": "0868c882",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "call executor_add_numpy(Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-0.6906, -0.9761],\n",
- " [ 0.9819, -0.1328]], device='cuda:0'), Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-1.3271, -1.8759],\n",
- " [-0.2897, 0.0392]], device='cuda:0'))\n",
- " call torch_to_numpy((Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-0.6906, -0.9761],\n",
- " [ 0.9819, -0.1328]], device='cuda:0'), Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-1.3271, -1.8759],\n",
- " [-0.2897, 0.0392]], device='cuda:0')))\n",
- " |<- torch_to_numpy = ([[-0.6905969 -0.97613984]\n",
- " [ 0.98193294 -0.13276565]], [[-1.3271405 -1.8758768 ]\n",
- " [-0.28966585 0.03916528]])\n",
- "\n",
- " call add_numpy([[-0.6905969 -0.97613984]\n",
- " [ 0.98193294 -0.13276565]], [[-1.3271405 -1.8758768 ]\n",
- " [-0.28966585 0.03916528]])\n",
- " |<- add_numpy = [[-2.0177374 -2.8520167 ]\n",
- " [ 0.69226706 -0.09360038]]\n",
- "\n",
- " call numpy_to_torch(([[-2.0177374 -2.8520167 ]\n",
- " [ 0.69226706 -0.09360038]]), cuda:0)\n",
- " |<- numpy_to_torch = (Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-2.0177, -2.8520],\n",
- " [ 0.6923, -0.0936]], device='cuda:0'))\n",
- "\n",
- "|<- executor_add_numpy = Tensor(shape=torch.Size([2, 2]), stride=(2, 1), dtype=torch.float32, device=cuda:0) with values tensor([[-2.0177, -2.8520],\n",
- " [ 0.6923, -0.0936]], device='cuda:0')\n",
- "\n"
- ]
- }
- ],
- "source": [
- "# Let's test whether the result is correct\n",
- "cfun_torch = thunder.compile(fun, executors_list=[thunder.executors.TORCH])\n",
- "expected = cfun_torch(a, b)\n",
- "actual = cfun(a, b)\n",
- "torch.testing.assert_close(expected, actual) # Should not raise an exception"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 20,
- "id": "f978b2de",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "[SampleInput args=(tensor([[-5.6039, 5.0201, -8.2948, -0.1738],\n",
- " [ 8.4915, -2.8353, -7.4601, -4.3015],\n",
- " [ 6.0777, -7.6420, 3.4135, 3.2371],\n",
- " [-0.8413, -1.7334, -1.0025, -0.7366]], device='cuda:0'), tensor([[ 4.5391, 1.5542, 7.9208, -1.3760],\n",
- " [-6.5864, 8.6491, 6.1823, -1.8481],\n",
- " [ 7.9385, -0.4884, 4.2281, 1.3158],\n",
- " [-4.6107, 3.5805, 3.1749, -4.5989]], device='cuda:0')) kwargs={}]\n"
- ]
- }
- ],
- "source": [
- "from thunder.tests.opinfos import add_opinfo\n",
- "\n",
- "sample = next(add_opinfo.sample_input_generator(add_opinfo, device=\"cuda\", dtype=torch.float32, requires_grad=False))\n",
- "print(sample)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 21,
- "id": "f07882f2",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "call checker_add_numpy(TensorProxy(name=a, shape=(4, 4), dtype=float32, device=cuda:0), TensorProxy(name=t0, shape=(4, 4), dtype=float32, device=cuda:0))\n",
- "|<- checker_add_numpy = False\n",
- "\n"
- ]
- }
- ],
- "source": [
- "# Let's test whether the result is correct\n",
- "expected = cfun_torch(*sample.args)\n",
- "actual = cfun(*sample.args)\n",
- "torch.testing.assert_close(expected, actual) # Should not raise an exception"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 22,
- "id": "057689f5",
- "metadata": {},
- "outputs": [
- {
- "data": {
- "text/plain": [
- "# Constructed by Delete Last Used\n",
- "# import torch as torch\n",
- "import torch\n",
- "\n",
- "@torch.no_grad()\n",
- "def fun(a, b):\n",
- " # a: \"cuda:0 f32[2, 2]\" \n",
- " # b: \"cuda:0 f32[2, 1]\" \n",
- " t1 = torch.mul(b, a) # t1: \"cuda:0 f32[2, 2]\"\n",
- " del [b]\n",
- " t2 = torch.add(a, t1) # t2: \"cuda:0 f32[2, 2]\"\n",
- " del [a, t1]\n",
- " return t2"
- ]
- },
- "execution_count": 22,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "# The order of executors matters today\n",
- "cfun_torch_first = thunder.jit(fun, executors=[thunder.executors.TORCH, \"custom_add_executor\"])\n",
- "cfun_torch_first(a, b)\n",
- "thunder.last_traces(cfun_torch_first)[-1]"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 23,
- "id": "2aca9db7",
- "metadata": {},
- "outputs": [],
- "source": [
- "# Let's try inputs that are not supported by our executor\n",
- "a = torch.randn(3, 2, device=\"cuda\", dtype=torch.float64)\n",
- "b = torch.randn(3, 1, device=\"cuda\", dtype=torch.float64)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 24,
- "id": "4b3e1589",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "call checker_add_numpy(TensorProxy(name=a, shape=(3, 2), dtype=float64, device=cuda:0), TensorProxy(name=t1, shape=(3, 2), dtype=float64, device=cuda:0))\n",
- "|<- checker_add_numpy = False\n",
- "\n"
- ]
- },
- {
- "data": {
- "text/plain": [
- "# Constructed by Delete Last Used\n",
- "# import torch as torch\n",
- "import torch\n",
- "\n",
- "@torch.no_grad()\n",
- "def fun(a, b):\n",
- " # a: \"cuda:0 f64[3, 2]\" \n",
- " # b: \"cuda:0 f64[3, 1]\" \n",
- " t1 = torch.mul(b, a) # t1: \"cuda:0 f64[3, 2]\"\n",
- " del [b]\n",
- " t2 = torch.add(a, t1) # t2: \"cuda:0 f64[3, 2]\"\n",
- " del [a, t1]\n",
- " return t2"
- ]
- },
- "execution_count": 24,
- "metadata": {},
- "output_type": "execute_result"
- }
- ],
- "source": [
- "# Let's see how our function is represented in the execution trace now with the new unsupported inputs\n",
- "cfun(a, b)\n",
- "thunder.last_traces(cfun)[-1]"
- ]
- },
- {
- "cell_type": "markdown",
- "id": "122ead11",
- "metadata": {},
- "source": [
- "That's it! We've created our first executor. The process is very similar for other existing operators. There are two ingridients that are required to create an executor:\n",
- "* `checker_fn` that checks whether the executor is applicable for a particular set of inputs (works with `TensorProxy` objects),\n",
- "* `implementation_fn` that implements the operator (works with regular PyTorch tensors)."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "aec61cf6",
- "metadata": {},
- "outputs": [],
- "source": []
- }
- ],
- "metadata": {
- "kernelspec": {
- "display_name": "Python 3 (ipykernel)",
- "language": "python",
- "name": "python3"
- },
- "language_info": {
- "codemirror_mode": {
- "name": "ipython",
- "version": 3
- },
- "file_extension": ".py",
- "mimetype": "text/x-python",
- "name": "python",
- "nbconvert_exporter": "python",
- "pygments_lexer": "ipython3",
- "version": "3.10.4"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/notebooks/dev_tutorials/extend.ipynb b/notebooks/dev_tutorials/extend.ipynb
index 8b2fbe9036..304f6d9f6f 100644
--- a/notebooks/dev_tutorials/extend.ipynb
+++ b/notebooks/dev_tutorials/extend.ipynb
@@ -70,7 +70,7 @@
"source": [
"# Our operator executor will use the \"multimul\" function as a new example operator.\n",
"# This function uses NumPy to perform two multiplications of four inputs.\n",
- "# This functions very contrived, but will be useful to illustrate the extend submodule's capabilities.\n",
+ "# This function's contrived, but will be useful to illustrate the extend submodule's capabilities.\n",
"def multimul_impl(\n",
" a: Number | torch.Tensor, \n",
" b: Number | torch.Tensor,\n",
diff --git a/notebooks/dev_tutorials/fsdp_tutorial.ipynb b/notebooks/dev_tutorials/fsdp_tutorial.ipynb
new file mode 100644
index 0000000000..a4f61b47c3
--- /dev/null
+++ b/notebooks/dev_tutorials/fsdp_tutorial.ipynb
@@ -0,0 +1,2025 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## FSDP Tutorial\n",
+ "\n",
+ "In this tutorial, we will walk through the implementation of Fully Sharded Data Parallel (FSDP) with Zero2 sharding strategy in `thunder`."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Introduction\n",
+ "\n",
+ "In recent times, the LLM models have grown so large that all the model parameters don't fit on a single GPU. To circumvent this problem, there are various strategies like Tensor Parallel, Pipeline Parallel, Fully Sharded Data Parallel, etc to train these large models. In this tutorial, we discuss and implement Zero2 strategy for Fully Sharded Data Parallel (FSDP).\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### What is Zero2 strategy for FSDP?\n",
+ "\n",
+ "In this strategy, we shard the model parameters across all the availabe GPUs. That is each GPU holds onto only a chunk of the parameter. During the forward pass, all GPUs call `all_gather` communication primitive to gather the parameters from other GPUs. Unlike Zero3 strategy which frees the parameter after forward pass, we save these unsharded parameters for backward pass. This is to save the overhead of extra communication. In the backward pass, we utilize the saved parameters and compute the gradients. Once the gradients are computed, we use `reduce_scatter` communication primitive to reduce (average) the gradients across all GPUs and scatter those gradients so that a given GPU holds only a chunk of gradient.\n",
+ "\n",
+ "For more information on FSDP, we recommend reading\n",
+ "\n",
+ "1. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel - [Link](https://arxiv.org/abs/2304.11277)\n",
+ "2. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models - [Link](https://arxiv.org/abs/1910.02054)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "#### Example Model\n",
+ "\n",
+ "For this example we will have a simple model `Linear(Tanh(Linear(x)))` which will be sharded over 2 GPUs\n",
+ "\n",
+ "**NOTE**: We are generating the abstract trace so we don't actually need a system with 2 GPUs for this. It is only required when we execute this trace."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import torch\n",
+ "import torch.distributed\n",
+ "import thunder\n",
+ "import thunder.distributed\n",
+ "from IPython.display import Code"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "device='cuda'\n",
+ "dim = 64\n",
+ "def create_model():\n",
+ " layers = [torch.nn.Linear(dim, dim, bias=False),\n",
+ " torch.nn.Tanh(),\n",
+ " torch.nn.Linear(dim, dim, bias=False)]\n",
+ " return torch.nn.Sequential(*layers).to(device)\n",
+ "\n",
+ "# Model\n",
+ "model = create_model()\n",
+ "# Input\n",
+ "x = torch.randn(dim, dim, device=device)\n",
+ "\n",
+ "\n",
+ "# we want to obtain a functional version of our model. The JIT does that internally and we reach into those\n",
+ "# internals here\n",
+ "thunder_model = thunder.jit(model)\n",
+ "cache_rec, i_, _ = thunder.compile_data(thunder_model).get_computation_and_inputs(x)\n",
+ "computation_trace = cache_rec.computation_traces[0]\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def wrap_as_highlighted_code(trace):\n",
+ " return Code(str(trace), language=\"python\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "We can show the functional version:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "# Constructed by Dead Code Elimination (took 0 milliseconds)\n",
+ "import torch\n",
+ "import torch.nn.functional\n",
+ "from thunder.executors.torchex import no_autocast\n",
+ "\n",
+ "@torch.no_grad()\n",
+ "@no_autocast()\n",
+ "def augmented_forward_fn(*args):\n",
+ " # args: "Collection" \n",
+ " t0, \\\n",
+ " t1, \\\n",
+ " t2, \\\n",
+ " = args\n",
+ " t3 = torch.nn.functional.linear(t0, t1, None) # t3: "cuda:0 f32[64, 64]"\n",
+ " # t3 = ltorch.linear(t0, t1, None) # t3: "cuda:0 f32[64, 64]"\n",
+ " # t3 = prims.linear(t0, t1, None) # t3: "cuda:0 f32[64, 64]"\n",
+ " [t4] = nvFusion0(t3)\n",
+ " # t4 = prims.tanh(t3) # t4: "cuda:0 f32[64, 64]"\n",
+ " t5 = torch.nn.functional.linear(t4, t2, None) # t5: "cuda:0 f32[64, 64]"\n",
+ " # t5 = ltorch.linear(t4, t2, None) # t5: "cuda:0 f32[64, 64]"\n",
+ " # t5 = prims.linear(t4, t2, None) # t5: "cuda:0 f32[64, 64]"\n",
+ " return {'output': t5, 'flat_args': [t0, t1, t2], 'flat_output': (t5,)}, ((t0, t2, t4), ())\n",
+ "# Constructed by Dead Code Elimination (took 0 milliseconds)\n",
+ "import thunder\n",
+ "import thunder.core.prims as prims\n",
+ "import thunder.torch as ltorch\n",
+ "import torch\n",
+ "import torch.nn.functional\n",
+ "from thunder.executors.torchex import no_autocast\n",
+ "\n",
+ "@torch.no_grad()\n",
+ "@no_autocast()\n",
+ "def augmented_forward_fn(*args):\n",
+ " # args: "Collection" \n",
+ " t0, \\\n",
+ " t1, \\\n",
+ " t2, \\\n",
+ " = args\n",
+ " t3 = ltorch.linear(t0, t1, None) # t3: "cuda:0 f32[64, 64]"\n",
+ " # t3 = ltorch.linear(t0, t1, None) # t3: "cuda:0 f32[64, 64]"\n",
+ " # t3 = prims.linear(t0, t1, None) # t3: "cuda:0 f32[64, 64]"\n",
+ " t4 = prims.tanh(t3) # t4: "cuda:0 f32[64, 64]"\n",
+ " t5 = ltorch.linear(t4, t2, None) # t5: "cuda:0 f32[64, 64]"\n",
+ " # t5 = ltorch.linear(t4, t2, None) # t5: "cuda:0 f32[64, 64]"\n",
+ " # t5 = prims.linear(t4, t2, None) # t5: "cuda:0 f32[64, 64]"\n",
+ " return {'output': t5, 'flat_args': [t0, t1, t2], 'flat_output': (t5,)}, ((t0, t2, t4), ())\n",
+ "# Constructed by Dead Code Elimination (took 0 milliseconds)\n",
+ "import thunder\n",
+ "import thunder.core.prims as prims\n",
+ "import thunder.distributed.prims\n",
+ "import thunder.torch as ltorch\n",
+ "import torch\n",
+ "from thunder.executors.torchex import no_autocast\n",
+ "\n",
+ "@torch.no_grad()\n",
+ "@no_autocast()\n",
+ "def model_with_syncs(x, *params):\n",
+ " # x: "cuda:0 f32[64, 64]" \n",
+ " # params: "Collection" \n",
+ " t0, \\\n",
+ " t1, \\\n",
+ " = params\n",
+ " t2 = thunder.distributed.prims.synchronize(t0, _torch_distributed_distributed_c10d_ProcessGroup_0) # t2: "cuda:0 f32[64, 64]"\n",
+ " t3 = thunder.distributed.prims.synchronize(t1, _torch_distributed_distributed_c10d_ProcessGroup_0) # t3: "cuda:0 f32[64, 64]"\n",
+ " t4 = ltorch.linear(x, t2, None) # t4: "cuda:0 f32[64, 64]"\n",
+ " # t4 = prims.linear(x, t2, None) # t4: "cuda:0 f32[64, 64]"\n",
+ " t5 = prims.tanh(t4) # t5: "cuda:0 f32[64, 64]"\n",
+ " t6 = ltorch.linear(t5, t3, None) # t6: "cuda:0 f32[64, 64]"\n",
+ " # t6 = prims.linear(t5, t3, None) # t6: "cuda:0 f32[64, 64]"\n",
+ " return ({'output': t6, 'flat_args': [x, t2, t3], 'flat_output': (t6,)}, ((x, t3, t5), ()))\n",
+ "# Constructed by Dead Code Elimination (took 1 milliseconds)\n",
+ "import thunder\n",
+ "import thunder.core.devices as devices\n",
+ "import thunder.core.dtypes as dtypes\n",
+ "import thunder.core.prims as prims\n",
+ "import thunder.distributed.prims\n",
+ "import thunder.torch as ltorch\n",
+ "import torch\n",
+ "from thunder.executors.torchex import no_autocast\n",
+ "\n",
+ "@torch.no_grad()\n",
+ "@no_autocast()\n",
+ "def _value_and_grad(*args):\n",
+ " # args: "Collection" \n",
+ " t0, \\\n",
+ " t1, \\\n",
+ " t2, \\\n",
+ " = args\n",
+ " t3 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t3: "cuda:0 f32[64, 64]"\n",
+ " t4 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t4: "cuda:0 f32[64, 64]"\n",
+ " t5 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t5: "cuda:0 f32[64, 64]"\n",
+ " t6 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t6: "cuda:0 f32[64, 64]"\n",
+ " t7 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t7: "cuda:0 f32[64, 64]"\n",
+ " t8 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t8: "cuda:0 f32[64, 64]"\n",
+ " t9 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t9: "cuda:0 f32[64, 64]"\n",
+ " t10 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t10: "cuda:0 f32[64, 64]"\n",
+ " p11 = thunder.distributed.prims.all_gather(t1, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p11: "FUTURE cuda:0 f32[64, 64]"\n",
+ " t12 = thunder.distributed.prims.wait(p11) # t12: "cuda:0 f32[64, 64]"\n",
+ " p13 = thunder.distributed.prims.all_gather(t2, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p13: "FUTURE cuda:0 f32[64, 64]"\n",
+ " t14 = thunder.distributed.prims.wait(p13) # t14: "cuda:0 f32[64, 64]"\n",
+ " t15 = prims.linear(t0, t12, None) # t15: "cuda:0 f32[64, 64]"\n",
+ " t16 = prims.tanh(t15) # t16: "cuda:0 f32[64, 64]"\n",
+ " t17 = prims.linear(t16, t14, None) # t17: "cuda:0 f32[64, 64]"\n",
+ " t18 = prims.add(t6, t7) # t18: "cuda:0 f32[64, 64]"\n",
+ " t19 = prims.add(t3, t8) # t19: "cuda:0 f32[64, 64]"\n",
+ " t20 = prims.add(t5, t9) # t20: "cuda:0 f32[64, 64]"\n",
+ " t21 = ltorch.reshape(t18, -1, 64) # t21: "cuda:0 f32[64, 64]"\n",
+ " # t21 = prims.reshape(t18, (64, 64)) # t21: "cuda:0 f32[64, 64]"\n",
+ " t22 = ltorch.matmul(t21, t14) # t22: "cuda:0 f32[64, 64]"\n",
+ " # t22 = prims.matmul(t21, t14) # t22: "cuda:0 f32[64, 64]"\n",
+ " t23 = ltorch.reshape(t18, -1, 64) # t23: "cuda:0 f32[64, 64]"\n",
+ " # t23 = prims.reshape(t18, (64, 64)) # t23: "cuda:0 f32[64, 64]"\n",
+ " t24 = prims.transpose(t23, (1, 0)) # t24: "cuda:0 f32[64, 64]"\n",
+ " t25 = ltorch.reshape(t16, -1, 64) # t25: "cuda:0 f32[64, 64]"\n",
+ " # t25 = prims.reshape(t16, (64, 64)) # t25: "cuda:0 f32[64, 64]"\n",
+ " t26 = ltorch.matmul(t24, t25) # t26: "cuda:0 f32[64, 64]"\n",
+ " # t26 = prims.matmul(t24, t25) # t26: "cuda:0 f32[64, 64]"\n",
+ " t27 = prims.add(t10, t22) # t27: "cuda:0 f32[64, 64]"\n",
+ " t28 = prims.add(t20, t26) # t28: "cuda:0 f32[64, 64]"\n",
+ " t29 = ltorch.mul(t16, t16) # t29: "cuda:0 f32[64, 64]"\n",
+ " # t29 = prims.mul(t16, t16) # t29: "cuda:0 f32[64, 64]"\n",
+ " t30 = ltorch.sub(1, t29, alpha=None) # t30: "cuda:0 f32[64, 64]"\n",
+ " # _ = prims.convert_element_type(1, float)\n",
+ " # t30 = prims.sub(1.0, t29) # t30: "cuda:0 f32[64, 64]"\n",
+ " t31 = ltorch.mul(t27, t30) # t31: "cuda:0 f32[64, 64]"\n",
+ " # t31 = prims.mul(t27, t30) # t31: "cuda:0 f32[64, 64]"\n",
+ " t32 = ltorch.reshape(t31, -1, 64) # t32: "cuda:0 f32[64, 64]"\n",
+ " # t32 = prims.reshape(t31, (64, 64)) # t32: "cuda:0 f32[64, 64]"\n",
+ " t33 = ltorch.matmul(t32, t12) # t33: "cuda:0 f32[64, 64]"\n",
+ " # t33 = prims.matmul(t32, t12) # t33: "cuda:0 f32[64, 64]"\n",
+ " t34 = ltorch.reshape(t31, -1, 64) # t34: "cuda:0 f32[64, 64]"\n",
+ " # t34 = prims.reshape(t31, (64, 64)) # t34: "cuda:0 f32[64, 64]"\n",
+ " t35 = prims.transpose(t34, (1, 0)) # t35: "cuda:0 f32[64, 64]"\n",
+ " t36 = ltorch.reshape(t0, -1, 64) # t36: "cuda:0 f32[64, 64]"\n",
+ " # t36 = prims.reshape(t0, (64, 64)) # t36: "cuda:0 f32[64, 64]"\n",
+ " t37 = ltorch.matmul(t35, t36) # t37: "cuda:0 f32[64, 64]"\n",
+ " # t37 = prims.matmul(t35, t36) # t37: "cuda:0 f32[64, 64]"\n",
+ " t38 = prims.add(t19, t33) # t38: "cuda:0 f32[64, 64]"\n",
+ " t39 = prims.add(t4, t37) # t39: "cuda:0 f32[64, 64]"\n",
+ " t40 = ltorch.true_divide(t28, 2) # t40: "cuda:0 f32[64, 64]"\n",
+ " # _ = prims.convert_element_type(2, float)\n",
+ " # t40 = prims.div(t28, 2.0) # t40: "cuda:0 f32[64, 64]"\n",
+ " p41 = thunder.distributed.prims.reduce_scatter(t40, _DistributedReduceOps_1, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p41: "FUTURE cuda:0 f32[32, 64]"\n",
+ " t42 = thunder.distributed.prims.wait(p41) # t42: "cuda:0 f32[32, 64]"\n",
+ " t43 = ltorch.true_divide(t39, 2) # t43: "cuda:0 f32[64, 64]"\n",
+ " # _ = prims.convert_element_type(2, float)\n",
+ " # t43 = prims.div(t39, 2.0) # t43: "cuda:0 f32[64, 64]"\n",
+ " p44 = thunder.distributed.prims.reduce_scatter(t43, _DistributedReduceOps_1, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p44: "FUTURE cuda:0 f32[32, 64]"\n",
+ " t45 = thunder.distributed.prims.wait(p44) # t45: "cuda:0 f32[32, 64]"\n",
+ " return (({'output': t17, 'flat_args': [t0, t12, t14], 'flat_output': (t17,)}, ((t0, t14, t16), ())), (t38, t45, t42))\n",
+ "# Constructed by Delete Last Used (took 0 milliseconds)\n",
+ "import torch\n",
+ "import torch.nn.functional\n",
+ "from thunder.executors.torchex import no_autocast\n",
+ "\n",
+ "@torch.no_grad()\n",
+ "@no_autocast()\n",
+ "def _value_and_grad(*args):\n",
+ " # args: "Collection" \n",
+ " t0, \\\n",
+ " t1, \\\n",
+ " t2, \\\n",
+ " = args\n",
+ " del args\n",
+ " t3 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t3: "cuda:0 f32[64, 64]"\n",
+ " # t3 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t3: "cuda:0 f32[64, 64]"\n",
+ " # t3 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t3: "cuda:0 f32[64, 64]"\n",
+ " t4 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t4: "cuda:0 f32[64, 64]"\n",
+ " # t4 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t4: "cuda:0 f32[64, 64]"\n",
+ " # t4 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t4: "cuda:0 f32[64, 64]"\n",
+ " t5 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t5: "cuda:0 f32[64, 64]"\n",
+ " # t5 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t5: "cuda:0 f32[64, 64]"\n",
+ " # t5 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t5: "cuda:0 f32[64, 64]"\n",
+ " t6 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t6: "cuda:0 f32[64, 64]"\n",
+ " # t6 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t6: "cuda:0 f32[64, 64]"\n",
+ " # t6 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t6: "cuda:0 f32[64, 64]"\n",
+ " t7 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t7: "cuda:0 f32[64, 64]"\n",
+ " # t7 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t7: "cuda:0 f32[64, 64]"\n",
+ " # t7 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t7: "cuda:0 f32[64, 64]"\n",
+ " t8 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t8: "cuda:0 f32[64, 64]"\n",
+ " # t8 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t8: "cuda:0 f32[64, 64]"\n",
+ " # t8 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t8: "cuda:0 f32[64, 64]"\n",
+ " t9 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t9: "cuda:0 f32[64, 64]"\n",
+ " # t9 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t9: "cuda:0 f32[64, 64]"\n",
+ " # t9 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t9: "cuda:0 f32[64, 64]"\n",
+ " t10 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t10: "cuda:0 f32[64, 64]"\n",
+ " # t10 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t10: "cuda:0 f32[64, 64]"\n",
+ " # t10 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t10: "cuda:0 f32[64, 64]"\n",
+ " p11 = torch_all_gather_prim_impl(t1, _torch_distributed_distributed_c10d_ProcessGroup_2, True) # p11: "FUTURE cuda:0 f32[64, 64]"\n",
+ " del t1\n",
+ " t12 = torch_wait_prim_impl(p11) # t12: "cuda:0 f32[64, 64]"\n",
+ " del p11\n",
+ " p13 = torch_all_gather_prim_impl(t2, _torch_distributed_distributed_c10d_ProcessGroup_2, True) # p13: "FUTURE cuda:0 f32[64, 64]"\n",
+ " del t2\n",
+ " t14 = torch_wait_prim_impl(p13) # t14: "cuda:0 f32[64, 64]"\n",
+ " del p13\n",
+ " t15 = torch.nn.functional.linear(t0, t12, None) # t15: "cuda:0 f32[64, 64]"\n",
+ " # t15 = ltorch.linear(t0, t12, None) # t15: "cuda:0 f32[64, 64]"\n",
+ " # t15 = prims.linear(t0, t12, None) # t15: "cuda:0 f32[64, 64]"\n",
+ " t16 = torch.tanh(t15) # t16: "cuda:0 f32[64, 64]"\n",
+ " # t16 = ltorch.tanh(t15) # t16: "cuda:0 f32[64, 64]"\n",
+ " # t16 = prims.tanh(t15) # t16: "cuda:0 f32[64, 64]"\n",
+ " del t15\n",
+ " t17 = torch.nn.functional.linear(t16, t14, None) # t17: "cuda:0 f32[64, 64]"\n",
+ " # t17 = ltorch.linear(t16, t14, None) # t17: "cuda:0 f32[64, 64]"\n",
+ " # t17 = prims.linear(t16, t14, None) # t17: "cuda:0 f32[64, 64]"\n",
+ " t18 = torch.add(t6, t7) # t18: "cuda:0 f32[64, 64]"\n",
+ " # t18 = ltorch.add(t6, t7, alpha=None) # t18: "cuda:0 f32[64, 64]"\n",
+ " # t18 = prims.add(t6, t7) # t18: "cuda:0 f32[64, 64]"\n",
+ " del t6, t7\n",
+ " t19 = torch.add(t3, t8) # t19: "cuda:0 f32[64, 64]"\n",
+ " # t19 = ltorch.add(t3, t8, alpha=None) # t19: "cuda:0 f32[64, 64]"\n",
+ " # t19 = prims.add(t3, t8) # t19: "cuda:0 f32[64, 64]"\n",
+ " del t3, t8\n",
+ " t20 = torch.add(t5, t9) # t20: "cuda:0 f32[64, 64]"\n",
+ " # t20 = ltorch.add(t5, t9, alpha=None) # t20: "cuda:0 f32[64, 64]"\n",
+ " # t20 = prims.add(t5, t9) # t20: "cuda:0 f32[64, 64]"\n",
+ " del t5, t9\n",
+ " t21 = torch.reshape(t18, (-1, 64)) # t21: "cuda:0 f32[64, 64]"\n",
+ " # t21 = ltorch.reshape(t18, (-1, 64)) # t21: "cuda:0 f32[64, 64]"\n",
+ " # t21 = prims.reshape(t18, (64, 64)) # t21: "cuda:0 f32[64, 64]"\n",
+ " t22 = torch.matmul(t21, t14) # t22: "cuda:0 f32[64, 64]"\n",
+ " # t22 = ltorch.matmul(t21, t14) # t22: "cuda:0 f32[64, 64]"\n",
+ " # t22 = prims.matmul(t21, t14) # t22: "cuda:0 f32[64, 64]"\n",
+ " del t21\n",
+ " t23 = torch.reshape(t18, (-1, 64)) # t23: "cuda:0 f32[64, 64]"\n",
+ " # t23 = ltorch.reshape(t18, (-1, 64)) # t23: "cuda:0 f32[64, 64]"\n",
+ " # t23 = prims.reshape(t18, (64, 64)) # t23: "cuda:0 f32[64, 64]"\n",
+ " del t18\n",
+ " t24 = torch.permute(t23, (1, 0)) # t24: "cuda:0 f32[64, 64]"\n",
+ " # t24 = ltorch.permute(t23, (1, 0)) # t24: "cuda:0 f32[64, 64]"\n",
+ " # t24 = prims.transpose(t23, (1, 0)) # t24: "cuda:0 f32[64, 64]"\n",
+ " del t23\n",
+ " t25 = torch.reshape(t16, (-1, 64)) # t25: "cuda:0 f32[64, 64]"\n",
+ " # t25 = ltorch.reshape(t16, (-1, 64)) # t25: "cuda:0 f32[64, 64]"\n",
+ " # t25 = prims.reshape(t16, (64, 64)) # t25: "cuda:0 f32[64, 64]"\n",
+ " t26 = torch.matmul(t24, t25) # t26: "cuda:0 f32[64, 64]"\n",
+ " # t26 = ltorch.matmul(t24, t25) # t26: "cuda:0 f32[64, 64]"\n",
+ " # t26 = prims.matmul(t24, t25) # t26: "cuda:0 f32[64, 64]"\n",
+ " del t24, t25\n",
+ " t27 = torch.add(t10, t22) # t27: "cuda:0 f32[64, 64]"\n",
+ " # t27 = ltorch.add(t10, t22, alpha=None) # t27: "cuda:0 f32[64, 64]"\n",
+ " # t27 = prims.add(t10, t22) # t27: "cuda:0 f32[64, 64]"\n",
+ " del t10, t22\n",
+ " t28 = torch.add(t20, t26) # t28: "cuda:0 f32[64, 64]"\n",
+ " # t28 = ltorch.add(t20, t26, alpha=None) # t28: "cuda:0 f32[64, 64]"\n",
+ " # t28 = prims.add(t20, t26) # t28: "cuda:0 f32[64, 64]"\n",
+ " del t20, t26\n",
+ " t29 = torch.mul(t16, t16) # t29: "cuda:0 f32[64, 64]"\n",
+ " # t29 = ltorch.mul(t16, t16) # t29: "cuda:0 f32[64, 64]"\n",
+ " # t29 = prims.mul(t16, t16) # t29: "cuda:0 f32[64, 64]"\n",
+ " t30 = torch.sub(1, t29) # t30: "cuda:0 f32[64, 64]"\n",
+ " # t30 = ltorch.sub(1, t29, alpha=None) # t30: "cuda:0 f32[64, 64]"\n",
+ " # _ = prims.convert_element_type(1, float)\n",
+ " # t30 = prims.sub(1.0, t29) # t30: "cuda:0 f32[64, 64]"\n",
+ " del t29\n",
+ " t31 = torch.mul(t27, t30) # t31: "cuda:0 f32[64, 64]"\n",
+ " # t31 = ltorch.mul(t27, t30) # t31: "cuda:0 f32[64, 64]"\n",
+ " # t31 = prims.mul(t27, t30) # t31: "cuda:0 f32[64, 64]"\n",
+ " del t27, t30\n",
+ " t32 = torch.reshape(t31, (-1, 64)) # t32: "cuda:0 f32[64, 64]"\n",
+ " # t32 = ltorch.reshape(t31, (-1, 64)) # t32: "cuda:0 f32[64, 64]"\n",
+ " # t32 = prims.reshape(t31, (64, 64)) # t32: "cuda:0 f32[64, 64]"\n",
+ " t33 = torch.matmul(t32, t12) # t33: "cuda:0 f32[64, 64]"\n",
+ " # t33 = ltorch.matmul(t32, t12) # t33: "cuda:0 f32[64, 64]"\n",
+ " # t33 = prims.matmul(t32, t12) # t33: "cuda:0 f32[64, 64]"\n",
+ " del t32\n",
+ " t34 = torch.reshape(t31, (-1, 64)) # t34: "cuda:0 f32[64, 64]"\n",
+ " # t34 = ltorch.reshape(t31, (-1, 64)) # t34: "cuda:0 f32[64, 64]"\n",
+ " # t34 = prims.reshape(t31, (64, 64)) # t34: "cuda:0 f32[64, 64]"\n",
+ " del t31\n",
+ " t35 = torch.permute(t34, (1, 0)) # t35: "cuda:0 f32[64, 64]"\n",
+ " # t35 = ltorch.permute(t34, (1, 0)) # t35: "cuda:0 f32[64, 64]"\n",
+ " # t35 = prims.transpose(t34, (1, 0)) # t35: "cuda:0 f32[64, 64]"\n",
+ " del t34\n",
+ " t36 = torch.reshape(t0, (-1, 64)) # t36: "cuda:0 f32[64, 64]"\n",
+ " # t36 = ltorch.reshape(t0, (-1, 64)) # t36: "cuda:0 f32[64, 64]"\n",
+ " # t36 = prims.reshape(t0, (64, 64)) # t36: "cuda:0 f32[64, 64]"\n",
+ " t37 = torch.matmul(t35, t36) # t37: "cuda:0 f32[64, 64]"\n",
+ " # t37 = ltorch.matmul(t35, t36) # t37: "cuda:0 f32[64, 64]"\n",
+ " # t37 = prims.matmul(t35, t36) # t37: "cuda:0 f32[64, 64]"\n",
+ " del t35, t36\n",
+ " t38 = torch.add(t19, t33) # t38: "cuda:0 f32[64, 64]"\n",
+ " # t38 = ltorch.add(t19, t33, alpha=None) # t38: "cuda:0 f32[64, 64]"\n",
+ " # t38 = prims.add(t19, t33) # t38: "cuda:0 f32[64, 64]"\n",
+ " del t19, t33\n",
+ " t39 = torch.add(t4, t37) # t39: "cuda:0 f32[64, 64]"\n",
+ " # t39 = ltorch.add(t4, t37, alpha=None) # t39: "cuda:0 f32[64, 64]"\n",
+ " # t39 = prims.add(t4, t37) # t39: "cuda:0 f32[64, 64]"\n",
+ " del t4, t37\n",
+ " t40 = torch.true_divide(t28, 2) # t40: "cuda:0 f32[64, 64]"\n",
+ " # t40 = ltorch.true_divide(t28, 2) # t40: "cuda:0 f32[64, 64]"\n",
+ " # _ = prims.convert_element_type(2, float)\n",
+ " # t40 = prims.div(t28, 2.0) # t40: "cuda:0 f32[64, 64]"\n",
+ " del t28\n",
+ " p41 = torch_reduce_scatter_prim_impl(t40, _DistributedReduceOps_3, _torch_distributed_distributed_c10d_ProcessGroup_2, True) # p41: "FUTURE cuda:0 f32[32, 64]"\n",
+ " del t40\n",
+ " t42 = torch_wait_prim_impl(p41) # t42: "cuda:0 f32[32, 64]"\n",
+ " del p41\n",
+ " t43 = torch.true_divide(t39, 2) # t43: "cuda:0 f32[64, 64]"\n",
+ " # t43 = ltorch.true_divide(t39, 2) # t43: "cuda:0 f32[64, 64]"\n",
+ " # _ = prims.convert_element_type(2, float)\n",
+ " # t43 = prims.div(t39, 2.0) # t43: "cuda:0 f32[64, 64]"\n",
+ " del t39\n",
+ " p44 = torch_reduce_scatter_prim_impl(t43, _DistributedReduceOps_3, _torch_distributed_distributed_c10d_ProcessGroup_2, True) # p44: "FUTURE cuda:0 f32[32, 64]"\n",
+ " del t43\n",
+ " t45 = torch_wait_prim_impl(p44) # t45: "cuda:0 f32[32, 64]"\n",
+ " del p44\n",
+ " return (({'output': t17, 'flat_args': [t0, t12, t14], 'flat_output': (t17,)}, ((t0, t14, t16), ())), (t38, t45, t42))\n",
+ "# Constructed by Dead Code Elimination (took 0 milliseconds)\n",
- "import thunder\n",
- "import thunder.distributed.prims\n",
- "import thunder.torch as ltorch\n",
- "import torch\n",
- "from thunder.executors.torchex import no_autocast\n",
- "\n",
- "@torch.no_grad()\n",
- "@no_autocast()\n",
- "def model_with_syncs(*params, x):\n",
- " # params \n",
- " # x \n",
- " t0, \\\n",
- " t1, \\\n",
- " = params\n",
- " t2 = thunder.distributed.prims.synchronize(t0, _torch_distributed_distributed_c10d_ProcessGroup_0) # t2\n",
- " t3 = thunder.distributed.prims.synchronize(t1, _torch_distributed_distributed_c10d_ProcessGroup_0) # t3\n",
- " t4 = ltorch.linear(x, t2, None) # t4\n",
- " # t4 = prims.linear(x, t2, None) # t4\n",
- " t5 = ltorch.tanh(t4) # t5\n",
- " # t5 = prims.tanh(t4) # t5\n",
- " t6 = ltorch.linear(t5, t3, None) # t6\n",
- " # t6 = prims.linear(t5, t3, None) # t6\n",
- " return t6\n",
- "
# Constructed by Dead Code Elimination (took 0 milliseconds)\n",
- "import thunder\n",
- "import thunder.core.devices as devices\n",
- "import thunder.core.dtypes as dtypes\n",
- "import thunder.core.prims as prims\n",
- "import thunder.distributed.prims\n",
- "import thunder.torch as ltorch\n",
- "import torch\n",
- "from thunder.executors.torchex import no_autocast\n",
- "\n",
- "@torch.no_grad()\n",
- "@no_autocast()\n",
- "def _value_and_grad(*args, **kwargs):\n",
- " # args \n",
- " # kwargs \n",
- " t0, \\\n",
- " t1, \\\n",
- " = args\n",
- " t2 = kwargs['x']\n",
- " t3 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t3\n",
- " p4 = thunder.distributed.prims.all_gather(t0, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p4\n",
- " t5 = thunder.distributed.prims.wait(p4) # t5\n",
- " p6 = thunder.distributed.prims.all_gather(t1, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p6\n",
- " t7 = thunder.distributed.prims.wait(p6) # t7\n",
- " t8 = prims.linear(t2, t5, None) # t8\n",
- " t9 = prims.tanh(t8) # t9\n",
- " t10 = prims.linear(t9, t7, None) # t10\n",
- " t11 = ltorch.reshape(t3, -1, 64) # t11\n",
- " # t11 = prims.reshape(t3, (64, 64)) # t11\n",
- " t12 = ltorch.matmul(t11, t7) # t12\n",
- " # t12 = prims.matmul(t11, t7) # t12\n",
- " t13 = ltorch.reshape(t3, -1, 64) # t13\n",
- " # t13 = prims.reshape(t3, (64, 64)) # t13\n",
- " t14 = prims.transpose(t13, (1, 0)) # t14\n",
- " t15 = ltorch.reshape(t9, -1, 64) # t15\n",
- " # t15 = prims.reshape(t9, (64, 64)) # t15\n",
- " t16 = ltorch.matmul(t14, t15) # t16\n",
- " # t16 = prims.matmul(t14, t15) # t16\n",
- " t17 = ltorch.mul(t9, t9) # t17\n",
- " # t17 = prims.mul(t9, t9) # t17\n",
- " t18 = ltorch.sub(1.0, t17, alpha=None) # t18\n",
- " # t18 = prims.sub(1.0, t17) # t18\n",
- " t19 = ltorch.mul(t12, t18) # t19\n",
- " # t19 = prims.mul(t12, t18) # t19\n",
- " t20 = ltorch.reshape(t19, -1, 64) # t20\n",
- " # t20 = prims.reshape(t19, (64, 64)) # t20\n",
- " t21 = ltorch.matmul(t20, t5) # t21\n",
- " # t21 = prims.matmul(t20, t5) # t21\n",
- " t22 = ltorch.reshape(t19, -1, 64) # t22\n",
- " # t22 = prims.reshape(t19, (64, 64)) # t22\n",
- " t23 = prims.transpose(t22, (1, 0)) # t23\n",
- " t24 = ltorch.reshape(t2, -1, 64) # t24\n",
- " # t24 = prims.reshape(t2, (64, 64)) # t24\n",
- " t25 = ltorch.matmul(t23, t24) # t25\n",
- " # t25 = prims.matmul(t23, t24) # t25\n",
- " t26 = ltorch.true_divide(t16, 2) # t26\n",
- " # _ = prims.convert_element_type(2, float)\n",
- " # t26 = prims.div(t16, 2.0) # t26\n",
- " p27 = thunder.distributed.prims.reduce_scatter(t26, _DistributedReduceOps_1, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p27\n",
- " t28 = thunder.distributed.prims.wait(p27) # t28\n",
- " t29 = ltorch.true_divide(t25, 2) # t29\n",
- " # _ = prims.convert_element_type(2, float)\n",
- " # t29 = prims.div(t25, 2.0) # t29\n",
- " p30 = thunder.distributed.prims.reduce_scatter(t29, _DistributedReduceOps_1, _torch_distributed_distributed_c10d_ProcessGroup_0, True) # p30\n",
- " t31 = thunder.distributed.prims.wait(p30) # t31\n",
- " return (t10, (t31, t28, {'x': t21}))\n",
- "# Constructed by Delete Last Used (took 0 milliseconds)\n",
- "import torch\n",
- "import torch.nn.functional\n",
- "from thunder.executors.torchex import no_autocast\n",
- "\n",
- "@torch.no_grad()\n",
- "@no_autocast()\n",
- "def _value_and_grad(*args, **kwargs):\n",
- " # args \n",
- " # kwargs \n",
- " t0, \\\n",
- " t1, \\\n",
- " = args\n",
- " del args\n",
- " t2 = kwargs['x']\n",
- " del kwargs\n",
- " t3 = torch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t3\n",
- " # t3 = ltorch.full((64, 64), 1, device=torch.device("cuda:0"), dtype=torch.float32) # t3\n",
- " # t3 = prims.full((64, 64), 1, device=devices.Device("cuda:0"), dtype=dtypes.float32) # t3\n",
- " p4 = torch_all_gather_prim_impl(t0, _torch_distributed_distributed_c10d_ProcessGroup_2, True) # p4\n",
- " del t0\n",
- " t5 = torch_wait_prim_impl(p4) # t5\n",
- " del p4\n",
- " p6 = torch_all_gather_prim_impl(t1, _torch_distributed_distributed_c10d_ProcessGroup_2, True) # p6\n",
- " del t1\n",
- " t7 = torch_wait_prim_impl(p6) # t7\n",
- " del p6\n",
- " t8 = torch.nn.functional.linear(t2, t5, None) # t8\n",
- " # t8 = ltorch.linear(t2, t5, None) # t8\n",
- " # t8 = prims.linear(t2, t5, None) # t8\n",
- " t9 = torch.tanh(t8) # t9\n",
- " # t9 = ltorch.tanh(t8) # t9\n",
- " # t9 = prims.tanh(t8) # t9\n",
- " del t8\n",
- " t10 = torch.nn.functional.linear(t9, t7, None) # t10\n",
- " # t10 = ltorch.linear(t9, t7, None) # t10\n",
- " # t10 = prims.linear(t9, t7, None) # t10\n",
- " t11 = torch.reshape(t3, (-1, 64)) # t11\n",
- " # t11 = ltorch.reshape(t3, (-1, 64)) # t11\n",
- " # t11 = prims.reshape(t3, (64, 64)) # t11\n",
- " t12 = torch.matmul(t11, t7) # t12\n",
- " # t12 = ltorch.matmul(t11, t7) # t12\n",
- " # t12 = prims.matmul(t11, t7) # t12\n",
- " del t11, t7\n",
- " t13 = torch.reshape(t3, (-1, 64)) # t13\n",
- " # t13 = ltorch.reshape(t3, (-1, 64)) # t13\n",
- " # t13 = prims.reshape(t3, (64, 64)) # t13\n",
- " del t3\n",
- " t14 = torch.permute(t13, (1, 0)) # t14\n",
- " # t14 = ltorch.permute(t13, (1, 0)) # t14\n",
- " # t14 = prims.transpose(t13, (1, 0)) # t14\n",
- " del t13\n",
- " t15 = torch.reshape(t9, (-1, 64)) # t15\n",
- " # t15 = ltorch.reshape(t9, (-1, 64)) # t15\n",
- " # t15 = prims.reshape(t9, (64, 64)) # t15\n",
- " t16 = torch.matmul(t14, t15) # t16\n",
- " # t16 = ltorch.matmul(t14, t15) # t16\n",
- " # t16 = prims.matmul(t14, t15) # t16\n",
- " del t14, t15\n",
- " t17 = torch.mul(t9, t9) # t17\n",
- " # t17 = ltorch.mul(t9, t9) # t17\n",
- " # t17 = prims.mul(t9, t9) # t17\n",
- " del t9\n",
- " t18 = torch.sub(1.0, t17) # t18\n",
- " # t18 = ltorch.sub(1.0, t17, alpha=None) # t18\n",
- " # t18 = prims.sub(1.0, t17) # t18\n",
- " del t17\n",
- " t19 = torch.mul(t12, t18) # t19\n",
- " # t19 = ltorch.mul(t12, t18) # t19\n",
- " # t19 = prims.mul(t12, t18) # t19\n",
- " del t12, t18\n",
- " t20 = torch.reshape(t19, (-1, 64)) # t20\n",
- " # t20 = ltorch.reshape(t19, (-1, 64)) # t20\n",
- " # t20 = prims.reshape(t19, (64, 64)) # t20\n",
- " t21 = torch.matmul(t20, t5) # t21\n",
- " # t21 = ltorch.matmul(t20, t5) # t21\n",
- " # t21 = prims.matmul(t20, t5) # t21\n",
- " del t20, t5\n",
- " t22 = torch.reshape(t19, (-1, 64)) # t22\n",
- " # t22 = ltorch.reshape(t19, (-1, 64)) # t22\n",
- " # t22 = prims.reshape(t19, (64, 64)) # t22\n",
- " del t19\n",
- " t23 = torch.permute(t22, (1, 0)) # t23\n",
- " # t23 = ltorch.permute(t22, (1, 0)) # t23\n",
- " # t23 = prims.transpose(t22, (1, 0)) # t23\n",
- " del t22\n",
- " t24 = torch.reshape(t2, (-1, 64)) # t24\n",
- " # t24 = ltorch.reshape(t2, (-1, 64)) # t24\n",
- " # t24 = prims.reshape(t2, (64, 64)) # t24\n",
- " del t2\n",
- " t25 = torch.matmul(t23, t24) # t25\n",
- " # t25 = ltorch.matmul(t23, t24) # t25\n",
- " # t25 = prims.matmul(t23, t24) # t25\n",
- " del t23, t24\n",
- " t26 = torch.true_divide(t16, 2) # t26\n",
- " # t26 = ltorch.true_divide(t16, 2) # t26\n",
- " # _ = prims.convert_element_type(2, float)\n",
- " # t26 = prims.div(t16, 2.0) # t26\n",
- " del t16\n",
- " p27 = torch_reduce_scatter_prim_impl(t26, _DistributedReduceOps_3, _torch_distributed_distributed_c10d_ProcessGroup_2, True) # p27\n",
- " del t26\n",
- " t28 = torch_wait_prim_impl(p27) # t28\n",
- " del p27\n",
- " t29 = torch.true_divide(t25, 2) # t29\n",
- " # t29 = ltorch.true_divide(t25, 2) # t29\n",
- " # _ = prims.convert_element_type(2, float)\n",
- " # t29 = prims.div(t25, 2.0) # t29\n",
- " del t25\n",
- " p30 = torch_reduce_scatter_prim_impl(t29, _DistributedReduceOps_3, _torch_distributed_distributed_c10d_ProcessGroup_2, True) # p30\n",
- " del t29\n",
- " t31 = torch_wait_prim_impl(p30) # t31\n",
- " del p30\n",
- " return (t10, (t31, t28, {'x': t21}))\n",
- " \n",
- "\n",
- "\n",
- " 2 0 LOAD_FAST 2 (layer_0) || 0) LOAD_FAST\n",
- " 2 LOAD_FAST 0 (x) || LOAD_FAST\n",
- " 4 LOAD_FAST 1 (mask) || LOAD_FAST\n",
- " 6 CALL_FUNCTION 2 || CALL_FUNCTION\n",
- " 8 STORE_FAST 0 (x) || STORE_FAST\n",
- " || LOAD_FAST\n",
- " 3 10 LOAD_FAST 3 (layer_1) || LOAD_FAST\n",
- " 12 LOAD_FAST 0 (x) || LOAD_FAST\n",
- " 14 LOAD_FAST 1 (mask) || LOAD_CONST\n",
- " 16 LOAD_CONST 0 (None) || IS_OP\n",
- " 18 IS_OP 1 || POP_JUMP_IF_FALSE\n",
- " 20 POP_JUMP_IF_FALSE 13 (to 26) || \n",
- " 22 LOAD_FAST 1 (mask) || 1) LOAD_FAST\n",
- " 24 JUMP_FORWARD 1 (to 28) || JUMP_FORWARD\n",
- " >> 26 LOAD_CONST 1 (1) || \n",
- " >> 28 CALL_FUNCTION 2 || 2) LOAD_CONST\n",
- " 30 STORE_FAST 0 (x) || JUMP_ABSOLUTE\n",
- " || \n",
- " 4 32 LOAD_FAST 0 (x) || 3) CALL_FUNCTION\n",
- " 34 LOAD_FAST 1 (mask) || STORE_FAST\n",
- " 36 LOAD_CONST 0 (None) || LOAD_FAST\n",
- " 38 IS_OP 0 || LOAD_FAST\n",
- " 40 BUILD_TUPLE 2 || LOAD_CONST\n",
- " 42 RETURN_VALUE || IS_OP\n",
- " || BUILD_TUPLE\n",
- " || RETURN_VALUE\n",
- " || \n"
- ]
- }
- ],
- "source": [
- "import dis\n",
- "import io\n",
- "import itertools\n",
- "\n",
- "print(disassembled.code, \"\\n\\n\")\n",
- "dis.dis(disassembled.code, file=(buffer := io.StringIO()))\n",
- "buffer.seek(0)\n",
- "dis_lines = buffer.read().splitlines(False)\n",
- "\n",
- "block_lines = []\n",
- "for idx, block in enumerate(disassembled.raw.values()):\n",
- " block_lines.extend(f\"{'' if idy else f'{idx})':<4}{instruction.opname}\" for idy, instruction in enumerate(block))\n",
- " block_lines.append(\"\")\n",
- "\n",
- "pad = max(len(l) for l in dis_lines)\n",
- "for dis_line, block_line in itertools.zip_longest(dis_lines, block_lines, fillvalue=\"\"):\n",
- " print(f\"{dis_line:<{pad + 10}} ||{' ' * 10}{block_line}\")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 7,
- "id": "dec0f942-0915-4cf5-a79a-9e37fd2f4e6f",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Block 0: [] => [layer_1, v0]\n",
- " LOAD[layer_0, x, mask]\n",
- " CALL_FUNCTION . . . . . . . . . . . . . . (layer_0, x, mask) -> v0\n",
- " STORE[x]\n",
- " LOAD[layer_1, x, mask, None: CONST]\n",
- " IS_OP . . . . . . . . . . . . . . . . . . (mask, None) -> v1\n",
- " POP_JUMP_IF_FALSE . . . . . . . . . . . . (v1) -> \n",
- " -> 1, 2(Jump)\n",
- "\n",
- "Block 1: [⓵ , ⓶ ] => [⓵ , ⓶ , mask]\n",
- " LOAD[mask]\n",
- " JUMP_FORWARD\n",
- " -> 3(Jump)\n",
- "\n",
- "Block 2: [⓵ , ⓶ ] => [⓵ , ⓶ , 1]\n",
- " LOAD[1: CONST]\n",
- " JUMP_ABSOLUTE*\n",
- " -> 3(Jump)\n",
- "\n",
- "Block 3: [⓵ , ⓶ , ⓷ ] => []\n",
- " CALL_FUNCTION . . . . . . . . . . . . . . (⓵ , ⓶ , ⓷ ) -> v0\n",
- " STORE[x]\n",
- " LOAD[x, mask, None: CONST]\n",
- " IS_OP . . . . . . . . . . . . . . . . . . (mask, None) -> v1\n",
- " BUILD_TUPLE . . . . . . . . . . . . . . . (v0, v1) -> v2\n",
- " RETURN_VALUE . . . . . . . . . . . . . . (v2) -> \n",
- "\n"
- ]
- }
- ],
- "source": [
- "print(parsed.summary)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 8,
- "id": "3a1cdbfe-996f-48c2-a5d4-b05389d2316d",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "CALL_FUNCTION (layer_0(LOCAL), x(LOCAL), mask(LOCAL)) -> (IntermediateValue(at 0x7f064f0c26b0),)\n",
- "IS_OP (mask(LOCAL), None(CONST)) -> (IntermediateValue(at 0x7f064f0c3070),)\n",
- "POP_JUMP_IF_FALSE (OutputRef(IS_OP, idx=0)) -> ()\n",
- "\n",
- "JUMP_FORWARD () -> ()\n",
- "\n",
- "JUMP_ABSOLUTE () -> ()\n",
- "\n",
- "CALL_FUNCTION (0(STACK), 1(STACK), 2(STACK)) -> (IntermediateValue(at 0x7f064f0f0040),)\n",
- "IS_OP (mask(LOCAL), None(CONST)) -> (IntermediateValue(at 0x7f064f0f0250),)\n",
- "BUILD_TUPLE (OutputRef(CALL_FUNCTION, idx=0), OutputRef(IS_OP, idx=0)) -> (IntermediateValue(at 0x7f064f0f0460),)\n",
- "RETURN_VALUE (OutputRef(BUILD_TUPLE, idx=0)) -> ()\n",
- "\n"
- ]
- }
- ],
- "source": [
- "def pretty_repr(x) -> str:\n",
- " if isinstance(x, parse.VariableKey):\n",
- " return f\"{x.identifier}({x.scope.name})\"\n",
- " if isinstance(x, values.OutputRef):\n",
- " return f\"OutputRef({x.instruction.opname}, idx={x.idx})\"\n",
- " return repr(x)\n",
- "\n",
- "for block, begin, end in parsed_symbolic.blocks:\n",
- " # At this point `begin` isn't very interesting as it's all just placeholders.\n",
- " for instruction, symbolic in block.items():\n",
- " \n",
- " inputs = \", \".join(pretty_repr(i) for i in symbolic.inputs.ordered)\n",
- " print(f\"{instruction.opname:<25} ({inputs}) -> {symbolic.outputs}\")\n",
- " print()"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 9,
- "id": "8daca4b6-31b3-4d8e-8be2-9744b44d8f4f",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Unlink\n",
- "MarkTuples\n",
- "AddTransitive\n",
- "Connect\n",
- "\n",
- "================================================================================\n",
- " ProtoBlock: 0x7f064f0f1cf0\n",
- " CALL_FUNCTION\n",
- " IS_OP\n",
- " POP_JUMP_IF_FALSE\n",
- "\n",
- "ProtoBlock: 0x7f064ef1e560\n",
- " JUMP_FORWARD\n",
- "\n",
- "ProtoBlock: 0x7f064ef81b10\n",
- " JUMP_ABSOLUTE\n",
- "\n",
- "ProtoBlock: 0x7f064ef9cfd0\n",
- " CALL_FUNCTION\n",
- " IS_OP\n",
- " BUILD_TUPLE\n",
- " RETURN_VALUE\n"
- ]
- }
- ],
- "source": [
- "for transform, _ in protograph_transforms:\n",
- " print(transform.__name__)\n",
- "\n",
- "print(f\"\\n{'=' * 80}\\n\", proto_graph)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 10,
- "id": "5c56fbcf-13db-4749-940b-ecbb0dfd0af5",
- "metadata": {},
- "outputs": [
- {
- "name": "stdout",
- "output_type": "stream",
- "text": [
- "Graph of\n",
- " Block (reached from [None])\n",
- " CALL_FUNCTION([, , ])\n",
- " IS_OP 1 (1)\n",
- " POP_JUMP_IF_FALSE 13 (26)\n",
- " Block (reached from [])\n",
- " JUMP_FORWARD 1 (28)\n",
- " Block (reached from [])\n",
- " JUMP_ABSOLUTE 14 (None)\n",
- " Block (reached from [, ])\n",
- " CALL_FUNCTION([, , ])\n",
- " IS_OP 0 (0)\n",
- " BUILD_TUPLE 2 (2)\n",
- " RETURN_VALUE None (None)\n"
- ]
- }
- ],
- "source": [
- "g = acquire_method(masked_apply)\n",
- "print(g)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "22ef39c9-04aa-40ff-acc2-1387a3f372f4",
- "metadata": {},
- "outputs": [],
- "source": []
- }
- ],
- "metadata": {
- "kernelspec": {
- "display_name": "Python 3",
- "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.11"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/thunder/core/script/parse/__init__.py b/thunder/core/script/parse/__init__.py
deleted file mode 100644
index c29adb91a6..0000000000
--- a/thunder/core/script/parse/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from thunder.core.script.parse.disassemble import *
-from thunder.core.script.parse.functionalize import *
-from thunder.core.script.parse.instructions import *
-from thunder.core.script.parse.stack_effect import *
-
-# This will be populated as parse-time narrowing is introduced.
-FORBIDDEN_INSTRUCTIONS = InstructionSet()
diff --git a/thunder/core/script/parse/disassemble.py b/thunder/core/script/parse/disassemble.py
deleted file mode 100644
index 3ae5ae1bbe..0000000000
--- a/thunder/core/script/parse/disassemble.py
+++ /dev/null
@@ -1,166 +0,0 @@
-"""Convert a `CodeType` object into a series of simple blocks."""
-from __future__ import annotations
-
-import dataclasses
-import dis
-import itertools
-from types import CodeType
-from typing import Any, NewType, TypeVar
-from collections.abc import Iterable, Mapping
-
-from thunder.core.script.parse import stack_effect
-from thunder.core.script.parse.instructions import * # There are a lot of constants, and it defines `__all__`
-
-
-__all__ = ("Disassembled", "ParseDetailInstruction", "EpilogueFixup", "Jump")
-
-BlockIdx = NewType("BlockIdx", int)
-Jump = NewType("Jump", bool)
-
-
-@dataclasses.dataclass
-class Disassembled:
- code: CodeType
-
- _StartIndex = NewType("_StartIndex", int)
- _RawBlocks = dict[_StartIndex, tuple[ThunderInstruction, ...]]
- raw: RawBlocks
-
- _Blocks = tuple[tuple["ThunderInstruction", ...], ...]
- blocks: _Blocks
-
- _Edges = tuple[tuple[BlockIdx, BlockIdx, Jump], ...]
- edges: _Edges
-
- @classmethod
- def make(cls, co: CodeType) -> Disassembled:
- raw_blocks, last_line_no = partition(co)
- blocks, edges = connect_blocks(consolidate_returns(raw_blocks))
- for instruction in itertools.chain(*blocks):
- instruction.line_no = getattr(instruction, "line_no", last_line_no)
- return cls(code=co, raw=raw_blocks, blocks=blocks, edges=edges)
-
-
-class ParseDetailInstruction(ThunderInstruction):
- """Allow us to distinguish instructions that are added during parsing."""
-
- pass
-
-
-class EpilogueFixup(ParseDetailInstruction):
- pass
-
-
-def compute_jump(instruction: ThunderInstruction, position: int) -> int | None:
- if instruction in ABSOLUTE_JUMP_INSTRUCTIONS:
- return instruction.oparg
-
- elif instruction in UNCONDITIONAL_BACKWARD:
- return position + 1 - instruction.oparg
-
- elif "BACKWARD" in instruction.opname:
- # TODO: POP_JUMP_BACKWARD_IF_... variants
- raise NotImplementedError(instruction.opname)
-
- elif instruction in RELATIVE_JUMP_INSTRUCTIONS:
- return position + 1 + instruction.oparg
-
- return None
-
-
-IntT = TypeVar("IntT", bound=int, covariant=True)
-StartIndex = NewType("StartIndex", int)
-RawBlocks = dict[StartIndex, tuple[ThunderInstruction, ...]]
-
-
-def get_free_key(x: Mapping[IntT, Any]) -> int:
- key = -len(x)
- while key in x:
- key -= 1
- return key
-
-
-def partition(co: CodeType) -> tuple[RawBlocks, int]:
- bytecode = tuple(ThunderInstruction(*i) for i in dis.get_instructions(co, first_line=0))
-
- # Determine the boundaries for the simple blocks.
- split_after = JUMP_INSTRUCTIONS | RETURN_INSTRUCTIONS
- follows_jump = itertools.chain([0], (int(i in split_after) for i in bytecode))
- new_block = (int(i or j.is_jump_target) for i, j in zip(follows_jump, bytecode))
-
- # Split the bytecode (and instruction number) into groups
- group_indices = tuple(itertools.accumulate(new_block))
- groups = itertools.groupby(enumerate(bytecode), key=lambda args: group_indices[args[0]])
-
- # Drop the group index, copy from the groupby iter, and unzip `enumerate`.
- groups = (zip(*tuple(i)) for _, i in groups)
- blocks: dict[StartIndex, list[ThunderInstruction]] = {
- StartIndex(start): list(block) for (start, *_), block in groups
- }
-
- # If the last instruction is not a jump or return (which means we split
- # because the next instruction was a jump target) then we need to tell
- # the current block how to advance.
- for start, block in blocks.items():
- if block[-1] not in split_after:
- next_start = StartIndex(start + len(block))
- assert bytecode[next_start].is_jump_target
- block.append(ParseDetailInstruction.make_jump_absolute(next_start))
-
- line_no = 1
- for instruction in itertools.chain(*[block for block in blocks.values()]):
- instruction.line_no = line_no = instruction.starts_line or line_no
-
- return {k: tuple(v) for k, v in blocks.items()}, line_no
-
-
-def consolidate_returns(blocks: RawBlocks) -> RawBlocks:
- def is_return(block: tuple[ThunderInstruction, ...]) -> bool:
- assert block and not any(i.opname == RETURN_VALUE for i in block[:-1])
- return block[-1].opname == RETURN_VALUE
-
- blocks = blocks.copy()
- return_blocks = {k: v for k, v in blocks.items() if is_return(v)}
- if len(return_blocks) > 1:
- new_return_start = StartIndex(get_free_key(blocks))
- for start, (*body, prior_return) in return_blocks.items():
- assert is_return((prior_return,)), prior_return
- blocks[start] = (*body, ParseDetailInstruction.make_jump_absolute(new_return_start))
- return_blocks = {new_return_start: (ParseDetailInstruction.make_return(is_jump_target=True),)}
-
- # Move return block to the end. This isn't always valid (since a block might
- # expect to fall through and reach it), but that will be resolved by the
- # sort in `ProtoGraph`'s ctor.
- blocks = {k: v for k, v in blocks.items() if k not in return_blocks}
- blocks.update(return_blocks)
- return blocks
-
-
-def connect_blocks(blocks: RawBlocks) -> tuple[Disassembled._Blocks, Disassembled._Edges]:
- def iter_raw_edges(blocks: RawBlocks) -> Iterable[tuple[StartIndex, StartIndex, Jump, int, int]]:
- for start, block in tuple(blocks.items()):
- raw_block_len = sum(not isinstance(i, ParseDetailInstruction) for i in block)
- *_, last_i = block
- if last_i in JUMP_INSTRUCTIONS:
- end = start + raw_block_len - 1
- _, (push_nojump, push_jump) = stack_effect.stack_effect_detail(last_i)
- if last_i not in UNCONDITIONAL_JUMP_INSTRUCTIONS:
- yield start, StartIndex(end + 1), Jump(False), max(push_jump - push_nojump, 0), last_i.line_no
-
- if (jump_offset := compute_jump(last_i, end)) is not None:
- yield start, StartIndex(jump_offset), Jump(True), max(push_nojump - push_jump, 0), last_i.line_no
-
- blocks = blocks.copy()
- edges: list[tuple[StartIndex, StartIndex, Jump]] = []
- for source, destination, jump, pop_suffix, line_no in iter_raw_edges(blocks):
- if pop_suffix:
- blocks[epilogue := StartIndex(get_free_key(blocks))] = (
- *(EpilogueFixup.make(POP_TOP, None, line_no=line_no) for _ in range(pop_suffix)),
- EpilogueFixup.make_jump_absolute(destination, line_no=line_no),
- )
- edges.extend(((source, epilogue, jump), (epilogue, destination, jump)))
- else:
- edges.append((source, destination, jump))
-
- to_idx = {k: BlockIdx(idx) for idx, k in enumerate(blocks.keys())}
- return tuple(blocks.values()), tuple((to_idx[source], to_idx[sink], jump) for source, sink, jump in edges)
diff --git a/thunder/core/script/parse/functionalize.py b/thunder/core/script/parse/functionalize.py
deleted file mode 100644
index 9664271918..0000000000
--- a/thunder/core/script/parse/functionalize.py
+++ /dev/null
@@ -1,287 +0,0 @@
-"""Replay the CPython stack machine to determine data flow within a simple block."""
-from __future__ import annotations
-
-import dataclasses
-import enum
-import inspect
-import itertools
-import marshal
-import textwrap
-from types import CodeType
-from typing import Any, NamedTuple, NewType
-
-import networkx as nx
-
-from thunder.core.script import algorithms
-from thunder.core.script.parse import disassemble, instructions, stack_effect
-from thunder.core.utils import safe_zip, FrozenDict, InferringDict
-
-__all__ = ("VariableScope", "VariableKey", "ParsedFunctional", "FunctionalizedBlock", "PlaceholderValue")
-
-
-class VariableScope(enum.Enum):
- CONST = enum.auto()
- LOCAL = enum.auto()
- NONLOCAL = enum.auto()
- GLOBAL = enum.auto()
- STACK = enum.auto()
-
-
-class VariableKey(NamedTuple):
- """Denotes the location of a variable.
- For example, `x = 5` assigns the variable stored in `VariableKey(5, VariableScope.CONST)`
- to the location `VariableKey("x", VariableScope.LOCAL)`. (Provided `x` is a local variable.)
- The type of `identifier` varies based on `scope`:
- `marshal`able VariableScope.CONST
- str VariableScope.LOCAL / NONLOCAL / GLOBAL
- int VariableScope.STACK
- Any VariableScope.BOUNDARY
- """
-
- identifier: Any
- scope: VariableScope
-
- def __repr__(self) -> str:
- return f"VariableKey({self.identifier}, scope={self.scope.name})"
-
- def __eq__(self, other: object) -> bool:
- return (
- isinstance(other, VariableKey)
- and self.scope == other.scope
- and type(self.identifier) is (_ := type(other.identifier)) # Conflict between `ruff` and `yesqa`
- and self.identifier == other.identifier
- )
-
- def __lt__(self, other: tuple[Any, ...]) -> bool:
- assert isinstance(other, VariableKey), (self, other)
- try:
- return (self.scope.value, self.identifier) < (other.scope.value, other.identifier)
- except TypeError:
- assert self.scope == other.scope, (self, other)
- if self.scope == VariableScope.CONST:
- # We prefer to use native ordering. However for unorderable types (e.g. CodeType)
- # `marshal` at least provides a consistent ordering.
- return marshal.dumps(self.identifier) < marshal.dumps(other.identifier)
- raise
-
- @property
- def is_const(self) -> bool:
- return self.scope == VariableScope.CONST
-
-
-def _compute_stack_offsets(disassembled: disassemble.Disassembled) -> tuple[int, ...]:
- # If we convert the stack indices to a common basis then we can ignore stack effect and
- # treat VariableScope.STACK variables just like any other local.
- G = algorithms.TypedDiGraph[disassemble.BlockIdx]((i, j) for i, j, _ in disassembled.edges)
- G.add_nodes_from(range(len(disassembled.blocks)))
- offsets: dict[disassemble.BlockIdx, int] = {i: 0 for i in G.nodes if not G.in_degree(i)} # type: ignore[misc]
- assert len(offsets) == 1, G
-
- for source, sink in nx.edge_dfs(G):
- net_stack_effect = 0
- for instruction in disassembled.blocks[source]:
- pop, push_by_branch = stack_effect.stack_effect_detail(instruction)
- net_stack_effect += max(push_by_branch) - pop
- expected = offsets[source] + net_stack_effect
- actual = offsets.setdefault(sink, expected)
- assert actual == expected, (actual, expected)
-
- assert all(v >= 0 for v in offsets.values()), offsets
- return tuple(offsets[disassemble.BlockIdx(i)] for i in range(len(disassembled.blocks)))
-
-
-LOAD_OPNAMES = FrozenDict[str, VariableScope](
- LOAD_CONST=VariableScope.CONST,
- LOAD_FAST=VariableScope.LOCAL,
- LOAD_DEREF=VariableScope.NONLOCAL,
- LOAD_CLOSURE=VariableScope.NONLOCAL,
- LOAD_GLOBAL=VariableScope.GLOBAL,
-)
-
-STORE_OPNAMES = FrozenDict[str, VariableScope](
- STORE_FAST=VariableScope.LOCAL,
- STORE_DEREF=VariableScope.NONLOCAL,
- STORE_GLOBAL=VariableScope.GLOBAL,
-)
-
-DEL_OPNAMES = FrozenDict[str, VariableScope](
- DELETE_FAST=VariableScope.LOCAL,
- DELETE_DEREF=VariableScope.NONLOCAL,
- DELETE_GLOBAL=VariableScope.GLOBAL,
-)
-
-PlaceholderValue = NewType("PlaceholderValue", str)
-Inputs = NewType("Inputs", tuple[PlaceholderValue, ...])
-Outputs = NewType("Outputs", tuple[PlaceholderValue, ...])
-BeginState = FrozenDict["VariableKey", PlaceholderValue]
-EndState = FrozenDict["VariableKey", PlaceholderValue | None]
-FunctionalNode = tuple[instructions.ThunderInstruction, Inputs, Outputs]
-FunctionalizedBlock = NewType("FunctionalizedBlock", tuple[tuple[FunctionalNode, ...], BeginState, EndState])
-
-
-@dataclasses.dataclass(frozen=True)
-class ParsedFunctional:
- blocks: tuple[FunctionalizedBlock, ...]
- provenance: disassemble.Disassembled
-
- @staticmethod
- def make(co: CodeType) -> ParsedFunctional:
- disassembled = disassemble.Disassembled.make(co)
- return ParsedFunctional(_functionalize_blocks(disassembled), disassembled)
-
- @property
- def summary(self) -> str:
- return _summarize(self)
-
-
-def _functionalize_blocks(disassembled: disassemble.Disassembled) -> tuple[FunctionalizedBlock, ...]:
- code = disassembled.code
- errors: list[str] = []
- if code.co_cellvars:
- errors.append(
- "Nonlocal variables are not supported but\n"
- f" {code.co_name}() defined in {code.co_filename}:{code.co_firstlineno}\n"
- f" defines nonlocal variable{'s' if len(code.co_cellvars) > 1 else ''}: {', '.join(code.co_cellvars)}"
- )
-
- def report_unsupported(msg: str, instruction: instructions.ThunderInstruction) -> None:
- source_lines, _ = inspect.getsourcelines(code)
- errors.append(
- f"{msg}{instruction} found\n"
- f" {code.co_name}() defined in {code.co_filename}:{code.co_firstlineno}\n"
- f" line {instruction.line_no + code.co_firstlineno}: {source_lines[instruction.line_no].rstrip()}"
- )
-
- name_arrays = FrozenDict[VariableScope, tuple[str, ...]](
- {
- VariableScope.CONST: code.co_consts,
- VariableScope.LOCAL: code.co_varnames,
- VariableScope.NONLOCAL: (*code.co_cellvars, *code.co_freevars),
- VariableScope.GLOBAL: code.co_names,
- }
- )
-
- def to_key(instruction: instructions.ThunderInstruction, scope: VariableScope) -> VariableKey:
- assert scope != VariableScope.STACK, "Indexing into the stack is not permitted."
- assert scope in name_arrays, f"Unknown variable scope: {scope}"
- if scope == VariableScope.NONLOCAL:
- report_unsupported("nonlocal variables are not supported but instruction = ", instruction)
- return VariableKey(name_arrays[scope][instruction.oparg], scope)
-
- def convert(block: tuple[instructions.ThunderInstruction, ...], stack_offset: int) -> FunctionalizedBlock:
- stack: list[PlaceholderValue] = [PlaceholderValue(f"Initial_stack_{i}") for i in range(stack_offset)]
- begin_variables = {VariableKey(idx, VariableScope.STACK): v for idx, v in enumerate(stack)}
- end_variables = InferringDict[VariableKey, PlaceholderValue | None](
- lambda key: begin_variables.setdefault(key, PlaceholderValue(f"Initial: ({key.identifier} {key.scope})"))
- )
-
- assert block
- functionalized: list[tuple[instructions.ThunderInstruction, Inputs, Outputs]] = []
- for idx, instruction in enumerate(block):
- # These are already reflected in the next opcode's argument
- if instruction.opname == instructions.EXTENDED_ARG:
- continue
-
- elif instruction in instructions.UNSAFE_OPCODES:
- # These are unsafe to run, but we should still be able to parse them.
- report_unsupported("Unsupported instruction = ", instruction)
-
- pop, push_by_branch = stack_effect.stack_effect_detail(instruction)
- push = max(push_by_branch)
-
- def assert_expected_stack_effects(pop_i: int, push_i: int) -> None:
- assert (pop, push) == (pop_i, push_i), f"{instruction=} {pop=} {push=}"
-
- # Peek at the stack to track variable mutations.
- if (store_scope := STORE_OPNAMES.get(instruction.opname)) is not None:
- assert_expected_stack_effects(1, 0)
- end_variables[to_key(instruction, store_scope)] = stack.pop()
-
- elif (del_scope := DEL_OPNAMES.get(instruction.opname)) is not None:
- assert_expected_stack_effects(1, 0)
- end_variables[to_key(instruction, del_scope)] = None
-
- elif (load_scope := LOAD_OPNAMES.get(instruction.opname)) is not None:
- assert_expected_stack_effects(0, 1)
- loaded = end_variables[load_key := to_key(instruction, load_scope)]
- assert loaded is not None, f"Access to deleted variable: {load_key}, {instruction}"
- stack.append(loaded)
-
- else:
- # We have already functionalized variable accesses, so we can prune loads and stores.
- inputs = tuple(stack.pop() for _ in range(pop))
- outputs = Outputs(tuple(PlaceholderValue(f"{idx}_{instruction.opname}__{idy}") for idy in range(push)))
- stack.extend(outputs)
- functionalized.append((instruction, Inputs(tuple(reversed(inputs))), outputs))
-
- end_stack = {VariableKey(idx, VariableScope.STACK): v for idx, v in enumerate(stack)}
- end_state: EndState = FrozenDict({**end_variables, **end_stack})
- return FunctionalizedBlock((tuple(functionalized), FrozenDict(begin_variables), end_state))
-
- stack_offsets = _compute_stack_offsets(disassembled)
- functionalized = tuple(convert(block, offset) for block, offset in safe_zip(disassembled.blocks, stack_offsets))
- if errors:
- raise RuntimeError("Preprocessing issues detected:\n" + textwrap.indent("\n\n".join(errors), " " * 4))
-
- return functionalized
-
-
-# =============================================================================
-# == Summary for debugging and testing ========================================
-# =============================================================================
-def _summarize(parsed: ParsedFunctional) -> str:
- # Clear identifiers for input stack values.
- to_symbol = FrozenDict[int, str](enumerate("⓵ ⓶ ⓷ ⓸ ⓹ ⓺ ⓻ ⓼ ⓽ ⓾ Ⓐ Ⓑ Ⓒ Ⓓ Ⓔ Ⓕ".split()))
-
- # Group output edges.
- grouped_edges: dict[int, str] = {
- source: ", ".join(f"{sink}{'(Jump)' if jump else ''}" for _, sink, jump in sinks)
- for source, sinks in itertools.groupby(parsed.provenance.edges, lambda e: e[0])
- }
-
- # Best effort to apply descriptive names.
- inputs_outputs = {}
- block_headers: list[str] = []
- for idx, (functionalized_block, begin, end) in enumerate(parsed.blocks):
- begin_stack = tuple(v for k, v in begin.items() if k.scope == VariableScope.STACK)
- stack_names: dict[str, str] = {v: to_symbol.get(idx, f"S{idx}") + "\u2009" for idx, v in enumerate(begin_stack)}
- names: dict[str, str] = {**{v: f"{k.identifier}" for k, v in begin.items()}, **stack_names}
- names.update({v: f"v{idx}" for idx, v in enumerate(itertools.chain(*[o for _, _, o in functionalized_block]))})
- for instruction, inputs, outputs in functionalized_block:
- inputs_outputs[instruction] = (tuple(names[i] for i in inputs), tuple(names[o] for o in outputs))
-
- end_stack = {k: v for k, v in end.items() if k.scope == VariableScope.STACK}
- assert tuple(k.identifier for k in end_stack) == tuple(range(len(end_stack))), end_stack
- end_stack_str = ", ".join(names[i or ""] for i in end_stack.values())
- block_headers.append(f"Block {idx}: [{', '.join(stack_names.values())}] => [{end_stack_str}]")
-
- # Group loads and stores.
- prefix = {opname: opname.split("_")[0] for opname in itertools.chain(STORE_OPNAMES, LOAD_OPNAMES, DEL_OPNAMES)}
- condensed: list[list[tuple[str, instructions.ThunderInstruction | None]]] = []
- for raw_block in parsed.provenance.blocks:
- condensed.append([])
- for prefix_or_i, group in itertools.groupby(raw_block, lambda i: prefix.get(i.opname, i)):
- if isinstance(prefix_or_i, str):
- name = ", ".join(f"{i.argval}: {i.opname[len(prefix_or_i) + 1:]}" for i in group)
- condensed[-1].append((f"{prefix_or_i}[{name.replace(': FAST', '')}]", None))
- else:
- opname = f"{prefix_or_i.opname}{'' if type(prefix_or_i) is instructions.ThunderInstruction else '*'}"
- condensed[-1].append((opname, prefix_or_i))
-
- # Write lines.
- lines: list[str] = []
- width = max(len(name) for name, _ in itertools.chain(*condensed))
- width = max(width, max(len(i) for i in block_headers)) + 5
- for idx, (condensed_block, (_, _, end)) in enumerate(safe_zip(condensed, parsed.blocks)):
- lines.append(block_headers[idx])
- for name, maybe_i in condensed_block:
- inputs, outputs = inputs_outputs.get(maybe_i, ((), ())) # type: ignore[assignment, arg-type]
- if inputs or outputs:
- name = f"{name} ".ljust(width, ".").replace("..", ". ")
- name = f"{name} ({', '.join(inputs)}) -> {', '.join(outputs)}"
- lines.append(f" {name}")
- if idx in grouped_edges:
- lines.append(f" -> {grouped_edges[idx]}")
- lines.append("")
-
- return "\n".join(lines)
diff --git a/thunder/core/script/parse/instructions.py b/thunder/core/script/parse/instructions.py
deleted file mode 100644
index 34fb7fed07..0000000000
--- a/thunder/core/script/parse/instructions.py
+++ /dev/null
@@ -1,117 +0,0 @@
-"""Extension of the builtin `dis` module."""
-from __future__ import annotations
-
-import dis
-from typing import Any
-
-from typing_extensions import Self
-
-from thunder.core.utils import _OrderedSet
-
-__all__ = (
- "ThunderInstruction",
- "InstructionSet",
- "JUMP_ABSOLUTE",
- "RETURN_VALUE",
- "POP_TOP",
- "EXTENDED_ARG",
- "UNCONDITIONAL_BACKWARD",
- "UNCONDITIONAL_JUMP_INSTRUCTIONS",
- "ABSOLUTE_JUMP_INSTRUCTIONS",
- "RELATIVE_JUMP_INSTRUCTIONS",
- "JUMP_INSTRUCTIONS",
- "RAISE_RETURN_INSTRUCTIONS",
- "RETURN_INSTRUCTIONS",
- "UNSAFE_OPCODES",
-)
-
-
-class ThunderInstruction(dis.Instruction):
- """Thin wrapper on top of dis.Instruction to implement thunder specific logic."""
-
- line_no: int
-
- def __hash__(self) -> int:
- # We sometimes want to use an instruction as a key so we can map back to nodes.
- # `dis.Instruction` is a named tuple and therefore implements recursive constituent
- # hashing which can lead to unwanted collisions. We instead override this behavior
- # to instead use identity hashing.
- return id(self)
-
- def __eq__(self, other: object) -> bool:
- return self is other
-
- @property
- def oparg(self) -> int:
- assert self.arg is not None, self
- return self.arg
-
- def modify_copy(self, **kwargs: Any) -> ThunderInstruction:
- assert type(self) is ThunderInstruction, self
- if "opname" in kwargs:
- kwargs.setdefault("opcode", dis.opmap.get(kwargs["opname"], -1))
- result = ThunderInstruction(**{**self._asdict(), **kwargs})
- result.line_no = self.line_no
- return result
-
- @classmethod
- def make(cls, opname: str, arg: int | None, line_no: int, **kwargs: Any) -> Self:
- ctor_kwargs = dict(
- opname=opname,
- opcode=dis.opmap.get(opname, -1),
- arg=arg,
- argval=None,
- argrepr="None",
- offset=-999,
- starts_line=None,
- is_jump_target=False,
- )
- ctor_kwargs.update(kwargs)
- result = cls(**ctor_kwargs) # type: ignore
- result.line_no = line_no
- return result
-
- @classmethod
- def make_jump_absolute(cls, arg: int, line_no: int = -1) -> ThunderInstruction:
- return cls.make(JUMP_ABSOLUTE, arg, argrepr=f"{arg}", line_no=line_no)
-
- @classmethod
- def make_return(cls, is_jump_target: bool, line_no: int = -1) -> ThunderInstruction:
- return cls.make(RETURN_VALUE, arg=None, argrepr="", is_jump_target=is_jump_target, line_no=line_no)
-
-
-class InstructionSet(_OrderedSet[str, int | ThunderInstruction]):
- """Convenience class for checking opcode properties."""
-
- def canonicalize(self, i: str | int | ThunderInstruction) -> str:
- if isinstance(i, str):
- return i
-
- elif isinstance(i, int):
- return dis.opname[i]
-
- else:
- assert isinstance(i, ThunderInstruction)
- return i.opname
-
-
-# Special opcodes
-JUMP_ABSOLUTE = "JUMP_ABSOLUTE"
-RETURN_VALUE = "RETURN_VALUE"
-POP_TOP = "POP_TOP"
-EXTENDED_ARG = "EXTENDED_ARG"
-
-
-UNCONDITIONAL_BACKWARD = InstructionSet(("JUMP_BACKWARD", "JUMP_BACKWARD_NO_INTERRUPT"))
-UNCONDITIONAL_JUMP_INSTRUCTIONS = InstructionSet((JUMP_ABSOLUTE, "JUMP_FORWARD", *UNCONDITIONAL_BACKWARD))
-
-ABSOLUTE_JUMP_INSTRUCTIONS = InstructionSet(dis.hasjabs)
-RELATIVE_JUMP_INSTRUCTIONS = InstructionSet(dis.hasjrel)
-JUMP_INSTRUCTIONS = InstructionSet((*dis.hasjabs, *dis.hasjrel, *UNCONDITIONAL_JUMP_INSTRUCTIONS))
-
-RAISE_RETURN_INSTRUCTIONS = InstructionSet(("RAISE_VARARGS", "RERAISE"))
-RETURN_INSTRUCTIONS = InstructionSet((RETURN_VALUE, *RAISE_RETURN_INSTRUCTIONS))
-
-
-# https://github.com/Lightning-AI/lightning-thunder/issues/1075
-UNSAFE_OPCODES = InstructionSet(("SETUP_WITH", "SETUP_FINALLY"))
diff --git a/thunder/core/script/parse/stack_effect.py b/thunder/core/script/parse/stack_effect.py
deleted file mode 100644
index 1ddeff412d..0000000000
--- a/thunder/core/script/parse/stack_effect.py
+++ /dev/null
@@ -1,234 +0,0 @@
-import dis
-import opcode
-import sys
-from typing import NewType, TypeAlias, TypeVar
-from collections.abc import Callable
-from collections.abc import Iterable
-
-from types import EllipsisType
-
-from thunder.core.utils import FrozenDict
-
-__all__ = ("stack_effect_detail", "fill_ellipses")
-
-T = TypeVar("T")
-Pop = NewType("Pop", int)
-Push = NewType("Push", int)
-StackEffect: TypeAlias = tuple[Pop, Push] | tuple[Pop, tuple[Push, Push]]
-
-# Aliases for common cases
-NoStackEffect = (Pop(0), Push(0))
-PushTOS = (Pop(0), Push(1))
-PopTOS = (Pop(1), Push(0))
-ReplaceTOS = (Pop(1), Push(1))
-BinaryOp = (Pop(2), Push(1))
-
-
-def make_function_detail(*args: int) -> Callable[[int], StackEffect]:
- return lambda oparg: (Pop(2 + sum((oparg & flag) != 0 for flag in args)), Push(1))
-
-
-def fill_ellipses(**kwargs: T | EllipsisType) -> Iterable[tuple[str, T]]:
- prior_effect: T | EllipsisType = Ellipsis
- for opname, effect in kwargs.items():
- if effect is Ellipsis:
- effect = prior_effect
- assert effect is not Ellipsis
- prior_effect = effect
- yield opname, effect
-
-
-__EFFECTS = dict[str, StackEffect | Callable[[int], StackEffect] | EllipsisType](
- NOP=NoStackEffect, # ∅ -> ∅
- EXTENDED_ARG=NoStackEffect,
- #
- # Stack manipulation
- POP_TOP=PopTOS, # A -> ∅
- ROT_TWO=(Pop(2), Push(2)), # A,B -> B,A
- ROT_THREE=(Pop(3), Push(3)), # A,B,C -> C,A,B
- ROT_FOUR=(Pop(4), Push(4)), # A,B,C,D -> D,A,B,C
- ROT_N=lambda oparg: (Pop(oparg), Push(oparg)), # A,B,...,Z -> Z,A,B,...
- DUP_TOP=(Pop(1), Push(2)), # A -> A,A
- DUP_TOP_TWO=(Pop(2), Push(4)), # A,B -> A,B,A,B
- UNPACK_SEQUENCE=lambda oparg: (Pop(1), Push(oparg)), # A -> B,C,...
- #
- # Jumps & return
- JUMP_FORWARD=NoStackEffect, # ∅ -> ∅
- JUMP_ABSOLUTE=...,
- POP_JUMP_IF_FALSE=PopTOS, # A -> ∅
- POP_JUMP_IF_TRUE=...,
- RETURN_VALUE=...,
- JUMP_IF_NOT_EXC_MATCH=BinaryOp, # A,B -> ∅
- #
- # Exceptions and context managers:
- POP_BLOCK=NoStackEffect, # ∅ -> ∅
- POP_EXCEPT=(Pop(3), Push(0)), # A, B, C -> ∅
- RERAISE=...,
- RAISE_VARARGS=lambda oparg: (Pop(oparg), Push(0)), # A,B,... -> ∅
- WITH_EXCEPT_START=(Pop(7), Push(8)), # ??!?
- LOAD_ASSERTION_ERROR=PushTOS, # ∅ -> A
- #
- # Variable manipulation
- LOAD_CONST=PushTOS, # ∅ -> A
- LOAD_FAST=...,
- LOAD_GLOBAL=...,
- LOAD_NAME=...,
- STORE_FAST=PopTOS, # A -> ∅
- STORE_GLOBAL=...,
- STORE_NAME=...,
- DELETE_FAST=NoStackEffect, # ∅ -> ∅
- DELETE_GLOBAL=...,
- DELETE_NAME=...,
- #
- # Attributes
- LOAD_METHOD=(Pop(1), Push(2)), # A -> B,A
- LOAD_ATTR=ReplaceTOS, # A -> B
- STORE_ATTR=(Pop(2), Push(0)), # A, B -> ∅
- DELETE_ATTR=PopTOS, # A -> ∅
- #
- # Closures
- LOAD_CLOSURE=PushTOS, # ∅ -> A
- LOAD_DEREF=...,
- LOAD_CLASSDEREF=...,
- STORE_DEREF=PopTOS, # A -> ∅
- DELETE_DEREF=NoStackEffect, # ∅ -> ∅
- #
- # Functions and calls A,B,... -> Z
- CALL_FUNCTION=lambda x: (Pop(x + 1), Push(1)),
- CALL_METHOD=lambda x: (Pop(x + 2), Push(1)),
- CALL_FUNCTION_KW=...,
- CALL_FUNCTION_EX=make_function_detail(0x01),
- MAKE_FUNCTION=make_function_detail(0x01, 0x02, 0x04, 0x08),
- #
- # Build containers A,B,... -> Z
- BUILD_TUPLE=lambda oparg: (Pop(oparg), Push(1)),
- BUILD_LIST=...,
- BUILD_SET=...,
- BUILD_STRING=...,
- BUILD_MAP=lambda oparg: (Pop(oparg * 2), Push(1)),
- BUILD_CONST_KEY_MAP=lambda x: (Pop(x + 1), Push(1)),
- LIST_TO_TUPLE=ReplaceTOS, # A -> B
- #
- # Insertion leaves container on the stack A,B -> A
- SET_ADD=BinaryOp,
- SET_UPDATE=...,
- LIST_APPEND=...,
- LIST_EXTEND=...,
- DICT_MERGE=...,
- DICT_UPDATE=...,
- MAP_ADD=(Pop(3), Push(1)), # A,B,C -> A
- COPY_DICT_WITHOUT_KEYS=(Pop(2), Push(2)), # A,B -> A,C (I am unsure...)
- #
- # Unary operators A -> B
- UNARY_POSITIVE=ReplaceTOS,
- UNARY_NEGATIVE=...,
- UNARY_NOT=...,
- UNARY_INVERT=...,
- #
- # Binary operators A,B -> C
- BINARY_POWER=BinaryOp,
- BINARY_MULTIPLY=...,
- BINARY_MATRIX_MULTIPLY=...,
- BINARY_MODULO=...,
- BINARY_ADD=...,
- BINARY_SUBTRACT=...,
- BINARY_SUBSCR=...,
- BINARY_FLOOR_DIVIDE=...,
- BINARY_TRUE_DIVIDE=...,
- INPLACE_FLOOR_DIVIDE=...,
- INPLACE_TRUE_DIVIDE=...,
- INPLACE_ADD=...,
- INPLACE_SUBTRACT=...,
- INPLACE_MULTIPLY=...,
- INPLACE_MATRIX_MULTIPLY=...,
- INPLACE_MODULO=...,
- BINARY_LSHIFT=...,
- BINARY_RSHIFT=...,
- BINARY_AND=...,
- BINARY_XOR=...,
- BINARY_OR=...,
- COMPARE_OP=...,
- IS_OP=...,
- CONTAINS_OP=...,
- #
- # Binary operators (inplace)
- # https://docs.python.org/3/reference/datamodel.html?highlight=iadd#object.__iadd__
- # "... and return the result (which could be, but does not have to be, self)."
- INPLACE_POWER=BinaryOp,
- INPLACE_LSHIFT=...,
- INPLACE_RSHIFT=...,
- INPLACE_AND=...,
- INPLACE_XOR=...,
- INPLACE_OR=...,
- #
- # Indexing operators
- STORE_SUBSCR=(Pop(3), Push(0)), # A,B,C -> ∅
- DELETE_SUBSCR=(Pop(2), Push(0)), # A,B -> ∅
- BUILD_SLICE=lambda x: (Pop(x), Push(1)), # A,B,... -> Z
- UNPACK_EX=lambda x: (Pop(1), Push((x & 0xFF) + (x >> 8) + 1)), # A -> B,C,...
- #
- # Iterators
- GET_ITER=ReplaceTOS, # A -> B
- GET_YIELD_FROM_ITER=ReplaceTOS,
- #
- # Misc.
- FORMAT_VALUE=lambda oparg: (Pop(1 + bool(oparg & 0x04)), Push(1)), # (A?),B -> C
- PRINT_EXPR=PopTOS, # A -> ∅
- IMPORT_STAR=...,
- LOAD_BUILD_CLASS=PushTOS,
- SETUP_ANNOTATIONS=NoStackEffect,
- GET_LEN=(Pop(1), Push(2)),
- IMPORT_NAME=BinaryOp,
- IMPORT_FROM=(Pop(1), Push(2)),
- MATCH_CLASS=(Pop(3), Push(1)),
- MATCH_MAPPING=(Pop(1), Push(2)),
- MATCH_SEQUENCE=...,
- MATCH_KEYS=(Pop(2), Push(3 + bool(sys.version_info < (3, 11)))),
- #
- # Jump dependent
- FOR_ITER=(Pop(1), (Push(2), Push(0))),
- SETUP_WITH=(Pop(1), (Push(2), Push(7))),
- SETUP_FINALLY=(Pop(0), (Push(0), Push(6))),
- SETUP_ASYNC_WITH=(Pop(0), (Push(0), Push(6))),
- #
- # NOTE: These instructions have been removed since they are extraneous special cases.
- # https://github.com/faster-cpython/ideas/issues/567
- # https://github.com/python/cpython/issues/102859
- JUMP_IF_TRUE_OR_POP=(Pop(1), (Push(0), Push(1))),
- JUMP_IF_FALSE_OR_POP=...,
- #
- # TODO(robieta, t-vi): Iterators and generators
- # "GEN_START": PopTOS, # Where does TOS for this come from?
- # "YIELD_VALUE": ReplaceTOS, # I think
- # "YIELD_FROM": (2, PushNew), # I am very unsure
- # "GET_AWAITABLE": (1, 1),
- # "BEFORE_ASYNC_WITH": (1, 2),
- # "GET_AITER": (1, 1),
- # "GET_ANEXT": (1, 2),
- # "END_ASYNC_FOR": (7, 0),
-)
-
-
-# Split so MyPy can type check `__EFFECTS` without having to go through `fill_ellipses`.
-_RAW_STACK_EFFECTS = FrozenDict[str, StackEffect | Callable[[int], StackEffect]](fill_ellipses(**__EFFECTS))
-del __EFFECTS
-
-
-def stack_effect_detail(instruction: dis.Instruction) -> tuple[Pop, tuple[Push, Push]]:
- assert isinstance(instruction, dis.Instruction), instruction
- if callable(effect := _RAW_STACK_EFFECTS[instruction.opname]):
- assert instruction.arg is not None
- effect = effect(instruction.arg)
-
- assert isinstance(effect, tuple) and len(effect) == 2 and isinstance(effect[0], int)
- if isinstance(effect[1], int):
- effect = (effect[0], (effect[1],) * 2)
-
- # Python exposes a method to compute stack effect, so while it's not part
- # of the public API we may as well use it to check our bookkeeping.
- pop, (push_nojump, push_jump) = effect
- for jump, push in ((False, push_nojump), (True, push_jump)):
- expected = opcode.stack_effect(instruction.opcode, instruction.arg, jump=jump)
- assert expected == push - pop, (expected, push, pop, jump)
-
- return Pop(pop), (Push(push_nojump), Push(push_jump))
diff --git a/thunder/core/script/passes.py b/thunder/core/script/passes.py
deleted file mode 100644
index a53581ec01..0000000000
--- a/thunder/core/script/passes.py
+++ /dev/null
@@ -1,932 +0,0 @@
-import dis
-import copy
-import inspect
-import opcode
-import sys
-import types
-from typing import Any, Dict, List, Tuple, Union
-from collections.abc import Callable
-from collections.abc import Hashable
-from contextvars import ContextVar
-
-import networkx as nx
-import torch # # aehem.
-
-import thunder
-from thunder.core.script.frontend import acquire_method, remove_unused_values
-from thunder.core.script.graph import (
- assert_block,
- assert_node,
- assert_value,
- Graph,
- Block,
- clone_blocks,
- _generate_raises,
- GraphObject,
- Node,
- PhiValue,
- replace_values,
- SourceInformation,
- _Undefined,
- Value,
- repr_source_location,
-)
-from thunder.core.script.instrumentation import verbose_error, record
-from thunder.core.script.parse import ThunderInstruction, JUMP_ABSOLUTE
-from thunder.core.script.python_ir_data import get_instruction, X_THUNDER_STORE_ATTR
-from thunder.torch import _torch_to_thunder_complete_map
-from thunder.core.script.noinline import NOINLINE_METHODS
-from thunder.core.utils import debug_asserts_enabled, debug_asserts_level, OrderedSet
-
-MAX_INLINE_ITERS = 50
-
-
-def split_block(gr: "Graph", bl: "Block", n: "Node") -> Block:
- # The admin involved:
- # - create a new "bottom block", the input block is the "top block"
- # - split the .nodes
- # - block_inputs of the top block and block_outputs of the bottom are the original
- # block_inputs and block_outputs
- # - scan all the node inputs and block_outputs of the lower part to see
- # which need to be block_inputs of the lower block and thus outputs of the top one
- # - define outputs of the "top block" to be the required inputs
- # - add the input PhiValues and replace the outputs of the top block with them in the
- # uses in the bottom block
- # - add unconditional jump from top to bottom part
-
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
- i = 0
- while i < len(gr.blocks) and gr.blocks[i] is not bl:
- i += 1
- assert i < len(gr.blocks), "block not found"
- j = 0
- while j < len(bl.nodes) and bl.nodes[j] is not n:
- j += 1
- assert j < len(bl.nodes), "node not found"
- nbl = Block()
- nbl.nodes = bl.nodes[j:]
- del bl.nodes[j:]
- old_block_outputs = bl.block_outputs
- nbl.block_outputs = OrderedSet()
- bl.block_outputs = OrderedSet()
- nbl.block_inputs = []
-
- bl_jump_node = Node(i=ThunderInstruction.make_jump_absolute(arg=None), inputs=[], outputs=[])
- bl_jump_node.jump_targets = [nbl]
- if bl.nodes:
- bl_jump_node.source_infos = copy.deepcopy(bl.nodes[-1].source_infos)
- else:
- bl_jump_node.source_infos = copy.deepcopy(nbl.nodes[0].source_infos)
- bl.nodes.append(bl_jump_node)
- nbl.jump_sources.append(bl_jump_node)
- nbl.graph = gr
- gr.blocks.insert(i + 1, nbl)
-
- potential_bl_outputs = {i for i in bl.block_inputs}
- for n in bl.nodes:
- for o in n.outputs:
- potential_bl_outputs.add(o)
- for i in bl.block_inputs:
- potential_bl_outputs.add(i)
- value_map: dict[GraphObject, GraphObject] = {}
-
- def get_or_create_phi(v: Value) -> Value:
- if v in value_map:
- return assert_value(value_map[v])
- if v.is_const or v.is_global:
- return v
- if v in potential_bl_outputs: # priority follow parent vs. phi_value?
- phi_value = PhiValue([v], [bl_jump_node], nbl)
- nbl.block_inputs.append(phi_value)
- bl.block_outputs.add(v)
- value_map[v] = phi_value
- return phi_value
- if v.parent is not None:
- # this adds v.parent to the value_map, so that is used
- # for the clone's parent
- get_or_create_phi(v.parent)
- v_new = v.clone(translation_dict=value_map)
- v_new.block = nbl
- return v_new
- raise ValueError(f"unknwn value {v}")
-
- for n in nbl.nodes:
- n.inputs = [get_or_create_phi(i) for i in n.inputs]
- for o in n.outputs:
- o.block = nbl
- value_map[o] = o
-
- for o in old_block_outputs:
- if o not in value_map:
- bl.block_outputs.add(o)
- else:
- assert value_map[o].block is nbl or (
- value_map[o].is_function_arg or value_map[o].is_global
- ), f"value {repr(o)} mapped to {repr(value_map[o])} has block {gr.blocks.index(value_map[o].block)} instead of {gr.blocks.index(nbl)}"
- nbl.block_outputs.add(value_map[o])
- if o is not value_map[o]:
- for pv in o.phi_values[:]:
- if pv.block is not nbl:
- pv.replace_value(o, value_map[o])
-
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
-
- return nbl
-
-
-@verbose_error
-def find_method_through_phi_parent(fn_value: Value) -> tuple[Value, list[str]]:
- Point = tuple[Value, tuple[str, ...]]
- to_process: list[Point] = [(v, ()) for v in fn_value.resolve()]
- edges: OrderedSet[tuple[Point, Point]] = OrderedSet(((fn_value, ()), i) for i in to_process)
- while to_process:
- v, attr = to_process.pop()
- destination = (v, attr)
- if (parent := v.parent) is not None and (name := v.name) is not None:
- destination = (parent, (name, *attr))
-
- elif (node := v.node) is not None and node.i.opname == "BINARY_SUBSCR" and node.inputs[1].is_const:
- destination = (node.inputs[0], (repr(node.inputs[1].value), *attr))
-
- for vi in destination[0].resolve():
- edge = ((v, attr), (vi, destination[1]))
- if edge not in edges:
- edges.add(edge)
- to_process.append(edge[1])
-
- G = nx.from_edgelist(edges, nx.DiGraph)
- G.remove_edges_from(nx.selfloop_edges(G))
- assert nx.is_connected(G.to_undirected())
- assert nx.is_directed_acyclic_graph(G)
-
- # A size one topological generation means all flow must pass through that node. Thus, the latest
- # generation with that property is the farthest we can resolve attributes.
- *_, (fn_value, attr_lookups) = (i for i, *other in nx.topological_generations(G) if not other)
- return fn_value, list(attr_lookups)
-
-
-def find_and_evaluate_method_through_phi_parent(v: Value) -> object | Callable:
- fn_parent_value, attr_lookups = find_method_through_phi_parent(v)
- if fn_parent_value.value is None:
- return None
- fn_value = fn_parent_value.value
- for al in attr_lookups:
- value = getattr(fn_value, al, _Undefined)
- if value is _Undefined:
- return _Undefined(fn_value, al)
- fn_value = value
- return fn_value
-
-
-class SkipInlineError(NotImplementedError):
- pass
-
-
-@record(delegate_to="n")
-def inline_method_call(gr: "Graph", n: "Node") -> None:
- gr.ensure_links()
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
- found_block = False
- for i_bl, bl in enumerate(gr.blocks):
- for i_n, n1 in enumerate(bl.nodes):
- if n1 is n: # is?
- found_block = True
- break
- if found_block:
- break
- assert found_block
- if n.i.opname == "CALL_METHOD":
- fn_value: Callable = find_and_evaluate_method_through_phi_parent(n.inputs[0]) # type: ignore
- assert not isinstance(fn_value, _Undefined)
- if fn_value is None:
- raise NotImplementedError("cannot inline non-explicit function")
-
- ## TODO: value for self arg in Method calls?
- ### in general: What is with callables here?
- if isinstance(fn_value, torch.nn.Module):
- mod1: object = fn_value
- value_for_self1 = n.inputs[0]
- fn_value = fn_value.forward
- elif isinstance(fn_value, types.MethodType):
- mod1 = fn_value.__self__
- value_for_self1 = n.inputs[1]
- else:
- mod1 = None
- value_for_self1 = None
-
- if inspect.isbuiltin(fn_value):
- raise NotImplementedError("cannot inline built-in (C-implemented) function")
- elif n.i.opname in {"CALL_FUNCTION", "CALL_FUNCTION_KW"}:
- fn_value = find_and_evaluate_method_through_phi_parent(n.inputs[0]) # type: ignore
- assert not isinstance(fn_value, _Undefined)
- if fn_value is None:
- raise NotImplementedError("cannot inline non-explicit function")
-
- if isinstance(fn_value, torch.nn.Module):
- mod1 = fn_value
- value_for_self1 = n.inputs[0]
- fn_value = fn_value.forward
- else:
- if isinstance(fn_value, types.FunctionType):
- mod1 = None
- value_for_self1 = None
- elif isinstance(fn_value, types.MethodType):
- mod1 = fn_value.__self__
- value_for_self1 = n.inputs[0].parent
- assert value_for_self1 is not None
- else:
- source_str = repr_source_location(gr, n.source_infos)
- raise NotImplementedError(f"inlining {fn_value} in instruction {n} at\n{source_str}")
- else:
- raise NotImplementedError(f"inlining {n}")
-
- # splitting must be done before replacing values, but this is changed even if we don't inline...
- nbl = split_block(gr, bl, bl.nodes[i_n + 1])
-
- gr1 = acquire_method(fn_value, module=mod1, mro_klass=gr.mro_klass if mod1 == gr.module else None)
- for gr1_n in gr1.nodes():
- assert gr1_n.source_infos
- have_generated = False
- for si in gr1_n.source_infos:
- si.gen_line_no = si.gen_line_no + len(gr.source_lines) + 1
- si.gen_end_line_no = si.gen_end_line_no + len(gr.source_lines) + 1
- # prepend
- gr1_n.source_infos[:0] = copy.deepcopy(n.source_infos)
- gr.source_lines.append("\n")
- gr.source_lines += gr1.source_lines
-
- if gr1.ismethod:
- sig1 = inspect.signature(gr1.method.__func__)
- else:
- sig1 = inspect.signature(gr1.method)
- # transform defaults
- sig1 = sig1.replace(
- parameters=[
- p
- if p.default is inspect._empty
- else p.replace(default=Value(name=p.name, typ=type(p.default), value=p.default, is_const=True))
- for p in sig1.parameters.values()
- ]
- )
-
- if gr1.ismethod:
- call_args = [value_for_self1]
- else:
- call_args = []
-
- if n.i.opname == "CALL_METHOD":
- call_args += n.inputs[2:]
- call_kwargs: dict[str, Any] = {}
- elif n.i.opname == "CALL_FUNCTION":
- call_args += n.inputs[1:]
- call_kwargs = {}
- elif n.i.opname == "CALL_FUNCTION_KW":
- assert n.inputs[-1].is_const
- num_kwargs = len(n.inputs[-1].value)
- call_kwargs = {k: v for k, v in zip(n.inputs[-1].value, n.inputs[-1 - num_kwargs : -1])}
- call_args += n.inputs[1 : -1 - num_kwargs]
- else:
- raise NotImplementedError()
-
- # TODO: catch and translate error messages, check types(?)
- bound_args = sig1.bind(*call_args, **call_kwargs)
- bound_args.apply_defaults()
-
- gr1_varargs = [n for n, p in sig1.parameters.items() if p.kind == p.kind.VAR_POSITIONAL]
- gr1_varkwargs = [n for n, p in sig1.parameters.items() if p.kind == p.kind.VAR_KEYWORD]
- ## TODO: TRANSLATE args (=tuple of Values) and kwargs (=dict str->Value) to a Value to something Value of ... (probably needs at least BUILD_TUPLE etc)
- if gr1_varargs or gr1_varkwargs:
- raise SkipInlineError("varargs and kwargs are currently not implemented")
-
- n1 = bl.nodes.pop(i_n)
- assert n1 is n
-
- # there should be exactly one
- (ret_bl,) = (bl for bl in gr1.blocks if len(bl.nodes) > 0 and bl.nodes[-1].i.opname == "RETURN_VALUE")
-
- ret_node = ret_bl.nodes[-1]
- ret_node.i = ThunderInstruction.make(
- JUMP_ABSOLUTE,
- arg=-1,
- argrepr="None",
- offset=ret_node.i.offset,
- starts_line=ret_node.i.starts_line,
- is_jump_target=ret_node.i.is_jump_target,
- line_no=ret_node.i.line_no,
- )
- bl.nodes[-1].jump_targets = [gr1.blocks[0]]
- assert len(gr1.blocks[0].jump_sources) == 1
- gr1.blocks[0].jump_sources = [bl.nodes[-1]]
- for pv in gr1.blocks[0].block_inputs:
- assert pv.jump_sources == [None]
- pv.jump_sources = [bl.nodes[-1]]
- ret_node.jump_targets = [nbl]
- nbl.jump_sources = [ret_node if js == bl.nodes[-1] else js for js in nbl.jump_sources]
- for pv in nbl.block_inputs:
- pv.jump_sources = [ret_node if js == bl.nodes[-1] else js for js in pv.jump_sources]
-
- for bl1 in gr1.blocks:
- bl1.graph = gr
- gr.blocks[i_bl + 1 : i_bl + 1] = gr1.blocks
-
- assert len(n.outputs) == 1
- inp_map = {p: bound_args.arguments[p.name] for p in gr1.local_variables_at_start if p.name in bound_args.arguments}
- if n.outputs[0] in bl.block_outputs: # it may legitimately happen that we don't use the output
- bl.block_outputs.remove(n.outputs[0]) # TODO: what with inplace!!
- bl.block_outputs.update(inp_map.values()) # Note: This includes default args
- gr.ensure_links()
- replace_values(gr1, inp_map)
-
- # output value
- rv = ret_node.inputs.pop()
- assert not ret_node.inputs
- (orv,) = n.outputs
- replace_values(gr, {orv: rv})
- ret_bl.block_outputs.add(rv)
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
-
-
-def inline_submodule_calls(gr: "Graph") -> bool:
- # inlines submodule calls
- # returns whether something has changed
- # TODO: recursively and not from nested structures (ModuleList etc.)
- changed = False
- gr.ensure_links()
- for bl in gr.blocks[:]:
- for n in bl.nodes[:]:
- if n.i.opname in {"CALL_METHOD", "CALL_FUNCTION", "CALL_FUNCTION_KW"}:
- fn_value = find_and_evaluate_method_through_phi_parent(n.inputs[0])
- if isinstance(fn_value, _Undefined):
- # TODO: We could insert a RAISE here if we then delete the return
- # value and all (direct or indirect) uses.
- methval = Value(
- value=_generate_raises(
- f"attribute error '{type(fn_value.value)}' object has no attribute '{fn_value.attr}'"
- ),
- is_const=True,
- )
- n.i = n.i.modify_copy(opname="CALL_FUNCTION", arg=0, opcode=None)
- n.inputs = [methval]
- if isinstance(fn_value, torch.nn.Module) or (
- inspect.ismethod(fn_value)
- and isinstance(fn_value.__self__, torch.nn.Module)
- and (fn_value not in NOINLINE_METHODS.get())
- ):
- inline_method_call(gr, n)
- changed = True
-
- return changed
-
-
-def strongly_inline_functions(gr: "Graph") -> None:
- for _ in range(MAX_INLINE_ITERS):
- loop = False
- gr.ensure_links()
- for bl in gr.blocks[:]:
- for n in bl.nodes[:]:
- if n.i.opname in {"CALL_METHOD", "CALL_FUNCTION", "CALL_FUNCTION_KW"}:
- fn_value = find_and_evaluate_method_through_phi_parent(n.inputs[0])
- if (
- fn_value is not None
- and not inspect.isbuiltin(fn_value)
- and isinstance(fn_value, types.FunctionType)
- and fn_value not in _torch_to_thunder_complete_map
- and fn_value not in NOINLINE_METHODS.get()
- ):
- ## handle methods or nn.Modules / other classes?
- try:
- inline_method_call(gr, n)
- loop = True
- except SkipInlineError:
- pass
- except RuntimeError as e:
- (msg,) = e.args
- source_str = repr_source_location(gr, n.source_infos)
- msg = f"{msg}\nwhile inlining:\n{source_str}"
- e.args = (msg,)
- raise e
- if not loop:
- return
-
- raise AssertionError(f"Inlining did not complete after {MAX_INLINE_ITERS} passes.")
-
-
-def torch_to_thunder(gr: "Graph", fallback: bool = False) -> None:
- """replaces calls to torch.foo functions with calls into thunder's torch language."""
-
- def fill_in_value(v: Value, seen: OrderedSet[Value]) -> None:
- if v in seen:
- return
- seen.add(v)
- parent = v.parent
- if parent is None and isinstance(v, PhiValue):
- for vv in v.values:
- fill_in_value(vv, seen)
- for vv in v.values[1:]:
- if vv.value is not v.values[0].value:
- return
- v.value = v.values[0].value
- if v.value is None and parent is not None:
- fill_in_value(parent, seen)
- if v.name is None and isinstance(v, PhiValue) and parent is not None and parent.name is not None:
- v.name = parent.name
- if v.value is None and parent is not None and parent.value is not None and v.name is not None:
- v.value = getattr(parent.value, v.name, None)
-
- for bl in gr.blocks:
- for n in bl.nodes:
- for idx, i in enumerate(n.inputs):
- done = False
- fill_in_value(i, OrderedSet())
- i_or_parent = i
- while (
- not isinstance(i_or_parent.value, Hashable)
- or i_or_parent.value not in _torch_to_thunder_complete_map
- ) and i_or_parent.parent is not None:
- i_or_parent = i_or_parent.parent
-
- if isinstance(i_or_parent.value, Hashable) and i_or_parent.value in _torch_to_thunder_complete_map:
- i_or_parent.value = _torch_to_thunder_complete_map[i.value]
- # we reinstantiate because we don't want a PhiValue here
- i_new = Value(
- value=i_or_parent.value,
- typ=type(i_or_parent.value),
- parent=None,
- is_const=True,
- is_global=False,
- name=i_or_parent.name,
- )
- n.inputs[idx] = i_new
- if n.i.opname == "CALL_METHOD" and idx == 0:
- # todo get others, too
- n.i = get_instruction(opname="CALL_FUNCTION", arg=n.i.arg)
- del n.inputs[1]
- done = True
-
- if (not done) and fallback: # fallback
- # todo: change name?, deeper nesting?
- if i.value == torch:
- i.value = thunder.langs.torch
- if i.parent is not None and i.parent.value == torch:
- i.parent.value = thunder.langs.torch
- assert i.name is not None
- i.value = getattr(thunder.langs.torch, i.name)
-
- # replace other things by checking against torch module (make dict at startup?)
- name = getattr(i.value, "__name__", None)
- tf = None
- if name is not None:
- tf = getattr(torch, name, None)
- if tf is not None and i.value == tf:
- i.value = getattr(thunder.langs.torch, name)
- i.is_global = False
- i.is_const = True
-
-
-def merge_two_blocks(gr: "Graph", bl1: "Block") -> None:
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
- jt = bl1.nodes[-1].jump_targets
- if len(jt) != 1:
- raise RuntimeError("can only fuse blocks with deterministic connection")
- bl2 = jt[0]
- if len(bl2.jump_sources) != 1 or bl2.jump_sources[0] != bl1.nodes[-1]:
- raise RuntimeError("second block to be fused must only have first block as jump source")
-
- replacements: dict[Value, Value] = {}
- for i in bl2.block_inputs:
- assert isinstance(i, PhiValue) and len(i.values) == 1, (i, getattr(i, "values", None))
- (iv,) = i.values
- if iv in bl1.block_outputs:
- replacements[i] = iv
- else:
- if i.jump_sources == [bl1.nodes[-1]]:
- i.jump_sources = [iv.block.nodes[-1]]
- bl1.block_inputs.append(i)
- i.block = bl1
-
- replace_values(bl2, replacements, follow_phi_values=True)
- # TODO: Should this happen automatically in replace_values?
- # Should we also replace values in bl1?
- for o in bl1.block_outputs:
- for pv in o.phi_values[:]:
- if pv in replacements:
- pv.remove_value(o)
- else:
- pv.jump_sources = [js if js != bl1.nodes[-1] else bl2.nodes[-1] for js in pv.jump_sources]
-
- bl1_jump = bl1.nodes[-1]
- bl2_jump = bl2.nodes[-1]
-
- bl1.block_outputs = OrderedSet(o for o in bl1.block_outputs if o.phi_values)
- bl1.block_outputs.update(bl2.block_outputs)
-
- bl1.nodes[-1:] = bl2.nodes
- gr.blocks.remove(bl2)
-
- gr.ensure_links()
-
- # fix jump sources in other blocks
- for bl in gr.blocks:
- for i in bl.block_inputs:
- i.jump_sources = [(bl2_jump if js is bl1_jump else js) for js in i.jump_sources]
-
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
-
-
-def merge_blocks_where_possible(gr: "Graph") -> None:
- i_bl = 0
- while i_bl < len(gr.blocks):
- bl1 = gr.blocks[i_bl]
- jt = bl1.nodes[-1].jump_targets
- if len(jt) == 1:
- bl2 = jt[0]
- else:
- bl2 = None
- if bl2 is not None and len(bl2.jump_sources) == 1 and bl2.jump_sources[0] == bl1.nodes[-1]:
- merge_two_blocks(gr, bl1)
- else:
- i_bl += 1
-
-
-def find_blocks_of_for(gr: "Graph", for_block: "Block") -> list[Block]:
- assert for_block.nodes[-1].i.opname == "FOR_ITER"
-
- blocks_of_for_loop = OrderedSet({for_block})
- currently_looking_at = set()
-
- def find_blocks_of_for_rec(for_block: "Block", start_block: "Block") -> bool:
- if for_block == start_block:
- return True
- if start_block in currently_looking_at:
- return False
- currently_looking_at.add(start_block)
- found = False
- for jt in start_block.nodes[-1].jump_targets:
- found |= find_blocks_of_for_rec(for_block, jt)
- currently_looking_at.remove(start_block)
- if found:
- blocks_of_for_loop.add(start_block)
- return found
-
- find_blocks_of_for_rec(for_block, for_block.nodes[-1].jump_targets[0])
- return list(blocks_of_for_loop)
-
-
-def unroll_for_over_modules(gr: "Graph", for_iter_node: "Node") -> None:
- gr.ensure_links()
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
- get_iter_node = for_iter_node.inputs[0].values[0].node
- assert get_iter_node.i.opname == "GET_ITER"
-
- iterated_module_list_parent, attr_lookups = find_method_through_phi_parent(get_iter_node.inputs[0])
- assert iterated_module_list_parent.value is not None
- iterated_module_list = iterated_module_list_parent.value
- for al in attr_lookups:
- iterated_module_list = getattr(iterated_module_list, al)
-
- # what about more complex things?
- assert isinstance(iterated_module_list, (torch.nn.Sequential, torch.nn.ModuleList))
-
- for_loop_len = len(iterated_module_list)
- for_iter_block = for_iter_node.block
- assert for_iter_block is not None
- get_iter_block = get_iter_node.block
-
- (iter_v,) = get_iter_node.outputs
- (iter_phi,) = for_iter_node.inputs
-
- assert isinstance(iter_phi, PhiValue)
- assert iter_v in iter_phi.values
-
- ### first we find the blocks of the for loop
- bls = find_blocks_of_for(gr, for_iter_block)
-
- jmp_nodes = {bl.nodes[-1] for bl in bls}
- assert all((v is iter_v or js in jmp_nodes) for v, js in zip(iter_phi.values, iter_phi.jump_sources))
-
- for_iter_node.i = get_instruction(opname="BINARY_SUBSCR", arg=None)
- iter_phi.remove_value(iter_v)
- assert len(iter_v.phi_values) == 0
- get_iter_block.block_outputs.remove(iter_v)
-
- get_iter_block.block_outputs.add(get_iter_node.inputs[0])
-
- seen = set()
-
- def delete_value_and_sources(v: Value) -> None:
- # check that it is possible?
- if v in seen:
- return
- seen.add(v)
- if isinstance(v, PhiValue):
- for vv, js in zip(v.values, v.jump_sources):
- delete_value_and_sources(vv)
- assert js is not None and js.block is not None
- js.block.block_outputs.remove(vv)
- v.block.block_inputs.remove(v)
-
- delete_value_and_sources(iter_phi)
- seq_phi = PhiValue(values=[get_iter_node.inputs[0]], jump_sources=[get_iter_block.nodes[-1]], block=for_iter_block)
- get_iter_block.nodes.remove(get_iter_node)
- for_iter_block.block_inputs.append(seq_phi)
-
- idx = Value(value=0, is_const=True)
- for_iter_node.inputs = [seq_phi, idx]
- for_iter_node.outputs = [for_iter_node.outputs[1]]
-
- for_iter_block_jmp = Node(i=get_instruction(opname="JUMP_ABSOLUTE", arg=None))
- for_iter_block_jmp.source_infos = copy.deepcopy(for_iter_node.source_infos)
- for_iter_block.nodes.append(for_iter_block_jmp)
- for_iter_block_jmp.jump_targets = [for_iter_node.jump_targets[0]]
- for_iter_node_exit_jump_target = for_iter_node.jump_targets[1]
- for_iter_node.jump_targets = []
- for_iter_block_jmp.jump_targets[0].jump_sources = [
- (js if js is not for_iter_node else for_iter_block_jmp)
- for js in for_iter_block_jmp.jump_targets[0].jump_sources
- ]
-
- exit_block = Block()
- gr.blocks.append(exit_block)
- exit_node = Node(i=get_instruction(opname="JUMP_ABSOLUTE", arg=None))
- exit_node.source_infos = copy.deepcopy(for_iter_node.source_infos)
- exit_node.jump_targets = [for_iter_node_exit_jump_target]
- target_after_iter = exit_node.jump_targets[0]
- exit_node.jump_targets[0].jump_sources = [
- (js if js is not for_iter_node else exit_node) for js in exit_node.jump_targets[0].jump_sources
- ]
- exit_block.nodes.append(exit_node)
- for i in for_iter_block.block_inputs:
- exit_block.block_inputs.append(PhiValue([], [], exit_block))
-
- unroll_blocks: list[tuple[list[Block], dict[GraphObject, GraphObject]]] = [(list(bls), {})]
- unroll_blocks += [clone_blocks(bls) for _ in range(1, for_loop_len)]
- for idx, (nbls, td) in enumerate(unroll_blocks):
- if idx > 0:
- gr.blocks += nbls
- v_idx = Value(value=idx, is_const=True)
- assert_node(td[for_iter_node]).inputs[1] = v_idx
- fin_o = assert_node(td[for_iter_node]).outputs[0]
- assert fin_o.name is not None
- fin_o.name += f"_{idx}"
- else:
- assert for_iter_node.outputs[0].name is not None
- for_iter_node.outputs[0].name += "_0"
-
- gr.ensure_links()
-
- fixup_data = []
- for idx, (nbls, td) in enumerate(unroll_blocks):
- if idx == 0:
- fib_i = for_iter_block
- jump_sources_to_fix = [js for js in for_iter_block.jump_sources if js is not get_iter_block.nodes[-1]]
- else:
- fib_i = assert_block(td[for_iter_block])
- jump_sources_to_fix = fib_i.jump_sources[:]
- if idx + 1 < len(unroll_blocks):
- _, td_next = unroll_blocks[idx + 1]
- fib_next = assert_block(td_next[for_iter_block])
- else:
- fib_next = exit_block
-
- fixup_data.append((fib_i, jump_sources_to_fix, fib_next, nbls))
-
- for idx_it, (fib_i, jump_sources_to_fix, fib_next, nbls) in enumerate(fixup_data):
- for js in jump_sources_to_fix:
- assert js is not None
- for idx, jt in enumerate(js.jump_targets):
- if jt == fib_i:
- js.set_jump_target(fib_next, idx=idx)
-
- for idx_i, i in enumerate(fib_i.block_inputs):
- if any((js.block in nbls) for js in i.jump_sources):
- ## if this is a variable updated in the loop:
- ## - instead of looping back, point the update to the phi value of the next block (or the exit block)
- ## - if idx > 0: remove external (before the loop) value
- for v, js in zip(i.values[:], i.jump_sources[:]):
- if js is not None and js.block not in nbls and idx_it > 0:
- i.remove_value(v)
-
- for idx_it, (fib_i, jump_sources_to_fix, fib_next, nbls) in enumerate(fixup_data):
- for idx_i, i in enumerate(fib_i.block_inputs):
- if any((js is not None and js.block in nbls) for js in i.jump_sources):
- for v, js in zip(i.values[:], i.jump_sources[:]):
- if js is not None and assert_block(assert_node(js).block) in nbls:
- i.remove_value(v)
- assert_block(fib_next).block_inputs[idx_i].add_missing_value(v, jump_source=js)
- if idx_it == 0:
- for pv in i.phi_values[:]:
- if pv.block is target_after_iter:
- pv.remove_value(i)
- pv.add_missing_value(exit_block.block_inputs[idx_i], jump_source=exit_node)
-
- for i in exit_block.block_inputs[:]:
- if i.phi_values:
- exit_block.block_outputs.add(i)
- else:
- assert isinstance(i, PhiValue)
- for v in i.values[:]:
- i.remove_value(v)
- exit_block.block_inputs.remove(i)
- if debug_asserts_enabled():
- thunder.core.script.graph.check_graph(gr)
-
-
-def find_and_unroll_for_loop(gr: "Graph") -> bool:
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
- gr.ensure_links()
-
- for bl in gr.blocks[:]:
- for n in bl.nodes[:]:
- if n.i.opname == "FOR_ITER":
- for_iter_node = n
- get_iter_node = for_iter_node.inputs[0].values[0].node
- if get_iter_node.i.opname == "GET_ITER":
- (
- iterated_module_list_parent,
- attr_lookups,
- ) = find_method_through_phi_parent(get_iter_node.inputs[0])
- if iterated_module_list_parent.value is None:
- continue
- iterated_module_list = iterated_module_list_parent.value
- for al in attr_lookups:
- iterated_module_list = getattr(iterated_module_list, al)
- # what about more complex things? in particular enumerate, but zip, ...
- if isinstance(iterated_module_list, (torch.nn.Sequential, torch.nn.ModuleList)):
- thunder.core.script.passes.unroll_for_over_modules(gr, for_iter_node)
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
- thunder.core.script.passes.merge_blocks_where_possible(gr)
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
- return True
- if debug_asserts_enabled():
- thunder.core.script.graph.check_graph(gr)
- return False
-
-
-def unroll_for_loops_and_inline_modules(gr: "Graph") -> None:
- if debug_asserts_level() > 1:
- thunder.core.script.graph.check_graph(gr)
- iterate = True
- while iterate:
- iterate = find_and_unroll_for_loop(gr)
- if not iterate:
- iterate = inline_submodule_calls(gr)
- if iterate:
- thunder.core.script.passes.merge_blocks_where_possible(gr)
-
-
-def module_to_function(gr: "Graph") -> tuple[list[str], list[torch.Tensor]]:
- attr_dict: dict[str, int] = {}
- attr_list: list[str] = []
- attr_values = []
- return_values: dict[str, Value] = {} # PhiValues in the return block
-
- if debug_asserts_enabled():
- thunder.core.script.graph.check_graph(gr)
-
- def functionalize_value_if_possible(i):
- # TODO: inefficient because it looks twice
- v = find_and_evaluate_method_through_phi_parent(i)
- # assert not isinstance(v, _Undefined), f"undefined: {v.value} {v.attr}"
- if isinstance(v, _Undefined):
- return Value(value=v, is_const=True)
- maybe_self, attrs = find_method_through_phi_parent(i)
-
- attr_string = ".".join(attrs)
- if maybe_self.value is gr.module and (isinstance(v, torch.Tensor) or (attr_string in return_values)):
- # the new attributes come directly after the self argument
- idx = attr_dict.setdefault(attr_string, len(attr_list) + 1)
- if idx == len(attr_list) + 1:
- func_arg = Value(name=attr_string, is_function_arg=True)
- gr.local_variables_at_start.insert(idx, func_arg)
- attr_list.append(attr_string)
- attr_values.append(v)
- gr.co_argcount += 1
- # we need a default argument to be able to put the things at the end (but this will have to change for *args, **kwargs anyway...
- # gr.func_defaults.append(None)
- if attr_string in return_values:
- return_values[attr_string].add_missing_value(func_arg)
- else:
- func_arg = gr.local_variables_at_start[idx]
-
- pvs = [pv for pv in func_arg.phi_values if pv.block is bl]
- if not pvs:
- pv = PhiValue([func_arg], [None], bl)
- bl.block_inputs.append(pv)
- else:
- (pv,) = pvs
- ## remove old input from phi_values etc?
- return pv
- if maybe_self.value is gr.module and (
- n.i.opname not in {"BINARY_SUBSCR"} and not isinstance(v, torch.nn.Module)
- ):
- ## inline to const...
- i.value = v
- i.typ = type(i.value)
- i.parent = None
- i.is_const = True
- i.is_global = False
- return None
- return None
-
- return_block = None
- for bl in gr.blocks:
- if bl.nodes[-1].i.opname == "RETURN_VALUE":
- assert return_block is None, "multiple return statements should not happen here"
- return_block = bl
- assert return_block is not None, "could not find return block"
-
- for bl in gr.blocks:
- for n in bl.nodes:
- if n.i.opname == "STORE_ATTR":
- v = find_and_evaluate_method_through_phi_parent(n.inputs[1])
- if isinstance(v, _Undefined):
- n.inputs[1] = Value(value=v, is_const=True)
- continue
- # assert not isinstance(v, _Undefined), f"undefined: {v.value} {v.attr}"
- maybe_self, attrs = find_method_through_phi_parent(n.inputs[1])
- attrs.append(n.i.argval)
- if maybe_self.value is gr.module:
- attr_string = ".".join(attrs)
- n.i = n.i.modify_copy(opname=X_THUNDER_STORE_ATTR, opcode=None, argval=attr_string)
- pv = return_values.get(attr_string)
- if pv is None:
- pv = PhiValue([], [], return_block)
- pv.name = attr_string
- return_values[attr_string] = pv
- return_block.block_inputs.append(pv)
- v = Value(node=n, name=attr_string, block=bl) # disambiguate?
- pv.add_missing_value(v, jump_source=bl.nodes[-1])
- n.outputs = [v]
- bl.block_outputs.add(v)
- del n.inputs[1]
-
- for bl in gr.blocks:
- for n in bl.nodes:
- if n.i.opname == "CALL_METHOD":
- if n.inputs[0].parent == n.inputs[1]:
- v = find_and_evaluate_method_through_phi_parent(n.inputs[0])
- if not isinstance(v, types.MethodType) or v.__self__ != find_and_evaluate_method_through_phi_parent(
- n.inputs[1]
- ):
- # this case (not a proper method call is usually handled in executing the LOAD_METHOD opcode)
- n.i = n.i.modify_copy(opname="CALL_FUNCTION", opcode=None)
- del n.inputs[1]
-
- for idx_i, i in enumerate(n.inputs):
- v = functionalize_value_if_possible(i)
- if v is not None:
- n.inputs[idx_i] = v
-
- bl.block_outputs = OrderedSet(
- [v if (v := functionalize_value_if_possible(o)) is not None else o for o in bl.block_outputs]
- )
-
- if return_values:
- bt_extra = Node(
- i=get_instruction(opname="BUILD_TUPLE", arg=1 + len(return_values)),
- source_infos=copy.deepcopy(return_block.nodes[-1].source_infos),
- )
- bt_extra.inputs = return_block.nodes[-1].inputs + list(return_values.values())
- v_tuple_extra = Value(node=bt_extra, block=return_block)
- bt_extra.outputs = [v_tuple_extra]
- return_block.nodes.insert(-1, bt_extra)
- return_block.nodes[-1].inputs = [v_tuple_extra]
-
- remove_unused_values(gr)
- if gr.local_variables_at_start[0].phi_values:
- gr.summary(print_lines=True)
- raise RuntimeError(
- """could not eliminate self argument
- this most likely means that you are setting attributes in forward or using them
- in an unexpected way that thunder does not yet support.
- The problem lies in (indirect) uses of V_0 in the graph above."""
- )
-
- # check to avoid assignments for both a.b and a.b.c
- sorted_keys = sorted(return_values.keys()) # this uses that '.' sorts before other things
- for i in range(len(sorted_keys) - 1):
- kbase = sorted_keys[i]
- knext = sorted_keys[i + 1]
- if knext.startswith(kbase) and knext[len(kbase)] == ".":
- # N.B. we know that knext is longer if kbase is a prefix so the knext[len(kbase)] above will not be out of bounds.
- raise RuntimeError(f"Assigning to members of assigned members ('{kbase}' and '{knext}') is not supported.")
-
- del gr.local_variables_at_start[0]
- gr.co_argcount -= 1
- if gr.co_posonlyargcount > 0:
- gr.co_posonlyargcount -= 1
-
- # thunder.core.script.graph.check_graph(gr)
- # gr.summary(print_lines=True)
-
- return attr_list, attr_values, list(return_values.keys())
diff --git a/thunder/core/script/protograph.py b/thunder/core/script/protograph.py
deleted file mode 100644
index 169e4bc60c..0000000000
--- a/thunder/core/script/protograph.py
+++ /dev/null
@@ -1,518 +0,0 @@
-from __future__ import annotations
-
-
-import abc
-import collections
-import dataclasses
-import functools
-import inspect
-import itertools
-from types import CodeType
-from typing import cast, overload, Any, Literal, NewType
-from collections.abc import Iterable, Iterator, Mapping
-
-from thunder.core.script import algorithms, instrumentation, parse, values
-from thunder.core.utils import debug_asserts_enabled, FrozenDict, OrderedSet
-
-__all__ = ("ProtoBlock", "ProtoGraph")
-
-# =============================================================================
-# == Inter-ProtoBlock abstract value flow =====================================
-# =============================================================================
-#
-# ProtoBlocks are weakly coupled by design. The `VariableKey` slots allow edges
-# to be deduced (e.g. `x` at the start of one block must be the same as `x` at
-# the end of the prior block), but there's no strong requirement. (And indeed,
-# the ProtoGraph immediately after parsing has all unconnected `AbstractRef`s
-# for input values.) Similarly, ProtoGraph serves only to record organize the
-# block topology, check invariants, and provide various helper methods.
-#
-# This weak coupling exists to facilitate graph rewrites and reduce the surface
-# area for self-inconsistent representation. By readily discarding (deduced)
-# information we don't need to carry invariants through complex passes; we can
-# simply decouple the graph, perform whatever local modifications we like, and
-# then reconnect everything. This representation is immutable (notwithstanding
-# a few implementation details), so "decouple" means emitting a new erased
-# graph. (Though simple value replacements can be done directly.)
-JumpTarget = NewType("JumpTarget", tuple["ProtoBlock", parse.Jump])
-Uses = NewType("Uses", OrderedSet[parse.VariableKey])
-
-
-@dataclasses.dataclass(frozen=True, eq=False)
-class ProtoBlock(instrumentation.InstrumentingBase): # type: ignore[misc,no-any-unimported]
- """Stores abstract data flow for a code block."""
-
- flow: values.IntraBlockFlow
- jump_targets: tuple[JumpTarget, ...] = dataclasses.field(default=(), init=False)
- uses: Uses = dataclasses.field(default_factory=lambda: Uses(OrderedSet()), init=False)
-
- def __repr__(self) -> str:
- ops = "\n".join(f" {i.opname}" for i, _ in self.flow.symbolic)
- return f"ProtoBlock: {hex(id(self))}\n{ops}"
-
- def __hash__(self) -> int:
- return id(self)
-
- def __post_init__(self) -> None:
- self.uses.update(self.flow.uses)
-
- def add_jump_target(self, other: ProtoBlock, jump: parse.Jump) -> None:
- """We need to add jump targets after all ProtoBlocks are initialized."""
-
- # Override `frozen=True` for this one limited use case.
- object.__setattr__(self, "jump_targets", self.jump_targets + ((other, jump),))
-
-
-@dataclasses.dataclass(frozen=True, eq=False)
-class ProtoGraph:
- protoblocks: tuple[ProtoBlock, ...]
- root: ProtoBlock
- parents: Mapping[ProtoBlock, tuple[ProtoBlock, ...]]
-
- Provenance = values.ParsedSymbolic | tuple[type["ProtoGraphTransform"], "ProtoGraph"]
- provenance: Provenance
-
- def __init__(self, protoblocks: Iterable[ProtoBlock], provenance: Provenance) -> None:
- G = algorithms.TypedDiGraph[ProtoBlock]()
- for protoblock in (protoblocks := tuple(protoblocks)):
- is_return = tuple(protoblock.flow.symbolic)[-1][0].opname == parse.RETURN_VALUE
- G.add_node(protoblock, is_return=is_return)
-
- for protoblock in protoblocks:
- for destination, jump in protoblock.jump_targets:
- G.add_edge(protoblock, destination, adjacent=not jump)
-
- assert protoblocks
- object.__setattr__(self, "protoblocks", tuple(algorithms.sort_adjacent(G)))
- assert len(G) == len(self.protoblocks) == len(protoblocks), (len(G), len(self.protoblocks), len(protoblocks))
-
- object.__setattr__(self, "root", self.protoblocks[0])
- root_stack = [(k, v) for k, v in self.root.flow.begin_state if k.scope == parse.VariableScope.STACK]
- assert not root_stack, f"Root block should not have stack inputs: {root_stack}"
-
- nodes = cast(Iterable[ProtoBlock], G.nodes) # For some reason mypy needs this.
- parents = {protoblock: tuple(G.predecessors(protoblock)) for protoblock in nodes}
- object.__setattr__(self, "parents", FrozenDict(parents))
- object.__setattr__(self, "provenance", provenance)
-
- @classmethod
- def from_code(cls, co: CodeType) -> ProtoGraph:
- """Given a method, disassemble it to a sequence of simple blocks."""
- parsed = values.ParsedSymbolic.make(parse.ParsedFunctional.make(co))
- protoblocks = tuple(
- ProtoBlock(values.IntraBlockFlow(symbolic, begin, end)) for symbolic, begin, end in parsed.blocks
- )
- for source, sink, jump in parsed.provenance.provenance.edges:
- protoblocks[source].add_jump_target(protoblocks[sink], jump)
-
- return cls(protoblocks, parsed)
-
- def __iter__(self) -> Iterator[ProtoBlock]:
- yield from self.protoblocks
-
- def __getitem__(self, index: int) -> ProtoBlock:
- return self.protoblocks[index]
-
- def __len__(self) -> int:
- return len(self.protoblocks)
-
- def __repr__(self) -> str:
- return "\n\n".join(repr(protoblock) for protoblock in self)
-
- @property
- def flat_flow(self) -> Iterable[tuple[parse.ThunderInstruction, values.Symbolic, values.Materialized]]:
- for protoblock in self:
- for instruction, symbolic in protoblock.flow.symbolic:
- yield instruction, symbolic, protoblock.flow.materialized[instruction]
-
- @property
- def is_linked(self) -> bool:
- # NOTE: `is_linked` is vacuously True for a single block graph.
- flat_begin = itertools.chain(*(i.flow._begin.values() for i in self if i is not self.root))
- return len(self) == 1 or any(not isinstance(i, values.AbstractRef) for i in flat_begin)
-
- def unlink(self) -> ProtoGraph:
- return Unlink(self).apply(or_default=True)
-
- def link(self) -> ProtoGraph:
- if result := ProtoGraphTransform.chain(self, AddTransitive, MatchStacks, Connect):
- assert AddTransitive(result).apply(or_default=False) is None
- assert (result or self).is_linked
- return result or self
-
- def debug_print_protoflows(self) -> None:
- """
- Print out the node_flow for each protoblock in the
- protograph, in a way that's nice to read and debug with.
- """
-
- counter = 0
- idxes: dict[values.AbstractValue, int] = {}
- for pb in self:
- for node in pb.flow.materialized.values():
- for val in itertools.chain(node.inputs.ordered, node.outputs):
- if val not in idxes.keys():
- idxes[val] = counter
- counter += 1
-
- def to_index_str(values: tuple[values.AbstractValue, ...]) -> str:
- indices = (str(idxes[v]) for v in values)
- return f"({', '.join(indices)})"
-
- for i, pb in enumerate(self):
- print(f"Protoblock {i}:")
- print(f"{'':>22}Inputs, Outputs")
- for instruction, node in pb.flow.materialized.items():
- print(f" {instruction.opname:>20}, {to_index_str(node.inputs.ordered)} -> {to_index_str(node.outputs)}")
- print("\n")
-
-
-# =============================================================================
-# == Graph transforms (Base classes) ==========================================
-# =============================================================================
-# ProtoGraphTransform
-# ReplaceProtoBlocks
-# ReplaceValues
-# CondenseValues
-# ReplaceSymbolic
-
-
-class ProtoGraphTransform(abc.ABC):
- """Handles mechanical portions of graph rewrites.
- The base case is unopinionated; it simply accepts whatever new ProtoGraph is
- emitted by `self._apply`. The primary feature it provides is checking.
-
- NOTE:
- The convention adopted is for the pass logic to produce `T | None`
- (e.g. `ProtoGraph | None`) where `None` signals that no change is
- applicable.
-
- Forbid Linked:
- A key invariant of ProtoGraph is that every AbstractValue has exactly
- **one** producer, which is set by the symbolic flow. (With the exception
- of `AbstractRef`s which are placeholders for an as-yet unspecified
- AbstractValue.) However, within a block there is a flat list of
- **concrete** values specifying the state at the start of the block.
-
- If one were to replace all instances of `X` in a ProtoGraph with `Y`,
- this invariant would be preserved. On the other hand, if one were to
- replace `X` with `Y` **only at the symbolic producer of `X`** then
- downstream blocks could still have `X` as a block input, despite the
- fact that `X` no longer has a producer. (Note that this is only a problem
- across blocks; within blocks the materialization pass respects the update
- and emits a consistent materialized state for the new ProtoBlock.)
-
- It is often convenient to simply rewrite the symbolic flow within a
- single ProtoBlock. In that case the correct procedure is to generate an
- unlinked ProtoGraph, perform the local rewrites, and then relink it.
- (Where the connection pass will handle reconciliation automatically.)
-
- Check idempotence:
- Nearly all passes are expected to be idempotent. This provides a good
- deal of free test coverage since it produces both a test case (the result
- of `self._apply`) and an expected result. (That `self._apply` returns
- `None`.) We perform this check many times in order to flush out
- non-deterministic passes. (Though the value is configurable if a pass is
- particularly expensive.)
-
- However, given the potential added start up latency and possibility of
- spurious failures this check is gated by `debug_asserts_enabled`, which
- defaults to `False`. (Except for unit tests.)
- """
-
- _forbid_linked: bool = False
- _kwargs: FrozenDict[str, Any] # Used to replay transform for `idempotent` check.
- _idempotent_repeats: int = 10
-
- @abc.abstractmethod
- def _apply(self) -> ProtoGraph | None:
- """Override this method to emit an (optional) new ProtoGraph."""
- ...
-
- def __new__(cls, *args: Any, **kwargs: Any) -> ProtoGraphTransform:
- self = super().__new__(cls)
- bound = inspect.signature(self.__class__.__init__).bind(None, *args, **kwargs).arguments
- bound.pop("self")
- bound.pop("proto_graph")
- self._kwargs = FrozenDict(bound)
- return self
-
- def __init__(self, proto_graph: ProtoGraph) -> None:
- assert not (self._forbid_linked and len(proto_graph) > 1 and proto_graph.is_linked), self
- assert isinstance(proto_graph, ProtoGraph)
- self._protograph = proto_graph
-
- @property
- def protograph(self) -> ProtoGraph:
- return self._protograph
-
- @overload
- def apply(self, or_default: Literal[False]) -> ProtoGraph | None:
- ...
-
- @overload
- def apply(self, or_default: Literal[True]) -> ProtoGraph:
- ...
-
- def apply(self, or_default: bool = False) -> ProtoGraph | None:
- result = self._apply()
- if debug_asserts_enabled():
- result_to_check = result or self.protograph
- for i in range(self._idempotent_repeats):
- assert self.__class__(proto_graph=result_to_check, **self._kwargs)._apply() is None, (i, self)
- return result or (self.protograph if or_default else None)
-
- @staticmethod
- def chain(proto_graph: ProtoGraph, *transforms: type[ProtoGraphTransform]) -> ProtoGraph | None:
- initial = proto_graph
- for transform in transforms:
- proto_graph = transform(proto_graph).apply(or_default=True)
- return None if proto_graph is initial else proto_graph
-
-
-class ReplaceProtoBlocks(ProtoGraphTransform):
- """Helper to replace individual ProtoBlocks while retaining the same ProtoGraph topology."""
-
- @abc.abstractmethod
- def apply_to_protoblock(self, protoblock: ProtoBlock) -> values.IntraBlockFlow | None:
- ...
-
- def post_apply(self, old: ProtoBlock, new: ProtoBlock) -> None:
- pass
-
- def _apply(self) -> ProtoGraph | None:
- # TODO(robieta): Right now block order is load bearing, so we have to preserve it.
- transformed = {i: self.apply_to_protoblock(i) for i in self.protograph}
-
- if any(transformed.values()):
- replacements = {i: ProtoBlock(flow or i.flow) for i, flow in transformed.items()}
- for old_protoblock, new_protoblock in replacements.items():
- self.post_apply(old_protoblock, new_protoblock)
- for old_target, is_jump in old_protoblock.jump_targets:
- new_protoblock.add_jump_target(replacements[old_target], is_jump)
- return ProtoGraph(replacements.values(), provenance=(self.__class__, self.protograph))
- return None
-
-
-class ReplaceValues(ReplaceProtoBlocks):
- """Copies the ProtoGraph with value replacements.
-
- NOTE: This is strictly a condensing transform, and this is only invertible
- (using another `ReplaceValues`) in trivial cases.
- """
-
- _retain_uses: bool = True
-
- @abc.abstractproperty
- def replace_map(self) -> values.ReplaceMap:
- ...
-
- @functools.cached_property
- def _replace_map(self) -> values.ReplaceMap:
- replace_map = self.replace_map
- assert not (invalid := [k for k in replace_map if isinstance(k, values.NonPyObject)]), invalid
- return FrozenDict(algorithms.flatten_map(replace_map))
-
- def apply_to_protoblock(self, protoblock: ProtoBlock) -> values.IntraBlockFlow | None:
- return protoblock.flow.substitute(self._replace_map)
-
- def post_apply(self, old: ProtoBlock, new: ProtoBlock) -> None:
- if self._retain_uses:
- new.uses.update(old.uses)
-
-
-class CondenseValues(ReplaceValues):
- ValueEdges = Iterable[tuple[values.AbstractValue, values.AbstractValue]]
-
- @abc.abstractproperty
- def edges(self) -> ValueEdges:
- ...
-
- @property
- def replace_map(self) -> values.ReplaceMap:
- replace_map: dict[values.AbstractValue, values.AbstractValue] = {}
- edges = itertools.chain(self.edges, self._phivalue_constituent_edges)
- for v, condensed in algorithms.compute_condense_map(edges).items():
- # Check invariants.
- assert condensed
- if not isinstance(v, values.AbstractPhiValue):
- invariants = ({c.identity for c in condensed} == {v.identity}, not isinstance(v, values.AbstractRef))
- assert all(invariants) or not any(invariants), (invariants, v, condensed)
-
- # `AbstractPhiValue._unpack_apply` will determine if we need an AbstractPhiValue.
- if (replacement := values.substitute_value(values.AbstractPhiValue(tuple(condensed)), {})) != v:
- replace_map[v] = replacement
-
- return FrozenDict(replace_map)
-
- @property
- def _phivalue_constituent_edges(self) -> ValueEdges:
- # AbstractPhiValues are somewhat unusual in that mismatches between blocks
- # are expected (that's sort of the point...) so we need to decompose them
- # so the condense pass doesn't get tripped up.
- for _, initial_ref in self.protograph.root.flow.begin_state:
- if isinstance(initial_ref, values.AbstractPhiValue):
- yield from ((constituent, initial_ref) for constituent in initial_ref.constituents)
-
-
-class ReplaceSymbolic(ReplaceProtoBlocks):
- _forbid_linked = True
-
- @abc.abstractmethod
- def apply_to_symbolic(
- self,
- instruction: parse.ThunderInstruction,
- symbolic: values.Symbolic,
- inputs: values.HybridMap[values.AbstractValue],
- ) -> values.Symbolic | None:
- ...
-
- def apply_to_protoblock(self, protoblock: ProtoBlock) -> values.IntraBlockFlow | None:
- flow_state = values.DigestFlow(protoblock.flow._begin)
- updated_symbolic: dict[parse.ThunderInstruction, values.Symbolic | None] = {}
- for i, symbolic in protoblock.flow.symbolic:
- updated_symbolic[i] = self.apply_to_symbolic(i, symbolic, symbolic.inputs.map(flow_state.get))
- _ = flow_state.next(i, updated_symbolic[i] or symbolic)
-
- if any(updated_symbolic.values()):
- new_symbolic = {k: v or protoblock.flow._symbolic[k] for k, v in updated_symbolic.items()}
- return dataclasses.replace(protoblock.flow, _symbolic=FrozenDict(new_symbolic))
- return None
-
-
-# =============================================================================
-# == Graph transforms (Applied) ===============================================
-# =============================================================================
-class Unlink(ReplaceProtoBlocks):
- def apply_to_protoblock(self, protoblock: ProtoBlock) -> values.IntraBlockFlow | None:
- if protoblock is not self.protograph.root:
- uses = (flow := protoblock.flow).uses.copy()
- end: values.Symbolic.EndState = FrozenDict({k: v for k, v in flow._end.items() if k != v})
- uses.update(v for v in end.values() if isinstance(v, parse.VariableKey) and not v.is_const)
- any_non_ref = any(not isinstance(i, values.AbstractRef) for i in flow._begin.values())
- if any_non_ref or len(end) < len(flow._end) or flow._begin.keys() ^ uses: # symmetric_difference
- begin: FrozenDict[parse.VariableKey, values.AbstractValue]
- begin = FrozenDict({k: values.AbstractRef(f"Unlink: {k}") for k in uses})
- return dataclasses.replace(protoblock.flow, _begin=begin, _end=end)
-
- return None
-
- def _apply(self) -> ProtoGraph | None:
- result = super()._apply()
- assert len(result or self.protograph) == 1 or not (result or self.protograph).is_linked, result
- return result
-
-
-class AddTransitive(ReplaceProtoBlocks):
- """Extend abstract value flows to include those needed by downstream blocks.
- This pass effectively functionalizes the abstract value flow by plumbing
- reads through parents as transitive dependencies. Note that we assume
- variables are only modified by `STORE_...` and `DELETE_...` instructions.
- This is not a sound assumption since opaque calls (`CALL_FUNCTION`,
- `CALL_METHOD`, etc.) could mutate global and nonlocal variables. This does
- not, however, pose an overall soundness problem because we can check for
- state mutations during inlining and rerun flow analysis.
- """
-
- def apply_to_protoblock(self, protoblock: ProtoBlock) -> values.IntraBlockFlow | None:
- flow = protoblock.flow
- end = FrozenDict({**{use: use for use in self.target_uses(protoblock, self.expanded_uses)}, **flow._end})
- if (missing := self.expanded_uses[protoblock].difference(protoblock.uses)) or (end != flow._end):
- begin = {**{k: values.AbstractRef("Transitive") for k in missing}, **flow._begin}
- return dataclasses.replace(flow, _begin=FrozenDict(begin), _end=FrozenDict(end))
- return None
-
- def post_apply(self, old: ProtoBlock, new: ProtoBlock) -> None:
- new.uses.update(self.expanded_uses[old])
-
- @functools.cached_property
- def expanded_uses(self) -> Mapping[ProtoBlock, Uses]:
- """Identify new transitive value dependencies.
- The process is more involved than simply checking for mismatches because
- adding a transitive value to a block may necessitate adding a transitive
- value to the prior block and so on.
- """
- uses = {protoblock: protoblock.uses.copy() for protoblock in self.protograph}
- blocks_to_process = collections.deque(uses.keys())
-
- while blocks_to_process:
- protoblock = blocks_to_process.popleft()
- target_uses = self.target_uses(protoblock, uses)
-
- # The reason we can ignore ALL `_OutputRef`s (including those that would index into a composite)
- # is that the (potential) composite's dependencies are already handled by `ProtoBlock._flow_uses`.
- transitive_uses = OrderedSet(
- source
- for use in target_uses
- if isinstance(source := protoblock.flow._end.get(use, use), parse.VariableKey)
- and source.scope != parse.VariableScope.CONST
- )
-
- if transitive_uses - uses[protoblock]:
- uses[protoblock].update(transitive_uses)
- blocks_to_process.extend(self.protograph.parents[protoblock])
-
- return FrozenDict(uses)
-
- @staticmethod
- def target_uses(protoblock: ProtoBlock, uses: Mapping[ProtoBlock, Uses] = FrozenDict()) -> Uses:
- flat_uses = itertools.chain(*(uses.get(target, target.uses) for target, _ in protoblock.jump_targets))
- return Uses(OrderedSet(use for use in flat_uses if use.scope != parse.VariableScope.CONST))
-
-
-class MatchStacks(ReplaceProtoBlocks):
- """Ensure stacks match across blocks.
-
- ProtoGraph doesn't rely on stack behavior (push, pop TOS, etc.), however it
- is still a good sanity check. (Which is why `Connect._inter_block_edges` asserts.)
- """
-
- def apply_to_protoblock(self, protoblock: ProtoBlock) -> values.IntraBlockFlow | None:
- upstream = OrderedSet[parse.VariableKey]()
- for parent in self.protograph.parents[protoblock]:
- upstream.update(k for k in parent.flow._end if k.scope == parse.VariableScope.STACK)
-
- if delta := upstream - protoblock.flow._begin:
- begin = {**protoblock.flow._begin, **{k: values.AbstractRef(f"Match stack: {k}") for k in delta}}
- return dataclasses.replace(protoblock.flow, _begin=FrozenDict(begin))
- return None
-
- def post_apply(self, old: ProtoBlock, new: ProtoBlock) -> None:
- new.uses.update(old.uses)
-
-
-class Connect(CondenseValues):
- @property
- def edges(self) -> CondenseValues.ValueEdges:
- yield from self._inter_block_edges(self.protograph)
- yield from self._graph_input_edges(self.protograph)
-
- @staticmethod
- def _graph_input_edges(proto_graph: ProtoGraph) -> CondenseValues.ValueEdges:
- for key, initial_ref in proto_graph.root.flow.begin_state:
- if isinstance(initial_ref.identity, values.ExternalRef):
- continue
-
- assert isinstance(initial_ref, values.AbstractRef), initial_ref
- assert key.scope not in (
- parse.VariableScope.CONST,
- parse.VariableScope.STACK,
- ), (key, proto_graph.root.flow._begin)
- yield values.CompositeValue().add_identity(values.ExternalRef(key)), initial_ref
-
- @staticmethod
- def _inter_block_edges(proto_graph: ProtoGraph) -> CondenseValues.ValueEdges:
- for protoblock in proto_graph:
- for child, _ in protoblock.jump_targets:
- outputs = dict(protoblock.flow.end_state)
- child_inputs = dict(child.flow.begin_state)
- for key, child_input in child_inputs.items():
- yield outputs.get(key, values.NonPyObject(values.NonPyObject.Tag.MISSING)), child_input
-
- # `AddTransitive` should ensure the stacks match.
- # (Except for return blocks which may discard the stack.)
- opname = tuple(child.flow.symbolic)[-1][0].opname
- if opname not in parse.RAISE_RETURN_INSTRUCTIONS:
- s_out = tuple(sorted(i.identifier for i in outputs if i.scope == parse.VariableScope.STACK))
- s_in = tuple(sorted(i.identifier for i in child_inputs if i.scope == parse.VariableScope.STACK))
- assert s_out == s_in, f"{s_out=} != {s_in=}, {opname}"
diff --git a/thunder/core/script/protograph_passes.py b/thunder/core/script/protograph_passes.py
deleted file mode 100644
index 657d6d6c8f..0000000000
--- a/thunder/core/script/protograph_passes.py
+++ /dev/null
@@ -1,74 +0,0 @@
-import dataclasses
-from collections.abc import Iterable
-
-from thunder.core.script import parse, values
-from thunder.core.script.protograph import ProtoGraph, ProtoGraphTransform, AddTransitive, ReplaceSymbolic
-from thunder.core.utils import debug_asserts_enabled
-
-ValueEdges = Iterable[tuple[values.AbstractValue, values.AbstractValue]]
-KNOWN_TUPLE = values.TraitName("__known_tuple")
-
-
-def _connect_protograph(proto_graph: "ProtoGraph") -> "ProtoGraph":
- proto_graph = proto_graph.link()
- assert AddTransitive(proto_graph).apply(or_default=False) is None
- for protoblock in proto_graph:
- for k, v in protoblock.flow.begin_state:
- assert not v.is_detail, (k, v)
- return proto_graph
-
-
-class MarkTuples(ReplaceSymbolic):
- def apply_to_symbolic(
- self,
- instruction: parse.ThunderInstruction,
- symbolic: values.Symbolic,
- _: values.HybridMap[values.AbstractValue],
- ) -> values.Symbolic | None:
- if instruction.opname == "BUILD_TUPLE" and isinstance(output := symbolic.outputs[0], values.IntermediateValue):
- assert len(symbolic.outputs) == 1, symbolic.outputs
- ordered = tuple(values.Reference(i) for i in range(-len(symbolic.inputs.ordered), 0))
- new_output = values.CompositeRef(ordered=ordered).add_named(KNOWN_TUPLE, values.ConstRef(True))
- return dataclasses.replace(symbolic, outputs=(new_output.add_identity(output),))
- return None
-
-
-class IndexTuples(ReplaceSymbolic):
- def apply_to_symbolic(
- self,
- instruction: parse.ThunderInstruction,
- symbolic: values.Symbolic,
- inputs: values.HybridMap[values.AbstractValue],
- ) -> values.Symbolic | None:
- replacement: values.Symbolic | None = None
- if instruction.opname == "BINARY_SUBSCR":
- to_index, index = inputs.ordered
- is_tuple = isinstance(to_index, values.CompositeValue) and to_index.get(KNOWN_TUPLE)
- index_key = index.key if isinstance(index, values.ExternalRef) and index.key.is_const else None
- if is_tuple and index_key and isinstance(idx := index_key.identifier, int):
- assert len(symbolic.outputs) == 1
- replacement = dataclasses.replace(symbolic, outputs=((values.Reference(0), values.Reference(idx)),))
-
- elif instruction.opname == "UNPACK_SEQUENCE":
- (to_unpack,) = inputs.ordered
- if isinstance(to_unpack, values.CompositeValue) and to_unpack.get(KNOWN_TUPLE):
- indices = (values.Reference(idx) for idx in range(-1, -len(symbolic.outputs) - 1, -1))
- outputs = tuple((values.Reference(0), idx) for idx in indices)
- replacement = dataclasses.replace(symbolic, outputs=outputs)
-
- elif instruction.opname == "UNPACK_EX":
- pass # TODO(apaz-cli): figure out indexing.
-
- return replacement if (replacement and replacement.outputs != symbolic.outputs) else None
-
-
-def _tuple_fold(proto_graph: ProtoGraph) -> ProtoGraph:
- """Replace tuple accesses (`BINARY_SUBSCR`, `UNPACK_SEQUENCE` instructions) with their members, if known."""
- return ProtoGraphTransform.chain(proto_graph, MarkTuples, IndexTuples) or proto_graph
-
-
-def apply_protograph_passes(protograph: ProtoGraph) -> ProtoGraph:
- protograph = _tuple_fold(protograph.unlink())
- protograph = _connect_protograph(protograph)
- assert AddTransitive(protograph).apply(or_default=False) is None
- return protograph
diff --git a/thunder/core/script/python_ir.py b/thunder/core/script/python_ir.py
deleted file mode 100644
index 2d66fbb77b..0000000000
--- a/thunder/core/script/python_ir.py
+++ /dev/null
@@ -1,507 +0,0 @@
-import collections
-import dis
-import inspect
-import sys
-import types
-from typing import Any, Dict, List, Optional, Tuple, Union
-from collections.abc import Callable
-from collections.abc import Hashable
-
-from thunder.core.script.graph import (
- assert_block,
- _generate_raises,
- Graph,
- GraphSummaryCallback,
- MROAwareObjectRef,
- Node,
- SourceInformation,
- Value,
- insert_before,
- insert_after,
- _Undefined,
-)
-from thunder.core.script.parse import RETURN_VALUE
-from thunder.core.script.python_ir_data import get_instruction, X_THUNDER_STORE_ATTR
-from thunder.core.utils import OrderedSet
-
-
-def undo_ssa(gr: "Graph") -> tuple[list[Value], list[str], list[str], list[Any]]:
- consts: list[Any] = []
- names: list[str] = []
-
- def get_value(v: Value, n: Node, inpidx: int | None = None) -> None:
- if n.i.opname == "CALL_METHOD" and inpidx == 1:
- bl = assert_block(n.block)
- idx = bl.nodes.index(n)
- if idx > 0 and bl.nodes[idx - 1].i.opname == "LOAD_METHOD":
- # if we just a LOAD_METHOD, that did put input 0 and 1 on the stack
- return
- else:
- # else the loading has been separated from the call, so we
- # switch to call LOAD_ATTR/CALL_FUNCTION instead
- n.i = n.i.modify_copy(opname="CALL_FUNCTION", opcode=None)
- return
- if isinstance(v.value, _Undefined):
- idx = len(consts)
- consts.append(
- _generate_raises(f"attribute error '{type(v.value.value)}' object has no attribute '{v.value.attr}'")
- )
- new_n = Node(
- i=get_instruction(opname="LOAD_CONST", arg=idx),
- outputs=[Value(value=consts[idx], is_const=True)],
- inputs=[],
- )
- new_n.inserted_for = n
- insert_before(new_n, n)
- new_n = Node(i=get_instruction(opname="CALL_FUNCTION", arg=0), outputs=[v], inputs=[consts[idx]])
- new_n.inserted_for = n
- insert_before(new_n, n)
- return
- if v.is_const:
- idx = len(consts)
- consts.append(v.value)
- new_n = Node(i=get_instruction(opname="LOAD_CONST", arg=idx), outputs=[v], inputs=[])
- new_n.inserted_for = n
- insert_before(new_n, n)
- elif isinstance(v.value, MROAwareObjectRef):
- # this works for attribs, but for methods? maybe have a pass eliminating/making explicit the super...
- get_value(v.value.obj, n)
- elif v.parent is not None:
- assert v.name is not None
- get_value(v.parent, n)
- if n.i.opname == "CALL_METHOD" and inpidx == 0:
- # print("###inputs", n.inputs, v, v in n.inputs)
- try:
- idx = names.index(v.name)
- except ValueError:
- idx = len(names)
- names.append(v.name)
- new_n = Node(
- i=get_instruction(opname="LOAD_METHOD", arg=idx),
- outputs=[v, v.parent],
- inputs=[v.parent],
- )
- new_n.inserted_for = n
- insert_before(new_n, n)
- elif n.i.opname == "LOAD_ATTR":
- # print("###load attr", n.outputs, n.i.argval)
- pass
- else:
- assert v.name is not None
- try:
- idx = names.index(v.name)
- except ValueError:
- idx = len(names)
- names.append(v.name)
- new_n = Node(
- i=get_instruction(opname="LOAD_ATTR", arg=idx),
- outputs=[v],
- inputs=[v.parent],
- )
- new_n.inserted_for = n
- insert_before(new_n, n)
- elif v.is_global and isinstance(v.value, Hashable) and v.value in __builtins__:
- # Builtins are unmarshallable and meant to be loaded globally. If they are
- # included in co_consts, the resulting function cannot go into a .pyc file.
- # Originally, the plan was to check if the value is a builtin by checking
- # if its type is "". However, this
- # turned out not to work since torch for some reason decided to set the
- # type of `torch.nn.functional.has_torch_function` to also be a builtin.
- if v.name not in names:
- names.append(v.name)
- idx = names.index(v.name)
- new_n = Node(i=get_instruction(opname="LOAD_GLOBAL", arg=idx), outputs=[v], inputs=[])
- new_n.inserted_for = n
- insert_before(new_n, n)
- elif v.is_global: # make binding the globals optional?
- if v.value not in consts:
- consts.append(v.value)
- idx = consts.index(v.value)
- new_n = Node(i=get_instruction(opname="LOAD_CONST", arg=idx), outputs=[v], inputs=[])
- new_n.inserted_for = n
- insert_before(new_n, n)
- else:
- idx = local_vars[v]
- # assert idx >= 0
- new_n = Node(i=get_instruction(opname="LOAD_FAST", arg=idx), outputs=[v], inputs=[])
- new_n.inserted_for = n
- insert_before(new_n, n)
-
- for bl in gr.blocks:
- for n in bl.nodes:
- n.block = bl
-
- local_vars: dict[Value, int] = {}
- lv_names: OrderedSet[str] = OrderedSet()
-
- def get_or_add_lv(v: Value, name: str | None = None) -> int:
- idx = local_vars.get(v)
- if idx is None:
- idx = len(local_vars)
- local_vars[v] = idx
-
- # handle name collisions...
- if name is None:
- name = v.name
-
- if name is None:
- name = f"_tmp_{idx}"
- else:
- name = name.replace(".", "_").replace("[", "").replace("]", "")
-
- if not name[:1].isalpha():
- name = "_" + name
- fullname = name
- suffix = 0
- while fullname in lv_names:
- suffix += 1
- fullname = f"{name}_{suffix}"
- lv_names.add(fullname)
- if v.name is None: # TODO: or do this always?
- v.name = fullname
- return idx
-
- nodes_to_skip = set()
-
- def store_phi_values(o: Value, o_idx: int, last_n: Node | None, cur_n: Node | None) -> Node | None:
- phi_values_in_processing = set()
-
- def store_phi_values_inner(o: Value, o_idx: int, last_n: Node | None) -> Node | None:
- if o in phi_values_in_processing:
- # avoid loops
- return last_n
- phi_values_in_processing.add(o)
- for v in o.phi_values:
- # TODO: refactor into general mechanism
- idx2 = get_or_add_lv(v)
- # last_n = store_phi_values_inner(v, o_idx, last_n)
- if o.is_const:
- if o.value not in consts:
- consts.append(o.value)
- o_idx = consts.index(o.value)
- new_n = Node(i=get_instruction(opname="LOAD_CONST", arg=o_idx), outputs=[o], inputs=[])
- new_n.inserted_for = cur_n
- else:
- new_n = Node(i=get_instruction(opname="LOAD_FAST", arg=o_idx), outputs=[o], inputs=[])
- new_n.inserted_for = cur_n
-
- nodes_to_skip.add(new_n)
- if last_n is None:
- insert_before(new_n, gr.blocks[0].nodes[0])
- else:
- insert_after(new_n, last_n)
- last_n = new_n
- new_n = Node(i=get_instruction(opname="STORE_FAST", arg=idx2), outputs=[], inputs=[o])
- new_n.inserted_for = cur_n
- nodes_to_skip.add(new_n)
- insert_after(new_n, last_n)
- last_n = new_n
- return last_n
-
- return store_phi_values_inner(o, o_idx, last_n)
-
- for v in gr.local_variables_at_start:
- if v is not None:
- get_or_add_lv(v)
-
- # inputs in phi values
- last_n = None
- # need to make a copy of the list because we're adding items to the list
- for idx, i in enumerate(tuple(local_vars.keys())):
- last_n = store_phi_values(i, idx, last_n, cur_n=None)
- for i in gr.blocks[0].block_inputs: # inlined parameters (partial) will be here
- for v, js in zip(i.values, i.jump_sources):
- if js is None and v.is_const:
- last_n = store_phi_values(v, None, last_n, cur_n=None)
- # print(i.values, i.jump_sources)
-
- names = []
-
- for bl in gr.blocks:
- jump_node = bl.nodes[-1]
- for n in bl.nodes[:]:
- processed_block_outputs = set()
- if n not in nodes_to_skip:
- n.inserted_for = n
- for inpidx, i in enumerate(n.inputs):
- get_value(i, n=n, inpidx=inpidx)
- last_n = n
- for o in n.outputs[::-1]:
- idx = get_or_add_lv(o)
- new_n = Node(
- i=get_instruction(opname="STORE_FAST", arg=idx),
- outputs=[],
- inputs=[o],
- )
- new_n.inserted_for = n
- assert last_n is not None
- insert_after(new_n, last_n)
- last_n = new_n
- if o in bl.block_outputs:
- processed_block_outputs.add(o)
- last_n = store_phi_values(o, idx, last_n, cur_n=n)
- if n.i.opname in ("STORE_ATTR", "IMPORT_NAME"): # STORE_ATTR for unknown objs
- # have a utility for this?
- try:
- idx = names.index(n.i.argval)
- except ValueError:
- idx = len(names)
- names.append(n.i.argval)
- n.i = n.i.modify_copy(arg=idx)
- if n.i.opname == X_THUNDER_STORE_ATTR:
- bl.nodes.remove(n)
- if bl.nodes[-1].i.opname != RETURN_VALUE: # TODO Should the return block have outputs (probably not)
- for o in bl.block_outputs:
- if o not in processed_block_outputs:
- get_value(o, n=jump_node) # before the jump
- idx = get_or_add_lv(o, name="bo")
- new_n = Node(
- i=get_instruction(opname="STORE_FAST", arg=idx),
- outputs=[],
- inputs=[o],
- )
- new_n.inserted_for = jump_node
- insert_before(new_n, n=jump_node)
- store_phi_values(o, idx, new_n, cur_n=jump_node)
-
- return list(local_vars.keys()), list(lv_names), names, consts
-
-
-# this function is taken from PyTorch Dynamo (c) 2022 by Facebook/Meta licensed
-# as per https://github.com/pytorch/pytorch/blob/master/LICENSE
-def linetable_writer(first_lineno: int) -> tuple[list[int], Callable, Callable]:
- """Used to create typing.CodeType.co_linetable See
- https://github.com/python/cpython/blob/main/Objects/lnotab_notes.txt This is the internal format of the line number
- table if Python >= 3.10."""
- assert sys.version_info >= (3, 9)
- linetable: list[int] = []
- lineno = first_lineno
- lineno_delta = 0
- byteno = 0
-
- def _update(byteno_delta: int, lineno_delta: int) -> None:
- while byteno_delta != 0 or lineno_delta != 0:
- byte_offset = max(0, min(byteno_delta, 254))
- line_offset = max(-127, min(lineno_delta, 127))
- assert byte_offset != 0 or line_offset != 0
- byteno_delta -= byte_offset
- lineno_delta -= line_offset
- linetable.extend((byte_offset, line_offset & 0xFF))
-
- def update(lineno_new: int, byteno_new: int) -> None:
- nonlocal lineno, lineno_delta, byteno
- byteno_delta = byteno_new - byteno
- byteno = byteno_new
- _update(byteno_delta, lineno_delta)
- lineno_delta = lineno_new - lineno
- lineno = lineno_new
-
- def end(total_bytes: int) -> None:
- _update(total_bytes - byteno, lineno_delta)
-
- return linetable, update, end
-
-
-def generate_function(gr: "Graph") -> Callable:
- orig_gr = gr
- gr, map_from_orig = gr.clone()
-
- local_vars, lv_names, names, consts = undo_ssa(gr)
- assert len(local_vars) == len(lv_names)
-
- NodeKey = Union[Node, tuple[Node, bool]]
- instruction_sizes: dict[NodeKey, int] = {}
-
- def build_address_map(end=False) -> dict[NodeKey, int]:
- # Key either (for jump nodes and jump=True)
- # or (, False) for non-jump in conditional jump
- address_map: dict[NodeKey, int] = {}
- ctr = 0
- for bl in gr.blocks:
- # assumes first block is function start
- for n in bl.nodes:
- address_map[n] = ctr
- ctr += instruction_sizes.get(n, 1)
- if len(n.jump_targets) == 2: # implicit unconditional jump
- ctr += instruction_sizes.get((n, False), 1)
- if end:
- address_map[n] = ctr - 1
- return address_map
-
- def make_bc() -> tuple[list[int], bool]:
- bc = []
-
- def write_extended_args(node_key: NodeKey, arg: int) -> bool:
- # returns if instruction size has changed
- instruction_size = instruction_sizes.get(node_key, 1)
- if arg > 0x_FF_FF_FF or instruction_size == 4:
- instruction_size = 4
- bc.append(dis.opmap["EXTENDED_ARG"])
- bc.append(arg >> 24)
- if arg > 0x_FF_FF or instruction_size >= 3:
- instruction_size = max(instruction_size, 3)
- bc.append(dis.opmap["EXTENDED_ARG"])
- bc.append((arg >> 16) & 0xFF)
- if arg > 0x_FF or instruction_size >= 2:
- instruction_size = max(instruction_size, 2)
- bc.append(dis.opmap["EXTENDED_ARG"])
- bc.append((arg >> 8) & 0xFF)
- else:
- instruction_size = 1
-
- if instruction_size != instruction_sizes.get(node_key, 1):
- instruction_sizes[node_key] = instruction_size
- return True
- return False
-
- changed_size = False
- line_no = None
- for bl in gr.blocks:
- jump_node = None
- for n in bl.nodes:
- opcode = n.i.opcode
- if opcode is None or opcode == -1: # Todo: opcode is typed int in ThunderInstruction, remove None here?
- opcode = dis.opmap[n.i.opname]
- assert opcode is not None, f"{n} has invalid opcode"
- # source range instead for 3.11?
- n_line_no = n.source_infos[-1].gen_line_no if n.source_infos else None
- if n_line_no is not None and n_line_no != line_no: # really, the last generated one...
- linetable_update(
- n_line_no + gr.source_start_line, address_map[n] * 2
- ) # byte offset for Python 3.10, too...
- line_no = n_line_no
- if opcode in dis.hasjabs:
- arg = address_map[n.jump_targets[-1].nodes[0]]
- elif opcode in dis.hasjrel:
- # TODO forward, backward
- arg = address_map[n.jump_targets[-1].nodes[0]] - address_map[n] - 1
- else:
- arg_ = n.i.arg
- arg = 0 if arg_ is None else arg_
-
- changed_size |= write_extended_args(n, arg)
-
- bc.append(opcode)
- bc.append(arg & 0x_FF)
- if len(n.jump_targets) > 1:
- jump_node = n
- if jump_node is not None:
- assert len(jump_node.jump_targets) == 2
- jarg = address_map[jump_node.jump_targets[0].nodes[0]]
- changed_size |= write_extended_args((jump_node, False), jarg)
- i = get_instruction(opname="JUMP_ABSOLUTE", arg=jarg & 0xFF)
- bc.append(i.opcode)
- assert i.arg is not None
- bc.append(i.arg)
- return bc, not changed_size
-
- done = False
- while not done:
- linetable, linetable_update, linetable_end = linetable_writer(gr.source_start_line)
- address_map = build_address_map()
- bc, done = make_bc()
-
- inserted_for = collections.defaultdict(list)
- end_address_map = build_address_map(end=True)
- for n in gr.nodes():
- inserted_for[getattr(n, "inserted_for", None)].append(end_address_map[n])
- for n in orig_gr.nodes():
- info = inserted_for[map_from_orig[n]]
- n.bytecode_range = (min(info), max(info)) if info else (None, None)
-
- linetable_end(len(bc))
- linetable_bytes = bytes(linetable)
- bc_bytes = bytes(bc)
-
- lv_at_start = [v for v in gr.local_variables_at_start if v is not None]
- co_argcount = gr.co_argcount
- co_posonlyargcount = gr.co_posonlyargcount
- co_kwonlyargcount = gr.co_kwonlyargcount
- co_nlocals = len(local_vars)
- # TODO: actually track the stack size when doing codegen (for optimizations)
- co_stacksize = max(max(len(n.inputs), len(n.outputs)) for n in gr.nodes())
- co_flags = gr.co_flags
- co_codestring = bc_bytes
- co_consts = tuple(consts)
- co_names = tuple(names)
- co_varnames = tuple(lv_names)
- co_filename = f""
- co_name = gr.co_name
- co_firstlineno = gr.source_start_line
- co_linetable = linetable_bytes # XXX
- co_freevars = ()
- co_cellvars = ()
-
- c = types.CodeType(
- co_argcount, # int
- co_posonlyargcount, # int
- co_kwonlyargcount, # int
- co_nlocals, # int
- co_stacksize, # int
- co_flags, # int
- co_codestring, # bytes
- co_consts, # tuple
- co_names, # tuple
- co_varnames, # tuple
- co_filename, # string
- co_name, # string
- co_firstlineno, # integer
- co_linetable, # bytes
- co_freevars, # tuple
- co_cellvars, # tuple
- )
-
- # types.FunctionType(code, globals, name=None, argdefs=None, closure=None)
- func = types.FunctionType(
- c,
- {
- "__builtins__": __builtins__,
- },
- argdefs=tuple(gr.func_defaults),
- )
- func.__kwdefaults__ = gr.func_kwdefaults
- func._gr = orig_gr
-
- # simple cache hack
- mtime = None # this signals that the cache should not be invalidated(!)
- lines = gr.source_lines
- size = len("".join(lines))
- inspect.linecache.cache[co_filename] = size, mtime, lines, co_filename
-
- try:
- _ = tuple(dis.get_instructions(func))
- except BaseException as e:
- raise RuntimeError("Unknown error generating callable") from e
-
- return func
-
-
-def annotated_dis(thunder_fn, print_lines=True):
- instructions = list(dis.get_instructions(thunder_fn))
- cur_pos = 0
-
- class Callback(GraphSummaryCallback):
- def node(self, n):
- nonlocal cur_pos
- before = []
- after = []
- begin_offset, end_offset = n.bytecode_range
- if begin_offset is not None:
- # the * 2 here and below is from Instruction.offset containing byte offsets and each
- # bytecode is 2 bytes (this is true at least for Python 3.8-3.12)
- while cur_pos < len(instructions) and instructions[cur_pos].offset < begin_offset * 2:
- before.append(instructions[cur_pos]._disassemble())
- cur_pos += 1
- if end_offset is not None:
- while cur_pos < len(instructions) and instructions[cur_pos].offset <= end_offset * 2:
- after.append(instructions[cur_pos]._disassemble())
- cur_pos += 1
- return before, after
-
- def finish(self):
- nonlocal cur_pos
- l = [i._disassemble() for i in instructions[cur_pos:]]
- cur_pos = len(instructions)
- return l
-
- return thunder_fn._gr.summary(print_lines=print_lines, callback=Callback())
diff --git a/thunder/core/script/python_ir_data.py b/thunder/core/script/python_ir_data.py
deleted file mode 100644
index c82eb168aa..0000000000
--- a/thunder/core/script/python_ir_data.py
+++ /dev/null
@@ -1,68 +0,0 @@
-import functools
-import sys
-from types import CodeType
-from typing import Union
-from collections.abc import Callable
-from collections.abc import Iterable
-
-from thunder.core.script import parse
-
-
-SUPPORTS_PREPROCESSING = (3, 9) <= sys.version_info < (3, 11)
-X_THUNDER_STORE_ATTR = "X_THUNDER_STORE_ATTR"
-
-
-# TODO(robieta): replace callsites.
-get_instruction = functools.partial(parse.ThunderInstruction.make, line_no=-1)
-
-
-def debug_compare_functions_print(diffs: dict[str, tuple[list, list]]):
- for k, (v1, v2) in diffs.items():
- if not (v1 is None and v2 is None):
- print(f"Differences in: {k}")
- print(f" CodeObject 1: {v1}")
- print(f" CodeObject 2: {v2}")
-
-
-def debug_compare_functions(
- code1: CodeType | Callable, code2: CodeType | Callable, *, show=False
-) -> dict[str, tuple[list, list]]:
- if not isinstance(code1, CodeType):
- code1 = code1.__code__
- if not isinstance(code2, CodeType):
- code2 = code2.__code__
-
- attrs = [
- "co_argcount",
- "co_kwonlyargcount",
- "co_nlocals",
- "co_stacksize",
- "co_flags",
- "co_consts",
- "co_names",
- "co_varnames",
- "co_filename",
- "co_name",
- "co_freevars",
- "co_cellvars",
- ]
-
- diffs = {}
- for attr in attrs:
- v1 = getattr(code1, attr)
- v2 = getattr(code2, attr)
-
- if v1 != v2:
- if isinstance(v1, dict) and isinstance(v2, dict):
- diffs[attr] = (v1 - v2, v2 - v1)
- if isinstance(v1, str) and isinstance(v2, str):
- diffs[attr] = (v1, v2)
- elif isinstance(v1, Iterable) and isinstance(v2, Iterable):
- diffs[attr] = (set(v1) - set(v2), set(v2) - set(v1))
- else:
- diffs[attr] = (v1, v2)
-
- if show:
- debug_compare_functions_print(diffs)
-
- return diffs
diff --git a/thunder/core/script/values/__init__.py b/thunder/core/script/values/__init__.py
deleted file mode 100644
index f1361cd857..0000000000
--- a/thunder/core/script/values/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from thunder.core.script.values.base import *
-from thunder.core.script.values.composite import *
-from thunder.core.script.values.materialization import *
-from thunder.core.script.values.symbolic import *
diff --git a/thunder/core/script/values/base.py b/thunder/core/script/values/base.py
deleted file mode 100644
index 1956c31050..0000000000
--- a/thunder/core/script/values/base.py
+++ /dev/null
@@ -1,200 +0,0 @@
-from __future__ import annotations
-
-import dataclasses
-import enum
-import textwrap
-from typing import overload, Any, Generic, NewType, TypeVar
-from collections.abc import Callable
-from collections.abc import Mapping
-
-from typing_extensions import Self
-
-from thunder.core.script import parse
-from thunder.core.utils import FrozenDict
-
-__all__ = (
- # HybridMap
- "Reference",
- "TraitName",
- "HybridMap",
- #
- # Values
- "AbstractValue",
- "AbstractRef",
- "NonPyObject",
- "IntermediateValue",
- "ExternalRef",
- #
- # Substitution
- "substitute_value",
- "ReplaceMap",
-)
-
-
-# =============================================================================
-# == Generic (hybrid tuple/dict) container ====================================
-# =============================================================================
-T = TypeVar("T")
-T1 = TypeVar("T1")
-Reference = NewType("Reference", int)
-TraitName = NewType("TraitName", str)
-
-
-@dataclasses.dataclass(frozen=True, eq=True)
-class HybridMap(Generic[T]):
- ordered: tuple[T, ...] = dataclasses.field(kw_only=True, default_factory=tuple)
- named: FrozenDict[TraitName, T] = dataclasses.field(kw_only=True, default_factory=FrozenDict)
-
- def __getitem__(self, key: Reference | TraitName) -> T:
- if isinstance(key, int):
- return self.ordered[key]
- elif isinstance(key, str):
- return self.named[key]
- raise TypeError(f"Invalid key: {key}")
-
- def __repr__(self) -> str:
- parts = [f"{self.__class__.__name__}("]
- if self.ordered:
- ordered = "\n".join(repr(i) for i in self.ordered)
- parts.append(f" ordered:\n{textwrap.indent(ordered, ' ' * 4)}")
-
- if self.named:
- named = "\n".join(f"{k}: {v}" for k, v in self.named.items())
- parts.append(f" named:\n{textwrap.indent(named, ' ' * 4)}")
-
- return "\n".join((*parts, ")"))
-
- @overload
- def map(self, f: Callable[[T], T]) -> Self:
- ...
-
- @overload
- def map(self, f: Callable[[T], T1]) -> HybridMap[T1]:
- ...
-
- def map(self, f: Any) -> Any:
- ordered = tuple(f(i) for i in self.ordered)
- named: FrozenDict[TraitName, T] = FrozenDict({k: f(v) for k, v in self.named.items()})
- return dataclasses.replace(self, ordered=ordered, named=named)
-
- def get(self, name: TraitName) -> T | None:
- return self.named.get(name)
-
- def add_named(self, name: TraitName, value: T) -> Self:
- named = dict(self.named)
- named.update({name: value}) # Preserve order.
- return dataclasses.replace(self, named=FrozenDict(named))
-
-
-# =============================================================================
-# == Simple value types =======================================================
-# =============================================================================
-class AbstractValue:
- """Represents a value during instruction parsing. (Prior to type binding.)"""
-
- __is_detail = True
-
- def __init_subclass__(cls, **kwargs: Any) -> None:
- cls.__is_detail = kwargs.pop("__is_detail", False) # `dataclasses` forces this into kwargs for some reason.
- super().__init_subclass__(**kwargs)
-
- def __copy__(self) -> AbstractValue:
- raise NotImplementedError
-
- @property
- def is_detail(self) -> bool:
- return self.__is_detail
-
- @property
- def identity(self) -> AbstractValue:
- """Analogous to `id(obj)`. For composites there is a layer of state management above the value itself.
-
- This is not suitable for equality checks (for example, mutation does not change
- an object's identity), but it is often the appropriate target for `isinstance` checks.
- """
- return self
-
- def _unpack_apply(self, _: ReplaceMap) -> AbstractValue:
- """Recursively update any constituent references in the abstract value."""
- return self
-
-
-@dataclasses.dataclass(frozen=True, eq=False)
-class AbstractRef(AbstractValue, __is_detail=True):
- """Placeholder value which will be resolved during parsing."""
-
- _debug_info: str = "N/A"
-
-
-@dataclasses.dataclass(frozen=True, eq=True)
-class NonPyObject(AbstractValue):
- """Singleton values used to signal some special interpreter state."""
-
- class Tag(enum.Enum):
- DELETED = enum.auto()
- MISSING = enum.auto()
- NULL = enum.auto()
-
- tag: Tag
-
- def __repr__(self) -> str:
- return self.tag.name
-
-
-class IntermediateValue(AbstractValue):
- """A (potentially) new value produced by an instruction."""
-
- def __repr__(self) -> str:
- return f"{self.__class__.__name__}(at {hex(id(self))})"
-
-
-@dataclasses.dataclass(frozen=True, eq=True)
-class ExternalRef(AbstractValue):
- """Reference values outside of the parsed code. (Arguments, constants, globals, etc.)"""
-
- key: parse.VariableKey
-
- def __repr__(self) -> str:
- if self.key.is_const:
- return f"Const({self.key.identifier})"
- return f"{self.__class__.__name__}({self.key.identifier}, {self.key.scope.name})"
-
-
-# =============================================================================
-# == Value substitution =======================================================
-# =============================================================================
-ReplaceMap = Mapping["AbstractValue", "AbstractValue"]
-
-
-@overload
-def substitute_value(v: AbstractValue, replace_map: ReplaceMap) -> AbstractValue:
- ...
-
-
-@overload
-def substitute_value(v: T, replace_map: ReplaceMap) -> T:
- ...
-
-
-def substitute_value(v: Any, replace_map: ReplaceMap) -> Any:
- """Find the replacement for `v`, and recursively substitute. (If applicable.)
-
- Some abstract values reference other abstract values. When we make substitution during
- graph transformations it is necessary to also consider replacement of an abstract
- value's constituents. Any subclass which must be unpacked in this manner should
- override `_unpack_apply`.
- """
- if not isinstance(v, AbstractValue):
- return v
-
- new_v = replace_map.get(v, v)
- if new_v != (x := replace_map.get(new_v, new_v)):
- msg = f"""
- `replace_map` may not contain chains.
- {v}
- {new_v}
- {x}
- See `flatten_map`."""
- raise ValueError(textwrap.dedent(msg))
-
- return new_v._unpack_apply(replace_map)
diff --git a/thunder/core/script/values/composite.py b/thunder/core/script/values/composite.py
deleted file mode 100644
index bc00dd5469..0000000000
--- a/thunder/core/script/values/composite.py
+++ /dev/null
@@ -1,177 +0,0 @@
-import abc
-import dataclasses
-import itertools
-from typing import Any, TypeVar
-from collections.abc import Iterable
-
-from typing_extensions import Self
-
-from thunder.core.script import parse
-from thunder.core.script.values import base, symbolic
-from thunder.core.utils import FrozenDict
-
-__all__ = ("InternalRef", "OrderedSlice", "CompositeValue", "CompositeRef", "AbstractPhiValue")
-
-T = TypeVar("T")
-
-
-# =============================================================================
-# == References ===============================================================
-# =============================================================================
-class InternalRef(base.AbstractValue, abc.ABC):
- @abc.abstractmethod
- def _resolve(self, inputs: base.HybridMap[base.AbstractValue]) -> base.AbstractValue:
- """Defines how to concretize itself."""
- ...
-
- @property
- def is_detail(self) -> bool:
- # All ref types are unsuitable for Graph binding.
- return True
-
- @classmethod
- def resolve(
- cls, output: symbolic.Symbolic.Output, *, inputs: base.HybridMap[base.AbstractValue] | None = None
- ) -> base.AbstractValue:
- inputs = base.HybridMap() if inputs is None else inputs
- if isinstance(output, (int, str)):
- return inputs[output]
-
- elif isinstance(output, symbolic.ConstRef):
- return base.ExternalRef(parse.VariableKey(output.identifier, parse.VariableScope.CONST))
-
- if isinstance(output, tuple):
- cls.validate_reference(output)
- result = inputs[output[0]]
- for idx in output[1:]:
- # We can only unpack a (possibly nested) composite.
- assert isinstance(result, base.HybridMap), result
- result = result[idx]
-
- assert isinstance(result, base.AbstractValue)
- return result
-
- elif isinstance(output, InternalRef):
- return output._resolve(inputs)
-
- return output
-
- @staticmethod
- def validate_reference(x: symbolic.NestedReference) -> None:
- x = (x,) if isinstance(x, (int, str)) else x
- assert isinstance(x, tuple) and x and all(isinstance(xi, (int, str)) for xi in x), x
-
-
-# =============================================================================
-# == Nesting ==================================================================
-# =============================================================================
-@dataclasses.dataclass(frozen=True, eq=True)
-class OrderedSlice:
- reference: symbolic.NestedReference
- slice: slice
-
- def __hash__(self) -> int:
- # `slice` isn't hashable until 3.12
- return hash((self.reference, self.slice.start, self.slice.stop, self.slice.step))
-
-
-@dataclasses.dataclass(frozen=True, eq=True, repr=False)
-class _Composite(base.AbstractValue, base.HybridMap[T], __is_detail=True):
- """Models an AbstractValue that references other (possibly also AbstractValue) state.
-
- Note: `ordered` and `named` should not contain cycles.
- """
-
- Identity = base.TraitName("__Thunder_Object_Identity")
-
- def _unpack_apply(self, replace_map: base.ReplaceMap) -> base.AbstractValue:
- new_self = self.map(lambda x: base.substitute_value(x, replace_map))
- assert isinstance(new_self, _Composite) # For mypy since we can't hint `_Composite[T] -> _Composite[T1]`
- return new_self
-
- def add_identity(self, identity: T) -> Self:
- return self.add_named(self.Identity, identity)
-
- # NOTE: We don't override `identity`. (Instead retaining `return self` from AbstractValue.)
- # This is because we generally won't know a good value, and passes should do their
- # own type checking. (And that checking should almost always be done on the materialized
- # value, not the symbolic reference.)
-
-
-@dataclasses.dataclass(frozen=True, eq=True, repr=False)
-class CompositeValue(_Composite[base.AbstractValue]):
- def __post_init__(self) -> None:
- assert all(isinstance(i, base.AbstractValue) for i in self.ordered)
- assert all(isinstance(i, base.AbstractValue) for i in self.named.values())
-
- @property
- def is_detail(self) -> bool:
- return any(i.is_detail for i in itertools.chain(self.ordered, self.named.values()))
-
- @property
- def identity(self) -> base.AbstractValue:
- return self.named.get(self.Identity, self)
-
-
-@dataclasses.dataclass(frozen=True, eq=True)
-class CompositeRef(InternalRef, _Composite[symbolic.Symbolic.Output | OrderedSlice]):
- def __post_init__(self) -> None:
- assert not any(isinstance(i, OrderedSlice) for i in self.named.values())
-
- def _resolve(self, inputs: base.HybridMap[base.AbstractValue]) -> CompositeValue:
- ordered: list[base.AbstractValue] = []
- for i in self.ordered:
- if isinstance(i, OrderedSlice):
- slice_target = self.resolve(i.reference, inputs=inputs) if i.reference else inputs
- assert isinstance(slice_target, base.HybridMap)
- ordered.extend(slice_target.ordered[i.slice])
- else:
- ordered.append(self.resolve(i, inputs=inputs))
-
- named: dict[base.TraitName, base.AbstractValue] = {}
- for k, v in self.named.items():
- assert not isinstance(v, OrderedSlice)
- named[k] = self.resolve(v, inputs=inputs)
-
- return CompositeValue(ordered=tuple(ordered), named=FrozenDict(named))
-
-
-# =============================================================================
-# == Unions ===================================================================
-# =============================================================================
-@dataclasses.dataclass(frozen=True, eq=True)
-class AbstractPhiValue(base.AbstractValue):
- constituents: tuple[base.AbstractValue, ...]
-
- def __post_init__(self) -> None:
- # Flatten nested PhiValues. e.g.
- # 𝜙[𝜙[A, B], 𝜙[A, C]] -> 𝜙[A, B, C]
- constituents = itertools.chain(*[self.flatten(i) for i in self.constituents])
-
- # Ensure a consistent order.
- constituents = tuple(v for _, v in sorted({hash(v): v for v in constituents}.items()))
- assert not any(isinstance(i, InternalRef) for i in constituents)
- object.__setattr__(self, "constituents", constituents)
-
- def __getitem__(self, _: Any) -> base.AbstractValue:
- # The semantics of indexing into an `AbstractPhiValue`` are not well defined:
- # - The order of `constituents` is arbitrary
- # - It's unclear if the desire is to select one constituent or create a new `AbstractPhiValue`
- # which indexes into each constituent.
- # If a concrete use case emerges we can tackle it; until then we refuse for safety.
-
- # TODO(robieta): Handle traits
- raise NotImplementedError
-
- def _unpack_apply(self, replace_map: base.ReplaceMap) -> base.AbstractValue:
- result = AbstractPhiValue(tuple(base.substitute_value(v, replace_map) for v in self.constituents))
- return result if len(result.constituents) > 1 else result.constituents[0]
-
- @classmethod
- def flatten(cls, v: base.AbstractValue) -> Iterable[base.AbstractValue]:
- constituents = [cls.flatten(i) for i in v.constituents] if isinstance(v, AbstractPhiValue) else [[v]]
- yield from itertools.chain(*constituents)
-
- @property
- def is_detail(self) -> bool:
- return any(i.is_detail for i in self.constituents)
diff --git a/thunder/core/script/values/materialization.py b/thunder/core/script/values/materialization.py
deleted file mode 100644
index e12e35d0e5..0000000000
--- a/thunder/core/script/values/materialization.py
+++ /dev/null
@@ -1,171 +0,0 @@
-from __future__ import annotations
-
-import dataclasses
-import functools
-import itertools
-from types import MappingProxyType
-from typing import Literal, TypeVar
-from collections.abc import Callable, Iterator
-
-from thunder.core.script import parse
-from thunder.core.script.values import base, composite, symbolic
-from thunder.core.utils import FrozenDict, OrderedSet
-from collections.abc import Iterable
-
-__all__ = ("Materialized", "DigestFlow", "IntraBlockFlow")
-T = TypeVar("T")
-
-
-# =============================================================================
-# == Intra-ProtoBlock abstract value flow =====================================
-# =============================================================================
-#
-# `ProtoBlocks` employ a dual representation, where node inputs and outputs can
-# be viewed as either a reference based DAG or a sequence of ops with concrete
-# `AbstractValue` inputs and outputs.
-#
-# At the boundaries of a ProtoBlock values have named (VariableKey) slots;
-# within the ProtoBlock there is no need for such slots (since there is no
-# control flow within a block and those named slots tell you how to build the
-# directed *cyclic* graph for the larger program) so they are stripped during
-# parsing.
-#
-# The inputs of a protoblock are stored as a map of `VariableKey -> AbstractValue`
-# and act as the intra-block DAG sources. The outputs are stored as references
-# since every ProtoBlock output must have a unique producer. (Either an input
-# or a node within the block.)
-#
-# The canonical representation for intra-block flow is "symbolic" (reference
-# based). If an `AbstractValue` appear in a symbolic node's outputs that
-# indicates that the node is that value's producer. Otherwise all inputs and
-# outputs are references: inputs reference either the begin state or the
-# outputs of a prior node while output references index into the node's inputs.
-#
-# When analyzing a graph we are generally interested in the concrete properties
-# of values; provenance is generally only important when connecting blocks and
-# performing rewrites. For these cases `IntraBlockFlow` generates a
-# "materialized" flow which resolves all references to `AbstractValue`s. The
-# symbolic representation is sufficient to emit the materialized representation,
-# but the reverse is not true.
-VarT = TypeVar("VarT", bound=base.AbstractValue, covariant=True)
-ConcreteState = FrozenDict[parse.VariableKey, base.AbstractValue]
-EndState = FrozenDict[parse.VariableKey, symbolic.Symbolic.Input]
-
-
-@dataclasses.dataclass(frozen=True, eq=False)
-class Materialized:
- """Flow element where all symbolic references have been resolved to concrete `AbstractValue`s."""
-
- inputs: base.HybridMap[base.AbstractValue]
- outputs: tuple[base.AbstractValue, ...]
-
-
-class DigestFlow:
- """One-shot helper for materializing a block."""
-
- GetT = Callable[[symbolic.Symbolic.Input], base.AbstractValue]
-
- def __init__(self, begin: ConcreteState) -> None:
- self._begin = begin
- self._result: dict[parse.ThunderInstruction, Materialized] = {}
-
- def next(self, instruction: parse.ThunderInstruction, symbolic: symbolic.Symbolic) -> Materialized:
- """Lookup the materialized node corresponding to a symbolic node."""
-
- # NB: `inputs_after_op` will be needed after we introduce mutations.
- inputs = inputs_after_op = symbolic.inputs.map(self.get)
- outputs = tuple(composite.InternalRef.resolve(o, inputs=inputs_after_op) for o in symbolic.outputs)
- assert all(isinstance(o, base.AbstractValue) for o in outputs), outputs
-
- self._result[instruction] = result = Materialized(inputs, outputs)
- return result
-
- def get(self, key: symbolic.Symbolic.Input) -> base.AbstractValue:
- """Resolve a Symbolic input based on the block state at that node."""
- result: base.AbstractValue
- if isinstance(key, base.NonPyObject.Tag):
- result = base.NonPyObject(key)
-
- elif isinstance(key, symbolic.OutputRef):
- inputs = base.HybridMap(ordered=self._result[key.instruction].outputs)
- result = composite.InternalRef.resolve(key.idx, inputs=inputs)
-
- else:
- assert isinstance(key, parse.VariableKey), key
- result = base.ExternalRef(key) if key.is_const else self._begin[key]
- return result
-
-
-@dataclasses.dataclass(frozen=True, eq=False)
-class IntraBlockFlow:
- _symbolic: FrozenDict[parse.ThunderInstruction, symbolic.Symbolic]
- _begin: ConcreteState
- _end: EndState
-
- StateIterT = Iterator[tuple[parse.VariableKey, base.AbstractValue]]
-
- def __post_init__(self) -> None:
- assert not (forbidden := tuple(i for i in self._symbolic if i in parse.FORBIDDEN_INSTRUCTIONS)), forbidden
- object.__setattr__(self, "_symbolic", FrozenDict(self._symbolic))
-
- missing = {i: base.AbstractRef("Inferred") for i in self.uses if i not in self._begin}
- object.__setattr__(self, "_begin", FrozenDict({**missing, **self._begin}))
- assert not any(k.is_const for k in self._begin), self._begin
- assert not any(isinstance(v, composite.InternalRef) for v in self._begin.values()), self._begin
-
- object.__setattr__(self, "_end", FrozenDict(self._end))
-
- @functools.cache
- def __getitem__(self, key: tuple[symbolic.Symbolic.Input, Literal[0, 1]]) -> base.AbstractValue:
- assert key[1] in (0, 1)
- return self._computed[1:][key[1]](key[0])
-
- @property
- def symbolic(self) -> Iterable[tuple[parse.ThunderInstruction, symbolic.Symbolic]]:
- yield from self._symbolic.items()
-
- @property
- def materialized(self) -> FrozenDict[parse.ThunderInstruction, Materialized]:
- return self._computed[0]
-
- @property
- def begin_state(self) -> StateIterT:
- yield from self._sort_and_filter_state(iter(self._begin.items()))
-
- @property
- def end_state(self) -> StateIterT:
- yield from self._sort_and_filter_state((k, self[v, 1]) for k, v in self._end.items())
-
- @staticmethod
- def _sort_and_filter_state(kv: StateIterT) -> StateIterT:
- yield from ((k, v) for k, v in sorted(kv) if not isinstance(v, base.NonPyObject))
-
- @property
- def uses(self) -> OrderedSet[parse.VariableKey]:
- assignment = (v for k, v in self._end.items() if isinstance(v, parse.VariableKey) and not v.is_const and k != v)
- return OrderedSet(itertools.chain(*(s.uses for _, s in self.symbolic), assignment))
-
- _Computed = tuple[
- FrozenDict[parse.ThunderInstruction, Materialized],
- DigestFlow.GetT, # Begin
- DigestFlow.GetT, # End
- ]
-
- @functools.cached_property
- def _computed(self) -> _Computed:
- flow_state = DigestFlow(self._begin)
- materialized_flow: FrozenDict[parse.ThunderInstruction, Materialized]
- materialized_flow = FrozenDict({i: flow_state.next(i, s) for i, s in self.symbolic}) # Populates `flow_state`
- return materialized_flow, DigestFlow(self._begin).get, flow_state.get
-
- def substitute(self, replace_map: base.ReplaceMap) -> IntraBlockFlow | None:
- """Replace `AbstractValue`s within the flow. (Block inputs and producer nodes.)"""
- replace_map_view = MappingProxyType(replace_map)
- new_symbolic: FrozenDict[parse.ThunderInstruction, symbolic.Symbolic]
- new_symbolic = FrozenDict({k: (s.substitute(replace_map_view)) for k, s in self.symbolic})
- begin = ConcreteState({k: base.substitute_value(v, replace_map_view) for k, v in self._begin.items()})
-
- # TODO(robieta): Check if a value is only present in `materialized` and error.
- if self._symbolic != new_symbolic or self._begin != begin:
- return dataclasses.replace(self, _symbolic=new_symbolic, _begin=begin)
- return None
diff --git a/thunder/core/script/values/symbolic.py b/thunder/core/script/values/symbolic.py
deleted file mode 100644
index 60bbaa08b6..0000000000
--- a/thunder/core/script/values/symbolic.py
+++ /dev/null
@@ -1,175 +0,0 @@
-"""Introduce references inside simple blocks."""
-from __future__ import annotations
-
-import dataclasses
-import itertools
-import sys
-from typing import Any, NamedTuple, TypeAlias
-from collections.abc import Callable, Iterable
-
-from typing_extensions import Self
-
-from thunder.core.script import parse
-from thunder.core.script.values import base
-from thunder.core.utils import FrozenDict, safe_zip
-
-__all__ = ("OutputRef", "ParsedSymbolic", "Symbolic", "NestedReference", "ConstRef")
-
-
-# =============================================================================
-# == Opcode-specific behavior =================================================
-# =============================================================================
-def rotate_N(oparg: int) -> tuple[int, ...]:
- return (-1,) + tuple(range(-oparg, -1))
-
-
-_AliasMask = tuple[int | None, ...]
-ALIAS_OPCODES = FrozenDict[str, _AliasMask | Callable[[int], _AliasMask]](
- parse.fill_ellipses(
- #
- # Stack manipulation
- ROT_N=rotate_N, # A,B,...,Z -> Z,A,B,...
- ROT_FOUR=rotate_N,
- ROT_THREE=rotate_N,
- ROT_TWO=rotate_N,
- DUP_TOP=(-1, -1), # A -> A,A
- DUP_TOP_TWO=(-2, -1) * 2, # A,B -> A,B,A,B
- #
- # Insertion leaves container on the stack A,B -> A
- SET_ADD=(-2,),
- SET_UPDATE=...,
- LIST_APPEND=...,
- LIST_EXTEND=...,
- DICT_MERGE=...,
- DICT_UPDATE=...,
- MAP_ADD=(-3,),
- COPY_DICT_WITHOUT_KEYS=(-2, None), # A,B -> A,C (I am unsure...)
- #
- # Misc.
- GET_LEN=(-1, None),
- MATCH_MAPPING=(-1, None),
- MATCH_SEQUENCE=...,
- MATCH_KEYS=(-1, -2, None) + () if sys.version_info >= (3, 11) else (None,),
- #
- # Jump dependent
- FOR_ITER=(-1, None),
- # NOTE: These instructions have been removed since they are extraneous special cases.
- # https://github.com/faster-cpython/ideas/issues/567
- # https://github.com/python/cpython/issues/102859
- JUMP_IF_TRUE_OR_POP=(-1,),
- JUMP_IF_FALSE_OR_POP=(-1,),
- #
- # This isn't actually correct. `LOAD_METHOD` will return either
- # A -> B, A
- # A -> B, NULL
- # However the `A | NULL` is only consumed by `CALL_METHOD`, so it's ok to use this alias.
- LOAD_METHOD=(None, -1), # A -> B,A
- )
-)
-
-
-# =============================================================================
-# == Symbolic flow ============================================================
-# =============================================================================
-IndexT: TypeAlias = base.Reference | base.TraitName
-NestedReference: TypeAlias = IndexT | tuple[IndexT, ...]
-
-
-@dataclasses.dataclass(frozen=True, eq=True)
-class OutputRef:
- """Identifies the producer of a value within a block."""
-
- instruction: parse.ThunderInstruction # Acts as a key for the producer Flow.
- idx: NestedReference # Indexes the producer's outputs.
-
-
-class ConstRef(NamedTuple):
- """Convenience wrapper to access `ExternalRef(VariableKey(..., CONST))` as a reference.
-
- This saves us from having to plumb constants through named inputs, since:
- A) `ExternalRef`s cannot appear in Symbolic outputs.
- (Since that implies a producer relationship which doesn't make sense.)
- B) Symbolic reference outputs must reference an input, which would mean an entry of
- `{some_random_name: VariableKey(..., CONST)}` would have to be added to named inputs
- which is tedious.
- """
-
- identifier: Any
-
-
-@dataclasses.dataclass(frozen=True, eq=True)
-class Symbolic:
- """Represents abstract flow immediately after functionalization."""
-
- # VariableKey: References the value of that variable at the start of the block
- # OutputRef: Reference values created by an earlier instruction within the block
- # SingletonValue.Tag: Reserved for special cases.
- Input = parse.VariableKey | OutputRef | base.NonPyObject.Tag
- inputs: base.HybridMap[Input]
-
- # NestedReference: Aliases the input at this position.
- # AbstractValue: New value created by this instruction
- Output = NestedReference | ConstRef | base.AbstractValue
- outputs: tuple[Output, ...]
-
- BeginState = FrozenDict[parse.VariableKey, base.AbstractValue]
- EndState = FrozenDict[parse.VariableKey, Input]
- Block = tuple[FrozenDict[parse.ThunderInstruction, "Symbolic"], BeginState, EndState]
-
- def __post_init__(self) -> None:
- # If an `AbstractValue` appears in `Symbolic.outputs` that implies that the symbolic
- # node in question is the value's producer. However it doesn't make sense for an external
- # value to be produced within the compiled function.
- assert not any(isinstance(o, base.ExternalRef) for o in self.outputs), self
-
- @property
- def uses(self) -> Iterable[parse.VariableKey]:
- """Block inputs used by this node.
-
- NOTE: This does not include values produced by an earlier node in the block.
- """
- for i in itertools.chain(self.inputs.ordered, self.inputs.named.values()):
- if isinstance(i, parse.VariableKey) and not i.is_const:
- yield i
-
- def substitute(self, replace_map: base.ReplaceMap) -> Self:
- outputs = tuple(base.substitute_value(o, replace_map) for o in self.outputs)
- return dataclasses.replace(self, outputs=outputs)
-
-
-# =============================================================================
-# == Conversion from functional representation ================================
-# =============================================================================
-@dataclasses.dataclass(frozen=True)
-class ParsedSymbolic:
- blocks: tuple[Symbolic.Block, ...]
- provenance: parse.ParsedFunctional
-
- @classmethod
- def make(cls, parsed: parse.ParsedFunctional) -> ParsedSymbolic:
- blocks: list[Symbolic.Block] = []
- for block, begin_state, end_state in parsed.blocks:
- # `functionalize_blocks` produces unique values, so provenance is unambiguous.
- producers: dict[parse.PlaceholderValue | None, Symbolic.Input] = {v: k for k, v in begin_state.items()}
- producers[None] = base.NonPyObject.Tag.DELETED
- assert len(producers) == len(begin_state) + 1, (producers, end_state)
-
- symbolic_blocks: dict[parse.ThunderInstruction, Symbolic] = {}
- for instruction, raw_inputs, raw_outputs in block:
- for idx, o in enumerate(raw_outputs):
- assert o not in producers
- producers[o] = OutputRef(instruction, base.Reference(idx))
-
- outputs: tuple[Symbolic.Output, ...] = tuple(base.IntermediateValue() for _ in raw_outputs)
- if alias := ALIAS_OPCODES.get(instruction.opname):
- mask = alias(len(outputs)) if callable(alias) else alias
- mask = (base.Reference(i) if i is not None else i for i in mask)
- outputs = tuple(o if o_mask is None else o_mask for o, o_mask in safe_zip(outputs, mask))
- inputs = base.HybridMap(ordered=tuple(producers[i] for i in raw_inputs))
- symbolic_blocks[instruction] = Symbolic(inputs, outputs)
-
- begin = {k: base.AbstractRef(v) for k, v in begin_state.items() if not k.is_const}
- end = {k: producers[v] for k, v in end_state.items() if not k.is_const}
- blocks.append((FrozenDict(symbolic_blocks), FrozenDict(begin), FrozenDict(end)))
-
- return cls(tuple(blocks), parsed)
diff --git a/thunder/core/symbol.py b/thunder/core/symbol.py
index 354e7ddc18..cb22553aa0 100644
--- a/thunder/core/symbol.py
+++ b/thunder/core/symbol.py
@@ -109,10 +109,10 @@ def default_python_printer(
# A symbol represents a function and how it can be transformed
# name is a string name for the operation
-# meta should use lightning.compile functions to evaluate the function;
-# it will be called with lightning.compile proxies
+# meta should use thunder.jit functions to evaluate the function;
+# it will be called with thunder.jit proxies
# id is an optional value to use when translating the function to executors
-# is_prim should be True if the Symbol represents a lightning.compile primitive
+# is_prim should be True if the Symbol represents a thunder.jit primitive
# python_printer is a function that will produce valid Python for calling the
# operation; this can usually be set to None, in which case the default python
# printer will be used for the Symbol. Symbols that control their own printing
@@ -196,14 +196,6 @@ def module(self) -> None | ModuleType:
result = inspect.getmodule(fn_)
return result
- # Properties used in transforms (defined later)
- # TODO https://github.com/Lightning-AI/lightning-thunder/issues/326
- # Remove this from here (think how symbols could be extended with transforms)
- # self.grad_defined = False
- # self.grad_ignored = False
- # self.grad_fwd = None
- # self.grad_bwd = None
-
def __repr__(self) -> str:
return f"[Symbol name={self.name}]"
@@ -313,7 +305,6 @@ def __post_init__(self):
# Constructs a new BoundSymbol with default values taken from this BoundSymbol
# Override values can be specified as kwargs
- # TODO https://github.com/Lightning-AI/lightning-thunder/issues/680
# Issue -- Provide a pattern for updating subsymbols when swapping outputs
# Maybe this can also just swap one set of symbols for another?
# Consider adding verification that the new and old output have the same metadata
diff --git a/thunder/core/trace.py b/thunder/core/trace.py
index 4e2bed9eba..88f1950585 100644
--- a/thunder/core/trace.py
+++ b/thunder/core/trace.py
@@ -19,7 +19,7 @@
from thunder.core.codeutils import ContextObject
-# TODO https://github.com/Lightning-AI/lightning-thunder/issues/327
+# TODO see issue "Improve TraceProvenance"
# Make this more interesting / printer better -- maybe let
# practitioners acquire the pass callable so they can replicate the pass?
# This class is intended to describe how the trace was constructed
@@ -36,7 +36,7 @@ def __repr__(self) -> str:
# TODO Should traces be BoundSymbols?
-# TODO https://github.com/Lightning-AI/lightning-thunder/issues/323
+# TODO issue "Create a mechanism for freezing TraceCtx objects"
# Add validation that a constant is never assigned to / reassigned
# Possibly separate the ideas of a trace -- a series of scopes containing bound symbols --
# and a TraceCtx, which can produce new traces
@@ -303,7 +303,7 @@ def python_ctx(self) -> dict:
return import_ctx
# TODO Account for multi-line signatures
- # TODO https://github.com/Lightning-AI/lightning-thunder/issues/324
+ # TODO issue "Add type annotations to Python function produced by traces"
# Consider extending the signature with type information, in particular the
# the type information of the return value might be interesting
def python(self, *, print_depth: int = 1) -> str:
@@ -395,7 +395,7 @@ def keyfn(class_or_module: type | ModuleType) -> str:
reset_tracectx(token)
# Returns a Python callable that executes the trace
- # TODO https://github.com/Lightning-AI/lightning-thunder/issues/323
+ # TODO issue "Create a mechanism for freezing TraceCtx objects"
# Create a mechanism for freezing traces and cache the compilation
def python_callable(self, *, global_dicts: None | dict = None) -> Callable:
python_str: str
diff --git a/thunder/core/transform_common.py b/thunder/core/transform_common.py
index 7b602084b6..fa74313bfe 100644
--- a/thunder/core/transform_common.py
+++ b/thunder/core/transform_common.py
@@ -134,9 +134,9 @@ def replace_redundant_inputs(
return new_bsyms
-# TODO(crcrpar): Implement a mechanism to keep track of supported ops that cannot be CSE'd.
-# For example, `uniform`, `dropout`, and `scaled_dot_product_attention`.
-# See: https://github.com/Lightning-AI/lightning-thunder/issues/671
+# These are ops that are not referentially transparent. We need to treat such
+# ops specially when optimizing; for example, CSE cannot coalesce two calls
+# into one for ops in this set.
NON_FUNCTIONAL_OPS: set[prims.PrimIDs | str] = {
prims.PrimIDs.UNIFORM,
"torch.uniform", # this doesn't exist as of the PR
diff --git a/thunder/core/transforms.py b/thunder/core/transforms.py
index 9922fede58..772e65a84d 100644
--- a/thunder/core/transforms.py
+++ b/thunder/core/transforms.py
@@ -24,6 +24,7 @@
NumberProxy,
Proxy,
TensorProxy,
+ FloatProxy,
variableify,
unvariableify,
CollectionProxy,
@@ -59,6 +60,7 @@
convolution,
)
from thunder.core.transform_common import dce
+from thunder.core.vjp_utils import make_aug_forward_and_backward
from thunder.extend import Executor
import thunder.torch as ltorch
@@ -538,24 +540,29 @@ def _flatten(bsym: BoundSymbol):
#
#
-# Functions related to functionalizing TOMs
+# Functions related to functionalizing ThunderOptimizedModules
#
# TODO Test with buffers
def populate_grads(grads: list[TensorProxy], tom: None | torch.nn.Module = None, args=None, kwargs=None) -> None:
idx: int = 0
- from thunder.common import ThunderOptimizedModule
-
- if tom is not None and isinstance(tom, ThunderOptimizedModule) and tom._additional_param_values is not None:
- for p in tom._additional_param_values:
- if p.requires_grad:
+ from thunder import ThunderModule, compile_data
+
+ if isinstance(tom, ThunderModule) or thunder.compile_data(tom).using_jit:
+ assert args is not None, "populate grad needs args (and possibly kwargs) to work with ThunderModules"
+ if kwargs is None:
+ kwargs = {}
+ _, computation_inputs, _ = compile_data(tom).get_computation_and_inputs(*args, **kwargs)
+ for p in computation_inputs:
+ if isinstance(p, torch.Tensor) and p.requires_grad:
# Supports grad accumulation (like when weight tying)
if p.grad is not None:
p.grad += grads[idx]
else:
p.grad = grads[idx]
idx += 1
+ return
# Short-circuits if there are no args or kwargs
if args is None and kwargs is None:
@@ -587,7 +594,6 @@ def clear_grads(module: torch.nn.Module) -> None:
b.grad = None
-from thunder.core.script.noinline import noinline
from thunder.core.interpreter import make_opaque
from thunder.core.langctxs import langctx, Languages
@@ -595,7 +601,7 @@ def clear_grads(module: torch.nn.Module) -> None:
# TODO RC1 Replace with langctx
def torchctx(fn):
_fn = langctx(Languages.TORCH)(fn)
- return make_opaque(noinline(_fn))
+ return make_opaque(_fn)
_grad_fn_map: dict[Any, Callable] = {}
@@ -664,8 +670,16 @@ def _convert_element_type_prim_grad(a: Number | TensorProxy, dtype: type | dtype
# NOTE prims.iota creates no grad associations
register_grad(pids.IOTA, prims.iota)
-# NOTE prims.uniform creates no grad associations
-register_grad(pids.UNIFORM, prims.uniform)
+
+def _uniform_grad(shape, minval, maxval, *, device, dtype):
+ fwd, saved = uniform_aug_fwd(shape, minval, maxval, device=device, dtype=dtype)
+ g = get_grad(fwd)
+ _, gminval, gmaxval = uniform_backward(*saved, g)
+ put_grads((minval, maxval), (gminval, gmaxval))
+ return fwd
+
+
+register_grad(pids.UNIFORM, _uniform_grad)
#
# Reshaping and permuting operator grads
@@ -863,12 +877,11 @@ def _abs_prim_grad(a: Number | TensorProxy) -> Number | TensorProxy:
register_grad(pids.ABS, _abs_prim_grad)
-@torchctx
def _cos_prim_grad(a: Number | TensorProxy) -> Number | TensorProxy:
- fwd = prims.abs(a)
+ fwd = prims.cos(a)
g = get_grad(fwd)
- put_grad(a, g * (-ltorch.sin(a)))
+ put_grad(a, g * (-prims.sin(a)))
return fwd
@@ -948,12 +961,11 @@ def _rsqrt_prim_grad(a: Number | TensorProxy, /) -> Number | TensorProxy:
register_grad(pids.RSQRT, _rsqrt_prim_grad)
-@torchctx
def _sin_prim_grad(a: Number | TensorProxy) -> Number | TensorProxy:
- fwd = prims.abs(a)
+ fwd = prims.sin(a)
g = get_grad(fwd)
- put_grad(a, g * ltorch.cos(a))
+ put_grad(a, g * prims.cos(a))
return fwd
@@ -1224,7 +1236,9 @@ def _embedding_prim_grad(
#
-def _get_gradfn(bsym: BoundSymbol, *, executors_list: Sequence[Any]) -> None | Callable:
+def _get_gradfn(bsym: BoundSymbol, *, executors_list: Sequence[Any] = tuple()) -> None | Callable:
+ cd = get_compile_data()
+ executors_list = cd.executors_list if cd is not None else executors_list
# Checks if the executor which has priority for this operation has a specific grad transform for it
for ex in executors_list:
if ex.can_execute_or_fuse(bsym):
@@ -1564,7 +1578,7 @@ def _selector(eligible_nodes: list[Node]) -> int:
return gradtrc
- # NOTE This is a kludge to indicate that we shouldn't support PyTorch's autograd because
+ # NOTE This is a kludge to indicate that we shouldn't use PyTorch's autograd because
# we're using our own autograd transform
cfn._using_grad_transform = True
@@ -1831,7 +1845,6 @@ def broadcast_in_dim_vmap(
) -> BatchedValue:
bdim = a.batch_dim
# TODO: remove this when shape and broadcast_dimensions become mandatory kwargs
- # See https://github.com/Lightning-AI/lightning-thunder/issues/181
shape, _ = safe_zip(*shape)
if len(broadcast_dimensions) > 0:
broadcast_dimensions, _ = safe_zip(*broadcast_dimensions)
@@ -2148,7 +2161,6 @@ def broadcast_in_dim_jvp(a: JVPDual, shape: tuple[JVPDual, ...], broadcast_dimen
x, xd = a
# TODO: shape and broadcast_dimensions should be tuples of ints
# but for now it's a tuple of JVPDuals
- # See https://github.com/Lightning-AI/lightning-thunder/issues/181
if len(shape) > 0 and isinstance(shape[0], JVPDual):
shape, _ = safe_zip(*shape)
if len(broadcast_dimensions) > 0 and isinstance(broadcast_dimensions[0], JVPDual):
@@ -2408,43 +2420,30 @@ def zeros_like(x):
# The augmented_primal function takes the primal values and returns the primal
# result and the residuals (saved values for the backward).
augmented_forward_impls = {
- prims.PrimIDs.ABS: lambda x: (prims.abs(x), (x,)),
prims.PrimIDs.ACOS: lambda x: (prims.acos(x), (x,)),
prims.PrimIDs.ACOSH: lambda x: (prims.acosh(x), (x,)),
- prims.PrimIDs.ADD: lambda x, y: (prims.add(x, y), tuple()),
prims.PrimIDs.ASIN: lambda x: (prims.asin(x), (x,)),
prims.PrimIDs.ASINH: lambda x: (prims.asinh(x), (x,)),
prims.PrimIDs.ATAN: lambda x: (prims.atan(x), (x,)),
prims.PrimIDs.ATANH: lambda x: (prims.atanh(x), (x,)),
prims.PrimIDs.ATAN2: lambda x, y: (prims.atan2(x, y), (x, y)),
- prims.PrimIDs.COS: lambda x: (prims.cos(x), (x,)),
prims.PrimIDs.COSH: lambda x: (prims.cosh(x), (x,)),
prims.PrimIDs.DIGAMMA: lambda x: (prims.digamma(x), (x,)),
- prims.PrimIDs.DIV: lambda x, y: (prims.div(x, y), (x, y)),
- prims.PrimIDs.ERF: lambda x: (prims.erf(x), (x,)),
prims.PrimIDs.ERFC: lambda x: (prims.erfc(x), (x,)),
prims.PrimIDs.ERFINV: lambda x: (prims.erfinv(x), (prims.erfinv(x),)),
prims.PrimIDs.ERFCINV: lambda x: (prims.erfcinv(x), (prims.erfcinv(x),)),
prims.PrimIDs.EXP2: lambda x: (prims.exp2(x), (prims.exp2(x),)),
prims.PrimIDs.EXPM1: lambda x: (prims.expm1(x), (prims.expm1(x),)),
prims.PrimIDs.LGAMMA: lambda x: (prims.lgamma(x), (x,)),
- prims.PrimIDs.MUL: lambda x, y: (prims.mul(x, y), (x, y)),
prims.PrimIDs.NDTRI: lambda x: (prims.ndtri(x), (prims.ndtri(x),)),
- prims.PrimIDs.SIN: lambda x: (prims.sin(x), (x,)),
prims.PrimIDs.SINH: lambda x: (prims.sinh(x), (x,)),
- prims.PrimIDs.SUB: lambda x, y: (prims.sub(x, y), tuple()),
prims.PrimIDs.SQRT: lambda x: (prims.sqrt(x), (prims.sqrt(x),)),
- prims.PrimIDs.EQ: lambda x, y: (prims.eq(x, y), (x, y)),
prims.PrimIDs.NE: lambda x, y: (prims.ne(x, y), (x, y)),
- prims.PrimIDs.GE: lambda x, y: (prims.ge(x, y), (x, y)),
prims.PrimIDs.GT: lambda x, y: (prims.gt(x, y), (x, y)),
prims.PrimIDs.LE: lambda x, y: (prims.le(x, y), (x, y)),
- prims.PrimIDs.LT: lambda x, y: (prims.lt(x, y), (x, y)),
- prims.PrimIDs.LOG: lambda x: (prims.log(x), (x,)),
prims.PrimIDs.LOG10: lambda x: (prims.log10(x), (x,)),
prims.PrimIDs.LOG1P: lambda x: (prims.log1p(x), (x,)),
prims.PrimIDs.LOG2: lambda x: (prims.log2(x), (x,)),
- prims.PrimIDs.NEG: lambda x: (prims.neg(x), tuple()),
prims.PrimIDs.ZETA: lambda x, y: (prims.zeta(x, y), (x, y)),
prims.PrimIDs.FMOD: lambda x, y: (prims.fmod(x, y), (x, y)),
}
@@ -2454,55 +2453,32 @@ def zeros_like(x):
# The backward function takes the residuals and cotangents and returns the
# vector-Jacobian products for each argument.
backward_impls = {
- prims.PrimIDs.ABS: lambda x, g: g * prims.sign(x),
prims.PrimIDs.ACOS: lambda x, g: -g / prims.sqrt(1.0 - x * x),
prims.PrimIDs.ACOSH: lambda x, g: g * prims.rsqrt(x * x - 1.0),
- prims.PrimIDs.ADD: lambda g: (g, g),
prims.PrimIDs.ASIN: lambda x, g: g / prims.sqrt(1.0 - x * x),
prims.PrimIDs.ASINH: lambda x, g: g * prims.rsqrt(1.0 + x * x),
prims.PrimIDs.ATAN: lambda x, g: g / (1.0 + x * x),
prims.PrimIDs.ATANH: lambda x, g: g / (1.0 - x * x),
- prims.PrimIDs.COS: lambda x, g: prims.mul(g, -prims.sin(x)),
prims.PrimIDs.COSH: lambda x, g: prims.mul(g, prims.sinh(x)),
- prims.PrimIDs.DIV: lambda x, y, g: (g / y, -g * x / (y**2)),
- prims.PrimIDs.ERF: lambda x, g: g * 2.0 / math.sqrt(math.pi) * prims.exp(-x * x),
prims.PrimIDs.ERFC: lambda x, g: -g * 2.0 / math.sqrt(math.pi) * prims.exp(-x * x),
prims.PrimIDs.ERFINV: lambda result, g: g * 0.5 * math.sqrt(math.pi) * prims.exp(result**2),
prims.PrimIDs.ERFCINV: lambda result, g: -g * 0.5 * math.sqrt(math.pi) * prims.exp(result**2),
prims.PrimIDs.EXP2: lambda result, g: g * result * math.log(2.0),
prims.PrimIDs.EXPM1: lambda result, g: g * (result + 1.0),
prims.PrimIDs.LGAMMA: lambda x, g: g * prims.digamma(x),
- prims.PrimIDs.MUL: lambda x, y, g: (g * y, g * x),
prims.PrimIDs.NDTRI: lambda result, g: g * prims.exp(0.5 * result**2) * math.sqrt(2.0 * math.pi),
- prims.PrimIDs.SIN: lambda x, g: prims.mul(g, prims.cos(x)),
prims.PrimIDs.SINH: lambda x, g: prims.mul(g, prims.cosh(x)),
- prims.PrimIDs.SUB: lambda g: (g, -g),
prims.PrimIDs.SQRT: lambda result, g: g / (2.0 * result),
- prims.PrimIDs.FULL: NoPullback(num_args=2),
- prims.PrimIDs.EQ: ZeroBackward(num_args=2),
prims.PrimIDs.NE: ZeroBackward(num_args=2),
- prims.PrimIDs.GE: ZeroBackward(num_args=2),
prims.PrimIDs.GT: ZeroBackward(num_args=2),
prims.PrimIDs.LE: ZeroBackward(num_args=2),
- prims.PrimIDs.LT: ZeroBackward(num_args=2),
- prims.PrimIDs.LOG: lambda x, g: g / x,
prims.PrimIDs.LOG10: lambda x, g: g / (x * 2.302585092994046),
prims.PrimIDs.LOG1P: lambda x, g: g / (x + 1),
prims.PrimIDs.LOG2: lambda x, g: g / (x * 0.6931471805599453),
- prims.PrimIDs.NEG: lambda g: -g,
prims.PrimIDs.FMOD: lambda x, y, g: (g, -g * prims.trunc(x / y)),
}
-@dataclass(**default_dataclass_params)
-class RuleInfo:
- checker: Callable
- rule: Callable
- fw_fallback: Callable
- bw_fallback: Callable
- executor: Executor
-
-
def register_augmented_forward(op):
"""Decorator to register an augmented forward implementation for a symbol.
@@ -2520,40 +2496,6 @@ def decorator(func):
return decorator
-def register_augmented_forward_with_checker(executor, op, checker, rule):
- """Decorator to register an augmented forward implementation for a symbol.
-
- Args:
- executor (Executor): Executor to which the rule applies.
- op (Ops): Symbol for which to register the augmented forward implementation.
- checker (Callable): Function that checks if the rule should be applied.
- rule (Callable): Function that applies the rule.
- """
- fw_fallback = augmented_forward_impls.get(op, None)
- bw_fallback = backward_impls.get(op, None)
- augmented_forward_impls[executor, op] = RuleInfo(checker, rule, fw_fallback, bw_fallback, executor)
-
-
-def deregister_augmented_forward_and_backward(op):
- """Deregisters an augmented forward implementation and a backward
- implementation for a symbol.
-
- Args:
- op (Ops): Symbol for which to deregister the augmented forward
- implementation and the backward implementation.
-
- Returns:
- None
- """
- # Restore the fallback implementation if it exists
- if isinstance(augmented_forward_impls[op], RuleInfo):
- backward_impls[op] = augmented_forward_impls[op].bw_fallback
- augmented_forward_impls[op] = augmented_forward_impls[op].fw_fallback
- else:
- del augmented_forward_impls[op]
- del backward_impls[op]
-
-
def register_backward(op):
"""Decorator to register a backward implementation for a symbol.
@@ -2623,29 +2565,6 @@ def polygamma_backward(n: int, a: Proxy, g):
return None, g * polygamma(n + 1, a)
-@register_augmented_forward(prims.PrimIDs.RSQRT)
-def rsqrt_augmented(x):
- """Augmented rsqrt operation.
-
- Args:
- x (Variable): input tensor.
-
- Returns:
- VJPDual: Primal and residuals.
- """
- primal = prims.rsqrt(x)
- residuals = (primal,)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.RSQRT)
-def rsqrt_backward(result, g):
- # An alternative derivation used by JAX is -0.5 * g * rsqrt(x) / x
- # where rsqrt(x) and x are saved for the backwards pass.
- # This derivation was selected because it avoids saving the input tensor.
- return -0.5 * g * result**3.0
-
-
@register_backward(prims.PrimIDs.ATAN2)
def atan2_backward(x, y, g):
alpha = 1.0 / (x * x + y * y)
@@ -2654,32 +2573,6 @@ def atan2_backward(x, y, g):
return grad_x, grad_y
-@register_augmented_forward(prims.PrimIDs.SUM)
-def sum_aug_fwd(x, dims):
- """Augmented sum operation.
-
- Args:
- x (Variable): Tensor to be summed.
- dims (Tuple[int, ...]): Dimensions to be summed.
-
- Returns:
- VJPDual: Primal and residuals.
- """
- primal = prims.sum(x, dims)
- residuals = (
- x.shape,
- dims,
- )
-
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.SUM)
-def sum_backward(x_shape, reduced_dims, g):
- # One return per positional argument of prims.sum
- return restore_reduced_dims(g, reduced_dims, x_shape), None
-
-
@register_augmented_forward(prims.PrimIDs.VAR)
def var_aug_fwd(a, dim, *, correction):
v = prims.var(a, dim, correction=correction)
@@ -2701,13 +2594,6 @@ def var_backward(a, dim, correction, v, g):
return (2 * g * (a - mean)) / normalization_scalar
-@register_augmented_forward(prims.PrimIDs.VAR_MEAN)
-def _var_mean_aug_fwd(a, dim, *, correction):
- v, m = prims.var_mean(a, dim, correction=correction)
-
- return (v, m), (a, dim, correction, m)
-
-
def n_elem_reduced(a_ndim, a_shape, dims):
dims = utils.canonicalize_dims(a_ndim, dims)
reduction_size = 1
@@ -2722,27 +2608,6 @@ def mean_backward(a_ndim, a_shape, dims, grad):
return restore_reduced_dims(grad, dims, a_shape) * mean_local_grad
-# TODO: fix division by zero when n_elem_reduced == 0 or when mean.numel == 0
-# by returning zeros_like(a) or similar.
-# TODO: fix grad when correction > n_elem_reduced.
-@register_backward(prims.PrimIDs.VAR_MEAN)
-def _var_mean_bwd(a, dim, correction, mean, grad_v, grad_m):
- n_elem_reduced = a.numel // mean.numel if a.numel != 0 else 1
-
- def mean_backward(a, dims, grad):
- mean_scale = 1.0 / n_elem_reduced
- grad = restore_reduced_dims(grad, dims, a.shape)
- return mean_scale * grad
-
- def var_backward(a, dims, correction, mean, grad):
- normalization_scalar = n_elem_reduced - correction
- grad = restore_reduced_dims(grad, dims, a.shape)
- mean = restore_reduced_dims(mean, dims, a.shape)
- return (2.0 * grad * (a - mean)) / normalization_scalar
-
- return var_backward(a, dim, correction, mean, grad_v) + mean_backward(a, dim, grad_m)
-
-
@register_augmented_forward(prims.PrimIDs.PAD)
def pad_aug_fwd(a, padding_value, padding_config):
return VJPDual((prims.pad(a, padding_value, padding_config),), (a, padding_config))
@@ -2812,42 +2677,15 @@ def grad_chooser_backward(primal, x, x_shape, reduced_dims, g):
return out
-register_backward(prims.PrimIDs.AMAX)(grad_chooser_backward)
register_backward(prims.PrimIDs.AMIN)(grad_chooser_backward)
-# TODO: exact same for amin, argmax, argmin
-@register_augmented_forward(prims.PrimIDs.AMAX)
-def amax_aug_fwd(x, dims):
- """Augmented amax operation.
-
- Args:
- x (Variable): Tensor to compute amax on.
- dims (Tuple[int, ...]): Dimensions to compute amax over.
-
- Returns:
- VJPDual: Primal and residuals.
- """
- primal = prims.amax(x, dims)
-
- residuals = (
- primal,
- x,
- x.shape,
- dims,
- )
-
- return VJPDual(primal, residuals)
-
-
@register_augmented_forward(prims.PrimIDs.AMIN)
def amin_aug_fwd(x, dims):
"""Augmented amin operation.
-
Args:
x (Variable): Tensor to compute amin on.
dims (Tuple[int, ...]): Dimensions to compute amin over.
-
Returns:
VJPDual: Primal and residuals.
"""
@@ -2863,26 +2701,6 @@ def amin_aug_fwd(x, dims):
return VJPDual(primal, residuals)
-@register_augmented_forward(prims.PrimIDs.EXP)
-def exp_aug_fwd(x):
- """Augmented exp operation.
-
- Args:
- x (Variable): Tensor to be exponentiated.
-
- Returns:
- VJPDual: Primal and residuals.
- """
- primal = prims.exp(x)
- residuals = (primal,)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.EXP)
-def exp_backward(result, g):
- return g * result
-
-
@register_augmented_forward(prims.PrimIDs.POW)
def pow_aug_fed(x, y):
"""Augmented the pow operation.
@@ -2927,106 +2745,11 @@ def tan_backward(result, g):
return g * (1 + result * result)
-@register_augmented_forward(prims.PrimIDs.TANH)
-def tanh_aug_fwd(x):
- """Augmented tanh operation.
-
- Args:
- x (Variable): Tensor to be passed to tanh.
-
- Returns:
- VJPDual: Primal and residuals.
- """
- primal = prims.tanh(x)
- residuals = (primal,)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.TANH)
-def tanh_backward(result, g):
- return g * (1.0 - result * result)
-
-
# NOTE: Jax uses np.argsort in its transpose vjp computation
def _argsort(seq):
return sorted(range(len(seq)), key=seq.__getitem__)
-@register_augmented_forward(prims.PrimIDs.TRANSPOSE)
-def transpose_aug_fwd(a, permutation):
- primal = prims.transpose(a, tuple(permutation))
- residuals = (permutation,)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.TRANSPOSE)
-def transpose_backward(permutation, g):
- undo = _argsort(permutation)
- return prims.transpose(g, tuple(undo))
-
-
-@register_augmented_forward(prims.PrimIDs.RESHAPE)
-def reshape_aug_fwd(a, shape):
- primal = prims.reshape(a, shape)
- residuals = (a.shape,)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.RESHAPE)
-def reshape_backward(orig_shape, g):
- return prims.reshape(g, orig_shape)
-
-
-@register_augmented_forward(prims.PrimIDs.SLICE)
-def slice_aug_fwd(a, start_indices, end_indices, strides):
- primal = prims.slice_prim(a, start_indices, end_indices, strides)
- residuals = (a.shape, start_indices, end_indices, strides)
- return VJPDual(primal, residuals)
-
-
-# Adapted from https://github.com/google/jax/blob/main/jax/_src/lax/slicing.py#L768
-@register_backward(prims.PrimIDs.SLICE)
-def slice_backward(shape, start_indices, end_indices, strides, g):
- padding = None
- if strides is None or np.all(np.equal(strides, 1)):
- padding = tuple(zip(start_indices, np.subtract(shape, end_indices), (0,) * len(start_indices)))
- else:
- real_limits = np.add(
- start_indices,
- np.where(np.equal(g.shape, 0), 0, np.add(1, np.multiply(np.subtract(g.shape, 1), strides))),
- )
- padding = tuple(zip(start_indices, np.subtract(shape, real_limits), np.subtract(strides, 1)))
-
- # We used NumPy arithmetics above, but the current infra expects Python ints.
- padding = tree_map(int, padding)
- result = prims.pad(g, const_as(0, g.dtype), padding)
-
- return result
-
-
-@register_augmented_forward(prims.PrimIDs.BROADCAST_IN_DIM)
-def broadcast_in_dim_aug_fwd(a: Proxy, shape: Sequence[int], broadcast_dimensions: Sequence[int]) -> VJPDual:
- primal = prims.broadcast_in_dim(a, shape, broadcast_dimensions)
- residuals = (a, shape, broadcast_dimensions)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.BROADCAST_IN_DIM)
-def broadcast_in_dim_backward(a, shape, broadcast_dimensions, g):
- from thunder.torch import sum
-
- # If g is None, then the primal was a constant and the pullback is zero.
- # TODO: implement None propagation in the VJP infrastructure so that we don't need to do this.
- if g is None:
- return None, None, None
- unit_dims = tuple(i for i, s in enumerate(a.shape) if s == 1)
- bcast_dims = tuple(b for i, b in enumerate(broadcast_dimensions) if i not in unit_dims)
- reduce_dims = tuple(s for i, s in enumerate(range(len(shape))) if i not in bcast_dims)
- g = sum(g, reduce_dims)
- g = unsqueeze(g, unit_dims)
- return g
-
-
@register_augmented_forward(prims.PrimIDs.DEVICE_PUT)
def device_put_aug_fwd(a: TensorProxy, device: Device) -> TensorProxy:
primal = prims.device_put(a, device)
@@ -3039,19 +2762,6 @@ def device_put_backward(orig_device, g):
return prims.device_put(g, orig_device), None
-@register_augmented_forward(prims.PrimIDs.CONVERT_ELEMENT_TYPE)
-def convert_element_type_aug_fwd(a: Proxy, dtype: dtypes.dtype) -> VJPDual:
- primal = prims.convert_element_type(a, dtype)
- residuals = (a.dtype if isinstance(a, TensorProxy) else (a.python_type if isinstance(a, NumberProxy) else type(a)),)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.CONVERT_ELEMENT_TYPE)
-def convert_element_type_backward(a_dtype, g):
- # perform cast back to input type during backward
- return prims.convert_element_type(g, a_dtype), None
-
-
@register_augmented_forward(prims.PrimIDs.CONVOLUTION)
def convolution_aug_fwd(
a: Proxy,
@@ -3253,32 +2963,6 @@ def pad_transpose_and_push_groups_into_batches(t):
return (input_grad, weight_grad, bias_grad)
-@register_augmented_forward("torch.nn.functional.cross_entropy")
-def cross_entropy_aug_fwd(
- input: Proxy,
- target: Proxy,
- weight=None,
- size_average=None,
- ignore_index=-100,
- reduce=None,
- reduction="mean",
- label_smoothing=0.0,
-) -> VJPDual:
- from thunder.torch import cross_entropy
-
- primal = cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
- residuals = (input, target, weight, reduction, ignore_index, label_smoothing)
- return VJPDual(primal, residuals)
-
-
-@register_backward("torch.nn.functional.cross_entropy")
-def cross_entropy_backward(input, target, weight, reduction, ignore_index, label_smoothing, g):
- from thunder.torch import cross_entropy_backward
-
- ginput = cross_entropy_backward(g, input, target, weight, reduction, ignore_index, label_smoothing)
- return ginput
-
-
@register_augmented_forward("torch.log_softmax")
def log_softmax_aug_fwd(input: TensorProxy, dim: int, *, dtype=None) -> VJPDual:
from thunder.torch import log_softmax
@@ -3408,64 +3092,6 @@ def softmax_backward(primal, dim, g):
return primal * (g - (primal * g).sum(dim, keepdim=True))
-@register_augmented_forward(prims.PrimIDs.MATMUL)
-def matmul_aug_fwd(a: TensorProxy, b: TensorProxy) -> VJPDual:
- primal = prims.matmul(a, b)
- residuals = (a, b)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.MATMUL)
-def matmul_backward(a, b, g):
- from thunder.torch import sum
-
- last_dim = (-1,)
- first_dim = (-2,)
- if a.ndim == 1 and b.ndim == 1:
- return g * b, g * a
-
- if b.ndim == 1:
- ga = unsqueeze(g, last_dim) @ unsqueeze(b, last_dim).mT
- gb = a.mT @ unsqueeze(g, last_dim)
- if g.ndim > 1:
- gb = squeeze(gb, last_dim)
- gb = sum(gb, tuple(range(gb.ndim - 1)))
- return ga, gb
-
- if a.ndim == 1:
- ga = unsqueeze(g, first_dim) @ b.mT
- if g.ndim > 1:
- ga = sum(ga, tuple(range(ga.ndim - 1)))
- gb = unsqueeze(a, first_dim).mT @ unsqueeze(g, first_dim)
- return ga, gb
-
- return g @ b.mT, a.mT @ g
-
-
-@register_augmented_forward(prims.PrimIDs.LINEAR)
-def linear_aug_fwd(a: TensorProxy, b: TensorProxy, c: TensorProxy | None) -> VJPDual:
- primal = prims.linear(a, b, c)
- residuals = (a, b, c)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.LINEAR)
-def linear_backward(a, b, c, g):
- from thunder.torch import matmul, sum
-
- first_dim = (-2,)
- ga = matmul(g.reshape(-1, g.shape[-1]), b).reshape(a.shape)
- if a.ndim == 1:
- gb = matmul(unsqueeze(g, first_dim).mT, unsqueeze(a, first_dim))
- else:
- gb = matmul(g.reshape(-1, g.shape[-1]).mT, a.reshape(-1, a.shape[-1]))
- assert list(gb.shape) == list(b.shape), f"linear_backward: {gb.shape} != {b.shape}"
- if c is None:
- return ga, gb, None
- gc = sum(g, tuple(range(g.ndim - 1))) if g.ndim > 1 else g
- return ga, gb, gc
-
-
def iter_bound_symbols(bound_symbols):
"""Iterate over bound symbols, skipping symbols that are not supported by
the transforms infrastructure.
@@ -3558,31 +3184,6 @@ def decomposed_fn_backward_rule(decomposed_fn, args, kwargs, saved_for_backward,
return result
-@register_augmented_forward(prims.PrimIDs.CAT)
-def cat_aug_fwd(tensors: list[TensorProxy], dim: int) -> VJPDual:
- primal = prims.cat(tensors, dim)
- residuals = (
- type(tensors),
- [t.shape[dim] for t in tensors],
- dim,
- )
-
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.CAT)
-def cat_backward(
- tensors_seq_type: type, tensor_dim_lens: list[int], dim: int, g: TensorProxy
-) -> tuple[Sequence[TensorProxy]]:
- grads = []
-
- slice_start = 0
- for dim_len in tensor_dim_lens:
- grads.append(slice_in_dim(g, slice_start, slice_start + dim_len, dim=dim))
- slice_start += dim_len
- return (tensors_seq_type(grads),)
-
-
@register_augmented_forward("torch.Tensor.contiguous")
@register_augmented_forward("torch.contiguous")
def contiguous_aug_fwd(x: TensorProxy, /, *, memory_format: torch.memory_format = torch.contiguous_format) -> VJPDual:
@@ -3599,25 +3200,11 @@ def contiguous_backward(*residuals_and_grad) -> TensorProxy:
return g
-@register_augmented_forward(prims.PrimIDs.WHERE)
-def where_aug_fwd(condition: TensorProxy, x: TensorProxy, y: TensorProxy) -> VJPDual:
- primal = prims.where(condition, x, y)
- residuals = (condition,)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.WHERE)
-def where_backward(condition, g):
- return prims.where(condition, g, 0.0), prims.where(condition, 0.0, g)
-
-
-@register_augmented_forward(prims.PrimIDs.RECIPROCAL)
def reciprocal_aug_fwd(a: TensorProxy) -> VJPDual:
primal = reciprocal(a)
return VJPDual(primal, (primal,))
-@register_backward(prims.PrimIDs.RECIPROCAL)
def reciprocal_backward(primal, g):
return -g * primal * primal
@@ -3631,38 +3218,6 @@ def reciprocal_joint_forward_backward_rule(a: TensorProxy) -> TensorProxy:
return result
-@register_augmented_forward(prims.PrimIDs.SQUEEZE)
-def squeeze_aug_fwd(a: TensorProxy, dims: Sequence[int]) -> VJPDual:
- primal = squeeze(a, dims)
- residuals = (dims,)
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.SQUEEZE)
-def squeeze_backward(dims: Sequence[int], g: TensorProxy) -> TensorProxy:
- return unsqueeze(g, dims)
-
-
-@register_augmented_forward(prims.PrimIDs.TAKE)
-def take_aug_fwd(x: TensorProxy, index: TensorProxy, dim: int) -> VJPDual:
- primal = prims.take(x, index, dim)
- residuals = (
- x.shape,
- x.device,
- x.dtype,
- index,
- dim,
- )
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.TAKE)
-def take_backward(
- shape: Sequence[int], device: Device, dtype: dtypes.dtype, index: TensorProxy, dim: int, g: TensorProxy
-):
- return prims.index_add(prims.full(shape, fill_value=0, device=device, dtype=dtype), index, g, dim)
-
-
@register_augmented_forward("torch.index_put")
def index_put_aug_fwd(
a: TensorProxy, /, indices: Sequence[TensorProxy], values: TensorProxy, accumulate: bool = False
@@ -3701,33 +3256,11 @@ def index_put_backward(indices: Sequence[TensorProxy], values: TensorProxy, accu
return clang.index_put(g, indices, ltorch.zeros_like(values), False), g_values
-@register_augmented_forward(prims.PrimIDs.TAKE_ALONG_AXIS)
-def take_along_axis_aug_fwd(x: TensorProxy, index: TensorProxy, dim: int) -> VJPDual:
- primal = prims.take_along_axis(x, index, dim)
- residuals = (
- x.shape,
- x.device,
- x.dtype,
- index,
- dim,
- )
- return VJPDual(primal, residuals)
-
-
-@register_backward(prims.PrimIDs.TAKE_ALONG_AXIS)
-def take_along_axis_backward(
- shape: Sequence[int], device: Device, dtype: dtypes.dtype, index: TensorProxy, dim: int, g: TensorProxy
-):
- return prims.scatter_add(prims.full(shape, fill_value=0, device=device, dtype=dtype), index, g, dim)
-
-
-@register_augmented_forward(prims.PrimIDs.UNIFORM)
def uniform_aug_fwd(shape, minval, maxval, *, device, dtype):
primal = prims.uniform(shape, minval, maxval, device=device, dtype=dtype)
return VJPDual(primal, (primal, minval, maxval))
-@register_backward(prims.PrimIDs.UNIFORM)
def uniform_backward(primal, minval, maxval, g):
# uniform is implemented as (maxval - minval) * uniform(shape, 0, 1) + minval
unscaled_primal = (primal - minval) / (maxval - minval)
@@ -3739,29 +3272,21 @@ def uniform_backward(primal, minval, maxval, g):
nondifferentiable_vjp_symbols = (prims.PrimIDs.BITWISE_AND, prims.PrimIDs.SIGNBIT, prims.PrimIDs.FULL)
-def get_executor_specific_aug_fwd_rule(symbol) -> RuleInfo | None:
- """Get executor specific augmented forward rule.
+def is_constant_for_vjp(symbol: prims.Symbol) -> bool:
+ """Check if a symbol is constant for the VJP transform.
Args:
- symbol (prims.Symbol): Symbol to get the rule for.
+ symbol (prims.Symbol): Symbol to check.
Returns:
- RuleInfo: Rule info for the symbol.
+ bool: True if the symbol is constant, False otherwise.
"""
- cd = get_compile_data()
- if cd is None:
- return None
-
- # Search for the executor specific rules. When there are multiple rules
- # for the same symbol, we use the left-most executor in the list (i.e.
- # the one with the highest priority) and we fallback to the next one if
- # the checker returns False.
- for executor in cd.executors_list:
- candidate = augmented_forward_impls.get((executor, symbol.sym.id))
- if isinstance(candidate, RuleInfo) and candidate.checker(*symbol.args, **symbol.kwargs):
- return candidate
-
- return None
+ are_all_args_non_differentiable = not any(isinstance(arg, (FloatProxy, TensorProxy)) for arg in symbol.flat_args)
+ return (
+ are_all_args_non_differentiable
+ or symbol.are_all_args_constant
+ or symbol.sym.id in nondifferentiable_vjp_symbols
+ )
def vjp_symbol_mapper(symbol: prims.Symbol, *args, **kwargs):
@@ -3776,7 +3301,7 @@ def vjp_symbol_mapper(symbol: prims.Symbol, *args, **kwargs):
Callable: A function that computes the VJP of the symbol.
"""
# Constant case
- if symbol.are_all_args_constant or symbol.sym.id in nondifferentiable_vjp_symbols:
+ if is_constant_for_vjp(symbol):
def vjp_impl_const(symbol, *args, **kwargs):
args, kwargs = tree_map(lambda x: x.primal if isinstance(x, VJPDual) else x, (args, kwargs))
@@ -3790,15 +3315,8 @@ def vjp_impl_const(symbol, *args, **kwargs):
# Normal case, we have a proxy tangent
vjp_impl = augmented_forward_impls.get(symbol.sym.id)
- vjp_impl = get_executor_specific_aug_fwd_rule(symbol) or vjp_impl
-
- if isinstance(vjp_impl, RuleInfo):
- # We should use this rule only if checker returns True for the current
- # symbol's arguments
- if vjp_impl.checker(*symbol.args, **symbol.kwargs):
- vjp_impl = vjp_impl.rule
- else:
- vjp_impl = vjp_impl.fw_fallback
+ if _get_gradfn(symbol) is not None:
+ vjp_impl, backward_fn = make_aug_forward_and_backward(symbol)
if vjp_impl is None:
# We could not find a VJP for this symbol, so we try to decompose it
@@ -3809,6 +3327,7 @@ def vjp_impl_const(symbol, *args, **kwargs):
# It could be a torch.dropout with 0.0 probability, so we skip it
if symbol.sym.id == "torch.nn.functional.dropout":
return None
+ print(f"VJP for {symbol} is not implemented")
raise NotImplementedError(f"VJP for {symbol.sym.id} is not implemented")
def _vjp_impl(*args, **kwargs):
@@ -3842,6 +3361,9 @@ def check_bsym_for_vjp(bsym):
if bsym.sym.id in backward_impls and bsym.sym.id in augmented_forward_impls:
return True
+ if bsym.sym.id in _grad_fn_map:
+ return True
+
# We could not find a VJP for this symbol, so we try to decompose it
# into sub-symbols and check if they are supported
if len(bsym.subsymbols) > 0 and not bsym.sym.is_prim:
@@ -3925,6 +3447,8 @@ def put_grad(v: Variable, val: Any) -> None:
elif isinstance(v, Sequence) and val is None:
# broadcast None to the right shape
safe_map(put_grad, v, [None] * len(v))
+ elif isinstance(v, Sequence) and isinstance(val, Sequence):
+ safe_map_flat(put_grad, v, val)
else:
# Skip writing to constants
pass
@@ -3943,7 +3467,7 @@ def put_grad(v: Variable, val: Any) -> None:
# Otherwise, we will need to rewrite the pullback functions
cotangents = tree_flatten(cotangents)[0]
residuals = forward_env[symbol_output[0].name].residuals
- if symbol.are_all_args_constant or symbol.sym.id in nondifferentiable_vjp_symbols:
+ if is_constant_for_vjp(symbol):
# We can skip the pullback if all the arguments are constant
continue
@@ -3955,17 +3479,15 @@ def put_grad(v: Variable, val: Any) -> None:
if symbol.sym.id == "torch.nn.functional.dropout" and not symbol.subsymbols:
# We can skip the pullback if the dropout probability is 0.0
# Assuming that the dropout symbol has the same output and argument
- # https://github.com/Lightning-AI/lightning-thunder/issues/906
assert symbol.output.name == symbol.args[0].name, "Dropout symbol has a different output and argument"
if symbol.args[1] == 0.0 or symbol.args[2] is False:
continue
backward = backward_impls.get(symbol.sym.id)
aug_forward = augmented_forward_impls.get(symbol.sym.id)
- aug_forward = get_executor_specific_aug_fwd_rule(symbol) or aug_forward
- if isinstance(aug_forward, RuleInfo):
- backward = backward_impls[aug_forward.executor, symbol.sym.id]
+ if _get_gradfn(symbol) is not None:
+ aug_forward, backward = make_aug_forward_and_backward(symbol)
if backward is None:
if len(symbol.subsymbols) > 0 and not isinstance(symbol.sym.id, prims.PrimIDs):
@@ -3980,9 +3502,11 @@ def put_grad(v: Variable, val: Any) -> None:
# If the backward returns a dict, we assume that it is a dict of
# forward arguments to the corresponding
# gradients/cotangents/adjoints/sensitivities.
+ used_names = set()
for i, (k, v) in enumerate(inspect.signature(aug_forward).parameters.items()):
if v.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD):
put_grad(symbol.args[i], result.get(k, None))
+ used_names.add(k)
# For developer convenience, we allow using the name from the
# forward meta in addition to the name from the augmented forward
@@ -3991,7 +3515,8 @@ def put_grad(v: Variable, val: Any) -> None:
# precedence.
for i, (k, v) in enumerate(inspect.signature(symbol.sym.meta).parameters.items()):
if v.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD):
- put_grad(symbol.args[i], result.get(k, None))
+ if k not in used_names:
+ put_grad(symbol.args[i], result.get(k, None))
continue
if not isinstance(result, Sequence):
@@ -4028,7 +3553,7 @@ def is_differentiable(arg):
result = tuple(next(iter_result) if is_differentiable(arg) else None for arg in symbol.args)
- # See https://github.com/Lightning-AI/lightning-thunder/issues/977.
+ # See "Backward impl for ops of the type Sequence[TensorProxy], ... -> ... results in None grads."
# This is a temporary workaround.
if symbol.sym.id in (prims.PrimIDs.CAT, "torch.cat", "torch.stack"):
safe_map_flat(put_grad, symbol.args, result)
@@ -4065,7 +3590,8 @@ def vjp_call_metafunc(detached: bool, primals, cotangents, trace: Trace, **kwarg
# TODO: Can't use a Symbol here because mixed executor sybsymbols seem to be
-# unsupported. See https://github.com/Lightning-AI/lightning-thunder/issues/1308
+# unsupported. See issue "Could not find an executor for bound symbol when its subsymbols
+# are not fully supported by a single executor"
vjp_call = partial(
vjp_call_metafunc, False
) # Symbol(id=Transforms.VjpOp, name="vjp_call", meta=partial(vjp_call_metafunc, False))
@@ -4190,7 +3716,7 @@ def unpacking_fn(saved_for_backward, cotangents):
# NOTE: Returning namedtuples from compiled functions doesn't work. See:
-# https://github.com/Lightning-AI/lightning-thunder/issues/881
+# "Allow returning namedtuples from compiled functions"
# Note [Grad forward output spec]
# If it did work it would be nice to use this namedtuple
# instead of the plain tuple or dict that we're using now.
@@ -4372,10 +3898,13 @@ def decorator(func):
def maybe_downcast_to(dtype, args):
allowed_downcast_types = (dtypes.float16, dtypes.bfloat16, dtypes.float32)
- if all(tree_map(lambda a: a.dtype in allowed_downcast_types, args)):
- return tree_map(lambda a: maybe_convert_to_dtype(a, dtype), args)
- else:
- return args
+
+ def map_fn(a):
+ if isinstance(a, TensorProxy) and a.dtype in allowed_downcast_types:
+ return maybe_convert_to_dtype(a, dtype)
+ return a
+
+ return tree_map(map_fn, args)
@register_autocast_rule("torch.matmul")
diff --git a/thunder/core/utils.py b/thunder/core/utils.py
index 9dfb95e767..1edebe70e4 100644
--- a/thunder/core/utils.py
+++ b/thunder/core/utils.py
@@ -812,8 +812,8 @@ def safe_map(f, *args):
def safe_map_flat(f, *args):
def convert_sequences_to_tuple(x):
- if isinstance(x, Sequence):
- return tuple(x)
+ if not isinstance(x, str) and isinstance(x, Sequence) and not isinstance(x, Proxy):
+ return tuple(convert_sequences_to_tuple(y) for y in x)
return x
args_flat_spec = safe_map(lambda x: tree_flatten(convert_sequences_to_tuple(x)), args)
diff --git a/thunder/core/vjp_utils.py b/thunder/core/vjp_utils.py
index 5f5a67febd..3c8128ffeb 100644
--- a/thunder/core/vjp_utils.py
+++ b/thunder/core/vjp_utils.py
@@ -1,4 +1,3 @@
-import copy
import inspect
from inspect import Parameter, Signature
from itertools import chain
@@ -13,6 +12,9 @@
from thunder.core.transform_common import dce
+_cache = {}
+
+
def make_aug_forward_and_backward(bsym: BoundSymbol) -> tuple[Callable, Callable]:
"""
Given a bound symbol, return a pair of forward and backward functions
@@ -33,13 +35,20 @@ def make_aug_forward_and_backward(bsym: BoundSymbol) -> tuple[Callable, Callable
A pair of forward and backward functions.
"""
import thunder
- from thunder.core.transforms import _grad_fn_map
+ from thunder.common import _make_cache_key
+ from thunder.core.transforms import _get_gradfn, eval_trace
- joint_forward_backward = _grad_fn_map.get(bsym.sym.id, None)
+ joint_forward_backward = _get_gradfn(bsym)
utils.check(
joint_forward_backward is not None,
lambda: f"Cannot generate forward and backward functions for {bsym.sym.name}",
)
+
+ key = (bsym.sym, subkey := _make_cache_key(bsym.args, bsym.kwargs))
+ cached_result = _cache.get(key, None) if subkey is not None else None
+ if cached_result is not None:
+ return cached_result
+
joint_trace = thunder.trace(inline_trace=False, use_dce=False)(joint_forward_backward, *bsym.args, **bsym.kwargs)
consumers = utils.consumers(joint_trace)
@@ -67,20 +76,31 @@ def find_backward_output(forward_input):
bw_outputs_args = tree_map(find_backward_output, joint_trace.args)
bw_outputs_kwargs = tree_map(find_backward_output, joint_trace.kwargs)
meta_parameters = inspect.signature(bsym.sym.meta).parameters
+ meta_parameters = {
+ name: param
+ for name, param in meta_parameters.items()
+ if param.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.POSITIONAL_ONLY)
+ }
bw_outputs = {name: bw_output for name, bw_output in utils.safe_zip(meta_parameters, bw_outputs_args)}
bw_outputs = bw_outputs | bw_outputs_kwargs
flat_bw_outputs, _ = tree_flatten(bw_outputs)
- backward_bsyms = utils.find_producer_symbols(joint_trace, flat_bw_outputs, bw_inputs)
- unpacking_ops = (
+ backward_bsyms = utils.find_producer_symbols(joint_trace, flat_bw_outputs, tree_flatten(bw_inputs)[0])
+ skip = (
prims.PrimIDs.UNPACK_EMPTY_DICT,
prims.PrimIDs.UNPACK_KEY,
prims.PrimIDs.UNPACK_SEQUENCE,
prims.PrimIDs.UNPACK_TRIVIAL,
+ prims.PrimIDs.GET_GRAD,
)
- backward_bsyms = [bsym for bsym in backward_bsyms if bsym.sym.id not in unpacking_ops]
+ backward_bsyms = [bsym for bsym in backward_bsyms if bsym.sym.id not in skip]
backward_bsyms.append(prims.python_return.bind(bw_outputs, output=()))
+ forward_input_proxies = tree_flatten((joint_trace.args, joint_trace.kwargs))[0]
+ forward_input_proxies = [arg for arg in forward_input_proxies if isinstance(arg, Proxy)]
+ forward_bsyms = utils.find_producer_symbols(joint_trace, tree_flatten(joint_trace.output)[0], forward_input_proxies)
+ backward_bsyms = [bsym for bsym in backward_bsyms if bsym not in forward_bsyms]
+
# Find required info from forward trace for backward trace
backward_producers = utils.producers(backward_bsyms)
saved_for_backward = []
@@ -91,7 +111,37 @@ def find_backward_output(forward_input):
if arg not in backward_producers and variableify(arg) not in map(variableify, tree_flatten(bw_inputs)[0]):
saved_for_backward.append(arg)
- backward_params = [Parameter(x.name, Parameter.POSITIONAL_OR_KEYWORD) for x in chain(saved_for_backward, bw_inputs)]
+ saved_for_backward = list({variableify(arg): arg for arg in saved_for_backward}.values())
+
+ # Augment forward trace to include saved_for_backward as output
+ augmented_forward_trace = from_trace(joint_trace)
+ augmented_forward_trace.bound_symbols = [
+ b for b in joint_trace.bound_symbols if b.sym.id not in (PrimIDs.PUT_GRAD, PrimIDs.GET_GRAD)
+ ]
+ return_bsym = augmented_forward_trace.bound_symbols[-1]
+ assert return_bsym.sym.id == PrimIDs.RETURN
+ augmented_forward_trace.bound_symbols[-1] = prims.python_return.bind(
+ (joint_trace.output, saved_for_backward), output=()
+ )
+ # Remove put/get grad and backward symbols from augmented forward trace
+ augmented_forward_trace = dce(augmented_forward_trace)
+
+ # Check if any of the bound symbols in the backward trace are also in the
+ # augmented forward trace
+ # If so, remove them from the backward trace
+ same_bsyms = set(augmented_forward_trace.bound_symbols) & set(backward_bsyms)
+ if same_bsyms:
+ backward_bsyms = [bsym for bsym in backward_bsyms if bsym not in same_bsyms]
+ additional_saved = [o for bsym in same_bsyms for o in bsym.flat_proxy_outs]
+ saved_for_backward += list({variableify(arg): arg for arg in additional_saved}.values())
+ augmented_forward_trace.bound_symbols[-1] = prims.python_return.bind(
+ (joint_trace.output, saved_for_backward), output=()
+ )
+
+ backward_params = [
+ Parameter(getattr(x, "name", f"arg{i}"), Parameter.POSITIONAL_OR_KEYWORD)
+ for i, x in enumerate(chain(saved_for_backward, bw_inputs))
+ ]
backward_signature = Signature(backward_params)
def backward_fn():
@@ -106,15 +156,15 @@ def backward_fn():
backward_trace.kwargs = {}
backward_trace.bound_symbols = backward_bsyms
- # Augment forward trace to include saved_for_backward as output
- augmented_forward_trace = from_trace(joint_trace)
- augmented_forward_trace.bound_symbols = copy.copy(joint_trace.bound_symbols)
- return_bsym = augmented_forward_trace.bound_symbols[-1]
- assert return_bsym.sym.id == PrimIDs.RETURN
- augmented_forward_trace.bound_symbols[-1] = prims.python_return.bind(
- (joint_trace.output, saved_for_backward), output=()
- )
- # Remove put/get grad from augmented forward trace
- augmented_forward_trace = dce(augmented_forward_trace)
+ # Creating new functions instead of using partial to avoid limitations in
+ # codeutils.get_siginfo
+ # https://github.com/Lightning-AI/lightning-thunder/blob/main/thunder/core/codeutils.py#L349-L353
+ def fw_fn(*args, **kwargs):
+ return eval_trace(augmented_forward_trace, *args, **kwargs)
+
+ def bw_fn(*args, **kwargs):
+ return eval_trace(backward_trace, *args, **kwargs)
+
+ _cache[key] = fw_fn, bw_fn
- return augmented_forward_trace.python_callable(), backward_trace.python_callable()
+ return fw_fn, bw_fn
diff --git a/thunder/distributed/__init__.py b/thunder/distributed/__init__.py
index ecb280149f..0a0eb943d9 100644
--- a/thunder/distributed/__init__.py
+++ b/thunder/distributed/__init__.py
@@ -68,9 +68,11 @@ def skip_data_parallel_grad_sync() -> None:
def _sync_grads(module: torch.nn.Module) -> None:
+ import thunder
+
params_with_grad = [p for p in module.parameters() if p.grad is not None]
grads = [p.grad for p in params_with_grad]
- process_group = module._lc_cd.process_group_for_ddp
+ process_group = thunder.compile_data(module).process_group_for_ddp
torch._foreach_div_(grads, process_group.size())
with tdist.distributed_c10d._coalescing_manager(group=process_group, async_ops=True) as cm:
for g in grads:
@@ -234,7 +236,7 @@ def main():
# Starts broadcasts
# TODO Make these broadcast asyncs
# TODO Perform up to two broadcasts at a time
- # https://github.com/Lightning-AI/lightning-thunder/issues/727
+ # See issue "Update ddp to use async broadcasts"
# TODO "Bucket" small tensors together before broadcasting
with torch.no_grad():
for param in model.parameters():
diff --git a/thunder/distributed/transforms/fsdp.py b/thunder/distributed/transforms/fsdp.py
index c07afce08e..ad92f5e6db 100644
--- a/thunder/distributed/transforms/fsdp.py
+++ b/thunder/distributed/transforms/fsdp.py
@@ -198,7 +198,7 @@ def maybe_swap_proxies_of_bsym_and_update_swap_map(bsym: BoundSymbol) -> bool:
lambda: f"{variableify(param)} not found in param set: {(variableify(p) for p in self.original_params)}",
)
if param not in self.param_to_bucket:
- # This path is hihly likely to be backward reduce-scatter bucketing:
+ # This path is highly likely to be backward reduce-scatter bucketing:
# when a param does not require grad, a trace could still have reduce-scatter
# and wait in its trace while the grad in the return statement is already
# replaced with `None`.
diff --git a/thunder/distributed/utils.py b/thunder/distributed/utils.py
index f34f98f3ad..9d9fa7bf42 100644
--- a/thunder/distributed/utils.py
+++ b/thunder/distributed/utils.py
@@ -52,7 +52,6 @@ def key(node: Node) -> int:
# TODO: Currently prefer the most memory-efficient way for ZeRO3,
-# https://github.com/Lightning-AI/lightning-thunder/issues/1925
# Need a strategy to balance the efficiency
# and memory usage in the future
def sort_waits_for_zero3(execution_trace):
diff --git a/thunder/examine/__init__.py b/thunder/examine/__init__.py
index 23661586bb..c75bc5f3b7 100644
--- a/thunder/examine/__init__.py
+++ b/thunder/examine/__init__.py
@@ -43,7 +43,7 @@ def __exit__(self, exc_type, exc_value, traceback):
# TODO Maybe have this print additional information and return more metadata?
-# TODO Accept kwargs for compile (like langctx)
+# TODO Accept kwargs for jit (like langctx)
# TODO Add profiling (or profiling option) to determine if we have a slowdown
# TODO If an error occurs, try to minify the program to produce a smaller sample to reproduce the error
def examine(fn: Callable, *args, show_call_stack: bool | int = False, **kwargs):
@@ -141,7 +141,7 @@ def examine(fn: Callable, *args, show_call_stack: bool | int = False, **kwargs):
return
- # Step 3 Attempts to compile the function using lightning.compile
+ # Step 3 Attempts to compile the function using thunder.jit
try:
cfn = thunder.jit(fn)
except Exception as e:
@@ -151,7 +151,7 @@ def examine(fn: Callable, *args, show_call_stack: bool | int = False, **kwargs):
)
raise e
- # Step 4 Attemps to execute the function using lightning.compile
+ # Step 4 Attempt to execute the function using thunder.jit
lc_result: Any
try:
lc_result = cfn(*args, **kwargs)
diff --git a/thunder/examine/memory_caculation.py b/thunder/examine/memory_caculation.py
index 45cc74a9a0..149349748c 100644
--- a/thunder/examine/memory_caculation.py
+++ b/thunder/examine/memory_caculation.py
@@ -22,7 +22,7 @@
"torch_wait_prim_impl",
)
-# A whitelist registry of symbols that require special memory calculation;
+# A registry of symbols that require special memory calculation;
# if not registered, the default memory calculation function is used.
memory_calculate_impls: dict[Symbol, Callable] = dict()
diff --git a/thunder/executors/apex_entropyex.py b/thunder/executors/apex_entropyex.py
index 7d5eec13fb..818199ad5b 100644
--- a/thunder/executors/apex_entropyex.py
+++ b/thunder/executors/apex_entropyex.py
@@ -10,11 +10,7 @@
from thunder.core.proxies import TensorProxy
from thunder.core.symbol import Symbol
from thunder.core.utils import check, same_shape
-from thunder.core.transforms import get_grad, put_grad, put_grads, mean_backward, sum_backward
-from thunder.core.transforms import (
- register_augmented_forward_with_checker,
- register_backward,
-)
+from thunder.core.transforms import get_grad, put_grad, put_grads, mean_backward, restore_reduced_dims
from thunder.extend import OperatorExecutor, register_executor
@@ -49,7 +45,8 @@ def apex_available() -> bool:
# TODO Consider performing the reduction as part of a traceable epilogue
-# See https://github.com/Lightning-AI/lightning-thunder/issues/1357
+# See "Update the apex cross entropy executor to put its reduction in a
+# traceable epilogue"
# NOTE Apex's cross entropy doesn't accept ignore_index >= 0, or the weight, size_average, or reduce parameters
def _apex_cross_entropy_impl(
a: torch.Tensor,
@@ -196,78 +193,6 @@ def _cross_entropy_checker(
return True
-# Check out
-# https://github.com/Lightning-AI/lightning-thunder/blob/main/dev_tutorials/thunder-add-vjp-rule.md
-# for a tutorial on how to add a VJP rule for any Symbol. We use our new
-# primitives to register a VJP rule for torch.nn.functional.cross_entropy. This
-# function is registered as the augmented forward rule for
-# torch.nn.functional.cross_entropy below
-def apex_cross_entropy_forward_rule(
- a,
- target,
- weight=None,
- size_average=None,
- ignore_index=-100,
- reduce=None,
- reduction="mean",
- label_smoothing=0.0,
-):
- loss, max_log_sum_exp = apex_xentropy(
- a,
- target=target,
- reduction=reduction,
- label_smoothing=label_smoothing,
- )
- primal = loss
- saved_for_backward = (a, target, max_log_sum_exp, reduction, label_smoothing)
- return primal, saved_for_backward
-
-
-register_augmented_forward_with_checker(
- apex_ex,
- ltorch.cross_entropy.id,
- _cross_entropy_checker,
- apex_cross_entropy_forward_rule,
-)
-
-
-# This function is the backward rule for torch.nn.functional.cross_entropy. It
-# accepts the primal output and saved_for_backward from the forward pass and
-# returns the backward output. The backward output is a tuple of the backward
-# output for each differentiable Tensor input to the forward pass. In this case,
-# the forward pass has 1 such input, so the backward output is a single Tensor.
-# This function is registered as the backward rule for
-# torch.nn.functional.cross_entropy
-@register_backward((apex_ex, ltorch.cross_entropy.id))
-def apex_cross_entropy_backward_rule(
- logits,
- labels,
- max_log_sum_exp,
- reduction,
- smoothing,
- grad,
-):
- from thunder.core.transforms import mean_backward, sum_backward
-
- if reduction == "mean":
- grad = mean_backward(max_log_sum_exp.ndim, max_log_sum_exp.shape, (0,), grad)
- elif reduction == "sum":
- grad = sum_backward(max_log_sum_exp.shape, (0,), grad)
- elif reduction == "none":
- pass
- else:
- raise ValueError(f"Invalid reduction: {reduction}")
-
- grad_logits = apex_xentropy_bwd(
- grad,
- logits,
- target=labels,
- max_log_sum_exp=max_log_sum_exp,
- label_smoothing=smoothing,
- )
- return grad_logits
-
-
# Translate calls from torch.nn.functional.cross_entropy to apex_xentropy (when the checker above returns True)
def _cross_entropy_transform(
a: TensorProxy,
@@ -304,7 +229,7 @@ def _apex_cross_entropy_grad(
if reduction == "mean":
g = mean_backward(max_log_sum_exp.ndim, max_log_sum_exp.shape, (0,), g)
elif reduction == "sum":
- g, _ = sum_backward(max_log_sum_exp.shape, (0,), g)
+ g = restore_reduced_dims(g, (0,), max_log_sum_exp.shape)
# NOTE Apex's xentropy bwd requires the grad computation to be performed in fp32
a_ = a.contiguous()
diff --git a/thunder/executors/cudnn_layernormex.py b/thunder/executors/cudnn_layernormex.py
index 1a84811278..b6f260cdd7 100644
--- a/thunder/executors/cudnn_layernormex.py
+++ b/thunder/executors/cudnn_layernormex.py
@@ -19,9 +19,7 @@ def cudnn_available() -> bool:
return CUDNN_AVAILABLE
-# WARNING: cudnn executor is experimental. Tests that use cudnn might fail.\n
-# Issue for tracking support: https://github.com/Lightning-AI/lightning-thunder/issues/880~
-
+# WARNING: cudnn layernorm executor is experimental. Tests that use cudnn might fail.
from dataclasses import dataclass
from functools import lru_cache
from typing import Union, Dict
diff --git a/thunder/executors/cudnnex.py b/thunder/executors/cudnnex.py
index 20759bb085..75494cff5a 100644
--- a/thunder/executors/cudnnex.py
+++ b/thunder/executors/cudnnex.py
@@ -21,9 +21,6 @@ def cudnn_available() -> bool:
return CUDNN_AVAILABLE
-# WARNING: cudnn executor is experimental. Tests that use cudnn might fail.\n
-# Issue for tracking support: https://github.com/Lightning-AI/lightning-thunder/issues/880~
-
from dataclasses import dataclass
from functools import lru_cache
from typing import Union, Dict
@@ -38,8 +35,6 @@ def cudnn_available() -> bool:
get_grad,
put_grad,
put_grads,
- register_augmented_forward_with_checker,
- register_backward,
)
from thunder.extend import OperatorExecutor, register_executor
import thunder.torch as ltorch
@@ -341,7 +336,11 @@ def _cudnn_sdpa_forward_checker(
if d % 8 != 0 or d > 128:
return False
- return True
+ is_backward_supported = _cudnn_sdpa_backward_checker(
+ query, key, value, attn_mask, dropout_p, is_causal, scale=scale
+ )
+
+ return True and is_backward_supported
@langctx("torch")
@@ -604,99 +603,6 @@ def cudnn_sdpa_bwd_impl(
)
-@langctx("torch")
-def cudnn_sdpa_aug_fw_rule_checker(
- query: TensorProxy,
- key: TensorProxy,
- value: TensorProxy,
- attn_mask: None | TensorProxy,
- dropout_p: float,
- is_causal: bool,
- *,
- scale: None | float,
-) -> bool:
- from thunder.core.compile_data import get_compile_data
-
- cd = get_compile_data()
- if cudnn_ex in cd.executors_list:
- is_forward_supported = _cudnn_sdpa_forward_checker(
- query, key, value, attn_mask, dropout_p, is_causal, scale=scale
- )
- is_backward_supported = _cudnn_sdpa_backward_checker(
- query, key, value, attn_mask, dropout_p, is_causal, scale=scale
- )
- return is_forward_supported and is_backward_supported
- return False
-
-
-def cudnn_sdpa_aug_fw_rule(
- query,
- key,
- value,
- attn_mask=None,
- dropout_p: float = 0.0,
- is_causal: bool = False,
- *,
- scale: float | None = None,
-):
- output, softmax_stats, seed, offset = cudnn_sdpa_fwd(
- query, key, value, attn_mask, dropout_p, is_causal, scale=scale
- )
- saved_for_backward = (
- query,
- key,
- value,
- attn_mask,
- dropout_p,
- is_causal,
- scale,
- output,
- softmax_stats,
- seed,
- offset,
- )
- return output, saved_for_backward
-
-
-register_augmented_forward_with_checker(
- cudnn_ex,
- "torch.nn.functional.scaled_dot_product_attention",
- cudnn_sdpa_aug_fw_rule_checker,
- cudnn_sdpa_aug_fw_rule,
-)
-
-
-@register_backward((cudnn_ex, "torch.nn.functional.scaled_dot_product_attention"))
-def cudnn_sdpa_backward_rule(
- query: Proxy,
- key: Proxy,
- value: Proxy,
- attn_mask: None | Proxy,
- dropout_p: float,
- is_causal: bool,
- scale: None | float,
- out: Proxy,
- softmax_stats: Proxy,
- seed: Proxy,
- offset: Proxy,
- grad_out: Proxy,
-):
- return cudnn_sdpa_bwd(
- grad_out,
- query,
- key,
- value,
- attn_mask,
- dropout_p,
- is_causal,
- out,
- softmax_stats,
- seed,
- offset,
- scale=scale,
- )
-
-
@langctx("torch")
def _cudnn_sdpa_transform(
query: TensorProxy,
@@ -729,7 +635,7 @@ def _cudnn_sdpa_grad(
)
g = get_grad(primal)
- grad_query, grad_key, grad_val, grad_attn_mask = cudnn_sdpa_bwd(
+ grads = cudnn_sdpa_bwd(
g,
query,
key,
@@ -743,6 +649,11 @@ def _cudnn_sdpa_grad(
offset,
scale=scale,
)
+ if attn_mask is None:
+ grad_query, grad_key, grad_val = grads
+ else:
+ grad_query, grad_key, grad_val, grad_attn_mask = grads
+
put_grads((query, key, value), (grad_query, grad_key, grad_val))
if attn_mask is not None:
put_grad(attn_mask, grad_attn_mask)
diff --git a/thunder/executors/nvfuserex.py b/thunder/executors/nvfuserex.py
index 3e25890dde..27d4ce721b 100644
--- a/thunder/executors/nvfuserex.py
+++ b/thunder/executors/nvfuserex.py
@@ -33,7 +33,6 @@ def required_nvfuser_version() -> LooseVersion:
return LooseVersion("0.0.1")
-# NOTE We require nvFuser version 0.0.1 or greater
def nvfuser_available() -> bool:
v = nvfuser_version()
return v is not None and v >= required_nvfuser_version()
diff --git a/thunder/executors/nvfuserex_impl.py b/thunder/executors/nvfuserex_impl.py
index 14fba83ffe..7a5ebfe5ae 100644
--- a/thunder/executors/nvfuserex_impl.py
+++ b/thunder/executors/nvfuserex_impl.py
@@ -206,7 +206,7 @@ def create_fd(
# NOTE nvFuser's default max length is 1024 operations at the time of this writing
# This arbitrarily increases it to 9999
# TODO Review splititng very large fusions or removing the max length restriction completely
- # See https://github.com/Lightning-AI/lightning-thunder/issues/901
+ # See "Very large nvFuser fusions hit max_length"
fd = FusionDefinition(max_length=9999)
with fd:
# NOTE Adding constants is disabled for the moment in favor of definining them inline
@@ -524,8 +524,7 @@ class nvFuserExecutor(FusionExecutor):
def __init__(self):
super().__init__("nvfuser", version=nvfuser.version())
- # TODO: Replace this with a query to current CompileData after
- # https://github.com/Lightning-AI/lightning-thunder/pull/1517 is merged
+ # TODO: Replace this with a query to a compile option
self._use_rematerialization = True
fuel_str = os.getenv("NVFUSER_OPTIMIZATION_FUEL")
@@ -762,8 +761,6 @@ def fusion_pass(self, trace: TraceCtx) -> TraceCtx:
# TODO has_cuda_input_or_output is too restrictive a check on what should be fused
# TODO check whether a function would output a CPU tensor? -- can nvFuser fuse such operations?
# ex. device_put to a CPU device from a CUDA device
- # (mruberry) I don't know if nvFuser even attempts to fuse any operation that can go
- # cross-device today
def _should_fuse(a: Node, b: Node):
def _can_fuse_node(n: Node):
# if already merged, then node can be fused
@@ -842,8 +839,7 @@ def _can_fuse_node(n: Node):
# Some of the operations might be better placed with its consumers (for
# example residual connection in transformer block). This pass moves
- # them to the consumer. See
- # https://github.com/Lightning-AI/lightning-thunder/issues/1520
+ # them to the consumer.
if self._use_rematerialization:
fusedtrace = rematerialize(fusedtrace)
@@ -1150,7 +1146,7 @@ def _pad_check(a: TensorProxy, padding_value: Number, padding_config: tuple[int,
# nvFuser's pad op requires pad_widths to be a sequence of Python numbers
# (lo_n, hi_n, lo_{n-1}, hi_{n-1}, ...) where dimensions are counted in reverse
# as shown, and dilation is not supported.
-# This is in constrant to lightning.compile's pad primitive, which specifies padding
+# This is in constrast to thunder.jit's pad primitive, which specifies padding
# and dilation as an ndim-length list of (lo, hi, dilation) triples.
# NOTE padding_value must be an nvConstant (or nvScalar?)
def pad(
@@ -1256,7 +1252,8 @@ def squeeze(a: TensorProxy, /, dims: Sequence[int], *, fd: FusionDefinition, lc_
# register_supported(PrimIDs.TAKE, take, _take_check)
# TAKE_ALONG_AXIS is currently disabled
-# See https://github.com/NVIDIA/Fuser/issues/458
+# There was an nvFuser bug that prevented this which is now fixed; we should
+# investigate re-enabling take_along_axis.
# # TODO Check that the nvFuser version is >= 0.0.10 when this operator was added
# def take_along_axis(a: TensorProxy, /, index: TensorProxy, dim: int, *, fd: FusionDefinition, lc_to_nv_map: dict) -> Any:
# nv_a = getnv(a, fd, lc_to_nv_map)
@@ -1709,12 +1706,6 @@ def div(a: TensorProxy | Number, b: TensorProxy | Number, *, fd: FusionDefinitio
nva = getnv(a, fd, lc_to_nv_map)
nvb = getnv(b, fd, lc_to_nv_map)
- # TODO nvFuser sometimes generates an innacurate result when dividing by a number
- # Remove this workaround once the issue is fixed
- # See: https://github.com/NVIDIA/Fuser/issues/160
- if isinstance(b, Number):
- return fd.ops.mul(nva, fd.ops.reciprocal(nvb))
-
# NOTE It's currently significantly faster for nvFuser to multiply the reciprocal than divide
# return fd.ops.div(nva, nvb)
return fd.ops.mul(nva, fd.ops.reciprocal(nvb))
diff --git a/thunder/executors/passes.py b/thunder/executors/passes.py
index ab6b35abba..8f1604e718 100644
--- a/thunder/executors/passes.py
+++ b/thunder/executors/passes.py
@@ -57,12 +57,13 @@ def preserve_bsym(bsym: BoundSymbol) -> Any:
# If the executor has an execution transform, it's called and True is returned
# If no executor can execute the BoundSymbol, False is returned
def visit_helper_(bsym: BoundSymbol) -> None | bool:
- if bsym.sym.executor is not None or bsym.sym.python_impl is not None:
+ if bsym.sym.python_impl is not None:
return None
ex: Executor
for ex in executors_list:
# TODO Consider allowing operator executors to claim portions of operations
+ # TODO Should FusionExecutors be allowed to claim bsym with bsym.sym.executor?
if (isinstance(ex, OperatorExecutor) and ex.can_execute(bsym)) or (
isinstance(ex, FusionExecutor) and ex.can_fuse(bsym)
):
@@ -87,6 +88,9 @@ def visit_helper_(bsym: BoundSymbol) -> None | bool:
safe_map_flat(update_swapmap, bsym.output, out)
return True
+ if bsym.sym.executor is not None:
+ return None
+
return False
def visit_(bsym: BoundSymbol) -> transforms.VISIT_TYPE:
@@ -158,8 +162,8 @@ def transform_for_execution(trace: TraceCtx, executors_list: Sequence[Executor])
return extrace
-# NOTE: See more details for motivation in the following issue:
-# https://github.com/Lightning-AI/lightning-thunder/issues/515
+# This is needed to ensure that subsymbol changes are reflected in the Python
+# code generator.
def _update_fusion_call_ctx(bsym: BoundSymbol) -> BoundSymbol:
"""Update the call_ctx information of the fusion BoundSymbol object.
diff --git a/thunder/executors/sdpaex.py b/thunder/executors/sdpaex.py
index 8c9aa4657c..a2ff22007b 100644
--- a/thunder/executors/sdpaex.py
+++ b/thunder/executors/sdpaex.py
@@ -17,8 +17,6 @@
get_grad,
put_grad,
put_grads,
- register_augmented_forward_with_checker,
- register_backward,
)
from thunder.extend import OperatorExecutor, register_executor
@@ -50,8 +48,8 @@ def ceil_div(a: int, b: int) -> int:
def _sdpa_pad_head_dimension(a: torch.Tensor) -> torch.Tensor:
head_size = a.shape[-1]
- # NOTE short-circuit path when we already have compatible head_size
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1505
+ # If the head is already a multiple of 8, then we don't need to pad. The
+ # pad op can be quite expensive in some cases.
if head_size % 8 == 0:
return a
padding_size = ceil_div(head_size, 8) * 8 - head_size
@@ -59,8 +57,7 @@ def _sdpa_pad_head_dimension(a: torch.Tensor) -> torch.Tensor:
def _sdpa_slice_head_dimension(a: torch.Tensor, head_size: int) -> torch.Tensor:
- # NOTE short-circuit path when we already have compatible head_size
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1505
+ # ditto pad_head_dimension: the slice can be expensive, so skip if possible.
if head_size % 8 == 0:
return a
return a[:, :, :, 0:head_size]
@@ -491,8 +488,8 @@ def _scaled_dot_product_attention_fused(
*,
scale: None | float = None,
):
- # NOTE Select fused sdpa using PyTorch eager mode selection behavior
- # See https://github.com/Lightning-AI/lightning-thunder/issues/622
+ # Figure out which SDPA to use. There are performance cliffs to the various
+ # implementations, and this makes the decision cognizant of those cliffs.
backend = _fused_sdp_choice(query, key, value, attn_mask, dropout_p, is_causal, scale)
utils.check(
@@ -530,8 +527,8 @@ def _scaled_dot_product_attention_grad(
*,
scale: None | float = None,
):
- # NOTE Select fused sdpa using PyTorch eager mode selection behavior
- # See https://github.com/Lightning-AI/lightning-thunder/issues/622
+ # Figure out which SDPA to use. There are performance cliffs to the various
+ # implementations, and this makes the decision cognizant of those cliffs.
backend = _fused_sdp_choice(query, key, value, attn_mask, dropout_p, is_causal, scale)
utils.check(
@@ -640,8 +637,9 @@ def _fused_sdp_choice(
is_causal = is_causal.value
if LooseVersion(torch.__version__) < LooseVersion("2.2.0"):
- # NOTE Select fused sdpa using PyTorch eager mode selection behavior
- # See https://github.com/Lightning-AI/lightning-thunder/issues/622
+ # Figure out which SDPA to use. There are performance cliffs to the
+ # various implementations, and this makes the decision cognizant of
+ # those cliffs.
backend = torch._fused_sdp_choice(
fake_query,
fake_key,
@@ -707,112 +705,3 @@ def _scaled_dot_product_attention_checker(
execution_transform=_scaled_dot_product_attention_fused,
grad_transform=_scaled_dot_product_attention_grad,
)
-
-
-def scaled_dot_product_attention_aug_fw(
- query: TensorProxy,
- key: TensorProxy,
- value: TensorProxy,
- attn_mask: TensorProxy | None,
- dropout_p: float = 0.0,
- is_causal: bool = False,
- *,
- scale: float | None = None,
-):
- # NOTE Select fused sdpa using PyTorch eager mode selection behavior
- # See https://github.com/Lightning-AI/lightning-thunder/issues/622
- backend = _fused_sdp_choice(query, key, value, attn_mask, dropout_p, is_causal, scale)
-
- utils.check(
- backend != SpdaBackend.ERROR,
- lambda: "Unable to find valid backend for scaled_dot_product_attention.",
- )
- utils.check(
- backend != SpdaBackend.MATH,
- lambda: "The fallback to sdpa thunder reference is not implemented.",
- exception_type=NotImplementedError,
- )
-
- tensor_args = (query, key, value)
- scalar_args = (dropout_p, is_causal)
- input_args = (*tensor_args, attn_mask, *scalar_args, scale)
- if backend == SpdaBackend.FLASH_ATTENTION:
- # Use flash attention kernel
- (primal, *remaining_results, debug_attn_mask) = sdpfa_gradfwd(*tensor_args, *scalar_args, scale=scale)
- # NOTE Remaining results contains [logsumexp, *flash_attn_only_residuals, *philox_residuals]
- residuals = (*input_args, primal, *remaining_results)
- return primal, residuals
- elif backend == SpdaBackend.MEMORY_EFFICIENT:
- # Use memory efficient kernel, which supports fp32 and attention mask arguments
- (primal, logsumexp, *philox_residuals) = sdpea_gradfwd(*tensor_args, attn_mask, *scalar_args, scale=scale)
- flash_attn_only_residuals = (None,) * 4
- residuals = (*input_args, primal, logsumexp, *flash_attn_only_residuals, *philox_residuals)
- return primal, residuals
-
-
-register_augmented_forward_with_checker(
- sdpa_ex,
- "torch.nn.functional.scaled_dot_product_attention",
- _scaled_dot_product_attention_checker,
- scaled_dot_product_attention_aug_fw,
-)
-
-
-@register_backward((sdpa_ex, "torch.nn.functional.scaled_dot_product_attention"))
-def scaled_dot_product_attention_backward(
- query: Proxy,
- key: Proxy,
- value: Proxy,
- attn_mask: None | Proxy,
- dropout_p: float,
- is_causal: bool,
- scale: None | float,
- out: Proxy,
- logsumexp: Proxy,
- cum_seq_q: None | Proxy,
- cum_seq_k: None | Proxy,
- max_q: None | int,
- max_k: None | int,
- philox_seed: Proxy,
- philox_offset: Proxy,
- grad_out: Proxy,
-):
- tensor_args = (query, key, value)
- scalar_args = (dropout_p, is_causal)
- flash_attention_args = (cum_seq_q, cum_seq_k, max_q, max_k)
- philox_args = (philox_seed, philox_offset)
- use_flash_attn = all(map(lambda a: a is not None, (cum_seq_q, cum_seq_k, max_q, max_k)))
- if use_flash_attn:
- (
- grad_query,
- grad_key,
- grad_val,
- ) = sdpfa_bwd(
- grad_out,
- *tensor_args,
- out,
- logsumexp,
- *flash_attention_args,
- *scalar_args,
- *philox_args,
- scale=scale,
- )
- # grad_attn_mask is None since it is not supported by flash_attention kernel
- return grad_query, grad_key, grad_val
- else:
- (
- grad_query,
- grad_key,
- grad_val,
- grad_attn_mask,
- ) = sdpea_bwd(
- grad_out,
- *tensor_args,
- attn_mask,
- out,
- logsumexp,
- *philox_args,
- *scalar_args,
- scale=scale,
- )
- return grad_query, grad_key, grad_val, grad_attn_mask
diff --git a/thunder/executors/torch_autograd.py b/thunder/executors/torch_autograd.py
index 2eb3f1f56b..93c0c2d874 100644
--- a/thunder/executors/torch_autograd.py
+++ b/thunder/executors/torch_autograd.py
@@ -161,7 +161,7 @@ def wrapper(*args, **kwargs):
return decorator
-def split_forward_backward(func, compile_data, compile_stats, /, *args, **kwargs):
+def split_forward_backward(computation_trc, compile_data, compile_stats, /, *args, **kwargs):
from thunder import trace
from thunder.executors.passes import transform_for_execution
from thunder.executors.passes import del_last_used
@@ -170,6 +170,17 @@ def split_forward_backward(func, compile_data, compile_stats, /, *args, **kwargs
from thunder.cudagraphs import CUDAGraphExecutor
from thunder.distributed.utils import sort_waits, sort_data_parallel_syncs, sort_waits_for_zero3
from thunder.distributed.transforms import FSDPCommBucketing
+ from thunder.core.transforms import eval_trace
+
+ # TODO: the trace->func->trace could likely be simplified (and look nicer)
+ # we cannot use python_callable() here, see the old repos 2458
+ if not isinstance(computation_trc, TraceCtx):
+ # for the legacy codepath
+ func = computation_trc
+ else:
+
+ def func(*args):
+ return eval_trace(computation_trc, *args)
utils.check(compile_data is not None, lambda: "`compile_data` is required")
@@ -178,6 +189,7 @@ def make_trace(func):
trace(compile_data=compile_data, inline_trace=False, insert_ddp_syncs=not compile_data.using_jit), func
)
+ computation_trc.kwargs = {}
# NOTE: This function is rather slow, so it's intended to be used
# behind a cache.
ba = signature(func).bind(*args, **kwargs)
@@ -193,6 +205,9 @@ def make_trace(func):
primal_trace = make_trace(func)(*args, **kwargs)
primal_trace = sort_data_parallel_syncs(primal_trace)
+ if compile_stats is not None:
+ compile_stats.last_traces.append(primal_trace)
+
# torch.autograd.Function doesn't support non-flat outputs, the
# grads wouldn't be propagated and backward receives None for each
# non-flat non-tensor output. The output must also be a flat tuple,
@@ -317,8 +332,7 @@ def make_trace(func):
bw_traces.append(bw_extrace)
if compile_stats is not None:
- compile_stats.primal_trace = primal_trace
- compile_stats.forward_last_traces = fw_traces
- compile_stats.backward_last_traces = bw_traces
+ compile_stats.last_traces += fw_traces
+ compile_stats.last_backward_traces += bw_traces
return fw_extrace, bw_extrace
diff --git a/thunder/executors/torch_compile.py b/thunder/executors/torch_compile.py
index 9d82e12526..a48635e878 100644
--- a/thunder/executors/torch_compile.py
+++ b/thunder/executors/torch_compile.py
@@ -72,7 +72,8 @@ def torch_interpreted_func(*args):
# _transform_for_operator_executor_execution implementation that need to be
# fixed first. One issue is that it doesn't maintain the ssa form of the
# trace, which is needed for all the passes to work correctly.
- # TODO: https://github.com/Lightning-AI/lightning-thunder/issues/1767
+ # TODO: issue "Try using _transform_for_operator_executor_execution for
+ # torch.compile executor"
torch_trace = trace(inline_trace=False)(torch_interpreted_func, *sorted_unique_inputs)
compiled_func = torch.compile(torch_trace.python_callable())
@@ -84,13 +85,13 @@ def compiled_func_wrapper(*args):
orig = getattr(torch._dynamo.eval_frame.guarded_backend_cache, "skip_backend_check_for_run_only_mode", None)
try:
# TODO: Remove this hack
- # This is a hack to get around the fact that for some reason Dynamo
- # doesn't recreate a guard for the compiled function called from the
- # backward thread. This is a problem because the guard is created
- # with the forward thread id, and the guard is not valid for the
- # backward thread. I couldn't come up with a small repro to file an
- # issue to PyTorch.
- # https://github.com/pytorch/pytorch/issues/114674
+ # Dynamo doesn't recreate a guard for the compiled function called
+ # from the backward thread. This is a problem because the guard is
+ # created with the forward thread ID, and the guard is not valid
+ # for the backward thread.
+ # Issue filed: https://github.com/pytorch/pytorch/issues/114674
+ # We should be able to remove this hack once we're sure that the
+ # above fix has propagated to all supported PyTorch releases.
torch._dynamo.eval_frame.guarded_backend_cache.skip_backend_check_for_run_only_mode = True
return compiled_func(*args)
finally:
diff --git a/thunder/executors/torchex.py b/thunder/executors/torchex.py
index 85cfd9c2ce..535fddd943 100644
--- a/thunder/executors/torchex.py
+++ b/thunder/executors/torchex.py
@@ -633,7 +633,7 @@ def _stride_order_prim_impl(a: torch.Tensor, order: Sequence[int]) -> torch.Tens
rsqrt = _register_torch_operation("rsqrt")
# # NOTE That PyTorch's "sgn" corresponds with the "sign" primitive
sgn = _register_torch_operation("sgn", like=ltorch.sign)
-# # NOTE torch.sign isn't bound here because lightning.compile always uses sgn
+# # NOTE torch.sign isn't bound here because thunder always uses sgn
# sign = _register_torch_operation("sign")
signbit = _register_torch_operation("signbit")
sin = _register_torch_operation("sin")
@@ -1095,8 +1095,8 @@ def _index_put_prim_transform(
return index_put(a, indices, values, accumulate)
-# NOTE torch.compile currently fails to compile scatter add in bfloat16
-# TODO RC1 Separate this into a torch.compile executor
+# NOTE torch.compile has a compilation issue with scatter add in bfloat16,
+# hence the special case here.
# NOTE The scatter add transforms must set the torch language context explicitly so the .to() method
# on tensors is resolved (alternatively they could explicitly call thunder.torch.to)
@langctx(Languages.TORCH)
@@ -1109,7 +1109,8 @@ def _scatter_add_prim_transform(a: TensorProxy, /, index: TensorProxy, value: Te
return scatter_add(a, dim, index, value)
-# NOTE torch.compile currently fails to compile scatter add in bfloat16
+# NOTE torch.compile has a compilation issue with scatter add in bfloat16,
+# hence the special case here.
@langctx(Languages.TORCH)
def _scatter_add_transform(a: TensorLike, /, dim: int, index: TensorLike, src: TensorLike) -> TensorLike:
# NOTE scatter_add does not participate in type promotion, so if a has the bfloat16 dtype, then so does src
@@ -1243,7 +1244,7 @@ def _cross_entropy_backward_impl(
)
# TODO Add support nll_loss_nd, weight tensor, and label_smoothing options.
- # See https://github.com/Lightning-AI/lightning-thunder/issues/704
+ # See issue "Add support for remaining cross_entropy_loss arguments."
utils.check(a.ndim <= 2 and target.ndim <= 1, lambda: f"multi-dimension cross-entropy is not supported.")
utils.check(weight is None, lambda: f"weight tensor argument is not supported.")
@@ -1601,7 +1602,7 @@ def _unpack_prim_impl(
) -> list[torch.Tensor]:
return torch._utils._unflatten_dense_tensors(buffer, tensors)
- # TODO(crcrpar): Make this compatible with the coming torch_compile executor as it's doing really well for cat and reshape.
+ # TODO(crcrpar): Make this compatible with the torch.compile executor as it's doing really well for cat and reshape.
# NOTE(crcrpar): why no caching/resue of buffer?
# This prim is only used by fsdp backward for now.
# Bucketing of reduce-scatter, i.e., creating a buffer for
@@ -1621,7 +1622,6 @@ def _unpack_prim_impl(
# To support individual copies from gradient to its bucket requires a mask or an arrayy of indices to achieve correct behavior.
# In PyTorch, the op for this is [`Tensor.index_copy_`](https://pytorch.org/docs/stable/generated/torch.Tensor.index_copy_.html) where even the index tensor needs to be on the same device as ``self`` and ``tensor``.
# So caching of the bucketing for fsdp backward would bloat up the memory consumption, which is the main reason this doesn't do any caching.
- # See https://github.com/Lightning-AI/lightning-thunder/pull/1669/commits/a942b87e88738ce94f874c21d4adc38749ff10d7#diff-c2fd275781ba0c4aa7eec811bebb7bf0b6ca52a236b510ce7dfbb831d4d9bb40R197-R233 for the potential implementation's clumisiness.
#
# example of two unsharded gradients of [4, 2] and [4], world size of 4:
# -------- ------
diff --git a/thunder/executors/transformer_engineex.py b/thunder/executors/transformer_engineex.py
index 91d32e7a88..ced5d8fdb1 100644
--- a/thunder/executors/transformer_engineex.py
+++ b/thunder/executors/transformer_engineex.py
@@ -19,10 +19,6 @@
import thunder.core.prims as prims
from thunder.core.proxies import TensorProxy, CollectionProxy
from thunder.core.symbol import Symbol
-from thunder.core.transforms import (
- register_augmented_forward_with_checker,
- register_backward,
-)
from thunder.extend import OperatorExecutor, register_executor
__all__ = [
@@ -411,15 +407,6 @@ def linear_forward_rule_checker(a: TensorProxy, w: TensorProxy, bias: None | Ten
return False
-register_augmented_forward_with_checker(
- transformer_engine_ex,
- prims.linear.id,
- linear_forward_rule_checker,
- linear_forwad_rule,
-)
-
-
-@register_backward((transformer_engine_ex, prims.linear.id))
def linear_backward_rule(a_shape, w_shape, b_shape, ctx_idx, grad):
return te_functional_linear_backward(grad, a_shape, w_shape, b_shape, ctx_idx)
@@ -429,9 +416,21 @@ def _linear_transform(a: TensorProxy, w: TensorProxy, b: TensorProxy) -> torch.T
return _create_fp8_linear_bound_symbol(a, w, b, is_grad_enabled=False)
+def _linear_grad(a: TensorProxy, w: TensorProxy, b: TensorProxy) -> TensorProxy:
+ out, saved_for_backward = linear_forwad_rule(a, w, b)
+ g = prims.get_grad(out)
+ ga, gw, gb = linear_backward_rule(*saved_for_backward, g)
+ prims.put_grad(a, ga)
+ prims.put_grad(w, gw)
+ if b is not None:
+ prims.put_grad(b, gb)
+ return out
+
+
# Registers the implementation for torch.nn.functional.linear
transformer_engine_ex.register_implementation(
prims.linear,
checker=_linear_checker,
execution_transform=_linear_transform,
+ grad_transform=_linear_grad,
)
diff --git a/thunder/extend/__init__.py b/thunder/extend/__init__.py
index 2d2ce1fd2a..2acdcf6427 100644
--- a/thunder/extend/__init__.py
+++ b/thunder/extend/__init__.py
@@ -25,6 +25,7 @@
"add_always_executor",
"remove_default_executor",
"remove_always_executor",
+ "register_lookaside",
]
@@ -208,6 +209,7 @@ def register_operator(
module: None | type | ModuleType = None,
fn: None | Callable = None,
bind_postprocess: None | Callable = None,
+ replaces: None | Callable = None,
python_printer: Callable = default_python_printer,
) -> Symbol:
assert (like is None) ^ (meta is None), "Expected one and only one of 'like' and 'meta' to be specified"
@@ -237,6 +239,9 @@ def _bind_postprocess(bsym: BoundSymbol) -> None:
)
self.opmap[name] = sym
+ if replaces is not None:
+ register_lookaside(replaces, sym)
+
return sym
def register_implementation(
@@ -381,3 +386,10 @@ def deregister_executor(ex: Hashable | Executor) -> None:
remove_always_executor(id)
remove_default_executor(id)
+
+
+def register_lookaside(function, symbol) -> None:
+ """register `symbol` as a lookaside for `function`"""
+ import thunder.core.jit_ext
+
+ thunder.core.jit_ext._general_jit_lookaside_map[function] = thunder.core.jit_ext.interpreter_needs_wrap(symbol)
diff --git a/thunder/numpy/__init__.py b/thunder/numpy/__init__.py
index b53e9f26f5..ed2e6474aa 100644
--- a/thunder/numpy/__init__.py
+++ b/thunder/numpy/__init__.py
@@ -11,9 +11,6 @@
from thunder.core.symbol import Symbol
import thunder.clang as clang
-# TODO RC1 Remove this
-from thunder.core.script.noinline import noinline
-
#
# NumPy operator definitions
@@ -28,7 +25,7 @@ def __init__(self, *, method_name: None | str = None):
def __call__(self, fn: Callable) -> Symbol:
_fn = langctx(Languages.NUMPY)(fn)
- _fn = noinline(_fn)
+ # TODO: register _fn as opaque with the interpreter or do this in jit_ext?
sym = Symbol(name=fn.__name__, meta=_fn)
if self.method_name is not None:
diff --git a/thunder/numpy/langctx.py b/thunder/numpy/langctx.py
index 0f6e445d8f..75961b6efb 100644
--- a/thunder/numpy/langctx.py
+++ b/thunder/numpy/langctx.py
@@ -23,7 +23,7 @@ def has_method(self, id: str) -> bool:
return id in _method_name_to_fn_map
def get_method(self, id: str, *args, **kwargs) -> Callable:
- # Note: concrete implmenetations should only raise AttributeError or
+ # Note: concrete implementations should only raise AttributeError or
# return None for "missing" methods as the proxies will
# route __getattr__ to here and hasattr relies on __getattr__
# throwing AttributeError (only) when the attribute does
diff --git a/thunder/core/script/__init__.py b/thunder/py.typed
similarity index 100%
rename from thunder/core/script/__init__.py
rename to thunder/py.typed
diff --git a/thunder/tests/distributed/test_ddp.py b/thunder/tests/distributed/test_ddp.py
index 4ae624b6db..c321a01394 100644
--- a/thunder/tests/distributed/test_ddp.py
+++ b/thunder/tests/distributed/test_ddp.py
@@ -134,7 +134,7 @@ def _run(cls, rank, test_name, file_name, pipe):
"DDP test requires CUDA and NCCL `torch.distributed` backend",
)
class CompileDDPTest(DataParallelTestCase):
- # Ref: https://github.com/Lightning-AI/lightning-thunder/issues/646
+ # Reference issue "Add an example of DDP(compile(model)) to tests"
def test_ddp_compile_module(self):
model = ToyModel().to(self.rank)
ddp_model = DDP(thunder.jit(model, device_ids=[self.rank]))
@@ -157,7 +157,7 @@ def test_ddp_compile_module(self):
last_loss = loss.detach().item()
assert init_loss > last_loss
- # Ref: https://github.com/Lightning-AI/lightning-thunder/issues/599
+ # Reference issue "[tracker] Support DistributedDataParallel"
def test_compile_ddp_module(self):
model = ToyModel().to(self.rank)
with self.assertRaisesRegex(
@@ -452,7 +452,7 @@ def test_ddp_grad_bucketing(self, executor, bucket_size_in_mb: int):
x = torch.ones((2, 12)).to(device)
cm(x).mean().backward()
- bwd_extrace = thunder.last_traces(cm)[1][-1]
+ bwd_extrace = thunder.last_backward_traces(cm)[-1]
bsym_sym_id_list = [bsym.sym.id for bsym in bwd_extrace.bound_symbols]
pack_syms = tuple(filter(lambda a: a == pack_prim_impl.id, bsym_sym_id_list))
unpack_syms = tuple(filter(lambda a: a == unpack_prim_impl.id, bsym_sym_id_list))
@@ -476,20 +476,24 @@ def test_rematerialize_all_gather(self):
m = ToyModel().to(device)
cm = thunder.jit(
fsdp(m, device=device, broadcast_from=0),
- interpretation=INTERPRETATION_OPTIONS.TRANSLATE_PYTHON,
)
x = torch.ones((2, 12), device=device)
cm(x).mean().backward()
- fwd_trc = thunder.last_traces(cm)[0][0]
- bwd_trc = thunder.last_traces(cm)[1][0]
+ (fwd_trc,) = (
+ t for t in thunder.last_traces(cm) if getattr(t.get_provenance(), "pss", "") == "Augmented forward pass"
+ )
+ bwd_trc = thunder.last_backward_traces(cm)[0]
from thunder.core.rematerialization import rematerialize_all_gather
result_fwd_trc, result_bwd_trc = rematerialize_all_gather(fwd_trc, bwd_trc)
# check the return statement in forward trace is updated
- sharded_param_names = ("t_net1_weight", "t_net2_weight")
- unshard_param_names = ("t5", "t16")
+ # TODO: this is not stable w.r.t. details of the processing, the sharded correspond to ("t_net1_weight", "t_net2_weight")
+ # in the original trace and are inputs to all_gather, the unshard are the outputs fo the corresponding wait
+ # If you fix this to be dynamically discerned, you'll be my hero.
+ sharded_param_names = ("t3", "t4")
+ unshard_param_names = ("t10", "t21")
result_saved_for_bwd = [x.name for x in fwd_trc.bound_symbols[-1].args[1][0]]
self.assertTrue(all(t not in sharded_param_names for t in result_saved_for_bwd))
self.assertTrue(all(t in result_saved_for_bwd for t in unshard_param_names))
@@ -648,7 +652,7 @@ def test_ddp_grad_parity_with_without_bucketing(self, executor):
else:
self.assertEqual(tuple(p.grad for p in cm.parameters() if p.grad is not None), gradients)
- # TODO(crcrpar): Add torch compile to executors_list once it's available.
+ # TODO(crcrpar): Add torch compile to executors_list
@common_utils.parametrize(
"executor,bucketing_strategy,fsdptype",
product(
@@ -826,8 +830,8 @@ def check_inflight_allgather_number(trc, n: int, is_bucket: bool):
loss.backward()
# get the trace before sorting
- fwd_trc = thunder.last_traces(cm)[0][-2]
- bwd_trc = thunder.last_traces(cm)[1][-2]
+ fwd_trc = thunder.last_traces(cm)[-2]
+ bwd_trc = thunder.last_backward_traces(cm)[-2]
from thunder.distributed.utils import limit_in_flight_allgathers
@@ -1101,7 +1105,7 @@ def _test_native_ddp_helper(input_data):
tdist.destroy_process_group(pg)
if rank == 0:
- bwd_extrace_sym_ids = [bsym.sym.id for bsym in thunder.last_traces(cmodel)[1][-1].bound_symbols]
+ bwd_extrace_sym_ids = [bsym.sym.id for bsym in thunder.last_backward_traces(cmodel)[-1].bound_symbols]
pack_unpack_update_bucket_view_found = (
"torch_pack_prim_impl" in bwd_extrace_sym_ids
and "torch_unpack_prim_impl" in bwd_extrace_sym_ids
diff --git a/thunder/tests/framework.py b/thunder/tests/framework.py
index ce237d9c97..dfdbeb8312 100644
--- a/thunder/tests/framework.py
+++ b/thunder/tests/framework.py
@@ -133,25 +133,11 @@ def executors_list(self) -> list[extend.Executor]:
@singledispatchmethod
def make_callable_legacy(self, fn, **kwargs):
- # TODO: an error is thrown for many functions because __code__ and
- # inspect.signature for wrapped functions is not matching.
- # KeyError: 'args'
- # thunder/core/script/frontend.py:125: KeyError
- # with disable_preprocessing=False
- # See: https://github.com/Lightning-AI/lightning-thunder/issues/386
- disable_preprocessing = kwargs.pop("disable_preprocessing", True)
- return thunder.compile(
- fn, executors_list=self.executors_list(), disable_preprocessing=disable_preprocessing, **kwargs
- )
+ assert kwargs.pop("disable_preprocessing", True)
+ return thunder.compile(fn, executors_list=self.executors_list(), disable_preprocessing=True, **kwargs)
@singledispatchmethod
def make_callable(self, fn, **kwargs):
- # TODO: an error is thrown for many functions because __code__ and
- # inspect.signature for wrapped functions is not matching.
- # KeyError: 'args'
- # thunder/core/script/frontend.py:125: KeyError
- # with disable_preprocessing=False
- # See: https://github.com/Lightning-AI/lightning-thunder/issues/386
return thunder.jit(fn, executors=self.executors_list(), **kwargs)
@make_callable.register
diff --git a/thunder/tests/lit_gpt_model.py b/thunder/tests/lit_gpt_model.py
index 32889c6cc5..57a85089bc 100644
--- a/thunder/tests/lit_gpt_model.py
+++ b/thunder/tests/lit_gpt_model.py
@@ -139,7 +139,7 @@ def forward(self, input_pos: torch.Tensor, k: torch.Tensor, v: torch.Tensor) ->
k = self.k.index_copy_(2, input_pos, k)
v = self.v.index_copy_(2, input_pos, v)
return k, v
- # THUNDER unsupported: https://github.com/Lightning-AI/lightning-thunder/issues/1145
+ # See issue: "Support more indexing operators (index_copy and index_add)"
k = self.k = torch.index_add(self.k, 2, input_pos, k)
v = self.v = torch.index_add(self.v, 2, input_pos, v)
# THUNDER bug: cannot return self.k, self.v here (may be cuda graphs related - no minimum repro)
diff --git a/thunder/tests/llama2_model.py b/thunder/tests/llama2_model.py
index bf70e56531..b8277cd757 100644
--- a/thunder/tests/llama2_model.py
+++ b/thunder/tests/llama2_model.py
@@ -64,7 +64,6 @@ def apply_rotary_emb(
xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)
# reshape freqs_cos and freqs_sin for broadcasting
- # https://github.com/Lightning-AI/lightning-thunder/issues/1106
a, b = freqs_cos.shape
freqs_cos = freqs_cos.view(1, a, 1, b)
freqs_sin = freqs_sin.view(1, a, 1, b)
@@ -247,7 +246,8 @@ def forward(self, tokens: torch.Tensor, targets: torch.Tensor | None = None) ->
if targets is not None:
# if we are given some desired targets also calculate the loss
logits = self.output(h)
- # https://github.com/Lightning-AI/lightning-thunder/issues/1108
+ # Workaround for issue "Unexpected KeyError when self attribute is
+ # set inside forward"
# self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
else:
# inference-time mini-optimization: only forward the output on the very last position
diff --git a/thunder/tests/nanogpt_model.py b/thunder/tests/nanogpt_model.py
index 128ac8c0f4..9d81b3abb8 100644
--- a/thunder/tests/nanogpt_model.py
+++ b/thunder/tests/nanogpt_model.py
@@ -69,8 +69,8 @@ def __init__(self, config):
# flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
self.flash = hasattr(torch.nn.functional, "scaled_dot_product_attention")
# NOTE: The original Karpathy's script hides bias registration behind a flag
- # but we don't do that here. We always register bias, because of preprocessing bug:
- # https://github.com/Lightning-AI/lightning-thunder/issues/605
+ # but we don't do that here. We always register bias due to a now-fixed
+ # bug in thunder.
# TODO: Move the bias registration to be happening `if not self.flash` once the bug is fixed.
# if not self.flash:
# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
@@ -173,8 +173,8 @@ def __init__(self, config):
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
# with weight tying when using torch.compile() some warnings get generated:
# "UserWarning: functional_call was passed multiple values for tied weights.
- # This behavior is deprecated and will be an error in future versions"
- # not 100% sure what this is, so far seems to be harmless. TODO investigate
+ # This behavior is deprecated and will be an error in future versions".
+ # So far this seems to be harmless. TODO investigate
self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
# init all weights
@@ -236,7 +236,6 @@ def forward(self, idx, targets=None):
# NOTE: Advanced indexing is not yet supported in Thunder
# RuntimeError: Advanced indexing currently only supports tensors as sequence elements
# inference-time mini-optimization: only forward the lm_head on the very last position
- # See https://github.com/Lightning-AI/lightning-thunder/issues/894
# logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
logits = self.lm_head(x)
loss = None
diff --git a/thunder/tests/opinfos.py b/thunder/tests/opinfos.py
index 7659d57fe2..17af749d37 100644
--- a/thunder/tests/opinfos.py
+++ b/thunder/tests/opinfos.py
@@ -770,8 +770,8 @@ def _abs_torch(x: torch.Tensor | Number):
# digamma is defined for all complex numbers EXCEPT negative integers and zero
digamma_opinfo = OpInfo(
clang.digamma,
- # NOTE: Restrict domain to avoid singularities because of
- # https://github.com/Lightning-AI/lightning-thunder/issues/1138
+ # NOTE: Restrict domain to avoid singularities because of issue
+ # "OpInfos do not use singularity_fn to produce "more stable" samples."
domain=(eps, math.inf),
# NOTE: digamma returns NaN for all negative integers. It returns -Inf when x = 0.
singularity_fn=lambda x: torch.where(x > 0, x, (x - torch.round(x))),
@@ -1079,7 +1079,7 @@ def _abs_torch(x: torch.Tensor | Number):
dtypes=(datatypes.float16, datatypes.complex32),
devicetypes=(devices.DeviceType.CPU,),
),
- # see https://github.com/csarofeen/pytorch/issues/2367
+ # Used to be an nvFuser bug here; TODO explore removing this xfail
DecorateInfo(
pytest.mark.xfail,
"test_core_vs_torch_consistency",
@@ -1179,8 +1179,7 @@ def _abs_torch(x: torch.Tensor | Number):
sample_input_generator=elementwise_unary_generator,
torch_reference=_elementwise_unary_torch(torch.sgn),
test_directives=(
- # TODO Need to add nvfuser specific support for complex sign
- # https://github.com/csarofeen/pytorch/issues/2492
+ # TODO nvFuser needs support for complex sign
DecorateInfo(
pytest.mark.xfail,
dtypes=(datatypes.complexfloating,),
@@ -1284,7 +1283,9 @@ def _abs_torch(x: torch.Tensor | Number):
sample_input_generator=elementwise_unary_generator,
torch_reference=_elementwise_unary_torch(torch.tan),
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2360
+ # TODO investigate nvFuser's implementation here; for complex datatypes
+ # nvFuser's tanh might be inaccurate, causing numerical mismatches, but
+ # also this concern is potentially stale in 03/2024.
DecorateInfo(
pytest.mark.xfail, "test_core_vs_torch_consistency", executors=("nvfuser",), dtypes=(datatypes.complex64,)
),
@@ -1305,7 +1306,9 @@ def _abs_torch(x: torch.Tensor | Number):
sample_input_generator=elementwise_unary_generator,
torch_reference=_elementwise_unary_torch(torch.tanh),
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2360
+ # TODO investigate nvFuser's implementation here; for complex datatypes
+ # nvFuser's tanh might be inaccurate, causing numerical mismatches, but
+ # also this concern is potentially stale in 03/2024.
DecorateInfo(
pytest.mark.xfail, "test_core_vs_torch_consistency", executors=("nvfuser",), dtypes=(datatypes.complex64,)
),
@@ -1358,7 +1361,9 @@ def _abs_torch(x: torch.Tensor | Number):
sample_input_generator=partial(elementwise_unary_generator, exclude_zero=True),
torch_reference=_elementwise_unary_torch(torch.log),
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2360
+ # TODO investigate nvFuser's implementation here; for complex datatypes
+ # nvFuser's tanh might be inaccurate, causing numerical mismatches, but
+ # also this concern is potentially stale in 03/2024.
DecorateInfo(
pytest.mark.xfail, "test_core_vs_torch_consistency", executors=("nvfuser",), dtypes=(datatypes.complex64,)
),
@@ -1379,7 +1384,9 @@ def _abs_torch(x: torch.Tensor | Number):
sample_input_generator=partial(elementwise_unary_generator, exclude_zero=True),
torch_reference=_elementwise_unary_torch(torch.log10),
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2360
+ # TODO investigate nvFuser's implementation here; for complex datatypes
+ # nvFuser's tanh might be inaccurate, causing numerical mismatches, but
+ # also this concern is potentially stale in 03/2024.
DecorateInfo(
pytest.mark.xfail, "test_core_vs_torch_consistency", executors=("nvfuser",), dtypes=(datatypes.complex64,)
),
@@ -1407,14 +1414,16 @@ def _abs_torch(x: torch.Tensor | Number):
sample_input_generator=elementwise_unary_generator,
torch_reference=_elementwise_unary_torch(torch.log1p),
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2360
+ # TODO investigate nvFuser's implementation here; for complex datatypes
+ # nvFuser's tanh might be inaccurate, causing numerical mismatches, but
+ # also this concern is potentially stale in 03/2024.
DecorateInfo(
pytest.mark.xfail,
"test_core_vs_torch_consistency",
executors=("nvfuser",),
dtypes=(datatypes.complexfloating,),
),
- # NOTE: Torch gives wrong result: https://github.com/pytorch/pytorch/issues/94333
+ # NOTE: Torch has an issue: https://github.com/pytorch/pytorch/issues/94333
DecorateInfo(
pytest.mark.skip,
"test_core_vs_torch_consistency",
@@ -1452,7 +1461,9 @@ def _abs_torch(x: torch.Tensor | Number):
sample_input_generator=partial(elementwise_unary_generator, exclude_zero=True),
torch_reference=_elementwise_unary_torch(torch.log2),
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2360
+ # TODO investigate nvFuser's implementation here; for complex datatypes
+ # nvFuser's tanh might be inaccurate, causing numerical mismatches, but
+ # also this concern is potentially stale in 03/2024.
DecorateInfo(
pytest.mark.xfail, "test_core_vs_torch_consistency", executors=("nvfuser",), dtypes=(datatypes.complex64,)
),
@@ -1556,7 +1567,7 @@ def relu6_error_generator(op, device, dtype=torch.float32, **kwargs):
dtypes=(datatypes.float16,),
devicetypes=(devices.DeviceType.CPU,),
),
- # TODO: https://github.com/Lightning-AI/lightning-thunder/issues/1444
+ # TODO: we might have a tolerance issue here with relu6.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -1577,7 +1588,7 @@ def relu6_error_generator(op, device, dtype=torch.float32, **kwargs):
"test_core_vs_torch_consistency",
dtypes=(datatypes.bool8,),
),
- # TODO: https://github.com/Lightning-AI/lightning-thunder/issues/1444
+ # TODO: we might have a tolerance issue here with relu6.
DecorateInfo(
pytest.mark.xfail(strict=False),
"test_vjp_correctness",
@@ -1613,7 +1624,7 @@ def selu_error_generator(op, device, dtype=torch.float32, **kwargs):
datatypes.bfloat16,
),
),
- # TODO: https://github.com/Lightning-AI/lightning-thunder/issues/1444
+ # TODO: we might have a tolerance issue here with relu6.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -1674,8 +1685,9 @@ def selu_error_generator(op, device, dtype=torch.float32, **kwargs):
devicetypes=(devices.DeviceType.CPU,),
active_if=LooseVersion(torch.__version__) < "1.13",
),
- # TODO: nvfuser needs to return copy for integer dtypes.
- # https://github.com/csarofeen/pytorch/issues/2499
+ # TODO: nvFuser does not define an integer trunc() and thus compilation
+ # fails. They should probably map integer trunc() to an identity op.
+ # Until they do, this test won't work for integer types.
DecorateInfo(
pytest.mark.xfail,
"test_core_vs_torch_consistency",
@@ -1742,7 +1754,8 @@ def elementwise_binary_generator(op, device, dtype, requires_grad, *, no_rhs_num
sample_input_generator=elementwise_binary_generator,
torch_reference=torch.add,
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2549
+ # See issue "broadcast_in_dim: The size of contiguity must equal to the
+ # number of non-broadcasting IterDomains"
DecorateInfo(
pytest.mark.skip,
"test_jvp_correctness",
@@ -1810,7 +1823,8 @@ def elementwise_binary_generator(op, device, dtype, requires_grad, *, no_rhs_num
sample_input_generator=elementwise_binary_generator,
torch_reference=torch.copysign,
test_directives=(
- # See https://github.com/Lightning-AI/lightning-thunder/issues/2218
+ # See issue: "flaky test:
+ # test_vjp_correctness_copysign_torch_cuda_float64 is flaky"
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -1827,8 +1841,7 @@ def elementwise_binary_generator(op, device, dtype, requires_grad, *, no_rhs_num
sample_input_generator=elementwise_comparison_generator,
torch_reference=torch.eq,
test_directives=(
- # There's a problem of reducing a tensor produced by full op
- # See https://github.com/NVIDIA/Fuser/issues/132
+ # TODO: enable this; there was a now-fixed nvFuser bug causing issues.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -1926,8 +1939,7 @@ def fmod_sample_input_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=elementwise_comparison_generator,
torch_reference=torch.ge,
test_directives=(
- # There's a problem of reducing a tensor produced by full op
- # See https://github.com/NVIDIA/Fuser/issues/132
+ # TODO: enable this; there was a now-fixed nvFuser bug causing issues.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -1969,8 +1981,7 @@ def fmod_sample_input_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=elementwise_comparison_generator,
torch_reference=torch.le,
test_directives=(
- # There's a problem of reducing a tensor produced by full op
- # See https://github.com/NVIDIA/Fuser/issues/132
+ # TODO: enable this; there was a now-fixed nvFuser bug causing issues.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -1987,8 +1998,7 @@ def fmod_sample_input_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=elementwise_comparison_generator,
torch_reference=torch.lt,
test_directives=(
- # There's a problem of reducing a tensor produced by full op
- # See https://github.com/NVIDIA/Fuser/issues/132
+ # TODO: enable this; there was a now-fixed nvFuser bug causing issues.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -2018,7 +2028,8 @@ def fmod_sample_input_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=elementwise_binary_generator,
torch_reference=torch.mul,
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2549
+ # See issue "broadcast_in_dim: The size of contiguity must equal to the
+ # number of non-broadcasting IterDomains"
DecorateInfo(
pytest.mark.skip,
"test_jvp_correctness",
@@ -2040,8 +2051,7 @@ def fmod_sample_input_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=elementwise_comparison_generator,
torch_reference=torch.ne,
test_directives=(
- # There's a problem of reducing a tensor produced by full op
- # See https://github.com/NVIDIA/Fuser/issues/132
+ # TODO: enable this; there was a now-fixed nvFuser bug causing issues.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -2066,15 +2076,15 @@ def fmod_sample_input_generator(op, device, dtype, requires_grad, **kwargs):
pytest.mark.skip,
dtypes=(datatypes.float16, datatypes.bfloat16),
),
- # See https://github.com/Lightning-AI/lightning-thunder/issues/972
- # PyTorch's nextafter may be causing CUDA illegal memory accesses
+ # TODO There was an issue with nextafter in PyTorch that should now be
+ # resolved; re-enable this and test.
DecorateInfo(
pytest.mark.skip,
"test_core_vs_torch_consistency",
devicetypes=(devices.DeviceType.CUDA,),
),
- # See https://github.com/Lightning-AI/lightning-thunder/issues/972
- # PyTorch's nextafter may be causing CUDA illegal memory accesses
+ # TODO There was an issue with nextafter in PyTorch that should now be
+ # resolved; re-enable this and test.
DecorateInfo(
pytest.mark.skip,
executors=("torch",),
@@ -2101,8 +2111,8 @@ def polygamma_sample_input_generator(op, device, dtype, requires_grad, *, no_rhs
polygamma_opinfo = OpInfo(
ltorch.polygamma,
- # NOTE: Restrict domain to avoid singularities because of
- # https://github.com/Lightning-AI/lightning-thunder/issues/1138
+ # NOTE: Restrict domain to avoid singularities. See issue "OpInfos do not
+ # use singularity_fn to produce "more stable" samples"
# NOTE: polygamma returns NaN, -Inf, or Inf for all negative integers.
domain=(eps, math.inf),
sample_input_generator=polygamma_sample_input_generator,
@@ -2155,7 +2165,7 @@ def pow_sample_input_generator(op, device, dtype, requires_grad, *, no_rhs_numbe
"test_core_vs_torch_consistency",
dtypes=(datatypes.complex32,),
),
- # See https://github.com/csarofeen/pytorch/issues/2361
+ # TODO For complex numbers we have some numerical consistency issues.
DecorateInfo(
pytest.mark.xfail,
"test_core_vs_torch_consistency",
@@ -2241,7 +2251,8 @@ def pow_sample_input_generator(op, device, dtype, requires_grad, *, no_rhs_numbe
),
# torch doesn't support bool true_divide
DecorateInfo(pytest.mark.xfail, "test_core_vs_torch_consistency", dtypes=(datatypes.bool8,)),
- # See https://github.com/csarofeen/pytorch/issues/2549
+ # See issue "broadcast_in_dim: The size of contiguity must equal to the
+ # number of non-broadcasting IterDomains"
DecorateInfo(
pytest.mark.skip,
"test_vjp_correctness",
@@ -2384,7 +2395,8 @@ def addcmul_addcdiv_sample_generator(op, device, dtype, requires_grad, **kwargs)
"test_core_vs_torch_consistency",
dtypes=(datatypes.exact,),
),
- # This test is flaky, see https://github.com/Lightning-AI/lightning-thunder/issues/2244
+ # See issue "flaky test:
+ # test_vjp_correctness_addcdiv_nvfuser_cuda_float64"
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -2510,8 +2522,7 @@ def clamp_sample_generator(op, device, dtype, requires_grad, **kwargs):
torch_reference=torch.clamp,
dtypes=(datatypes.signedinteger, datatypes.unsignedinteger, datatypes.floating),
test_directives=(
- # This test is flaky
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1992
+ # see issue "test_vjp_correctness_clamp_nvfuser_cuda_float64 is flaky"
DecorateInfo(
pytest.mark.skip,
"test_vjp_correctness",
@@ -2715,7 +2726,8 @@ def broadcast_in_dim_error_generator(op, device, **kwargs):
pytest.mark.xfail,
"test_errors",
),
- # See https://github.com/csarofeen/pytorch/issues/2549
+ # See issue "broadcast_in_dim: The size of contiguity must equal to the number of
+ # non-broadcasting IterDomains"
DecorateInfo(
pytest.mark.skip,
"test_jvp_correctness",
@@ -2831,6 +2843,11 @@ def diagonal_sample_generator(op, device, dtype, requires_grad, **kwargs):
ltorch.diagonal,
sample_input_generator=diagonal_sample_generator,
torch_reference=torch.diagonal,
+ test_directives=(
+ # thunder.torch.diagonal meta function is not correctly implemented for
+ # input case ((1, 2, 0, 3), -1, 0, -1)
+ DecorateInfo(pytest.mark.xfail(strict=True), "test_vjp_correctness"),
+ ),
)
shape_ops.append(diagonal_opinfo)
@@ -3214,8 +3231,7 @@ def pad_sample_generator(op, device, dtype, requires_grad, **kwargs):
# Versions of above examples but with padding between elements set to 0
((2, 2), ((1, 1, 0), (-1, 2, 0))),
((2, 0, 3), ((1, 0, 0), (1, 1, 0), (0, 0, 0))),
- # See https://github.com/Lightning-AI/lightning-thunder/issues/415
- # The PyTorch lowering does not handle this case properly
+ # See issue "PyTorch pad prim lowering handles out-of-bands negative padding incorrectly"
# ((7, 5), ((0, 0, 0), (-6, 2, 0))),
((5, 7), ((0, 0, 0), (-6, 2, 0))),
((3, 2, 5), ((-2, 1, 0), (1, -1, 0), (-1, 3, 0))), # negative pad in all 3 dims
@@ -3242,12 +3258,12 @@ def _jax_pad(a, padding_value, padding_config):
executors=("torch",),
dtypes=(datatypes.complexfloating,),
),
- # See issue https://github.com/Lightning-AI/lightning-thunder/issues/2053
+ # See issue "pad+nvFuser: wrong results when applied to 1-numel inputs"
DecorateInfo(
pytest.mark.xfail,
executors=("nvfuser",),
),
- # See issue https://github.com/Lightning-AI/lightning-thunder/issues/2053
+ # See issue "pad+nvFuser: wrong results when applied to 1-numel inputs"
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -3318,7 +3334,7 @@ def pad_torch_error_generator(op, device, dtype=torch.float32, **kwargs):
# TODO: only remove these cases when the executor is nvfuser
-# FIXME: Zero-dim cases are skipped due to https://github.com/csarofeen/pytorch/issues/2383
+# TODO: zero-dim cases had a bug, now fixed; re-enable.
# FIXME: tensors with no elements are skipped because of no nvfuser support
def reshape_sample_generator(op, device, dtype, requires_grad, **kwargs):
make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
@@ -3438,8 +3454,8 @@ def slice_in_dim_sample_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=slice_in_dim_sample_generator,
jax_reference=jax.lax.slice_in_dim if JAX_AVAILABLE else None,
test_directives=(
- # nvfuser executor doesn't support pad correctly
- # See https://github.com/Lightning-AI/lightning-thunder/issues/285
+ # TODO: nvfuser executor didn't support pad correctly, but now it should.
+ # Test and re-enable.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -3450,8 +3466,8 @@ def slice_in_dim_sample_generator(op, device, dtype, requires_grad, **kwargs):
shape_ops.append(slice_in_dim)
-# TODO https://github.com/Lightning-AI/lightning-thunder/issues/416
-# Add strides and slicing outside tensor boundaries
+# See issue "Slice prim samples need strides and slicing beyond tensor
+# boundaries"
def slice_prim_sample_generator(op, device, dtype, requires_grad, **kwargs):
make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
@@ -3473,8 +3489,8 @@ def slice_prim_sample_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=slice_prim_sample_generator,
jax_reference=jax.lax.slice if JAX_AVAILABLE else None,
test_directives=(
- # nvfuser executor doesn't support pad correctly
- # See https://github.com/Lightning-AI/lightning-thunder/issues/285
+ # TODO: nvfuser executor didn't support pad correctly, but now it should.
+ # Test and re-enable.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -3535,7 +3551,7 @@ def split_sample_generator(op, device, dtype, requires_grad, **kwargs):
((4, 6, 7), 3, -1),
((4, 6, 7), 9, 1),
((4, 6, 7), (1, 2, 1, 2), 1),
- # TODO https://github.com/Lightning-AI/lightning-thunder/issues/420
+ # See issue "nvFuser split test failure"
# ((4, 6, 7), (3, 1, 2, 0, 0, 1), -1),
((4, 4, 12), 4, 2),
)
@@ -3760,8 +3776,8 @@ def tensor_split_sample_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=tensor_split_sample_generator,
torch_reference=torch.tensor_split,
test_directives=(
- # nvfuser executor doesn't support pad correctly
- # See https://github.com/Lightning-AI/lightning-thunder/issues/285
+ # TODO: nvfuser executor didn't support pad correctly, but now it should.
+ # Test and re-enable.
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -3888,6 +3904,72 @@ def torch_permute_reference(a, *dims):
shape_ops.append(permute_opinfo)
+def t_sample_generator(op, device, dtype, requires_grad, **kwargs):
+ make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
+
+ # shape
+ cases = (
+ (),
+ (1),
+ (4),
+ (4, 5),
+ )
+
+ for shape in cases:
+ yield SampleInput(make(shape))
+
+
+def t_error_generator(op, device, dtype=torch.float32, **kwargs):
+ make = partial(make_tensor, device=device, dtype=dtype)
+
+ # shape, error type, error message
+ cases = (
+ ((4, 5, 6), RuntimeError, r"t\(\) expects a tensor with <= 2 dimensions, but self is 3D"),
+ (
+ (4, 5, 6, 7),
+ RuntimeError,
+ r"t\(\) expects a tensor with <= 2 dimensions, but self is 4D",
+ ),
+ )
+
+ for shape, err_type, err_msg in cases:
+ yield SampleInput(make(shape)), err_type, err_msg
+
+
+t_opinfo = OpInfo(
+ ltorch.t,
+ sample_input_generator=t_sample_generator,
+ error_input_generator=t_error_generator,
+ torch_reference=lambda x: torch.Tensor.t(x),
+)
+shape_ops.append(t_opinfo)
+
+
+def reverse_dims_T_sample_generator(op, device, dtype, requires_grad, **kwargs):
+ make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
+
+ # shape
+ cases = (
+ (),
+ (1),
+ (4),
+ (4, 5),
+ (4, 5, 6),
+ (4, 5, 6, 7),
+ )
+
+ for shape in cases:
+ yield SampleInput(make(shape))
+
+
+reverse_dims_T_opinfo = OpInfo(
+ ltorch.reverse_dims_T,
+ sample_input_generator=reverse_dims_T_sample_generator,
+ torch_reference=lambda x: x.T,
+)
+shape_ops.append(reverse_dims_T_opinfo)
+
+
def matrix_transpose_sample_generator(op, device, dtype, requires_grad, **kwargs):
make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
@@ -4183,7 +4265,8 @@ def unsqueeze_sample_generator(op, device, dtype, requires_grad, **kwargs):
sample_input_generator=unsqueeze_sample_generator,
jax_reference=jax.lax.expand_dims if JAX_AVAILABLE else None,
test_directives=(
- # See https://github.com/csarofeen/pytorch/issues/2549
+ # See issue "broadcast_in_dim: The size of contiguity must equal to the
+ # number of non-broadcasting IterDomains"
DecorateInfo(
pytest.mark.skip,
"test_jvp_correctness",
@@ -4398,7 +4481,8 @@ def _replace_random_percentage(a: torch.Tensor, value: Number, percentage: float
dtypes=(datatypes.complex32,),
devicetypes=(devices.DeviceType.CPU,),
),
- # See https://github.com/csarofeen/pytorch/issues/2369
+ # nvFuser had issues with complex reductions, now fixed; TODO re-enable
+ # this test.
DecorateInfo(
pytest.mark.xfail,
dtypes=(datatypes.complexfloating,),
@@ -4448,7 +4532,8 @@ def var_sample_generator(op, device, dtype, requires_grad):
dtypes=(datatypes.complex32,),
devicetypes=(devices.DeviceType.CPU, devices.DeviceType.CUDA),
),
- # See https://github.com/csarofeen/pytorch/issues/2369
+ # nvFuser had issues with complex reductions, now fixed; TODO re-enable
+ # this test.
DecorateInfo(
pytest.mark.xfail,
dtypes=(datatypes.complexfloating,),
@@ -4498,8 +4583,7 @@ def var_sample_generator(op, device, dtype, requires_grad):
# Complex var is not supported yet
dtypes=(datatypes.floating,),
test_directives=(
- # TODO FIXME nvFuser fails to compile var_mean for these tests
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1438
+ # See issue "nvFuser fails to compile some var_mean tests"
DecorateInfo(
pytest.mark.xfail,
"test_core_vs_torch_consistency",
@@ -5145,7 +5229,8 @@ def einsum_error_generator(op, device, **kwargs):
supports_grad=True,
# TODO: test all integer types and figure out their dtype.
dtypes=(datatypes.float32, datatypes.float64),
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1643.
+ # See issue "Disabled einsum tests might hide potential issues in our
+ # testing/op implementations"
# Testing only float32, float64 now.
# types=(datatypes.int64, datatypes.floating),
# domain=(-1, +1),
@@ -6023,7 +6108,7 @@ def group_norm_error_generator(op, device, **kwargs):
dtypes=(datatypes.float16, datatypes.bfloat16),
devicetypes=(devices.DeviceType.CUDA,),
),
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1405
+ # This should be fixed now; TODO re-enable, test
DecorateInfo(
pytest.mark.xfail,
executors=("nvfuser",),
@@ -6411,7 +6496,8 @@ def embedding_sample_generator(op, device, dtype, requires_grad, **kwargs):
dtypes=(datatypes.floating, datatypes.complexfloating),
test_directives=(
# TODO Investigate these discrepancies -- some dtype x executor configurations seem to be fine
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1387
+ # See issue "phantom grad's embedding computation is divergent from
+ # PyTorch's"
DecorateInfo(
custom_comparator(partial(assert_close, atol=1, rtol=2)),
"test_phantom_grad_vs_torch_consistency",
@@ -6749,9 +6835,8 @@ def cross_entropy_reference_generator(op, device, dtype, requires_grad, **kwargs
# TODO Enable cross entropy bwd weight support
-# see https://github.com/Lightning-AI/lightning-thunder/issues/834
# TODO Enable test cases after adding support nll_loss_nd, weight tensor, and label_smoothing options.
-# See https://github.com/Lightning-AI/lightning-thunder/issues/704
+# TODO see issue "Add support for remaining cross_entropy_loss arguments"
def cross_entropy_sample_generator(op, device, dtype, requires_grad, **kwargs):
make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
@@ -7098,7 +7183,7 @@ def interpolate_error_generator(op, device, dtype=torch.float32, **kwargs):
dtypes=(datatypes.float16,),
devicetypes=(devices.DeviceType.CPU,),
),
- # https://github.com/Lightning-AI/lightning-thunder/issues/1032
+ # This should be fixed now; TODO re-enable and test
DecorateInfo(
pytest.mark.xfail,
"test_vjp_correctness",
@@ -7115,7 +7200,8 @@ def interpolate_error_generator(op, device, dtype=torch.float32, **kwargs):
prob_distr_ops = []
-# multinomial testing is currently disabled due to https://github.com/Lightning-AI/lightning-thunder/issues/2258
+# multinomial testing is currently disabled due to issue "randomness: enable
+# PyTorch generators for operations like multinomial"
# def multinomial_sample_generator(op, device, dtype, requires_grad, **kwargs):
# make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
diff --git a/thunder/tests/test_core.py b/thunder/tests/test_core.py
index f7dfbd1267..54b27f4964 100644
--- a/thunder/tests/test_core.py
+++ b/thunder/tests/test_core.py
@@ -532,7 +532,7 @@ def foo(tup):
@instantiate(dtypes=NOTHING)
def test_type_promotion_tensors(executor, device, _):
if executor == TorchExecutor:
- pytest.xfail("https://github.com/Lightning-AI/lightning-thunder/issues/406")
+ pytest.xfail('see issue "vmap of sum doesn\'t work when dims are passed as a keyword argument"')
def foo(a, b):
return a + b
@@ -590,7 +590,7 @@ def bar(a, b, c):
@instantiate(dtypes=NOTHING)
def test_type_promotion_numbers_and_tensors(executor, device, _):
if executor == TorchExecutor:
- pytest.xfail("https://github.com/Lightning-AI/lightning-thunder/issues/406")
+ pytest.xfail('See issue "Type promotion with the torchexecutor and elementwise operations is incorrect"')
def foo(a, b, c):
return a + b + c
@@ -1130,7 +1130,8 @@ def test_detached_trace(executor, device: str, _):
def test_normalized_args_prims_sum(executor, device: str, dtype: dtypes.dtype):
# This test verifies that the recorded trace for a call to prims.sum
# has its positional and keyword arguments normalized to the same form.
- # See: https://github.com/Lightning-AI/lightning-thunder/issues/195
+ # See issue "vmap of sum doesn't work when dims are passed as a keyword
+ # argument"
a = make_tensor((2, 2), device=device, dtype=ltorch.to_torch_dtype(dtype))
def func_dim_posarg(x):
@@ -1221,7 +1222,8 @@ def foo(x):
assert str(trace).count("Testing") == 1
-# Check for https://github.com/Lightning-AI/lightning-thunder/issues/471
+# Check to verify the issue in "KeyError thrown in thunder.executor.utils.Region
+# when None is passed in as input".
@instantiate(dtypes=(thunder.float32,))
def test_argument_of_none(executor, device, dtype):
from thunder.executors.utils import Region
@@ -1683,7 +1685,7 @@ def test_transforms_vmap_axis_size(executor, device, _):
@instantiate(
dtypes=NOTHING,
- decorators=(pytest.mark.xfail(reason="https://github.com/Lightning-AI/lightning-thunder/issues/2118"),),
+ decorators=(pytest.mark.xfail(reason='issue "flaky test: test_transforms_vjp_{2_1, 1_2}_nvfuser_cuda_None"'),),
)
def test_transforms_vjp_1_2(executor, device, _):
from thunder.core.transforms import vjp
@@ -1790,7 +1792,7 @@ def func(x):
@instantiate(
dtypes=NOTHING,
- decorators=(pytest.mark.xfail(reason="https://github.com/Lightning-AI/lightning-thunder/issues/2118"),),
+ decorators=(pytest.mark.xfail(reason='issue "flaky test: test_transforms_vjp_{2_1, 1_2}_nvfuser_cuda_None"'),),
)
def test_transforms_vjp_2_1(executor, device, _):
from thunder.core.transforms import vjp
@@ -1830,7 +1832,8 @@ def func_2_1(x, y):
# executors=(
# nvFuserExecutor,
# # TODO: Enable Torch executor once the issue with sum is fixed
-# # See: https://github.com/Lightning-AI/lightning-thunder/issues/438
+# # See issue "Different behavior of sum(tensor, ()) for nvFuser and
+# # Torch executor"
# ),
# )
# def test_transforms_vmap_inline_value_and_grad(executor, device, _):
@@ -1950,8 +1953,8 @@ def f(a):
assert "thunder.computation" in excinfo.traceback[-1].path
-# TODO Add nvFuser support (https://github.com/Lightning-AI/lightning-thunder/issues/809)
-# TODO Make these OpInfo tests (https://github.com/Lightning-AI/lightning-thunder/issues/810)
+# TODO See issue "Add contiguous and clang.stride_order OpInfos that check stride
+# consistency with PyTorch"
@instantiate(
dtypes=NOTHING,
executors=(TorchExecutor,),
@@ -2191,7 +2194,8 @@ def func(qkv):
@instantiate(dtypes=NOTHING)
def test_no_passthrough_symbol(executor, device, _):
# A test case for the situation reported in
- # https://github.com/Lightning-AI/lightning-thunder/issues/1131
+ # "backward trace contains symbols not present in forward that cause
+ # NotImplementedError"
# When an operation simply passes through its input, we should not
# add it to the trace.
diff --git a/thunder/tests/test_cudnn_executor.py b/thunder/tests/test_cudnn_executor.py
index f90d2a231c..4128e02914 100644
--- a/thunder/tests/test_cudnn_executor.py
+++ b/thunder/tests/test_cudnn_executor.py
@@ -29,46 +29,47 @@ def grad_scaled_dot_product_attention_reference_generator(op, device, dtype, req
from thunder.tests.opinfos import SampleInput
# TODO: cudnnex seems to produce large mismatches against reference when tensor initialized from the wider default range of [-9,9]
- # https://github.com/Lightning-AI/lightning-thunder/issues/1871
+ # See issue "cuDNN SDPA backward might return NaNs for inputs with absolute
+ # value more than certain threshold"
make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=-0.5, high=0.5)
n_head = 2
- N = 8
+ N = 8 # batch size
# TODO: multiple of 8 seems to produce NaNs
- L = random.randint(1, 10) * 64
+ L = random.randint(1, 10) * 64 # query's sequence length
alignment_factor = 8
- S = random.randint(1, 10) * alignment_factor
- E = random.randint(8, 16) * alignment_factor
- Ev = random.randint(8, 16) * alignment_factor
+ S = random.randint(1, 10) * alignment_factor # key/value's sequence length
+ E = random.randint(8, 16) * alignment_factor # query/key's embedding size
+ Ev = random.randint(8, 16) * alignment_factor # value's embedding size
# 4-dim (multiheaded) causal cases
q, k, v = make(N, n_head, L, E), make(N, n_head, S, E), make(N, n_head, S, Ev)
- yield SampleInput(q, k, v, attn_mask := None, dropout_p := 0.0, is_causal := True)
+ yield SampleInput(q, k, v, None, dropout_p=0.0, is_causal=True)
# TODO: cudnnex seems to have a few mismatches. Will be enabled in a later PR.
# Non-contiguous input tensor case
nq = make(N, n_head, L, E).permute(0, 1, 3, 2)
nk = make(N, n_head, L, E).permute(0, 1, 3, 2)
nv = make(N, n_head, L, E).permute(0, 1, 3, 2)
- yield SampleInput(nq, nk, nv, attn_mask := None, dropout_p := 0.0, is_causal := False)
+ yield SampleInput(nq, nk, nv, None, dropout_p=0.0, is_causal=False)
# Test the scale factor which was added in torch 2.1
if LooseVersion(torch.__version__) >= LooseVersion("2.1.0"):
q, k, v = make(N, n_head, L, E), make(N, n_head, S, E), make(N, n_head, S, Ev)
- yield SampleInput(q, k, v, attn_mask := None, dropout_p := 0.0, is_causal := False, scale=0.123)
+ yield SampleInput(q, k, v, None, dropout_p=0.0, is_causal=False, scale=0.123)
# TODO: cudnnex only support of grad_attn_mask with batch dim 1 and both sequence lenghts divisible by 64. Release 9.0.1 will relax this constraint.
# Additive attn_mask
q, k, v = make(N, n_head, L, E), make(N, n_head, S, E), make(N, n_head, S, Ev)
additive_attn_mask = make((1, n_head, L, S), dtype=q.dtype).tril()
- yield SampleInput(q, k, v, attn_mask := additive_attn_mask, is_causal=False)
+ yield SampleInput(q, k, v, additive_attn_mask, is_causal=False)
# Boolean attn_mask
q, k, v = make(N, n_head, L, E), make(N, n_head, S, E), make(N, n_head, S, Ev)
bool_attn_mask = make((1, n_head, L, S), dtype=torch.bool, low=1, high=1, requires_grad=False).tril()
- yield SampleInput(q, k, v, attn_mask := bool_attn_mask, is_causal=False)
+ yield SampleInput(q, k, v, bool_attn_mask, is_causal=False)
grad_sdpa_cudnn_opinfo = OpInfo(
@@ -87,10 +88,6 @@ def grad_scaled_dot_product_attention_reference_generator(op, device, dtype, req
)
-# WARNING: cudnn executor is experimental. Tests that use cudnn might fail.\n
-# Issue for tracking support: https://github.com/Lightning-AI/lightning-thunder/issues/880
-# NOTE This test modifies the global executor map, so it technically should not
-# be run in parallel with other tests
@requiresCUDA
def test_cudnn_sdpa():
# expect sdpa to fail for 8.9.2 and below
@@ -113,10 +110,8 @@ def test_cudnn_sdpa():
query = 1 * (torch.randn(shape_Q, dtype=thunder.torch.to_torch_dtype(dtype), device="cuda") - 0.5)
key = 2 * (torch.randn(shape_K, dtype=thunder.torch.to_torch_dtype(dtype), device="cuda") - 0.5)
value = 3 * (torch.randn(shape_V, dtype=thunder.torch.to_torch_dtype(dtype), device="cuda") - 0.5)
- is_causal = False
- attn_mask = torch.randn(
- s_q, s_kv, requires_grad=False, device="cuda", dtype=thunder.torch.to_torch_dtype(dtype)
- )
+ is_causal = True
+ attn_mask = None
expected = torch.nn.functional.scaled_dot_product_attention(
query, key, value, is_causal=is_causal, attn_mask=attn_mask
@@ -127,7 +122,7 @@ def test(query, key, value, is_causal=False, attn_mask=None):
query, key, value, is_causal=is_causal, attn_mask=attn_mask
)
- ctest = thunder.compile(test, executors_list=[cudnn_ex])
+ ctest = thunder.jit(test, executors=[cudnn_ex])
actual = ctest(query, key, value, is_causal=is_causal, attn_mask=attn_mask)
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=1e-2)
last_trace = thunder.last_traces(ctest)[-1]
@@ -157,8 +152,6 @@ def snippet_torch_consistency(op, torch_op, sample):
assert_close(thunder_result, torch_result, equal_nan=True, atol=0.0625, rtol=5e-2)
-# WARNING: cudnn executor is experimental. Tests that use cudnn might fail.\n
-# Issue for tracking support: https://github.com/Lightning-AI/lightning-thunder/issues/880
# TODO Make it easier for executors to write tests like this, including writing them out-of-tree
# TODO The executor passed below is just a "dummy" that actually gets ignored -- we should provide
# a way to use decorators like @ops without a particular executor
@@ -184,7 +177,7 @@ def test_cudnn_vs_torch_consistency(op, device, dtype, *_):
pytest.xfail("Only interleaved layout is supported pre 8.9.2.")
for sample in op.reference_inputs(device, dtype, requires_grad=False):
- cfn = thunder.compile(op_name_to_fn[op.name], executors_list=[cudnn_ex, cudnn_layernorm_ex])
+ cfn = thunder.jit(op_name_to_fn[op.name], executors=[cudnn_ex, cudnn_layernorm_ex])
result = run_snippet(
snippet_torch_consistency,
@@ -207,17 +200,19 @@ def test_cudnn_vs_torch_consistency(op, device, dtype, *_):
supported_devicetypes=(devices.DeviceType.CUDA,),
)
def test_vjp_correctness_sdpa_cudnnex_manual(op, device, dtype, executor, comp):
- ran_atleast_one = False
for sample in op.reference_inputs(device, dtype, requires_grad=True):
- from thunder.executors.cudnnex import cudnn_ex
-
# Enforce tensor arguments are contiguous for torch reference
contiguous_args = list(map(lambda a: a.contiguous() if isinstance(a, torch.Tensor) else a, sample.args))
# query, key, value
grad_inputs = list(contiguous_args[:3])
- if (attn_mask := sample.args[3]) is not None and attn_mask.requires_grad:
- grad_inputs.append(attn_mask)
+ if (attn_mask := sample.args[3]) is not None:
+ if attn_mask.requires_grad:
+ grad_inputs.append(attn_mask)
+ # TODO(#2470): With cudnn frontend 1.1 and A100, this test hits
+ # RuntimeError when `attn_mask` is provided: `[cudnn_frontend]
+ # Error: No execution plans built successfully`.
+ continue
# Compute vjp result using PyTorch
expect_out = op.torch_reference(*contiguous_args, **sample.kwargs)
@@ -235,17 +230,10 @@ def test_vjp_correctness_sdpa_cudnnex_manual(op, device, dtype, executor, comp):
executors_list=executor.executors_list() + [cudnn_ex],
)
- try:
- actual_out, actual_grad = cfoo(filtered_args, (v,))
- except Exception as e:
- continue
+ actual_out, actual_grad = cfoo(filtered_args, (v,))
comp(actual_out, expect_out, atol=1e-2, rtol=1e-2)
# compare gradients of query, key, value, and attn_mask
for eg, ag in zip(expected_grad, actual_grad):
comp(eg, ag, atol=2e-1, rtol=2e-2)
-
- ran_atleast_one = True
-
- assert ran_atleast_one == True
diff --git a/thunder/tests/test_elementwise.py b/thunder/tests/test_elementwise.py
index a76db9e673..af5dcfb58b 100644
--- a/thunder/tests/test_elementwise.py
+++ b/thunder/tests/test_elementwise.py
@@ -25,7 +25,7 @@ def test_elementwise_dunder_operations_on_numbers(executor, device, dtype):
# (math.floor, (bool, int, float)),
# (operator.inv, (bool, int)),
# (operator.neg, (bool, int, float, complex)),
- # # TODO https://github.com/Lightning-AI/lightning-thunder/issues/713
+ # # TODO see issue "Implement positive operations"
# # operator.pos,
# (builtins.round, (bool, int, float)),
# (math.trunc, (bool, int, float)),
@@ -65,8 +65,8 @@ def foo(a):
assert_close(actual, expected)
-# TODO Test operator and method variants using OpInfos
-# See https://github.com/Lightning-AI/lightning-thunder/issues/710
+# TODO: see issue "Test operator and method variants of operations using
+# OpInfos"
@instantiate(dtypes=(thunder.float32,))
def test_core_tensor_methods(executor, device, dtype):
def foo(a, b, c, d):
@@ -126,7 +126,8 @@ def test_where(executor, device, dtype):
torch_result = torch_fn(pred, i64, -2.3)
assert_close(thunder_result, torch_result)
- # TODO Fix https://github.com/Lightning-AI/lightning-thunder/issues/711
+ # TODO fix issue "Currently nvFuser tensor x float operations result in
+ # float64 results"
# float x int
# thunder_result = thunder_fn(pred, 3., 5)
# torch_result = torch_fn(pred, 3., 5)
diff --git a/thunder/tests/test_examine_memory.py b/thunder/tests/test_examine_memory.py
index 1d01bef894..13bd05360b 100644
--- a/thunder/tests/test_examine_memory.py
+++ b/thunder/tests/test_examine_memory.py
@@ -63,12 +63,13 @@ def bar(a, b): # [4] [2,2]
with runtime_allocated_memory(device):
cbar(a, b)
- traces = thunder.last_traces(cbar)
- fwd_extrace = traces[0][-1]
+ fw_traces = thunder.last_traces(cbar)
+ fwd_extrace = fw_traces[-1]
max_mem_fwd = get_alloc_memory(fwd_extrace)
assert max_mem_fwd[0] == 144
assert sum(max_mem_fwd[1].values()) == get_return_memory(fwd_extrace.bound_symbols[-1]) # 144
- bw_extrace = traces[1][-1]
+ bw_traces = thunder.last_backward_traces(cbar)
+ bw_extrace = bw_traces[-1]
max_mem_bw = get_alloc_memory(bw_extrace)
assert max_mem_bw[0] == 144
assert sum(max_mem_bw[1].values()) == get_return_memory(bw_extrace.bound_symbols[-1]) # 32
@@ -137,9 +138,8 @@ def test_nanogpt_block(executor, device, dtype):
result = cblock(inp)
with runtime_allocated_memory(device):
result.backward(torch.ones_like(result))
- traces = thunder.last_traces(cblock)
- fw_extrace = traces[0][-1]
- bw_extrace = traces[1][-1]
+ fw_extrace = thunder.last_traces(cblock)[-1]
+ bw_extrace = thunder.last_backward_traces(cblock)[-1]
fw_alloc_mem = get_alloc_memory(fw_extrace)
bw_alloc_mem = get_alloc_memory(bw_extrace)
diff --git a/thunder/tests/test_extend.py b/thunder/tests/test_extend.py
index 9e77f8f496..b7dd300d12 100644
--- a/thunder/tests/test_extend.py
+++ b/thunder/tests/test_extend.py
@@ -131,3 +131,57 @@ def test_get_all_executors_includes_all_native_executors():
if torch.cuda.is_available():
expected.update({"nvfuser"})
assert actual == expected
+
+
+def test_register_implementation_custom_op():
+ myex = OperatorExecutor("myex", version="0.1")
+ register_executor(myex)
+
+ def official_add(a, b):
+ return a + b
+
+ def _myadd(a, b):
+ return a + b
+
+ myadd1 = myex.register_operator("myadd1", like=_myadd, fn=_myadd, replaces=official_add)
+ myadd2 = myex.register_operator("myadd2", like=_myadd, fn=_myadd)
+
+ def fn(a, b):
+ return official_add(a, b)
+
+ cfn = thunder.jit(fn, executors=[myex])
+
+ a = torch.randn(2, 2)
+ b = torch.randn(2, 2)
+
+ res = cfn(a, b)
+
+ assert "myadd1" in str(thunder.last_traces(cfn)[-1])
+
+ def myadd_trafo(a, b):
+ return myadd2(a, b)
+
+ def myadd_grad_trafo(a, b):
+ res = myadd2(a, b)
+ grad_res = get_grad(res)
+ put_grads((a, b), (grad_res, grad_res))
+ return res
+
+ myex.register_implementation(myadd1, execution_transform=myadd_trafo, grad_transform=myadd_grad_trafo)
+
+ cfn = thunder.jit(fn, executors=[myex])
+ res = cfn(a, b)
+
+ s = str(thunder.last_traces(cfn)[-1])
+ assert "myadd2" in s and "myadd1" not in s
+
+ a.requires_grad_()
+
+ res = cfn(a, b)
+
+ s = str(thunder.last_traces(cfn)[-1])
+ assert "myadd2" in s and "myadd1" not in s
+
+ a.requires_grad_()
+
+ deregister_executor(myex)
diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py
index 931c365fb9..3a44a25dbc 100644
--- a/thunder/tests/test_grad.py
+++ b/thunder/tests/test_grad.py
@@ -26,7 +26,7 @@
# TODO: Move this to thunder.tests.opinfos
op_skip = {
- # See https://github.com/Lightning-AI/lightning-thunder/issues/226
+ # See issue "Support closures of torch.Tensor"
# TODO: AttributeError: 'Tensor' object has no attribute 'true_dtype'
"masked_fill",
# TODO: RuntimeError: Expected index=tensor([2, 3, 2, 0, 3, 1, 0, 2],
@@ -514,7 +514,7 @@ def test_vjp_correctness_sdpa_manual(op, device, dtype, executor, comp):
vjp(filtered_op),
disable_torch_autograd_support=True,
disable_preprocessing=True,
- executors_list=executor.executors_list() + [sdpa_ex],
+ executors_list=[sdpa_ex, *executor.executors_list()],
)(filtered_args, (v,))
comp(actual_out, expect_out)
@@ -635,7 +635,8 @@ def bar(a, b):
dtypes=NOTHING,
)
def test_convert_element_type_with_float(executor, device, _):
- # Verifies a fix for https://github.com/Lightning-AI/lightning-thunder/issues/537
+ # Verifies the fix for "grad transform hits error: AttributeError: 'float'
+ # object has no attribute 'dtype'"
from thunder.core.transforms import value_and_grad
a = make_tensor([5], dtype=torch.float32, device=device)
@@ -708,7 +709,9 @@ def sincos_backward(sin_x, cos_x, g1, g2):
assert trace.output[0] == trace.bound_symbols[4].output
-# TODO: Fix flaky test https://github.com/Lightning-AI/lightning-thunder/issues/1919
+# TODO: see issue
+# "thunder/tests/test_grad.py::test_torch_autograd_saved_tensors_memory_release
+# is flaky"
@pytest.mark.xfail(strict=False, reason="This test is flaky")
@requiresCUDA
def test_torch_autograd_saved_tensors_memory_release():
@@ -822,6 +825,42 @@ def fun_bw(a, b, g):
torch.testing.assert_close(actual_bw, expected_bw)
+@instantiate(
+ dtypes=NOTHING,
+)
+def test_make_aug_forward_and_backward_var_mean(executor, device, _):
+ # This test checks that the split of the joint forward/backward function for
+ # var_mean correctly puts the forward part into the augmented forward
+ # function and the backward part into the backward function without
+ # overlapping symbols.
+ from thunder.core.vjp_utils import make_aug_forward_and_backward
+ from thunder.core.prims import var_mean
+
+ def fun(a):
+ return var_mean(a, (0,), correction=1)
+
+ x = torch.tensor((2, 2), device=device, dtype=torch.float32)
+
+ trace = thunder.trace()(fun, x)
+ var_mean_bsym = trace.bound_symbols[-2]
+ assert var_mean_bsym.sym.name == "var_mean"
+
+ aug_fw, bw = make_aug_forward_and_backward(var_mean_bsym)
+ aug_fw = executor.make_callable(aug_fw)
+ out, saved = aug_fw(x, (0,), correction=1)
+ bw = executor.make_callable(bw)
+ _ = bw(*saved, *out)
+ bw_trace = thunder.last_traces(bw)[0]
+ assert "var_mean" not in (s.sym.name for s in bw_trace.bound_symbols)
+
+
+def test_no_duplicate_backward_registered():
+ from thunder.core.transforms import backward_impls, _grad_fn_map
+
+ same_keys = set(_grad_fn_map.keys()).intersection(set(backward_impls.keys()))
+ assert not same_keys, f"Duplicate keys: {same_keys}"
+
+
@instantiate(
dtypes=NOTHING,
)
@@ -870,8 +909,7 @@ def func(a):
def test_torch_autograd_crazy_collections_in_and_out(executor, device, dtype):
from thunder.executors.torch_autograd import thunder_backward
- # Borrowed from
- # https://github.com/Lightning-AI/lightning-thunder/blob/3401475ee47d5a732b6b4d5dcbd88afcd9bed81d/thunder/tests/test_core.py#L117
+ # Borrowed from `test_crazy_collections_in_and_out`.
def foo(a, b, c, *, ka, kb, kc):
d = {
5: 2,
@@ -971,19 +1009,16 @@ def test_torch_autograd_module_get_compile_stats(executor, device, _):
out.backward(g)
compile_stats = compile_stats(lc)
- primal_trace = compile_stats.primal_trace
- forward_traces = compile_stats.forward_last_traces
- backward_traces = compile_stats.backward_last_traces
+ forward_traces = compile_stats.last_traces
+ backward_traces = compile_stats.last_backward_traces
assert isinstance(forward_traces, list)
assert len(forward_traces) >= 1
assert isinstance(backward_traces, list)
assert len(backward_traces) >= 1
- assert isinstance(primal_trace, TraceCtx)
- fw_bw_traces = thunder.last_traces(lc)
- assert isinstance(fw_bw_traces, tuple)
- assert len(fw_bw_traces) == 2
- assert fw_bw_traces[0] == forward_traces
- assert fw_bw_traces[1] == backward_traces
+ fw_traces = thunder.last_traces(lc)
+ bw_traces = thunder.last_backward_traces(lc)
+ assert fw_traces == forward_traces
+ assert bw_traces == backward_traces
@instantiate(
@@ -1250,11 +1285,11 @@ def test_populate_grads_mlp(executor, device, dtype):
clear_grads(model)
- tom = executor.make_callable_legacy(model, disable_preprocessing=False)
+ tom = executor.make_callable(model)
tom_grad = grad(tom)
thunder_grads = tom_grad(x)
- populate_grads(thunder_grads, tom)
+ populate_grads(thunder_grads, tom, args=(x,))
thunder_grads = extract_grads(tom)
assert_close(torch_grads, thunder_grads, atol=1e-3, rtol=1e-5)
@@ -1277,11 +1312,11 @@ def test_populate_grads_csa(executor, device, dtype):
clear_grads(model)
- tom = executor.make_callable_legacy(model, disable_preprocessing=False)
+ tom = executor.make_callable(model)
tom_grad = grad(tom)
thunder_grads = tom_grad(x)
- populate_grads(thunder_grads, tom)
+ populate_grads(thunder_grads, tom, args=[x])
thunder_grads = extract_grads(tom)
assert_close(torch_grads, thunder_grads, atol=1e-2, rtol=1e-2)
@@ -1304,11 +1339,11 @@ def test_populate_grads_block(executor, device, dtype):
clear_grads(model)
- tom = executor.make_callable_legacy(model, disable_preprocessing=False)
+ tom = executor.make_callable(model)
tom_grad = grad(tom)
thunder_grads = tom_grad(x)
- populate_grads(thunder_grads, tom)
+ populate_grads(thunder_grads, tom, args=[x])
thunder_grads = extract_grads(tom)
assert_close(torch_grads, thunder_grads, atol=1e-2, rtol=1e-2)
@@ -1340,7 +1375,7 @@ def test_populate_grads_nanogpt(executor, device, dtype):
clear_grads(model)
- tom = executor.make_callable_legacy(model, disable_preprocessing=False)
+ tom = executor.make_callable(model)
def grad_specifier(out) -> None:
logits, loss = out
@@ -1349,7 +1384,7 @@ def grad_specifier(out) -> None:
tom_grad = grad(tom, grad_specifier=grad_specifier)
thunder_grads = tom_grad(x, targets)
- populate_grads(thunder_grads, tom)
+ populate_grads(thunder_grads, tom, args=[x, targets])
thunder_grads = extract_grads(tom)
assert_close(torch_grads, thunder_grads, atol=1e-2, rtol=1e-2)
diff --git a/thunder/tests/test_interpreter.py b/thunder/tests/test_interpreter.py
index e113438f93..447647b8bb 100644
--- a/thunder/tests/test_interpreter.py
+++ b/thunder/tests/test_interpreter.py
@@ -1010,7 +1010,8 @@ def foo(x):
# }
-# See https://github.com/Lightning-AI/lightning-thunder/issues/2078
+# Test for issue "jit: passing jitted functions as arguments to jitted
+# functions fails."
def test_reduce_jitted_reduce_fn(jit):
import functools
@@ -1482,7 +1483,7 @@ def foo():
assert jfoo() is True
-@pytest.mark.xfail(reason="https://github.com/Lightning-AI/lightning-thunder/issues/1824")
+@pytest.mark.xfail(reason='"exec() and eval() lookaside ignores locals()"')
def test_exec_import_star(jit):
# Assert that we can actually generate the instruction
to_exec = "from itertools import *"
@@ -2606,8 +2607,8 @@ def test_displayhook(jit):
import io
import code
- # TODO: Implement the lookaside for exec(). Under he hood, `code.InteractiveInterpreter().runsource('5;6;7')``
- # just compiles the string and calls exec(), plus a little bit of irrelevant error handling.
+ # TODO: Implement the lookaside for exec(). Under the hood, `code.InteractiveInterpreter().runsource('5;6;7')``
+ # just compiles the string and calls exec(), plus a little bit of error handling.
# I'm not entirely convinced that the PRINT_EVAL is going through our system at the moment, but
# it for sure would with an exec() lookaside. I'm also not sure what makes InteractiveInterpreter
# interactive. It isn't *actually* in interactive mode. So, why is PRINT_EXPR in the interpreted
@@ -2616,7 +2617,7 @@ def test_displayhook(jit):
py_redirect = io.StringIO()
with redirect_stdout(py_redirect):
# Avoid clobbering this interpreter's display hook, and ensure it's interactive.
- # Why is this necessary? I'm not sure.
+ # Why is this necessary?
interpreter = code.InteractiveInterpreter()
def smt(s):
diff --git a/thunder/tests/test_jit_functional.py b/thunder/tests/test_jit_functional.py
index e65ed69aec..7fece51f72 100644
--- a/thunder/tests/test_jit_functional.py
+++ b/thunder/tests/test_jit_functional.py
@@ -292,7 +292,7 @@ def test_binary_ops_compare_numbers():
def test_binary_ops_int_numbers():
- # Issue https://github.com/Lightning-AI/lightning-thunder/issues/594 for more ops
+ # TODO: see issue "Implement logical and arithmetic left and right shifts"
# "<<", ">>",
int_ops = ["+", "&", "//", "*", "%", "|", "**", "-", "/", "^"]
@@ -1574,7 +1574,7 @@ def foo(a, b):
assert_close(expected, actual)
-@pytest.mark.xfail(reason="https://github.com/Lightning-AI/lightning-thunder/issues/2191")
+@pytest.mark.xfail(reason='issue: "jit-eager: allow sets as a return value"')
def test_return_set():
def foo(a, b):
return {a, b}
@@ -2453,7 +2453,7 @@ def foo():
jfoo()
-@pytest.mark.xfail(reason="https://github.com/Lightning-AI/lightning-thunder/issues/2184")
+@pytest.mark.xfail(reason='issue: "sharp edges: loading closures"')
def test_input_closure_sharp_edge():
x = 5
@@ -2486,7 +2486,7 @@ def _test_fn_global_no_sharp_edge_fn():
return 7
-@pytest.mark.xfail(reason="https://github.com/Lightning-AI/lightning-thunder/issues/2189")
+@pytest.mark.xfail(reason='issue: "sharp edge: allow function and module loads"')
def test_fn_global_no_sharp_edge():
def foo(x):
return x + _test_fn_global_no_sharp_edge_fn()
diff --git a/thunder/tests/test_jit_general.py b/thunder/tests/test_jit_general.py
index 2c591d3680..28fc8a54fc 100644
--- a/thunder/tests/test_jit_general.py
+++ b/thunder/tests/test_jit_general.py
@@ -23,6 +23,7 @@
import thunder.core.prims as prims
from thunder import pytorch_executor, nvfuser_executor
from thunder.executors.sdpaex import sdpa_ex
+from thunder.core.jit_ext import JITSharpEdgeError
#
@@ -49,6 +50,25 @@ def skipif_not_pytorch_2_1(f):
)(f)
+def test_jitting_through_opaque_torch_symbols_sharp_edge():
+ def no_sharp_edge(x):
+ # randn_like is in ltorch
+ return torch.randn_like(x)
+
+ def sharp_edge(x):
+ # rand_like is not yet in ltroch
+ return torch.rand_like(x)
+
+ x = torch.rand(1)
+
+ jno_sharp_edge = thunder.jit(no_sharp_edge, sharp_edges="error")
+ jno_sharp_edge(x)
+
+ jsharp_edge = thunder.jit(sharp_edge, sharp_edges="error")
+ with pytest.raises(JITSharpEdgeError):
+ jsharp_edge(x)
+
+
def test_binary_add_tensors():
def foo(a, b):
return a + b
@@ -330,7 +350,7 @@ def foo(a, b):
jfoo = thunder.jit(foo)
# TODO Add test for bool
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1990
+ # see issue "Binary addition on booleans should promote to an integer"
cases = (
(2, 3),
(2.1, 3.4),
@@ -377,7 +397,7 @@ def foo(a, b):
jfoo = thunder.jit(foo)
# TODO Add test for bool
- # See https://github.com/Lightning-AI/lightning-thunder/issues/1990
+ # see issue "Binary addition on booleans should promote to an integer"
cases = (
(2, 3),
(2.1, 3.4),
@@ -394,7 +414,7 @@ def foo(a, b):
_test_add_global_global = 2
-@pytest.mark.xfail(reason="https://github.com/Lightning-AI/lightning-thunder/issues/1935", raises=BaseException)
+@pytest.mark.xfail(reason='"disallow global reads and writes (temporarily)"', raises=BaseException)
def test_global_fails():
def foo():
return _test_add_global_global
@@ -405,7 +425,10 @@ def foo():
jfoo()
-@pytest.mark.xfail(reason="https://github.com/Lightning-AI/lightning-thunder/issues/1936", raises=BaseException)
+@pytest.mark.xfail(
+ reason='"Raise an error when a program attempts to write to a nonlocal that was captured from outside the interpreter"',
+ raises=BaseException,
+)
def test_nonlocal_outside_interpreter_fails():
def foo():
x = 3
diff --git a/thunder/tests/test_networks.py b/thunder/tests/test_networks.py
index 971201bcba..15760f2948 100644
--- a/thunder/tests/test_networks.py
+++ b/thunder/tests/test_networks.py
@@ -41,7 +41,7 @@ def test_nanogpt_complete(executor, device, dtype):
# TODO Investigate grad inconsistency
# TODO: Add float16 and bfloat16 comparison tests here and to all other tests in
# this file.
-# https://github.com/Lightning-AI/lightning-thunder/issues/907
+# See issue "Add half precision dtype tests to test_networks.py"
@instantiate(dtypes=(thunder.float32,))
def test_nanogpt_complete_autograd(executor, device, dtype):
tdtype = ttorch.to_torch_dtype(dtype)
diff --git a/thunder/tests/test_nvfuser.py b/thunder/tests/test_nvfuser.py
index 3cbe631bb9..cdeac92a29 100644
--- a/thunder/tests/test_nvfuser.py
+++ b/thunder/tests/test_nvfuser.py
@@ -315,7 +315,7 @@ def func(w, x, y, z):
@instantiate(dtypes=NOTHING, devicetypes=(devices.DeviceType.CUDA,), executors=(nvFuserExecutor,))
def test_cse_rematerialization(executor, device, _):
- # Unit test for https://github.com/Lightning-AI/lightning-thunder/issues/2046
+ # Unit test for "llama2.c example failed with bookend disabled."
from thunder.tests.llama2_model import Transformer, ModelArgs
from thunder.core.pytree import tree_flatten
@@ -338,10 +338,10 @@ def test_cse_rematerialization(executor, device, _):
x = torch.randint(0, vocab_size, (batch_size, max_seq_len), dtype=torch.int64, device=device)
y = torch.randint(0, vocab_size, (batch_size, max_seq_len), dtype=torch.int64, device=device)
- compiled_func = thunder.compile(
+ compiled_func = thunder.jit(
model.eval(),
- disable_torch_autograd_support=True,
- executors_list=executor.executors_list(),
+ disable_torch_autograd=True,
+ executors=executor.executors_list(),
nv_enable_bookend=False,
)
compiled_func(x, y)
@@ -356,11 +356,11 @@ def test_cse_rematerialization(executor, device, _):
assert len(fusion_bsyms) == 11
# fusion groups 1 and 6 correspond with the apply_rotary_emb function
# Nvfuser with recomputation should use precomputed cos and sin values.
- assert len(fusion_bsyms[1].args) == len(fusion_bsyms[6].args)
- assert fusion_bsyms[1].args[0].name == "freqs_cos"
- assert fusion_bsyms[1].args[1].name == "freqs_sin"
- assert fusion_bsyms[6].args[0].name == "freqs_cos"
- assert fusion_bsyms[6].args[1].name == "freqs_sin"
+ assert len(fusion_bsyms[1].args) == len(fusion_bsyms[7].args)
+ assert fusion_bsyms[1].subsymbols[0].output.name == "freqs_cos"
+ assert fusion_bsyms[1].subsymbols[1].output.name == "freqs_sin"
+ assert fusion_bsyms[7].subsymbols[0].output.name == "freqs_cos"
+ assert fusion_bsyms[7].subsymbols[1].output.name == "freqs_sin"
# Tests that two separated nvFuser regions can be merged when they don't depend
@@ -614,8 +614,7 @@ def func(x: torch.Tensor, s: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
dtypes=NOTHING,
executors=(
nvFuserExecutor,
- # NOTE torch executor does not have bookend optimization.
- # See comment: https://github.com/Lightning-AI/lightning-thunder/issues/571#issuecomment-1610778432
+ # NOTE We might want to do transpose bookend optimization for other executors than nvFuser.
),
)
def test_bookend_meta_optimization(executor, device, _):
diff --git a/thunder/tests/test_nvfuser_remat.py b/thunder/tests/test_nvfuser_remat.py
index 0f8560278b..b94403fcf3 100644
--- a/thunder/tests/test_nvfuser_remat.py
+++ b/thunder/tests/test_nvfuser_remat.py
@@ -310,8 +310,8 @@ def test_find_cut_dropout(executor, device, _):
ext_producer_outputs = find_external_producer_outputs(utils.consumers(trace), (), producer, consumer)
cut = find_cut(ext_producer_outputs, producer, consumer)
# Note t5 is the boolean mask for dropout. It should be chosen over the t6
- # that is the float32 mask. See this issue for the original problem:
- # https://github.com/Lightning-AI/lightning-thunder/issues/706
+ # that is the float32 mask. See this issue: "The Recomputation Algorithm on
+ # Dropout choses a float32 mask to save"
assert cut == ("t0", "t5", "t9")
diff --git a/thunder/tests/test_ops.py b/thunder/tests/test_ops.py
index b9d2fed286..867ca35a76 100644
--- a/thunder/tests/test_ops.py
+++ b/thunder/tests/test_ops.py
@@ -37,7 +37,7 @@ def snippet_torch_consistency(op: OpInfo, torch_op, sample: SampleInput, comp: C
thunder_result = op(*sample.args, **sample.kwargs)
torch_result = torch_op(*sample.args, **sample.kwargs)
- # TODO Review how lightning.compile returns Exception information
+ # TODO Review how thunder.jit returns Exception information
if isinstance(thunder_result, Exception):
raise thunder_result
diff --git a/thunder/tests/test_transformer_engine_executor.py b/thunder/tests/test_transformer_engine_executor.py
index 60aa6f64ea..41b6f83e36 100644
--- a/thunder/tests/test_transformer_engine_executor.py
+++ b/thunder/tests/test_transformer_engine_executor.py
@@ -70,7 +70,8 @@ def fn(x, w1, w2):
assert_close(w2.grad, te_linear2.weight.grad)
# Verifies te_linear was called
- forward_trace, backward_trace = thunder.last_traces(cfn)
+ forward_trace = thunder.last_traces(cfn)
+ backward_trace = thunder.last_backward_traces(cfn)
assert any(bsym.sym.name.startswith("te_linear") for bsym in forward_trace[-1].bound_symbols)
assert any(bsym.sym.name.startswith("te_functional_linear_backward") for bsym in backward_trace[-1].bound_symbols)
@@ -180,6 +181,6 @@ def foo(x, w):
)
cfunc(x, w)
- fwd_traces, _ = thunder.last_traces(cfunc)
+ fwd_traces = thunder.last_traces(cfunc)
# Verify that we have replaced `prims.linear` with `te_linear`
assert any(bsym.sym.name.startswith("te_linear") for bsym in fwd_traces[-1].bound_symbols)
diff --git a/thunder/torch/__init__.py b/thunder/torch/__init__.py
index a4388d9903..efb6319d83 100644
--- a/thunder/torch/__init__.py
+++ b/thunder/torch/__init__.py
@@ -12,7 +12,7 @@
import opt_einsum
-# Initialies the language context
+# Initializes the language context
from thunder.torch.langctx import register_method
import thunder.clang as clang
@@ -29,6 +29,7 @@
from thunder.core.symbol import Symbol
from thunder.core.transforms import register_grad, put_grads
from thunder.core.prims import get_grad, put_grad
+from thunder.core.baseutils import run_once
__all__ = [
"is_available",
@@ -224,8 +225,9 @@ def _parse_to_device_and_dtype(
dtype = to_dtype(dtype)
# Case 1 -- tensor first
else:
- # See https://github.com/Lightning-AI/lightning-thunder/issues/317
- # It'd be nice to write torch.Tensor here instead of TensorProxy
+ # It'd be nice to write torch.Tensor here instead of TensorProxy.
+ # See issue "Translate isinstance(a, torch.Tensor) calls so that
+ # TensorProxies can pass as torch.Tensors"
utils.check_type(tensor_dtype_or_device, TensorProxy)
device_ = tensor_dtype_or_device.device if device is None else to_device(device)
dtype_ = tensor_dtype_or_device.true_dtype if dtype is None else to_dtype(dtype)
@@ -413,7 +415,8 @@ def multinomial(
) -> TensorLike:
utils.check(out is None, lambda: "Non-None out is not supported", NotImplementedError)
- # See https://github.com/Lightning-AI/lightning-thunder/issues/2258
+ # See issue "randomness: enable PyTorch generators for operations like
+ # multinomial"
utils.check(
generator is None, lambda: f"multinomial does not yet support specifying a generator", NotImplementedError
)
@@ -430,7 +433,7 @@ def multinomial(
# TODO Maybe update this to return an offset of how far to advance the seed to acquire new values
-# See https://github.com/Lightning-AI/lightning-thunder/issues/1360
+# See issue "Maybe return offset from thunder.torch.uniform_philox"
@torchsymbol(is_method=False, id="torch.uniform_philox")
def uniform_philox(
shape: Sequence[int],
@@ -884,6 +887,33 @@ def squeeze(a: TensorLike, /, dim: None | int | Sequence[int] = None) -> TensorL
return clang.squeeze(a, dims)
+@torchsymbol(torch.t, is_method=True)
+def t(a: TensorLike, /) -> TensorLike:
+ utils.check(
+ a.ndim <= 2,
+ lambda: f"t() expects a tensor with <= 2 dimensions, but self is {a.ndim}D",
+ RuntimeError,
+ )
+ return prims.transpose(a, (1, 0)) if a.ndim == 2 else a
+
+
+@run_once
+def warn_ndim_not_2():
+ warnings.warn(
+ "The use of `x.T` on tensors of dimension other than 2 to reverse their shape is deprecated and will throw an error in a future release."
+ "Consider `x.mT` to transpose batches of matrices or `x.permute(*torch.arange(x.ndim - 1, -1, -1))` to reverse the dimensions of a tensor."
+ )
+
+
+def reverse_dims_T(a: TensorLike, /) -> TensorLike:
+ if a.ndim != 2:
+ warn_ndim_not_2()
+ return a if a.ndim < 2 else prims.transpose(a, tuple(reversed(range(a.ndim))))
+
+
+register_method("T", reverse_dims_T)
+
+
# TODO Add type annotations
# See https://pytorch.org/docs/master/generated/torch.tensor_split.html
@torchsymbol(torch.tensor_split, is_method=True)
@@ -1362,9 +1392,9 @@ def zeta(a, b, /):
# For calculate op1(a, op2(value, op2(b, c))) by promoting all input tensors at once
# NOTE use this explicit type promotion because a direct combination of add/mul will have a redundant cast,
-# which may lead to accuracy problems, see:
-# https://github.com/Lightning-AI/lightning-thunder/pull/1155#discussion_r1342653591 for details
-# TODO remove this when the optimization pass is ready: https://github.com/Lightning-AI/lightning-thunder/issues/1178
+# which may lead to accuracy problems.
+# TODO remove after issue "Redundant cast removal could be performed through metadata-only
+# operations, like broadcasting" is resolved
def addcmul_addcdiv_helper(
a, b, c, op1, op2, *, value=None, type_promotion_kind=utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT
):
@@ -2840,8 +2870,9 @@ def _avg_pool_helper(
# Dimensionality of the kernel.
kernel_numel = reduce(operator.mul, kernel_size, 1)
- # nn.functional.avg_pool does not have `divisor_override` for some reason.
- # TODO: seems like an oversight from PyTorch and/or 1d case is very niche.
+ # nn.functional.avg_pool does not have `divisor_override`.
+ # TODO: look into PyTorch side; is this behavior deliberate? Could be that
+ # 1D case is niche.
# If needed, handle it with checks and transforms. For now unconditionally
# override value with kernel_numel.
if divisor_override is None or dim == 1:
@@ -3003,7 +3034,8 @@ def _dropout_helper(a, p):
# TODO Add annotations, make not a prim
# The backward decomposition of cross_entropy cannot be efficiently fused, so we have this cross_entropy_backward
# primitive. Executors can override the primitive using internal implementations.
-# See https://github.com/Lightning-AI/lightning-thunder/issues/660
+# See issue "Cross_entropy is decomposed for backward but the decomposition is
+# not fusible currently"
@torchsymbol("cross_entropy_backward", id="cross_entropy_backward", is_prim=True)
def cross_entropy_backward(g, a, /, target, weight, reduction, ignore_index, label_smoothing):
return TensorProxy(like=g, shape=a.shape)
@@ -3571,7 +3603,8 @@ def log_softmax(a: TensorLike, /, dim: int, *, dtype: None | dtypeLike = None) -
# TODO Update annotations and consider moving to torchex
# We improve the efficiency of cross_entropy backward decomposition by adding the log_softmax_backward
# and nll_loss_backward primitives. Executors can override the primitives using internal implementations.
-# See https://github.com/Lightning-AI/lightning-thunder/issues/660
+# See issue "Cross_entropy is decomposed for backward but the decomposition is
+# not fusible currently"
@torchsymbol("log_softmax_backward", id="log_softmax_backward")
def log_softmax_backward(g: TensorProxy, /, output: TensorProxy, dim: int, dtype: dtypeLike) -> TensorLike:
dtype: dtypes.dtype = to_dtype(dtype)
@@ -3817,7 +3850,7 @@ def softmax(a: TensorLike, /, dim: int, *, dtype: None | dtypeLike = None) -> Te
if torch.distributed.is_available():
DistributedReduceOpLike = str | torch.distributed.ReduceOp | dist_prims.DistributedReduceOps
- # string name, PyTorch enum value, lightning.compile enum value
+ # string name, PyTorch enum value, thunder.jit enum value
_reduceop_triples = (("sum", torch.distributed.ReduceOp.SUM, dist_prims.DistributedReduceOps.SUM),)
def to_thunder_distributed_reduce_op(op: DistributedReduceOpLike | None):
+ 
+
+