[PyBindings] Pass float8 tensors through DLPack

extension/pybindings/README.md

@@ -206,10 +206,11 @@ Other dtypes NumPy lacks, such as float8, are shown the same way, as integers of
 the same width. NumPy converts that view as integers, so
 `np.asarray(out, dtype=np.float32)` gives the bit patterns as numbers, not the
 bfloat16 values. DLPack consumers, such as `torch.from_dlpack(out)`, receive the
-original bfloat16 dtype. A bfloat16 result passed back in is read through DLPack
-too, so it keeps its dtype, shape, and dim order. A bfloat16 input does not
-accept 16-bit integers, so that integer data is never read as bfloat16 by
-mistake; pass the result itself or raw bytes.
+original bfloat16 or float8 dtype. A bfloat16 or float8 result passed back in is
+read through DLPack too, so it keeps its dtype, shape, and dim order. A bfloat16
+input does not accept 16-bit integers, so that integer data is never read as
+bfloat16 by mistake; pass the result itself or raw bytes. A float8 input does not
+accept 8-bit integers either; pass the result itself or a float8 tensor.
 
 ```python
 import cupy as cp

extension/pybindings/pybindings.cpp

@@ -1709,8 +1709,8 @@ class BorrowedInput final {
 };
 
 #ifndef USE_ATEN_LIB
-// Asks for a DLPack 1.0 capsule, which can mark memory read-only, and falls
-// back to a legacy capsule for producers that predate it.
+// Asks for a versioned DLPack capsule, which can mark memory read-only, and
+// falls back to a legacy capsule for producers that predate versioned capsules.
 py::object dlpack_capsule(
     const py::object& producer,
     const py::object& stream) {
@@ -1720,7 +1720,8 @@ py::object dlpack_capsule(
       arguments["stream"] = stream;
     }
     if (versioned) {
-      arguments["max_version"] = py::make_tuple(1, 0);
+      arguments["max_version"] =
+          py::make_tuple(dlpack::kVersion.major, dlpack::kVersion.minor);
     }
     return producer.attr("__dlpack__")(**arguments);
   };
@@ -1990,6 +1991,20 @@ class DLPackTensorView final {
     if (dtype.code == dlpack::DataTypeCode::BFloat && dtype.bits == 16) {
       return ScalarType::BFloat16;
     }
+    if (dtype.code == dlpack::DataTypeCode::Float8_e4m3fn && dtype.bits == 8) {
+      return ScalarType::Float8_e4m3fn;
+    }
+    if (dtype.code == dlpack::DataTypeCode::Float8_e4m3fnuz &&
+        dtype.bits == 8) {
+      return ScalarType::Float8_e4m3fnuz;
+    }
+    if (dtype.code == dlpack::DataTypeCode::Float8_e5m2 && dtype.bits == 8) {
+      return ScalarType::Float8_e5m2;
+    }
+    if (dtype.code == dlpack::DataTypeCode::Float8_e5m2fnuz &&
+        dtype.bits == 8) {
+      return ScalarType::Float8_e5m2fnuz;
+    }
     if (dtype.code == dlpack::DataTypeCode::Complex && dtype.bits == 32) {
       return ScalarType::ComplexHalf;
     }

extension/pybindings/pybindings_dlpack.h

@@ -31,6 +31,11 @@ enum DataTypeCode : uint8_t {
   Complex = 5,
   BFloat = 4,
   Bool = 6,
+  // Since DLPack 1.1.
+  Float8_e4m3fn = 10,
+  Float8_e4m3fnuz = 11,
+  Float8_e5m2 = 12,
+  Float8_e5m2fnuz = 13,
 };
 
 struct Device {
@@ -65,6 +70,10 @@ struct Version {
   uint32_t minor;
 };
 
+// The version these bindings export and ask producers for: 1.1 added the
+// float8 codes.
+constexpr Version kVersion{1, 1};
+
 // DLPack 1.0 capsules carry this struct, named "dltensor_versioned".
 struct ManagedTensorVersioned {
   Version version;

extension/pybindings/pybindings_executorch_result.h

@@ -549,7 +549,7 @@ class PyExecuTorchResult final
           0};
       managed.manager_ctx = this;
       managed.deleter = release_dlpack_holder;
-      versioned.version = {1, 0};
+      versioned.version = dlpack::kVersion;
       versioned.manager_ctx = this;
       versioned.deleter = release_versioned_dlpack_holder;
       versioned.dl_tensor = managed.dl_tensor;
@@ -654,6 +654,14 @@ class PyExecuTorchResult final
         return {dlpack::DataTypeCode::UInt, 32, 1};
       case ScalarType::UInt64:
         return {dlpack::DataTypeCode::UInt, 64, 1};
+      case ScalarType::Float8_e4m3fn:
+        return {dlpack::DataTypeCode::Float8_e4m3fn, 8, 1};
+      case ScalarType::Float8_e4m3fnuz:
+        return {dlpack::DataTypeCode::Float8_e4m3fnuz, 8, 1};
+      case ScalarType::Float8_e5m2:
+        return {dlpack::DataTypeCode::Float8_e5m2, 8, 1};
+      case ScalarType::Float8_e5m2fnuz:
+        return {dlpack::DataTypeCode::Float8_e5m2fnuz, 8, 1};
       default:
         throw py::value_error(
             "ExecuTorch result dtype cannot be represented by DLPack");

extension/pybindings/test/test_pybindings.py

@@ -2291,6 +2291,113 @@ def forward(self, x):
             torch.equal(output.view(torch.float16), value.view(torch.float16))
         )
 
+    def test_float8_dlpack_input(self):
+        if self.kernel_mode != "portable":
+            self.skipTest("only portable bindings read DLPack inputs themselves")
+
+        class Identity(torch.nn.Module):
+            def forward(self, x):
+                return x
+
+        for dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+            value = torch.tensor([1.0, -2.0, 0.5, 0.0]).to(dtype)
+            program = to_edge(export(Identity(), (value,))).to_executorch(
+                ExecutorchBackendConfig(
+                    memory_planning_pass=MemoryPlanningPass(
+                        alloc_graph_input=False, alloc_graph_output=False
+                    )
+                )
+            )
+            method = self.load_prog_fn(program.buffer).load_method("forward")
+            # The wrapper takes no max_version, so it gets a legacy capsule; raw
+            # capsules also cover the versioned kind.
+            for kind, given in (
+                ("wrapper", DLPackOnly(value)),
+                ("legacy capsule", value.__dlpack__()),
+                ("versioned capsule", value.__dlpack__(max_version=(1, 1))),
+            ):
+                with self.subTest(dtype=dtype, input=kind):
+                    output = torch.from_dlpack(method(given)[0])
+                    self.assertEqual(output.dtype, dtype)
+                    self.assertTrue(
+                        torch.equal(output.view(torch.uint8), value.view(torch.uint8))
+                    )
+            if dtype == torch.float8_e4m3fn:
+                # exir cannot export fnuz dtypes, so their codes are checked by
+                # the dtype a float8_e4m3fn method reports when it refuses them,
+                # read from a versioned capsule.
+                for fnuz in (torch.float8_e4m3fnuz, torch.float8_e5m2fnuz):
+                    name = str(fnuz).removeprefix("torch.").capitalize()
+                    with self.subTest(dtype=fnuz):
+                        with self.assertRaisesRegex(
+                            ValueError,
+                            f"has dtype {name} but the method expects Float8_e4m3fn",
+                        ):
+                            method(value.to(fnuz).__dlpack__(max_version=(1, 1)))
+
+    def test_float8_result_round_trips_through_dlpack(self):
+        if self.kernel_mode != "portable":
+            self.skipTest("only portable bindings return ExecuTorchResult outputs")
+
+        class Identity(torch.nn.Module):
+            def forward(self, x):
+                return x
+
+        config = ExecutorchBackendConfig(
+            memory_planning_pass=MemoryPlanningPass(
+                alloc_graph_input=False, alloc_graph_output=False
+            )
+        )
+        script = """
+import importlib
+import sys
+import torch
+
+runtime = importlib.import_module(sys.argv[1])
+native = sys.modules[runtime.ExecuTorchResult.__module__]
+assert native._tensor_output == "executorch"
+for path, name in zip(sys.argv[2::2], sys.argv[3::2]):
+    dtype = getattr(torch, name)
+    value = torch.tensor([1.0, -2.0, 0.5, 0.0]).to(dtype)
+    with open(path, "rb") as source:
+        data = source.read()
+    method = runtime._load_program_from_buffer(data).load_method("forward")
+    output = method(value)[0]
+    assert isinstance(output, runtime.ExecuTorchResult), type(output)
+    tensor = torch.from_dlpack(output)
+    assert tensor.dtype == dtype, (name, tensor.dtype)
+    assert torch.equal(tensor.view(torch.uint8), value.view(torch.uint8)), name
+    roundtrip = torch.from_dlpack(method(output)[0])
+    assert roundtrip.dtype == dtype, (name, roundtrip.dtype)
+    assert torch.equal(roundtrip.view(torch.uint8), value.view(torch.uint8)), name
+print("float8 results round trip through DLPack")
+"""
+        arguments = [self.runtime.__name__]
+        with tempfile.TemporaryDirectory() as directory:
+            for dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+                name = str(dtype).removeprefix("torch.")
+                value = torch.tensor([1.0, -2.0, 0.5, 0.0]).to(dtype)
+                program = to_edge(export(Identity(), (value,))).to_executorch(config)
+                path = os.path.join(directory, f"{name}.pte")
+                with open(path, "wb") as source:
+                    source.write(program.buffer)
+                arguments += [path, name]
+            environment = os.environ.copy()
+            environment["EXECUTORCH_PYBINDINGS_TENSOR_OUTPUT"] = "executorch"
+            completed = subprocess.run(
+                [sys.executable, "-c", script, *arguments],
+                env=environment,
+                capture_output=True,
+                text=True,
+                timeout=60,
+            )
+        self.assertEqual(
+            completed.returncode,
+            0,
+            f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}",
+        )
+        self.assertIn("float8 results round trip through DLPack", completed.stdout)
+
     def test_unsupported_buffer_format_is_rejected(self):
         exported_program, _ = create_program(ModuleAddSingleInput())
         executorch_module = self.load_fn(exported_program.buffer)

extension/pybindings/test/test_pybindings_cuda_dlpack.py

@@ -1188,6 +1188,39 @@ def test_bfloat16_host_result_moves_to_the_device(self):
         with self.assertRaisesRegex(ValueError, "already on a device"):
             self.bindings.ExecuTorchResult.from_host(on_device, self.device.index)
 
+    def test_float8_host_bits_run_through_a_method(self):
+        # Python buffers have no float8 format, so the values come as bytes. exir
+        # cannot export fnuz dtypes, so only the others run through a method.
+        for dtype in (
+            torch.float8_e4m3fn,
+            torch.float8_e4m3fnuz,
+            torch.float8_e5m2,
+            torch.float8_e5m2fnuz,
+        ):
+            name = str(dtype).removeprefix("torch.")
+            with self.subTest(dtype=name):
+                value = torch.tensor([[1.0, -2.0, 0.5], [3.0, 0.0, -0.25]]).to(dtype)
+                host = memoryview(
+                    bytes(value.view(torch.uint8).flatten().tolist())
+                ).cast("B", value.shape)
+                uploaded = self.bindings.ExecuTorchResult.from_host(
+                    host, self.device.index, dtype=name
+                )
+                self.assertEqual(uploaded.dtype, name)
+                expected = value.to(self.device).view(torch.uint8)
+                as_torch = torch.from_dlpack(uploaded)
+                self.assertEqual(as_torch.dtype, dtype)
+                torch.testing.assert_close(as_torch.view(torch.uint8), expected)
+                if name.endswith("fnuz"):
+                    continue
+                data = self._device_identity(self.device.index, value)
+                method = self.bindings._load_program_from_buffer(data).load_method(
+                    "forward"
+                )
+                output = self._tensor(method(uploaded)[0])
+                self.assertEqual(output.dtype, dtype)
+                torch.testing.assert_close(output.view(torch.uint8), expected)
+
     def test_bfloat16_host_bits_move_to_the_device(self):
         # Host data without torch has no bfloat16 buffer format, so it comes as
         # 16-bit unsigned integers and from_host is told the dtype.

extension/pybindings/test/test_pybindings_no_aten.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import ctypes
 import gc
 import importlib
 import importlib.util
@@ -340,6 +341,16 @@ def test_dlpack_export_keywords(self) -> None:
         self.assertIn(
             '"dltensor_versioned"', repr(result.__dlpack__(max_version=(1, 0)))
         )
+        # Versioned capsules say 1.1, the DLPack version that added the float8
+        # codes they can carry.
+        get_pointer = ctypes.PYFUNCTYPE(
+            ctypes.c_void_p, ctypes.py_object, ctypes.c_char_p
+        )(("PyCapsule_GetPointer", ctypes.pythonapi))
+        capsule = result.__dlpack__(max_version=(1, 1))
+        version = (ctypes.c_uint32 * 2).from_address(
+            get_pointer(capsule, b"dltensor_versioned")
+        )
+        self.assertEqual(tuple(version), (1, 1))
         self.assertIn('"dltensor"', repr(result.__dlpack__()))
         with self.assertRaises(BufferError):
             result.__dlpack__(dl_device=(2, 0), copy=False)
@@ -367,7 +378,8 @@ def test_raw_dlpack_capsule_inputs(self) -> None:
                     output = executor(capsules)[0]
                     np.testing.assert_array_equal(np.from_dlpack(output), [3.0])
 
-        # A producer that knows DLPack 1.0 is asked for a versioned capsule.
+        # A producer that takes max_version is asked for a versioned capsule, up
+        # to DLPack 1.1.
         class VersionedProducer:
             def __init__(self, array: np.ndarray) -> None:
                 self.array = array
@@ -384,7 +396,7 @@ def __dlpack__(self, stream=None, max_version=None):
         np.testing.assert_array_equal(
             np.from_dlpack(method((first, producer))[0]), [3.0]
         )
-        self.assertEqual(producer.requests, [(1, 0)])
+        self.assertEqual(producer.requests, [(1, 1)])
         # A capsule gives its tensor away once.
         used = first.__dlpack__()
         method((used, second))