diff --git a/extension/pybindings/README.md b/extension/pybindings/README.md index 7e1ba466e88..71fc589f0f7 100644 --- a/extension/pybindings/README.md +++ b/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 diff --git a/extension/pybindings/pybindings.cpp b/extension/pybindings/pybindings.cpp index 071c8c10657..4af3d9ef921 100644 --- a/extension/pybindings/pybindings.cpp +++ b/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; } diff --git a/extension/pybindings/pybindings_dlpack.h b/extension/pybindings/pybindings_dlpack.h index 3e73fe78812..638745358a6 100644 --- a/extension/pybindings/pybindings_dlpack.h +++ b/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; diff --git a/extension/pybindings/pybindings_executorch_result.h b/extension/pybindings/pybindings_executorch_result.h index 6a3db374d15..6c5ab0253a8 100644 --- a/extension/pybindings/pybindings_executorch_result.h +++ b/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"); diff --git a/extension/pybindings/test/test_pybindings.py b/extension/pybindings/test/test_pybindings.py index 2a63c4edf53..fab69ed7379 100644 --- a/extension/pybindings/test/test_pybindings.py +++ b/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) diff --git a/extension/pybindings/test/test_pybindings_cuda_dlpack.py b/extension/pybindings/test/test_pybindings_cuda_dlpack.py index bea6d7ccf84..ef9b2f13c3c 100644 --- a/extension/pybindings/test/test_pybindings_cuda_dlpack.py +++ b/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. diff --git a/extension/pybindings/test/test_pybindings_no_aten.py b/extension/pybindings/test/test_pybindings_no_aten.py index a8ee126b4ff..4c51ba41950 100644 --- a/extension/pybindings/test/test_pybindings_no_aten.py +++ b/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))