From c910e46ebf883204db577af95ef1abcbe696420c Mon Sep 17 00:00:00 2001 From: Vikram Chandrashekar Date: Tue, 14 Jul 2026 14:26:01 -0700 Subject: [PATCH] Add FLOOR_MOD (floored modulo) pointwise mode Adds PointwiseMode_t::FLOOR_MOD, mapping to the backend CUDNN_POINTWISE_FLOOR_MOD enum (cuDNN 9.26.0+). FLOOR_MOD is floored modulo (y = a - floor(a / b) * b; the result follows the sign of the divisor), complementing the existing truncated MOD (C fmod/%; result follows the sign of the dividend); the two differ only for operands of opposite sign. Exposed as graph.floor_mod in the Python bindings, with a demo notebook contrasting mod (torch.fmod) and floor_mod (torch.remainder). The references to the backend CUDNN_POINTWISE_FLOOR_MOD enum are guarded by CUDNN_VERSION >= 92600, so the header-only frontend still compiles cleanly against older cuDNN releases (there floor_mod resolves to CUDNN_STATUS_NOT_SUPPORTED). Adds unit-test coverage for floor_mod. Signed-off-by: Vikram Chandrashekar --- include/cudnn_frontend_Operation.h | 1 + include/cudnn_frontend_utils.h | 14 ++ python/cudnn/_pygraph.py | 2 +- python/pygraph/pointwise.cpp | 18 ++ samples/README.md | 3 + samples/python/71_pointwise_mod.ipynb | 280 ++++++++++++++++++++++++++ test/python/test_pointwise.py | 86 ++++++++ 7 files changed, 403 insertions(+), 1 deletion(-) create mode 100644 samples/python/71_pointwise_mod.ipynb create mode 100644 test/python/test_pointwise.py diff --git a/include/cudnn_frontend_Operation.h b/include/cudnn_frontend_Operation.h index fa9dbba20..a3fceba6b 100644 --- a/include/cudnn_frontend_Operation.h +++ b/include/cudnn_frontend_Operation.h @@ -2930,6 +2930,7 @@ class OperationBuilder_v8 { (m_operation.pointwise_mode == PointwiseMode_t::LOG) || (m_operation.pointwise_mode == PointwiseMode_t::NEG) || (m_operation.pointwise_mode == PointwiseMode_t::MOD) || + (m_operation.pointwise_mode == PointwiseMode_t::FLOOR_MOD) || (m_operation.pointwise_mode == PointwiseMode_t::POW) || (m_operation.pointwise_mode == PointwiseMode_t::ABS) || (m_operation.pointwise_mode == PointwiseMode_t::CEIL) || diff --git a/include/cudnn_frontend_utils.h b/include/cudnn_frontend_utils.h index a7bbd06fb..48517ac32 100644 --- a/include/cudnn_frontend_utils.h +++ b/include/cudnn_frontend_utils.h @@ -556,6 +556,7 @@ enum class PointwiseMode_t { LOG, NEG, MOD, + FLOOR_MOD, POW, ABS, CEIL, @@ -611,6 +612,7 @@ NLOHMANN_JSON_SERIALIZE_ENUM(PointwiseMode_t, {PointwiseMode_t::LOG, "LOG"}, {PointwiseMode_t::NEG, "NEG"}, {PointwiseMode_t::MOD, "MOD"}, + {PointwiseMode_t::FLOOR_MOD, "FLOOR_MOD"}, {PointwiseMode_t::POW, "POW"}, {PointwiseMode_t::ABS, "ABS"}, {PointwiseMode_t::CEIL, "CEIL"}, @@ -818,6 +820,7 @@ get_pointwise_mode_port_count(PointwiseMode_t const& mode) { case PointwiseMode_t::MIN: case PointwiseMode_t::MAX: case PointwiseMode_t::MOD: + case PointwiseMode_t::FLOOR_MOD: case PointwiseMode_t::RELU_BWD: case PointwiseMode_t::TANH_BWD: case PointwiseMode_t::SIGMOID_BWD: @@ -1343,6 +1346,13 @@ convert_to_cudnn_type(cudnn_frontend::PointwiseMode_t const mode, cudnnPointwise case PointwiseMode_t::MOD: cudnn_mode = CUDNN_POINTWISE_MOD; return cudnnStatus_t::CUDNN_STATUS_SUCCESS; + case PointwiseMode_t::FLOOR_MOD: +#if (CUDNN_VERSION >= 92600) + cudnn_mode = CUDNN_POINTWISE_FLOOR_MOD; + return cudnnStatus_t::CUDNN_STATUS_SUCCESS; +#else + return cudnnStatus_t::CUDNN_STATUS_NOT_SUPPORTED; +#endif case PointwiseMode_t::POW: cudnn_mode = CUDNN_POINTWISE_POW; return cudnnStatus_t::CUDNN_STATUS_SUCCESS; @@ -2279,6 +2289,10 @@ convert_from_cudnn_type(cudnnPointwiseMode_t const cudnn_mode) { return PointwiseMode_t::NEG; case CUDNN_POINTWISE_MOD: return PointwiseMode_t::MOD; +#if (CUDNN_VERSION >= 92600) + case CUDNN_POINTWISE_FLOOR_MOD: + return PointwiseMode_t::FLOOR_MOD; +#endif case CUDNN_POINTWISE_POW: return PointwiseMode_t::POW; case CUDNN_POINTWISE_ABS: diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index 4993ed1dd..cf30cbb31 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -543,7 +543,7 @@ def matmul( }, # binary **{op: ("a", "b") for op in ("add", "add_square", "div", "logical_and", "logical_or", "mul", "sub")}, - **{op: ("input0", "input1") for op in ("max", "min", "mod", "pow")}, + **{op: ("input0", "input1") for op in ("max", "min", "mod", "floor_mod", "pow")}, **{op: ("input", "comparison") for op in ("cmp_eq", "cmp_ge", "cmp_gt", "cmp_le", "cmp_lt", "cmp_neq")}, "bias": ("input", "bias"), "scale": ("input", "scale"), diff --git a/python/pygraph/pointwise.cpp b/python/pygraph/pointwise.cpp index f8d3d4592..10fe96fc7 100644 --- a/python/pygraph/pointwise.cpp +++ b/python/pygraph/pointwise.cpp @@ -753,6 +753,24 @@ init_pygraph_pointwise_submodule(py::class_& m) { Returns: cudnn_tensor: The result of pointwise floating-point remainder of the input0 tensor's division by the input1 tensor )pbdoc"); + m.def("floor_mod", + &PyGraph::pointwise_binary, + py::arg("input0"), + py::arg("input1"), + py::arg_v("compute_data_type", cudnn_frontend::DataType_t::NOT_SET), + py::arg_v("name", ""), + R"pbdoc( + Compute the floored element-wise modulo (remainder) of two tensors: y = a - floor(a / b) * b. The result follows the sign of the divisor b. This differs from mod (truncated remainder, C fmod/%, result follows the sign of the dividend) only when a and b have opposite signs. + + Args: + input0 (cudnn_tensor): The dividend tensor a. + input1 (cudnn_tensor): The divisor tensor b. + compute_data_type (Optional[cudnn.data_type]): The data type for computation. Default is NOT_SET. + name (Optional[str]): A name for the operation to be performed. + + Returns: + cudnn_tensor: The result of the floored modulo operation. + )pbdoc"); m.def("pow", &PyGraph::pointwise_binary, py::arg("input0"), diff --git a/samples/README.md b/samples/README.md index 83a21dffd..66cde25fd 100644 --- a/samples/README.md +++ b/samples/README.md @@ -20,6 +20,9 @@ Samples leveraging FE's Python interface are located in [samples/python](python/ * [53_sdpa](python/53_sdpa_decode_with_paged_caches.ipynb) Shows how to run scaled dot product attention (decode phase) where the K and V caches are stored in non contiguous memory. +* [71_pointwise_mod](python/71_pointwise_mod.ipynb) + Demonstrates the two flavors of pointwise modulo: truncated (mod / torch.fmod) vs floored (floor_mod / torch.remainder). + ## C++ Interface Samples Samples leveraging FE's C++ interface are located in [samples/cpp](cpp/). diff --git a/samples/python/71_pointwise_mod.ipynb b/samples/python/71_pointwise_mod.ipynb new file mode 100644 index 000000000..919e6497e --- /dev/null +++ b/samples/python/71_pointwise_mod.ipynb @@ -0,0 +1,280 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Pointwise Modulo: truncated (`mod`) vs floored (`floor_mod`)\n", + "\n", + "cuDNN's pointwise API provides **two flavors of modulo (remainder)**. They produce the same result\n", + "whenever the dividend `a` and the divisor `b` have the **same sign**, and differ **only** when `a`\n", + "and `b` have **opposite signs**.\n", + "\n", + "## `graph.mod(a, b)` — truncated remainder\n", + "\n", + "Truncated modulo, matching C's `fmod` / `%` and **`torch.fmod(a, b)`**. The quotient is rounded\n", + "toward zero, so the result takes the **sign of the dividend `a`** (backend enum\n", + "`CUDNN_POINTWISE_MOD`).\n", + "\n", + "## `graph.floor_mod(a, b)` — floored remainder\n", + "\n", + "Floored modulo, defined as\n", + "\n", + "$$y = a - \\lfloor a / b \\rfloor \\cdot b$$\n", + "\n", + "matching Python's `%` and **`torch.remainder(a, b)`**. The quotient is rounded toward negative\n", + "infinity, so the result takes the **sign of the divisor `b`** (backend enum\n", + "`CUDNN_POINTWISE_FLOOR_MOD`).\n", + "\n", + "## Worked example\n", + "\n", + "For `a = [-7, 7, -7, 7]` and `b = [3, 3, -3, -3]`:\n", + "\n", + "| `a` | `b` | `mod` = `torch.fmod(a, b)` (sign of `a`) | `floor_mod` = `torch.remainder(a, b)` (sign of `b`) |\n", + "|----:|----:|----------------------------------------:|----------------------------------------------------:|\n", + "| -7 | 3 | -1 | 2 |\n", + "| 7 | 3 | 1 | 1 |\n", + "| -7 | -3 | -1 | -1 |\n", + "| 7 | -3 | 1 | -2 |\n", + "\n", + "The first and last rows have opposite-sign operands, so `mod` and `floor_mod` disagree there; the\n", + "middle two rows share a sign, so they agree." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites and Setup\n", + "\n", + "This notebook requires an NVIDIA GPU and a cuDNN build whose backend provides the floored-modulo\n", + "pointwise mode `CUDNN_POINTWISE_FLOOR_MOD` (**cuDNN 9.26+**). On older backends the `floor_mod`\n", + "cell fails at build time; the truncated `mod` cell works on any backend that supports\n", + "`CUDNN_POINTWISE_MOD`.\n", + "\n", + "If the backend does not support `floor_mod`, the `floor_mod` cells below catch `cudnnGraphNotSupportedError`\n", + "and are skipped gracefully instead of crashing the notebook.\n", + "\n", + "**Environment setup:**\n", + "\n", + "- **Option A — pip install:**\n", + " ```bash\n", + " pip install nvidia-cudnn-frontend\n", + " ```\n", + "- **Option B — set paths manually:**\n", + " ```bash\n", + " export LD_LIBRARY_PATH=/path/to/cudnn/lib:${LD_LIBRARY_PATH}\n", + " export PYTHONPATH=/path/to/cudnn_frontend/build_python:${PYTHONPATH}\n", + " ```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# get_ipython().system('pip install nvidia-cudnn-cu13')\n", + "# get_ipython().system('pip install nvidia-cudnn-frontend')\n", + "# get_ipython().system('pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "warnings.filterwarnings(\"ignore\", message=\"Failed to initialize NumPy\")\n", + "\n", + "import torch\n", + "import cudnn\n", + "\n", + "# floor_mod requires a cuDNN backend that provides CUDNN_POINTWISE_FLOOR_MOD (cuDNN 9.26+).\n", + "print(\"cuDNN backend version:\", cudnn.backend_version())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Illustrative operands, chosen so mod and floor_mod disagree on the opposite-sign entries.\n", + "# a and b share a sign in columns 1 and 2 (agree), and have opposite signs in columns 0 and 3 (differ).\n", + "device = \"cuda\"\n", + "\n", + "a_gpu = torch.tensor([[-7.0, 7.0, -7.0, 7.0]], dtype=torch.float32, device=device)\n", + "b_gpu = torch.tensor([[3.0, 3.0, -3.0, -3.0]], dtype=torch.float32, device=device)\n", + "\n", + "# One cuDNN handle is created here and reused for both graphs below.\n", + "handle = cudnn.create_handle()\n", + "\n", + "print(f\"a: {a_gpu.tolist()}\")\n", + "print(f\"b: {b_gpu.tolist()}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Truncated modulo: y = mod(a, b). Result takes the sign of the dividend a (matches torch.fmod).\n", + "mod_graph = cudnn.pygraph(\n", + " handle=handle,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + ")\n", + "\n", + "a = mod_graph.tensor_like(a_gpu)\n", + "b = mod_graph.tensor_like(b_gpu)\n", + "\n", + "y = mod_graph.mod(a, b, compute_data_type=cudnn.data_type.FLOAT, name=\"mod\")\n", + "y.set_output(True).set_data_type(cudnn.data_type.FLOAT)\n", + "\n", + "mod_graph.validate()\n", + "mod_graph.build_operation_graph()\n", + "mod_graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", + "mod_graph.check_support()\n", + "mod_graph.build_plans()\n", + "\n", + "mod_gpu = torch.empty_like(a_gpu)\n", + "workspace = torch.empty(mod_graph.get_workspace_size(), device=device, dtype=torch.uint8)\n", + "\n", + "mod_graph.execute({a: a_gpu, b: b_gpu, y: mod_gpu}, workspace, handle=handle)\n", + "torch.cuda.synchronize()\n", + "\n", + "print(f\"mod result: {mod_gpu.tolist()}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# torch.fmod is the truncated remainder (sign of the dividend) -- the reference for graph.mod.\n", + "ref_trunc = torch.fmod(a_gpu, b_gpu)\n", + "\n", + "print(f\"cuDNN mod: {mod_gpu.tolist()}\")\n", + "print(f\"torch.fmod: {ref_trunc.tolist()}\")\n", + "assert torch.allclose(mod_gpu, ref_trunc), f\"mod mismatch: {mod_gpu} vs {ref_trunc}\"\n", + "print(\"PASSED: cuDNN mod matches torch.fmod (truncated remainder).\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Floored modulo: y = a - floor(a / b) * b. Result takes the sign of the divisor b (matches torch.remainder).\n", + "# On a cuDNN version / GPU that does not support floor_mod, the build raises\n", + "# cudnnGraphNotSupportedError, which we catch to skip gracefully.\n", + "try:\n", + " floor_mod_graph = cudnn.pygraph(\n", + " handle=handle,\n", + " intermediate_data_type=cudnn.data_type.FLOAT,\n", + " compute_data_type=cudnn.data_type.FLOAT,\n", + " )\n", + "\n", + " a = floor_mod_graph.tensor_like(a_gpu)\n", + " b = floor_mod_graph.tensor_like(b_gpu)\n", + "\n", + " y = floor_mod_graph.floor_mod(a, b, compute_data_type=cudnn.data_type.FLOAT, name=\"floor_mod\")\n", + " y.set_output(True).set_data_type(cudnn.data_type.FLOAT)\n", + "\n", + " floor_mod_graph.validate()\n", + " floor_mod_graph.build_operation_graph()\n", + " floor_mod_graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])\n", + " floor_mod_graph.check_support()\n", + " floor_mod_graph.build_plans()\n", + "\n", + " floor_mod_gpu = torch.empty_like(a_gpu)\n", + " workspace = torch.empty(floor_mod_graph.get_workspace_size(), device=device, dtype=torch.uint8)\n", + "\n", + " floor_mod_graph.execute({a: a_gpu, b: b_gpu, y: floor_mod_gpu}, workspace, handle=handle)\n", + " torch.cuda.synchronize()\n", + "\n", + " print(f\"floor_mod result: {floor_mod_gpu.tolist()}\")\n", + "except cudnn.cudnnGraphNotSupportedError:\n", + " floor_mod_gpu = None\n", + " print(\"floor_mod not supported on this GPU / cuDNN version (requires cuDNN 9.26+) — skipping.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# torch.remainder is the floored remainder (sign of the divisor) -- the reference for graph.floor_mod.\n", + "if floor_mod_gpu is not None:\n", + " ref_floor = torch.remainder(a_gpu, b_gpu)\n", + "\n", + " print(f\"cuDNN floor_mod: {floor_mod_gpu.tolist()}\")\n", + " print(f\"torch.remainder: {ref_floor.tolist()}\")\n", + " assert torch.allclose(floor_mod_gpu, ref_floor), f\"floor_mod mismatch: {floor_mod_gpu} vs {ref_floor}\"\n", + " print(\"PASSED: cuDNN floor_mod matches torch.remainder (floored remainder).\")\n", + "else:\n", + " print(\"Skipped floor_mod reference check (op unsupported here).\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Side-by-side: mod and floor_mod agree on same-sign operands and differ on opposite-sign operands.\n", + "if floor_mod_gpu is not None:\n", + " a_list = a_gpu.flatten().tolist()\n", + " b_list = b_gpu.flatten().tolist()\n", + " mod_list = mod_gpu.flatten().tolist()\n", + " floor_list = floor_mod_gpu.flatten().tolist()\n", + "\n", + " print(f\"{'a':>4} {'b':>4} {'mod (torch.fmod)':>18} {'floor_mod (torch.remainder)':>28} note\")\n", + " print(\"-\" * 66)\n", + " for av, bv, mv, fv in zip(a_list, b_list, mod_list, floor_list):\n", + " note = \"differ\" if abs(mv - fv) > 1e-6 else \"agree\"\n", + " print(f\"{av:4.0f} {bv:4.0f} {mv:18.0f} {fv:28.0f} {note}\")\n", + "\n", + " print(\"\\nThe two flavors differ exactly on the opposite-sign rows (a=-7, b=3 and a=7, b=-3).\")\n", + "else:\n", + " print(\"floor_mod was not supported on this GPU / cuDNN version, so only the mod flavor ran above; skipping the side-by-side comparison.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "cudnn.destroy_handle(handle)" + ] + } + ], + "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.14.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/test/python/test_pointwise.py b/test/python/test_pointwise.py new file mode 100644 index 000000000..37e131fc6 --- /dev/null +++ b/test/python/test_pointwise.py @@ -0,0 +1,86 @@ +"""Functional coverage for the binary modulo pointwise ops. + +Each op is lowered into a single-node cuDNN graph that is built and executed on +the GPU, then compared against a torch reference: + + * mod -> truncated remainder (C fmod/%; the result follows the sign of + the dividend), reference torch.fmod + * floor_mod -> floored remainder (the result follows the sign of the divisor), + reference torch.remainder + +floor_mod maps to the backend CUDNN_POINTWISE_FLOOR_MOD enum, which was added in +cuDNN 9.26.0. On older backends the frontend resolves it to +CUDNN_STATUS_NOT_SUPPORTED, so the graph build declines. The floor_mod case is +therefore version-gated and additionally wrapped in the standard +cudnnGraphNotSupportedError skip, so it is waived (never failed) on pre-9.26 +cuDNN or on hardware/configs that cannot serve it. +""" + +import cudnn +import pytest +import torch + +# First cuDNN backend version exposing CUDNN_POINTWISE_FLOOR_MOD. +FLOOR_MOD_MIN_BACKEND_VERSION = 92600 + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "op_name, torch_reference", + [ + ("mod", torch.fmod), # truncated: result follows the sign of the dividend + ("floor_mod", torch.remainder), # floored: result follows the sign of the divisor + ], +) +def test_pointwise_modulo(op_name, torch_reference, cudnn_handle): + if cudnn_handle is None: + pytest.skip("cuDNN backend not available") + + if op_name == "floor_mod" and cudnn.backend_version() < FLOOR_MOD_MIN_BACKEND_VERSION: + pytest.skip("floor_mod (CUDNN_POINTWISE_FLOOR_MOD) requires cuDNN 9.26.0 or newer") + + # A mixed-sign dividend with a positive divisor is exactly where truncated (mod) + # and floored (floor_mod) remainder disagree, so this genuinely exercises the + # floored semantics rather than a case both modes share. Integer-valued fp32 + # keeps a - trunc/floor(a / b) * b exact, so no tolerance slack is needed. + a_gpu = torch.tensor([-7.0, -3.0, -1.0, 2.0, 5.0, 8.0, -6.0, 9.0], device="cuda", dtype=torch.float32).reshape(2, 4) + b_gpu = torch.full_like(a_gpu, 4.0) + + stream = torch.cuda.current_stream().cuda_stream + cudnn.set_stream(handle=cudnn_handle, stream=stream) + + graph = cudnn.pygraph( + io_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + handle=cudnn_handle, + ) + + a = graph.tensor_like(a_gpu) + b = graph.tensor_like(b_gpu) + + c = getattr(graph, op_name)(a, b) + c.set_output(True).set_data_type(cudnn.data_type.FLOAT) + + graph.validate() + + try: + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + except cudnn.cudnnGraphNotSupportedError as e: + pytest.skip(f"TEST WAIVED: {op_name} not supported on this backend/config. {e}") + + graph.build_plans(cudnn.build_plan_policy.HEURISTICS_CHOICE) + + c_expected = torch_reference(a_gpu, b_gpu) + c_actual = torch.zeros_like(c_expected) + + workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8) + graph.execute({a: a_gpu, b: b_gpu, c: c_actual}, workspace, handle=cudnn_handle) + + torch.cuda.synchronize() + torch.testing.assert_close(c_expected, c_actual) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])