diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for keel-abi + +## 0.1.0.0 -- 2026-08-18 + +* First release. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Zhe Zhang + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,30 @@ +# keel-abi + +Hand-written `Storable` bindings for the two frozen C ABIs of the +data-science world: + +- the Apache Arrow **C Data Interface** and **C Stream Interface** + (`ArrowSchema` / `ArrowArray` / `ArrowArrayStream`), and +- **DLPack**'s `DLManagedTensorVersioned`, + +in both directions: import (consume a foreign producer's structs) and +export (produce structs for a foreign consumer, e.g. pyarrow or numpy). + +The shipped library contains no C sources and depends only on `base`. +The struct layouts are written out by hand for 64-bit pointers — the +library refuses to build on any other architecture — and a +test-suite-only C file of `_Static_assert(offsetof(...))` checks fails +CI if the hand layouts ever disagree with a real C compiler. + +The managed layer (`Keel.Abi.Arrow`, `Keel.Abi.DLPack`) enforces the +protocols' ownership rules: release callbacks run exactly once, +double-release is a no-op, exceptions never escape into the foreign +caller, and consumption is exception-safe under `bracket`/`finally`. +The raw layer (`*.Raw`) is exported for callers who need the bare +structs. + +Round-trips against real pyarrow (arrays and streams, both directions) +and numpy (DLPack, both directions) are covered by the test suites. + +Part of the keel workspace — see the +[project repository](https://github.com/skymanbp/keel). diff --git a/keel-abi.cabal b/keel-abi.cabal new file mode 100644 --- /dev/null +++ b/keel-abi.cabal @@ -0,0 +1,97 @@ +cabal-version: 3.0 +name: keel-abi +version: 0.1.0.0 +synopsis: Arrow C Data/Stream Interface and DLPack, no cbits +description: + Hand-written @Storable@ bindings for the two frozen C ABIs of the + data-science world: the Apache Arrow C Data \/ C Stream Interface and + DLPack's @DLManagedTensorVersioned@ — import and export directions. + . + The shipped library contains no C sources and needs no C toolchain + beyond what GHC itself bundles: the structs are frozen ABIs, so their + layouts are written out by hand for 64-bit pointers — the library + refuses to build on any other architecture — and a /test-suite-only/ + C file of @_Static_assert(offsetof(...))@ checks fails CI if the + hand layouts ever disagree with a real C compiler. +license: MIT +license-file: LICENSE +author: Zhe Zhang +maintainer: Zhe Zhang +category: Foreign, Data +build-type: Simple +homepage: https://github.com/skymanbp/keel +bug-reports: https://github.com/skymanbp/keel/issues +extra-doc-files: + CHANGELOG.md + README.md +tested-with: GHC ==9.10.3 || ==9.12.4 || ==9.14.1 + +source-repository head + type: git + location: https://github.com/skymanbp/keel + +common warnings + ghc-options: -Wall + default-language: GHC2021 + +library + import: warnings + -- the hand-written struct offsets assume sizeof(void*) == 8; on a + -- 32-bit target they would compile silently and corrupt memory, so + -- refuse to build there instead + if !(arch(x86_64) || arch(aarch64)) + buildable: False + exposed-modules: + Keel.Abi.Arrow + Keel.Abi.Arrow.Raw + Keel.Abi.DLPack + Keel.Abi.DLPack.Raw + hs-source-dirs: src + build-depends: base >=4.20 && <4.23 + +test-suite keel-abi-layout + import: warnings + type: exitcode-stdio-1.0 + main-is: Layout.hs + hs-source-dirs: test + c-sources: test/cbits/layout_gate.c + cc-options: -std=c11 + build-depends: + base, + keel-abi + +test-suite keel-abi-managed + import: warnings + type: exitcode-stdio-1.0 + main-is: Managed.hs + hs-source-dirs: test + ghc-options: -with-rtsopts=-T + build-depends: + base, + keel-abi + +test-suite keel-abi-pyarrow + import: warnings + type: exitcode-stdio-1.0 + main-is: PyArrow.hs + other-modules: PyEmbed + hs-source-dirs: test + ghc-options: -threaded + build-depends: + base, + keel-abi >=0.1 && <0.2, + keel-dyn >=0.1 && <0.2, + process >=1.6 && <1.7 + +test-suite keel-abi-dlpack + import: warnings + type: exitcode-stdio-1.0 + main-is: DLPackNumpy.hs + other-modules: PyEmbed + hs-source-dirs: test + ghc-options: -threaded + build-depends: + base, + keel-abi >=0.1 && <0.2, + keel-dyn >=0.1 && <0.2, + process >=1.6 && <1.7 diff --git a/src/Keel/Abi/Arrow.hs b/src/Keel/Abi/Arrow.hs new file mode 100644 --- /dev/null +++ b/src/Keel/Abi/Arrow.hs @@ -0,0 +1,256 @@ +-- | Managed ownership over the raw Arrow C Data Interface structs. +-- +-- The C Data Interface protocol: the /consumer/ allocates struct +-- storage and passes it to the producer to fill; whoever ends up holding +-- a filled struct must call its release callback exactly once. This +-- module packages both sides: +-- +-- __Consumer (import)__: 'withArrowSchemaImport' \/ 'withArrowArrayImport' +-- allocate zeroed storage, hand it to your action (pass the pointer to +-- the foreign producer, then read the filled struct), and guarantee the +-- release callback runs on exit — including when the action throws. +-- A zeroed, never-filled struct has a null release member and releasing +-- it is a no-op, so allocation and release stay balanced no matter what +-- the producer did. The 'mallocArrowSchemaImport' variants give the same +-- storage as a 'ForeignPtr' whose finalizer releases at GC time, for +-- structs that must outlive a lexical scope. +-- +-- __Producer (export)__: 'exportArrowSchema' \/ 'exportArrowArray' fill a +-- consumer-provided struct from a template record plus a cleanup action. +-- The installed release callback is one process-wide trampoline that +-- runs the cleanup exactly once, frees its 'StablePtr', and nulls the +-- struct's release member per spec — no per-export @\"wrapper\"@ +-- 'FunPtr' is created, so there is nothing to free from inside its own +-- invocation (the classic exporter hazard). +module Keel.Abi.Arrow + ( -- * Consumer side: allocate, let a producer fill, release + withArrowSchemaImport + , withArrowArrayImport + , withArrowArrayStreamImport + , mallocArrowSchemaImport + , mallocArrowArrayImport + , mallocArrowArrayStreamImport + + -- * Producer side: fill a consumer's struct + , exportArrowSchema + , exportArrowArray + , ArrowStreamProducer (..) + , exportArrowArrayStream + ) where + +import Control.Exception (SomeException, bracket, finally, mask_, try) +import Foreign.C.String (CString) +import Foreign.C.Types (CInt (..)) +import Foreign.Concurrent qualified as FC +import Foreign.ForeignPtr (ForeignPtr) +import Foreign.Marshal.Alloc (callocBytes, free) +import Foreign.Ptr (FunPtr, Ptr, freeHaskellFunPtr, nullFunPtr, nullPtr) +import Foreign.StablePtr + ( StablePtr + , castPtrToStablePtr + , castStablePtrToPtr + , deRefStablePtr + , freeStablePtr + , newStablePtr + ) +import Foreign.Storable (Storable, peek, poke, sizeOf) +import System.IO.Unsafe (unsafePerformIO) + +import Keel.Abi.Arrow.Raw + +-- --------------------------------------------------------------------- +-- Consumer side + +withImport :: forall s a. Storable s => (Ptr s -> IO ()) -> (Ptr s -> IO a) -> IO a +withImport releaseIt act = + bracket + (callocBytes (sizeOf (undefined :: s))) + (\p -> releaseIt p `finally` free p) + act + +-- | Zeroed 'ArrowSchema' storage for a producer to fill; released (if +-- filled) and freed on exit, exception-safe. +withArrowSchemaImport :: (Ptr ArrowSchema -> IO a) -> IO a +withArrowSchemaImport = withImport releaseArrowSchema + +-- | Zeroed 'ArrowArray' storage for a producer to fill; released (if +-- filled) and freed on exit, exception-safe. +withArrowArrayImport :: (Ptr ArrowArray -> IO a) -> IO a +withArrowArrayImport = withImport releaseArrowArray + +-- | Zeroed 'ArrowArrayStream' storage for a producer to fill; released +-- (if filled) and freed on exit, exception-safe. +withArrowArrayStreamImport :: (Ptr ArrowArrayStream -> IO a) -> IO a +withArrowArrayStreamImport = withImport releaseArrowArrayStream + +-- masked: an async exception between the allocation and the finalizer +-- registration would leak the block +mallocImport :: forall s. Storable s => (Ptr s -> IO ()) -> IO (ForeignPtr s) +mallocImport releaseIt = mask_ $ do + p <- callocBytes (sizeOf (undefined :: s)) + FC.newForeignPtr p (releaseIt p `finally` free p) + +-- | Like 'withArrowSchemaImport' but GC-managed: the release callback +-- (if the struct was filled) and the storage are reclaimed by the +-- finalizer. Prefer the @with@ variant when lifetime is lexical — +-- finalizers give no promptness guarantee. +mallocArrowSchemaImport :: IO (ForeignPtr ArrowSchema) +mallocArrowSchemaImport = mallocImport releaseArrowSchema + +-- | GC-managed 'ArrowArray' import target; see 'mallocArrowSchemaImport'. +mallocArrowArrayImport :: IO (ForeignPtr ArrowArray) +mallocArrowArrayImport = mallocImport releaseArrowArray + +-- | GC-managed 'ArrowArrayStream' import target; see +-- 'mallocArrowSchemaImport'. +mallocArrowArrayStreamImport :: IO (ForeignPtr ArrowArrayStream) +mallocArrowArrayStreamImport = mallocImport releaseArrowArrayStream + +-- --------------------------------------------------------------------- +-- Producer side + +foreign import ccall "wrapper" + wrapSchemaRelease :: (Ptr ArrowSchema -> IO ()) -> IO (FunPtr (Ptr ArrowSchema -> IO ())) + +foreign import ccall "wrapper" + wrapArrayRelease :: (Ptr ArrowArray -> IO ()) -> IO (FunPtr (Ptr ArrowArray -> IO ())) + +-- One process-wide trampoline per struct kind, never freed. It reads the +-- cleanup action out of private_data, runs it once, and marks the struct +-- released. Consumers may call release from any OS thread; a threaded +-- RTS handles that, a non-threaded one blocks the call until the RTS is +-- idle (standard foreign-export semantics). + +{-# NOINLINE schemaReleaseTrampoline #-} +schemaReleaseTrampoline :: FunPtr (Ptr ArrowSchema -> IO ()) +schemaReleaseTrampoline = unsafePerformIO . wrapSchemaRelease $ \p -> do + s <- peek p + -- null the release member BEFORE the cleanup runs: the consumer then + -- sees the spec-mandated released state no matter what the cleanup does + poke p s { schemaRelease = nullFunPtr, schemaPrivateData = nullPtr } + runCleanup (schemaPrivateData s) + +{-# NOINLINE arrayReleaseTrampoline #-} +arrayReleaseTrampoline :: FunPtr (Ptr ArrowArray -> IO ()) +arrayReleaseTrampoline = unsafePerformIO . wrapArrayRelease $ \p -> do + a <- peek p + poke p a { arrayRelease = nullFunPtr, arrayPrivateData = nullPtr } + runCleanup (arrayPrivateData a) + +-- Runs the carried cleanup and frees its StablePtr. The cleanup runs +-- under 'try' with the exception dropped: this executes inside a +-- callback invoked by foreign code, and a Haskell exception escaping +-- into a C caller is undefined behaviour (in practice it aborts the +-- process) — the C Data Interface expects release callbacks not to fail. +runCleanup :: Ptr () -> IO () +runCleanup pd = do + let sp = castPtrToStablePtr pd :: StablePtr (IO ()) + cleanup <- deRefStablePtr sp + _ <- try @SomeException cleanup + freeStablePtr sp + +-- | Fill @out@ as an exported schema: every field is taken from the +-- template except @release@\/@private_data@, which are overwritten with +-- the trampoline and the cleanup action. The cleanup must free whatever +-- the template's pointers own (format\/name\/metadata strings, children, +-- dictionary) and runs exactly once, from whichever thread the consumer +-- releases on. It must not throw: a thrown exception is caught and +-- discarded — the C caller of the release callback cannot receive it. +exportArrowSchema :: Ptr ArrowSchema -> ArrowSchema -> IO () -> IO () +exportArrowSchema out template cleanup = mask_ $ do + sp <- newStablePtr cleanup + poke out + template + { schemaRelease = schemaReleaseTrampoline + , schemaPrivateData = castStablePtrToPtr sp + } + +-- | Fill @out@ as an exported array; see 'exportArrowSchema'. The +-- cleanup must keep the buffers alive until it runs and then free them +-- (typically: 'Foreign.ForeignPtr.touchForeignPtr' captures, or explicit +-- 'free's of malloc'd buffers plus the buffer-pointer table). +exportArrowArray :: Ptr ArrowArray -> ArrowArray -> IO () -> IO () +exportArrowArray out template cleanup = mask_ $ do + sp <- newStablePtr cleanup + poke out + template + { arrayRelease = arrayReleaseTrampoline + , arrayPrivateData = castStablePtrToPtr sp + } + +-- --------------------------------------------------------------------- +-- Producer side: streams + +-- | A Haskell implementation of an exported 'ArrowArrayStream'. The +-- consumer's struct pointer is dropped from each signature — callbacks +-- are closures, so carry state by capture, not through @private_data@. +data ArrowStreamProducer = ArrowStreamProducer + { producerGetSchema :: Ptr ArrowSchema -> IO CInt + -- ^ Fill the out-schema (e.g. via 'exportArrowSchema'); return 0, + -- or an errno-style code on failure. A thrown exception is caught + -- and reported to the consumer as @EIO@ (5). + , producerGetNext :: Ptr ArrowArray -> IO CInt + -- ^ Fill the out-array with the next chunk, or zero the whole + -- struct (null release member) to signal end-of-stream; return 0, + -- or an errno-style code on failure. A thrown exception is caught + -- and reported to the consumer as @EIO@ (5). + , producerGetLastError :: IO CString + -- ^ Description of the last error, or 'nullPtr'. The string must + -- stay valid until the next stream call. A thrown exception is + -- caught and reported as 'nullPtr'. + , producerCleanup :: IO () + -- ^ Runs exactly once when the consumer releases the stream. Must + -- not throw: a thrown exception is caught and discarded. + } + +foreign import ccall "wrapper" + wrapStreamGetSchema :: StreamGetSchemaFn -> IO (FunPtr StreamGetSchemaFn) + +foreign import ccall "wrapper" + wrapStreamGetNext :: StreamGetNextFn -> IO (FunPtr StreamGetNextFn) + +foreign import ccall "wrapper" + wrapStreamGetLastError :: StreamGetLastErrorFn -> IO (FunPtr StreamGetLastErrorFn) + +{-# NOINLINE streamReleaseTrampoline #-} +streamReleaseTrampoline :: FunPtr (Ptr ArrowArrayStream -> IO ()) +streamReleaseTrampoline = unsafePerformIO . wrapStreamRelease $ \p -> do + s <- peek p + poke p s { streamRelease = nullFunPtr, streamPrivateData = nullPtr } + runCleanup (streamPrivateData s) + +foreign import ccall "wrapper" + wrapStreamRelease :: (Ptr ArrowArrayStream -> IO ()) -> IO (FunPtr (Ptr ArrowArrayStream -> IO ())) + +-- These callbacks execute inside a call from foreign code, where a +-- Haskell exception must not escape (undefined behaviour in the C +-- caller) — it is caught and mapped to the value the protocol can +-- carry: an errno-style code, or a null error string. +guardErrno :: IO CInt -> IO CInt +guardErrno act = either (\(_ :: SomeException) -> 5 {- EIO -}) id <$> try act + +guardLastError :: IO CString -> IO CString +guardLastError act = either (\(_ :: SomeException) -> nullPtr) id <$> try act + +-- | Fill @out@ as an exported stream backed by the producer's Haskell +-- callbacks. The release callback (the shared trampoline again) frees +-- the three callback 'FunPtr's — none of them is the one executing — +-- and then runs 'producerCleanup'. +exportArrowArrayStream :: Ptr ArrowArrayStream -> ArrowStreamProducer -> IO () +exportArrowArrayStream out producer = mask_ $ do + gsF <- wrapStreamGetSchema (\_self o -> guardErrno (producerGetSchema producer o)) + gnF <- wrapStreamGetNext (\_self o -> guardErrno (producerGetNext producer o)) + geF <- wrapStreamGetLastError (\_self -> guardLastError (producerGetLastError producer)) + sp <- newStablePtr $ do + freeHaskellFunPtr gsF + freeHaskellFunPtr gnF + freeHaskellFunPtr geF + producerCleanup producer + poke out + ArrowArrayStream + { streamGetSchema = gsF + , streamGetNext = gnF + , streamGetLastError = geF + , streamRelease = streamReleaseTrampoline + , streamPrivateData = castStablePtrToPtr sp + } diff --git a/src/Keel/Abi/Arrow/Raw.hs b/src/Keel/Abi/Arrow/Raw.hs new file mode 100644 --- /dev/null +++ b/src/Keel/Abi/Arrow/Raw.hs @@ -0,0 +1,344 @@ +-- | The Apache Arrow C Data Interface and C Stream Interface structs, +-- 1:1 and unmanaged. +-- +-- Layouts are written out by hand against the frozen ABI +-- (), so this +-- module needs no headers and the library ships no C sources. The +-- test-suite layout gate (@test\/cbits\/layout_gate.c@) re-derives every +-- offset with a real C compiler and fails the build on any disagreement; +-- 'arrowSchemaLayout' & friends exist for that gate. +-- +-- Everything here is raw: 'Ptr'-level records, manual ownership. The +-- release-callback discipline of the spec applies — consumers must call +-- 'releaseArrowSchema' \/ 'releaseArrowArray' exactly once when done, and +-- a moved (released) struct has a null 'schemaRelease' \/ 'arrayRelease'. +-- Managed import\/export lives one layer up, in @Keel.Abi.Arrow@. +module Keel.Abi.Arrow.Raw + ( -- * ArrowSchema + ArrowSchema (..) + , emptyArrowSchema + , releaseArrowSchema + , arrowFlagDictionaryOrdered + , arrowFlagNullable + , arrowFlagMapKeysSorted + + -- * ArrowArray + , ArrowArray (..) + , emptyArrowArray + , releaseArrowArray + + -- * ArrowArrayStream + , ArrowArrayStream (..) + , StreamGetSchemaFn + , StreamGetNextFn + , StreamGetLastErrorFn + , callStreamGetSchema + , callStreamGetNext + , callStreamGetLastError + , releaseArrowArrayStream + + -- * Layout tables (consumed by the test-suite layout gate) + , arrowSchemaLayout + , arrowArrayLayout + , arrowArrayStreamLayout + ) where + +import Control.Monad (unless) +import Data.Int (Int64) +import Foreign.C.String (CString) +import Foreign.C.Types (CInt (..)) +import Foreign.Ptr (FunPtr, Ptr, nullFunPtr, nullPtr) +import Foreign.Storable (Storable (..)) + +-- --------------------------------------------------------------------- +-- ArrowSchema + +-- | @struct ArrowSchema@. Field semantics (format strings, ownership) +-- are exactly those of the C Data Interface spec. +data ArrowSchema = ArrowSchema + { schemaFormat :: CString + , schemaName :: CString + , schemaMetadata :: CString + , schemaFlags :: Int64 + , schemaNChildren :: Int64 + , schemaChildren :: Ptr (Ptr ArrowSchema) + , schemaDictionary :: Ptr ArrowSchema + , schemaRelease :: FunPtr (Ptr ArrowSchema -> IO ()) + , schemaPrivateData :: Ptr () + } + +-- | The @ARROW_FLAG_*@ bits of 'schemaFlags'. +arrowFlagDictionaryOrdered, arrowFlagNullable, arrowFlagMapKeysSorted :: Int64 +arrowFlagDictionaryOrdered = 1 +arrowFlagNullable = 2 +arrowFlagMapKeysSorted = 4 + +oSchemaFormat, oSchemaName, oSchemaMetadata, oSchemaFlags, oSchemaNChildren, + oSchemaChildren, oSchemaDictionary, oSchemaRelease, oSchemaPrivateData, + szArrowSchema :: Int +oSchemaFormat = 0 +oSchemaName = 8 +oSchemaMetadata = 16 +oSchemaFlags = 24 +oSchemaNChildren = 32 +oSchemaChildren = 40 +oSchemaDictionary = 48 +oSchemaRelease = 56 +oSchemaPrivateData = 64 +szArrowSchema = 72 + +-- | @(sizeof, [(field, offset)])@ in declaration order — compared +-- verbatim against C @offsetof@ by the layout gate. +arrowSchemaLayout :: (Int, [(String, Int)]) +arrowSchemaLayout = + ( szArrowSchema + , [ ("format", oSchemaFormat) + , ("name", oSchemaName) + , ("metadata", oSchemaMetadata) + , ("flags", oSchemaFlags) + , ("n_children", oSchemaNChildren) + , ("children", oSchemaChildren) + , ("dictionary", oSchemaDictionary) + , ("release", oSchemaRelease) + , ("private_data", oSchemaPrivateData) + ] + ) + +instance Storable ArrowSchema where + sizeOf _ = szArrowSchema + alignment _ = 8 + peek p = + ArrowSchema + <$> peekByteOff p oSchemaFormat + <*> peekByteOff p oSchemaName + <*> peekByteOff p oSchemaMetadata + <*> peekByteOff p oSchemaFlags + <*> peekByteOff p oSchemaNChildren + <*> peekByteOff p oSchemaChildren + <*> peekByteOff p oSchemaDictionary + <*> peekByteOff p oSchemaRelease + <*> peekByteOff p oSchemaPrivateData + poke p s = do + pokeByteOff p oSchemaFormat (schemaFormat s) + pokeByteOff p oSchemaName (schemaName s) + pokeByteOff p oSchemaMetadata (schemaMetadata s) + pokeByteOff p oSchemaFlags (schemaFlags s) + pokeByteOff p oSchemaNChildren (schemaNChildren s) + pokeByteOff p oSchemaChildren (schemaChildren s) + pokeByteOff p oSchemaDictionary (schemaDictionary s) + pokeByteOff p oSchemaRelease (schemaRelease s) + pokeByteOff p oSchemaPrivateData (schemaPrivateData s) + +-- | All-null\/zero template — the starting point for building an export, +-- and the spec's representation of a released\/moved struct. +emptyArrowSchema :: ArrowSchema +emptyArrowSchema = + ArrowSchema + { schemaFormat = nullPtr + , schemaName = nullPtr + , schemaMetadata = nullPtr + , schemaFlags = 0 + , schemaNChildren = 0 + , schemaChildren = nullPtr + , schemaDictionary = nullPtr + , schemaRelease = nullFunPtr + , schemaPrivateData = nullPtr + } + +foreign import ccall "dynamic" + callSchemaRelease :: FunPtr (Ptr ArrowSchema -> IO ()) -> Ptr ArrowSchema -> IO () + +-- | Call the struct's release callback unless it is already null (i.e. +-- the struct was moved or released before). Idempotent per the spec: +-- the callback itself nulls the release member. +releaseArrowSchema :: Ptr ArrowSchema -> IO () +releaseArrowSchema p = do + fp <- peekByteOff p oSchemaRelease + unless (fp == nullFunPtr) (callSchemaRelease fp p) + +-- --------------------------------------------------------------------- +-- ArrowArray + +-- | @struct ArrowArray@. Buffer count and meaning follow the format +-- string of the corresponding 'ArrowSchema'. +data ArrowArray = ArrowArray + { arrayLength :: Int64 + , arrayNullCount :: Int64 + , arrayOffset :: Int64 + , arrayNBuffers :: Int64 + , arrayNChildren :: Int64 + , arrayBuffers :: Ptr (Ptr ()) + , arrayChildren :: Ptr (Ptr ArrowArray) + , arrayDictionary :: Ptr ArrowArray + , arrayRelease :: FunPtr (Ptr ArrowArray -> IO ()) + , arrayPrivateData :: Ptr () + } + +oArrayLength, oArrayNullCount, oArrayOffset, oArrayNBuffers, oArrayNChildren, + oArrayBuffers, oArrayChildren, oArrayDictionary, oArrayRelease, + oArrayPrivateData, szArrowArray :: Int +oArrayLength = 0 +oArrayNullCount = 8 +oArrayOffset = 16 +oArrayNBuffers = 24 +oArrayNChildren = 32 +oArrayBuffers = 40 +oArrayChildren = 48 +oArrayDictionary = 56 +oArrayRelease = 64 +oArrayPrivateData = 72 +szArrowArray = 80 + +-- | @(sizeof, [(field, offset)])@ in declaration order. +arrowArrayLayout :: (Int, [(String, Int)]) +arrowArrayLayout = + ( szArrowArray + , [ ("length", oArrayLength) + , ("null_count", oArrayNullCount) + , ("offset", oArrayOffset) + , ("n_buffers", oArrayNBuffers) + , ("n_children", oArrayNChildren) + , ("buffers", oArrayBuffers) + , ("children", oArrayChildren) + , ("dictionary", oArrayDictionary) + , ("release", oArrayRelease) + , ("private_data", oArrayPrivateData) + ] + ) + +instance Storable ArrowArray where + sizeOf _ = szArrowArray + alignment _ = 8 + peek p = + ArrowArray + <$> peekByteOff p oArrayLength + <*> peekByteOff p oArrayNullCount + <*> peekByteOff p oArrayOffset + <*> peekByteOff p oArrayNBuffers + <*> peekByteOff p oArrayNChildren + <*> peekByteOff p oArrayBuffers + <*> peekByteOff p oArrayChildren + <*> peekByteOff p oArrayDictionary + <*> peekByteOff p oArrayRelease + <*> peekByteOff p oArrayPrivateData + poke p a = do + pokeByteOff p oArrayLength (arrayLength a) + pokeByteOff p oArrayNullCount (arrayNullCount a) + pokeByteOff p oArrayOffset (arrayOffset a) + pokeByteOff p oArrayNBuffers (arrayNBuffers a) + pokeByteOff p oArrayNChildren (arrayNChildren a) + pokeByteOff p oArrayBuffers (arrayBuffers a) + pokeByteOff p oArrayChildren (arrayChildren a) + pokeByteOff p oArrayDictionary (arrayDictionary a) + pokeByteOff p oArrayRelease (arrayRelease a) + pokeByteOff p oArrayPrivateData (arrayPrivateData a) + +-- | All-null\/zero template — the starting point for building an export, +-- and the spec's representation of a released\/moved struct (an +-- end-of-stream @get_next@ writes exactly this). +emptyArrowArray :: ArrowArray +emptyArrowArray = + ArrowArray + { arrayLength = 0 + , arrayNullCount = 0 + , arrayOffset = 0 + , arrayNBuffers = 0 + , arrayNChildren = 0 + , arrayBuffers = nullPtr + , arrayChildren = nullPtr + , arrayDictionary = nullPtr + , arrayRelease = nullFunPtr + , arrayPrivateData = nullPtr + } + +foreign import ccall "dynamic" + callArrayRelease :: FunPtr (Ptr ArrowArray -> IO ()) -> Ptr ArrowArray -> IO () + +-- | Call the struct's release callback unless it is already null. +releaseArrowArray :: Ptr ArrowArray -> IO () +releaseArrowArray p = do + fp <- peekByteOff p oArrayRelease + unless (fp == nullFunPtr) (callArrayRelease fp p) + +-- --------------------------------------------------------------------- +-- ArrowArrayStream + +-- | @get_schema@: fills the out-struct, returns 0 or an errno-style code. +type StreamGetSchemaFn = Ptr ArrowArrayStream -> Ptr ArrowSchema -> IO CInt + +-- | @get_next@: fills the out-struct, or marks end-of-stream by leaving +-- its release member null. Returns 0 or an errno-style code. +type StreamGetNextFn = Ptr ArrowArrayStream -> Ptr ArrowArray -> IO CInt + +-- | @get_last_error@: description of the last error, or null. +type StreamGetLastErrorFn = Ptr ArrowArrayStream -> IO CString + +-- | @struct ArrowArrayStream@ — all members are function pointers plus +-- @private_data@; drive it with the @call*@ wrappers below. +data ArrowArrayStream = ArrowArrayStream + { streamGetSchema :: FunPtr StreamGetSchemaFn + , streamGetNext :: FunPtr StreamGetNextFn + , streamGetLastError :: FunPtr StreamGetLastErrorFn + , streamRelease :: FunPtr (Ptr ArrowArrayStream -> IO ()) + , streamPrivateData :: Ptr () + } + +oStreamGetSchema, oStreamGetNext, oStreamGetLastError, oStreamRelease, + oStreamPrivateData, szArrowArrayStream :: Int +oStreamGetSchema = 0 +oStreamGetNext = 8 +oStreamGetLastError = 16 +oStreamRelease = 24 +oStreamPrivateData = 32 +szArrowArrayStream = 40 + +-- | @(sizeof, [(field, offset)])@ in declaration order. +arrowArrayStreamLayout :: (Int, [(String, Int)]) +arrowArrayStreamLayout = + ( szArrowArrayStream + , [ ("get_schema", oStreamGetSchema) + , ("get_next", oStreamGetNext) + , ("get_last_error", oStreamGetLastError) + , ("release", oStreamRelease) + , ("private_data", oStreamPrivateData) + ] + ) + +instance Storable ArrowArrayStream where + sizeOf _ = szArrowArrayStream + alignment _ = 8 + peek p = + ArrowArrayStream + <$> peekByteOff p oStreamGetSchema + <*> peekByteOff p oStreamGetNext + <*> peekByteOff p oStreamGetLastError + <*> peekByteOff p oStreamRelease + <*> peekByteOff p oStreamPrivateData + poke p s = do + pokeByteOff p oStreamGetSchema (streamGetSchema s) + pokeByteOff p oStreamGetNext (streamGetNext s) + pokeByteOff p oStreamGetLastError (streamGetLastError s) + pokeByteOff p oStreamRelease (streamRelease s) + pokeByteOff p oStreamPrivateData (streamPrivateData s) + +-- | Invoke a stream's @get_schema@ member (pass the struct's own +-- pointer as the first argument). +foreign import ccall "dynamic" + callStreamGetSchema :: FunPtr StreamGetSchemaFn -> StreamGetSchemaFn + +-- | Invoke a stream's @get_next@ member. +foreign import ccall "dynamic" + callStreamGetNext :: FunPtr StreamGetNextFn -> StreamGetNextFn + +-- | Invoke a stream's @get_last_error@ member. +foreign import ccall "dynamic" + callStreamGetLastError :: FunPtr StreamGetLastErrorFn -> StreamGetLastErrorFn + +foreign import ccall "dynamic" + callStreamRelease :: FunPtr (Ptr ArrowArrayStream -> IO ()) -> Ptr ArrowArrayStream -> IO () + +-- | Call the stream's release callback unless it is already null. +releaseArrowArrayStream :: Ptr ArrowArrayStream -> IO () +releaseArrowArrayStream p = do + fp <- peekByteOff p oStreamRelease + unless (fp == nullFunPtr) (callStreamRelease fp p) diff --git a/src/Keel/Abi/DLPack.hs b/src/Keel/Abi/DLPack.hs new file mode 100644 --- /dev/null +++ b/src/Keel/Abi/DLPack.hs @@ -0,0 +1,124 @@ +-- | Managed ownership over DLPack's versioned exchange tensor. +-- +-- The DLPack contract: the producer hands over a +-- 'DLManagedTensorVersioned' (usually inside a @dltensor_versioned@ +-- PyCapsule); the consumer checks the major version, uses the data, and +-- calls the deleter exactly once. This module packages both sides the +-- same way "Keel.Abi.Arrow" does for Arrow structs: consumption under +-- 'Control.Exception.finally', production through a process-wide deleter +-- trampoline plus a 'Foreign.StablePtr.StablePtr'-carried cleanup. +module Keel.Abi.DLPack + ( AbiError (..) + + -- * Consumer side + , consumeManagedTensor + + -- * Producer side + , newManagedTensor + ) where + +import Control.Exception (Exception, SomeException, finally, mask_, onException, throwIO, try) +import Data.Int (Int64) +import Data.Word (Word32, Word64) +import Foreign.Marshal.Alloc (free, mallocBytes) +import Foreign.Ptr (FunPtr, Ptr, nullPtr, plusPtr) +import Foreign.StablePtr + ( StablePtr + , castPtrToStablePtr + , castStablePtrToPtr + , deRefStablePtr + , freeStablePtr + , newStablePtr + ) +import Foreign.Storable (peek, poke, pokeElemOff, sizeOf) +import System.IO.Unsafe (unsafePerformIO) + +import Keel.Abi.DLPack.Raw + +-- | Failure modes of the exchange protocol itself. +newtype AbiError = DLPackMajorUnsupported Word32 + -- ^ The producer filled the struct under a DLPack major version + -- newer than these bindings ('dlpackMajorVersion') understand. + deriving (Eq, Show) + +instance Exception AbiError + +-- | Take ownership of a produced tensor: run the action on the peeked +-- struct, then invoke the deleter — also when the action throws. If the +-- producer's major version is newer than 'dlpackMajorVersion', the +-- tensor is deleted unused (the version\/deleter prologue is stable +-- across majors by design) and 'DLPackMajorUnsupported' is thrown. +consumeManagedTensor + :: Ptr DLManagedTensorVersioned + -> (DLManagedTensorVersioned -> IO a) + -> IO a +consumeManagedTensor p act = do + m <- peek p + let major = dlverMajor (mtvVersion m) + if major > dlpackMajorVersion + then do + callTensorDeleter p + throwIO (DLPackMajorUnsupported major) + else act m `finally` callTensorDeleter p + +foreign import ccall "wrapper" + wrapDeleter + :: (Ptr DLManagedTensorVersioned -> IO ()) + -> IO (FunPtr (Ptr DLManagedTensorVersioned -> IO ())) + +-- One process-wide deleter, never freed: runs the cleanup carried in +-- manager_ctx, then frees the struct block itself (the deleter deletes +-- @self@ per spec). The cleanup runs under 'try' with the exception +-- dropped: the deleter is invoked by foreign code, and a Haskell +-- exception escaping into a C caller is undefined behaviour. +{-# NOINLINE deleterTrampoline #-} +deleterTrampoline :: FunPtr (Ptr DLManagedTensorVersioned -> IO ()) +deleterTrampoline = unsafePerformIO . wrapDeleter $ \p -> do + m <- peek p + let sp = castPtrToStablePtr (mtvManagerCtx m) :: StablePtr (IO ()) + cleanup <- deRefStablePtr sp + _ <- try @SomeException cleanup + freeStablePtr sp + free p + +-- | Allocate and fill a 'DLManagedTensorVersioned' for handoff to a +-- consumer. The shape array lives in the same allocation as the struct; +-- the tensor is CPU-device, compact row-major (null strides), zero byte +-- offset, version 'dlpackMajorVersion'.'dlpackMinorVersion'. The cleanup +-- runs exactly once — from the consumer's deleter call, on whatever +-- thread that happens — and must free\/unpin the data buffer; the +-- struct block frees itself afterwards. It must not throw: a thrown +-- exception is caught and discarded (the C caller cannot receive it). +newManagedTensor + :: DLDataType + -> [Int64] -- ^ shape (row-major, compact) + -> Ptr () -- ^ data + -> Word64 -- ^ flags ('dlpackFlagReadOnly' \/ 'dlpackFlagIsCopied' \/ 0) + -> IO () -- ^ cleanup, owns the data buffer + -> IO (Ptr DLManagedTensorVersioned) +newManagedTensor dt shape dat flags cleanup = mask_ $ do + let ndim = length shape + structSz = sizeOf (undefined :: DLManagedTensorVersioned) + p <- mallocBytes (structSz + ndim * 8) + flip onException (free p) $ do + let shapeP = p `plusPtr` structSz + mapM_ (uncurry (pokeElemOff shapeP)) (zip [0 ..] shape) + sp <- newStablePtr cleanup + poke p + DLManagedTensorVersioned + { mtvVersion = DLPackVersion dlpackMajorVersion dlpackMinorVersion + , mtvManagerCtx = castStablePtrToPtr sp + , mtvDeleter = deleterTrampoline + , mtvFlags = flags + , mtvTensor = + DLTensor + { dltData = dat + , dltDevice = DLDevice kDLCPU 0 + , dltNDim = fromIntegral ndim + , dltDType = dt + , dltShape = shapeP + , dltStrides = nullPtr + , dltByteOffset = 0 + } + } + pure p diff --git a/src/Keel/Abi/DLPack/Raw.hs b/src/Keel/Abi/DLPack/Raw.hs new file mode 100644 --- /dev/null +++ b/src/Keel/Abi/DLPack/Raw.hs @@ -0,0 +1,315 @@ +-- | The DLPack tensor-exchange structs, 1:1 and unmanaged. +-- +-- Targets DLPack v1.x: the exchanged object is 'DLManagedTensorVersioned' +-- (the pre-1.0 unversioned @DLManagedTensor@ is deliberately not bound). +-- Layouts are hand-written against @dlpack.h@ +-- () and verified by the test-suite +-- layout gate; the shipped library has no C sources. +-- +-- Ownership follows the DLPack contract: the consumer of a +-- 'DLManagedTensorVersioned' calls 'callTensorDeleter' exactly once when +-- done; the producer keeps everything the tensor points at alive until +-- then via 'mtvManagerCtx'. +module Keel.Abi.DLPack.Raw + ( -- * Version + DLPackVersion (..) + , dlpackMajorVersion + , dlpackMinorVersion + + -- * Device + , DLDevice (..) + , kDLCPU + , kDLCUDA + , kDLCUDAHost + , kDLOpenCL + , kDLVulkan + , kDLMetal + , kDLVPI + , kDLROCM + + -- * Data type + , DLDataType (..) + , kDLInt + , kDLUInt + , kDLFloat + , kDLOpaqueHandle + , kDLBfloat + , kDLComplex + , kDLBool + + -- * Tensor + , DLTensor (..) + , DLManagedTensorVersioned (..) + , dlpackFlagReadOnly + , dlpackFlagIsCopied + , callTensorDeleter + + -- * Layout tables (consumed by the test-suite layout gate) + , dlPackVersionLayout + , dlDeviceLayout + , dlDataTypeLayout + , dlTensorLayout + , dlManagedTensorVersionedLayout + ) where + +import Control.Monad (unless) +import Data.Int (Int32, Int64) +import Data.Word (Word16, Word32, Word64, Word8) +import Foreign.Ptr (FunPtr, Ptr, nullFunPtr) +import Foreign.Storable (Storable (..)) + +-- --------------------------------------------------------------------- +-- DLPackVersion + +-- | @DLPackVersion@ — the ABI version the producer filled the struct +-- with. Consumers must check 'dlverMajor' against 'dlpackMajorVersion'. +data DLPackVersion = DLPackVersion + { dlverMajor :: Word32 + , dlverMinor :: Word32 + } + deriving (Eq, Show) + +-- | The DLPack version these bindings target (v1.1). +dlpackMajorVersion, dlpackMinorVersion :: Word32 +dlpackMajorVersion = 1 +dlpackMinorVersion = 1 + +oVerMajor, oVerMinor, szDLPackVersion :: Int +oVerMajor = 0 +oVerMinor = 4 +szDLPackVersion = 8 + +-- | @(sizeof, [(field, offset)])@ in declaration order. +dlPackVersionLayout :: (Int, [(String, Int)]) +dlPackVersionLayout = + (szDLPackVersion, [("major", oVerMajor), ("minor", oVerMinor)]) + +instance Storable DLPackVersion where + sizeOf _ = szDLPackVersion + alignment _ = 4 + peek p = DLPackVersion <$> peekByteOff p oVerMajor <*> peekByteOff p oVerMinor + poke p v = do + pokeByteOff p oVerMajor (dlverMajor v) + pokeByteOff p oVerMinor (dlverMinor v) + +-- --------------------------------------------------------------------- +-- DLDevice + +-- | @DLDevice@. 'dldevType' is the @DLDeviceType@ enum (a C @int@). +data DLDevice = DLDevice + { dldevType :: Int32 + , dldevId :: Int32 + } + deriving (Eq, Show) + +-- | @DLDeviceType@ values (the ones keel can ever produce or consume; +-- the full enum is larger but frozen upstream). +kDLCPU, kDLCUDA, kDLCUDAHost, kDLOpenCL, kDLVulkan, kDLMetal, kDLVPI, kDLROCM :: Int32 +kDLCPU = 1 +kDLCUDA = 2 +kDLCUDAHost = 3 +kDLOpenCL = 4 +kDLVulkan = 7 +kDLMetal = 8 +kDLVPI = 9 +kDLROCM = 10 + +oDevType, oDevId, szDLDevice :: Int +oDevType = 0 +oDevId = 4 +szDLDevice = 8 + +-- | @(sizeof, [(field, offset)])@ in declaration order. +dlDeviceLayout :: (Int, [(String, Int)]) +dlDeviceLayout = (szDLDevice, [("device_type", oDevType), ("device_id", oDevId)]) + +instance Storable DLDevice where + sizeOf _ = szDLDevice + alignment _ = 4 + peek p = DLDevice <$> peekByteOff p oDevType <*> peekByteOff p oDevId + poke p d = do + pokeByteOff p oDevType (dldevType d) + pokeByteOff p oDevId (dldevId d) + +-- --------------------------------------------------------------------- +-- DLDataType + +-- | @DLDataType@: type code, bit width, vector lanes (1 for scalars). +-- Example: @DLDataType kDLFloat 64 1@ is a C @double@. +data DLDataType = DLDataType + { dldtCode :: Word8 + , dldtBits :: Word8 + , dldtLanes :: Word16 + } + deriving (Eq, Show) + +-- | @DLDataTypeCode@ values. +kDLInt, kDLUInt, kDLFloat, kDLOpaqueHandle, kDLBfloat, kDLComplex, kDLBool :: Word8 +kDLInt = 0 +kDLUInt = 1 +kDLFloat = 2 +kDLOpaqueHandle = 3 +kDLBfloat = 4 +kDLComplex = 5 +kDLBool = 6 + +oDtCode, oDtBits, oDtLanes, szDLDataType :: Int +oDtCode = 0 +oDtBits = 1 +oDtLanes = 2 +szDLDataType = 4 + +-- | @(sizeof, [(field, offset)])@ in declaration order. +dlDataTypeLayout :: (Int, [(String, Int)]) +dlDataTypeLayout = + (szDLDataType, [("code", oDtCode), ("bits", oDtBits), ("lanes", oDtLanes)]) + +instance Storable DLDataType where + sizeOf _ = szDLDataType + alignment _ = 2 + peek p = + DLDataType + <$> peekByteOff p oDtCode + <*> peekByteOff p oDtBits + <*> peekByteOff p oDtLanes + poke p t = do + pokeByteOff p oDtCode (dldtCode t) + pokeByteOff p oDtBits (dldtBits t) + pokeByteOff p oDtLanes (dldtLanes t) + +-- --------------------------------------------------------------------- +-- DLTensor + +-- | @DLTensor@ — a borrowed view; owns nothing. 'dltShape' (and +-- 'dltStrides' when non-null) point at @ndim@ @int64_t@s owned by the +-- producer. Null 'dltStrides' means compact row-major. Strides are in +-- /elements/, not bytes. +data DLTensor = DLTensor + { dltData :: Ptr () + , dltDevice :: DLDevice + , dltNDim :: Int32 + , dltDType :: DLDataType + , dltShape :: Ptr Int64 + , dltStrides :: Ptr Int64 + , dltByteOffset :: Word64 + } + +oTData, oTDevice, oTNDim, oTDType, oTShape, oTStrides, oTByteOffset, + szDLTensor :: Int +oTData = 0 +oTDevice = 8 +oTNDim = 16 +oTDType = 20 +oTShape = 24 +oTStrides = 32 +oTByteOffset = 40 +szDLTensor = 48 + +-- | @(sizeof, [(field, offset)])@ in declaration order. +dlTensorLayout :: (Int, [(String, Int)]) +dlTensorLayout = + ( szDLTensor + , [ ("data", oTData) + , ("device", oTDevice) + , ("ndim", oTNDim) + , ("dtype", oTDType) + , ("shape", oTShape) + , ("strides", oTStrides) + , ("byte_offset", oTByteOffset) + ] + ) + +instance Storable DLTensor where + sizeOf _ = szDLTensor + alignment _ = 8 + peek p = + DLTensor + <$> peekByteOff p oTData + <*> peekByteOff p oTDevice + <*> peekByteOff p oTNDim + <*> peekByteOff p oTDType + <*> peekByteOff p oTShape + <*> peekByteOff p oTStrides + <*> peekByteOff p oTByteOffset + poke p t = do + pokeByteOff p oTData (dltData t) + pokeByteOff p oTDevice (dltDevice t) + pokeByteOff p oTNDim (dltNDim t) + pokeByteOff p oTDType (dltDType t) + pokeByteOff p oTShape (dltShape t) + pokeByteOff p oTStrides (dltStrides t) + pokeByteOff p oTByteOffset (dltByteOffset t) + +-- --------------------------------------------------------------------- +-- DLManagedTensorVersioned + +-- | @DLManagedTensorVersioned@ — the owned exchange object of DLPack +-- v1.x. The consumer calls 'callTensorDeleter' exactly once when done. +data DLManagedTensorVersioned = DLManagedTensorVersioned + { mtvVersion :: DLPackVersion + , mtvManagerCtx :: Ptr () + , mtvDeleter :: FunPtr (Ptr DLManagedTensorVersioned -> IO ()) + , mtvFlags :: Word64 + , mtvTensor :: DLTensor + } + +-- | @DLPACK_FLAG_BITMASK_READ_ONLY@: the consumer must not write through +-- 'dltData'. +dlpackFlagReadOnly :: Word64 +dlpackFlagReadOnly = 1 + +-- | @DLPACK_FLAG_BITMASK_IS_COPIED@: the tensor is a copy, not a view. +dlpackFlagIsCopied :: Word64 +dlpackFlagIsCopied = 2 + +oMtvVersion, oMtvManagerCtx, oMtvDeleter, oMtvFlags, oMtvTensor, + szDLManagedTensorVersioned :: Int +oMtvVersion = 0 +oMtvManagerCtx = 8 +oMtvDeleter = 16 +oMtvFlags = 24 +oMtvTensor = 32 +szDLManagedTensorVersioned = 80 + +-- | @(sizeof, [(field, offset)])@ in declaration order. +dlManagedTensorVersionedLayout :: (Int, [(String, Int)]) +dlManagedTensorVersionedLayout = + ( szDLManagedTensorVersioned + , [ ("version", oMtvVersion) + , ("manager_ctx", oMtvManagerCtx) + , ("deleter", oMtvDeleter) + , ("flags", oMtvFlags) + , ("dl_tensor", oMtvTensor) + ] + ) + +instance Storable DLManagedTensorVersioned where + sizeOf _ = szDLManagedTensorVersioned + alignment _ = 8 + peek p = + DLManagedTensorVersioned + <$> peekByteOff p oMtvVersion + <*> peekByteOff p oMtvManagerCtx + <*> peekByteOff p oMtvDeleter + <*> peekByteOff p oMtvFlags + <*> peekByteOff p oMtvTensor + poke p m = do + pokeByteOff p oMtvVersion (mtvVersion m) + pokeByteOff p oMtvManagerCtx (mtvManagerCtx m) + pokeByteOff p oMtvDeleter (mtvDeleter m) + pokeByteOff p oMtvFlags (mtvFlags m) + pokeByteOff p oMtvTensor (mtvTensor m) + +foreign import ccall "dynamic" + callDeleter + :: FunPtr (Ptr DLManagedTensorVersioned -> IO ()) + -> Ptr DLManagedTensorVersioned + -> IO () + +-- | Invoke the tensor's deleter — the consumer-side "I am done" call. +-- A null deleter (legal per spec: the producer has nothing to free) is +-- a no-op. +callTensorDeleter :: Ptr DLManagedTensorVersioned -> IO () +callTensorDeleter p = do + fp <- peekByteOff p oMtvDeleter + unless (fp == nullFunPtr) (callDeleter fp p) diff --git a/test/DLPackNumpy.hs b/test/DLPackNumpy.hs new file mode 100644 --- /dev/null +++ b/test/DLPackNumpy.hs @@ -0,0 +1,113 @@ +-- | Conformance test: exchange DLPack v1.x versioned tensors with numpy, +-- in-process, in both directions. +-- +-- CPython is loaded at run time through keel-dyn (see "PyEmbed"). +-- +-- * import direction: numpy exports @arange(6, float64)@ as a +-- @dltensor_versioned@ PyCapsule; the script writes the raw pointer +-- into a Haskell-allocated slot and renames the capsule to +-- @used_dltensor_versioned@ (ownership transferred to us); Haskell +-- verifies version\/dtype\/shape\/values under 'consumeManagedTensor', +-- whose exit calls numpy's deleter; +-- * export direction: Haskell builds a 2x3 float64 tensor with +-- 'newManagedTensor'; a minimal producer class hands the capsule to +-- @np.from_dlpack@; Python verifies shape and values and drops the +-- array — numpy's consumer must call our deleter trampoline, which +-- runs our cleanup. +-- +-- Requires numpy >= 2.1 (DLPack 1.0 protocol support); older or absent +-- numpy SKIPs unless @KEEL_ABI_REQUIRE_NUMPY@ is set (CI sets it). +module Main (main) where + +import Data.IORef (newIORef, readIORef, writeIORef) +import Data.Int (Int64) +import Foreign.Marshal.Alloc (callocBytes, free, mallocBytes) +import Foreign.Ptr (Ptr, castPtr, nullPtr) +import Foreign.Storable (peek, peekElemOff, pokeElemOff) + +import Keel.Abi.DLPack +import Keel.Abi.DLPack.Raw +import PyEmbed + +main :: IO () +main = withEmbeddedPython "numpy" "KEEL_ABI_REQUIRE_NUMPY" $ \runScript -> do + importDirection runScript + exportDirection runScript + putStrLn "keel-abi-dlpack: tensors round-tripped both directions against numpy" + +-- --------------------------------------------------------------------- +-- Direction 1: numpy -> Haskell + +importDirection :: RunScript -> IO () +importDirection runScript = do + slot <- callocBytes 8 :: IO (Ptr (Ptr DLManagedTensorVersioned)) + runScript "dlpack-import" $ + "import numpy as np, ctypes\n\ + \a = np.arange(6, dtype=np.float64)\n\ + \cap = a.__dlpack__(max_version=(1, 1))\n\ + \ctypes.pythonapi.PyCapsule_GetPointer.restype = ctypes.c_void_p\n\ + \ctypes.pythonapi.PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]\n\ + \p = ctypes.pythonapi.PyCapsule_GetPointer(cap, b'dltensor_versioned')\n\ + \ctypes.pythonapi.PyCapsule_SetName.argtypes = [ctypes.py_object, ctypes.c_char_p]\n\ + \ctypes.pythonapi.PyCapsule_SetName(cap, b'used_dltensor_versioned')\n\ + \ctypes.c_void_p.from_address(" <> addr slot <> ").value = p\n" + mtp <- peek slot + expect (mtp /= nullPtr) "python wrote no tensor pointer" + + vals <- consumeManagedTensor mtp $ \m -> do + expect (dlverMajor (mtvVersion m) == dlpackMajorVersion) + ("producer major version: " <> show (dlverMajor (mtvVersion m))) + let t = mtvTensor m + expect (dldtCode (dltDType t) == kDLFloat) "dtype code not float" + expect (dldtBits (dltDType t) == 64) "dtype bits not 64" + expect (dldtLanes (dltDType t) == 1) "dtype lanes not 1" + expect (dldevType (dltDevice t) == kDLCPU) "device not CPU" + expect (dltNDim t == 1) ("ndim: " <> show (dltNDim t)) + sh <- peekElemOff (dltShape t) 0 + expect (sh == 6) ("shape[0]: " <> show sh) + -- strides: null means compact; a non-null [1] is equivalent for 1-D + strideOk <- + if dltStrides t == nullPtr + then pure True + else (== (1 :: Int64)) <$> peekElemOff (dltStrides t) 0 + expect strideOk "strides neither null nor [1]" + expect (dltByteOffset t == 0) "byte_offset nonzero" + mapM (peekElemOff (castPtr (dltData t) :: Ptr Double)) [0 .. 5] + expect (vals == [0, 1, 2, 3, 4, 5]) ("values: " <> show vals) + free slot + +-- --------------------------------------------------------------------- +-- Direction 2: Haskell -> numpy + +exportDirection :: RunScript -> IO () +exportDirection runScript = do + cleanupRan <- newIORef False + buf <- mallocBytes (6 * 8) :: IO (Ptr Double) + mapM_ (uncurry (pokeElemOff buf)) (zip [0 ..] [1 .. 6]) + mt <- + newManagedTensor + (DLDataType kDLFloat 64 1) + [2, 3] + (castPtr buf) + 0 + (free buf >> writeIORef cleanupRan True) + + runScript "dlpack-export" $ + "import numpy as np, ctypes, gc\n\ + \ctypes.pythonapi.PyCapsule_New.restype = ctypes.py_object\n\ + \ctypes.pythonapi.PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]\n\ + \class _KeelTensor:\n\ + \ def __init__(self, p):\n\ + \ self._p = p\n\ + \ def __dlpack__(self, max_version=None, dl_device=None, copy=None):\n\ + \ return ctypes.pythonapi.PyCapsule_New(ctypes.c_void_p(self._p), b'dltensor_versioned', None)\n\ + \ def __dlpack_device__(self):\n\ + \ return (1, 0)\n\ + \b = np.from_dlpack(_KeelTensor(" <> addr mt <> "))\n\ + \if b.shape != (2, 3) or b.tolist() != [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]:\n\ + \ raise RuntimeError('dlpack import mismatch: %r' % (b.tolist(),))\n\ + \del b\n\ + \gc.collect()\n" + + done <- readIORef cleanupRan + expect done "numpy never called our deleter (cleanup did not run)" diff --git a/test/Layout.hs b/test/Layout.hs new file mode 100644 --- /dev/null +++ b/test/Layout.hs @@ -0,0 +1,80 @@ +-- | Layout gate: compare every offset\/size in the Haskell layout tables +-- against a real C compiler's @offsetof@\/@sizeof@ (computed in +-- @cbits\/layout_gate.c@, which also carries compile-time +-- @_Static_assert@s for the same numbers). +module Main (main) where + +import Control.Monad (forM, unless) +import Foreign.C.Types (CSize (..)) +import Foreign.Marshal.Array (allocaArray, peekArray) +import Foreign.Ptr (Ptr) +import System.Exit (exitFailure) + +import Keel.Abi.Arrow.Raw +import Keel.Abi.DLPack.Raw + +foreign import ccall unsafe "keel_layout_ArrowSchema" + c_layout_ArrowSchema :: Ptr CSize -> IO CSize + +foreign import ccall unsafe "keel_layout_ArrowArray" + c_layout_ArrowArray :: Ptr CSize -> IO CSize + +foreign import ccall unsafe "keel_layout_ArrowArrayStream" + c_layout_ArrowArrayStream :: Ptr CSize -> IO CSize + +foreign import ccall unsafe "keel_layout_DLPackVersion" + c_layout_DLPackVersion :: Ptr CSize -> IO CSize + +foreign import ccall unsafe "keel_layout_DLDevice" + c_layout_DLDevice :: Ptr CSize -> IO CSize + +foreign import ccall unsafe "keel_layout_DLDataType" + c_layout_DLDataType :: Ptr CSize -> IO CSize + +foreign import ccall unsafe "keel_layout_DLTensor" + c_layout_DLTensor :: Ptr CSize -> IO CSize + +foreign import ccall unsafe "keel_layout_DLManagedTensorVersioned" + c_layout_DLManagedTensorVersioned :: Ptr CSize -> IO CSize + +-- | Returns the list of mismatch descriptions (empty = pass). +checkStruct + :: String + -> (Int, [(String, Int)]) + -> (Ptr CSize -> IO CSize) + -> IO [String] +checkStruct structName (hsSize, fields) probe = + allocaArray (length fields) $ \out -> do + cSize <- probe out + cOffs <- peekArray (length fields) out + let sizeErrs = + [ structName <> ": sizeof C=" <> show cSize <> " hs=" <> show hsSize + | fromIntegral cSize /= hsSize + ] + fieldErrs = + [ structName <> "." <> fname + <> ": offsetof C=" <> show cOff <> " hs=" <> show hsOff + | ((fname, hsOff), cOff) <- zip fields cOffs + , fromIntegral cOff /= hsOff + ] + pure (sizeErrs <> fieldErrs) + +main :: IO () +main = do + errs <- fmap concat . forM checks $ \(nm, layout, probe) -> + checkStruct nm layout probe + unless (null errs) $ do + mapM_ putStrLn errs + exitFailure + putStrLn ("keel-abi: layout gate passed (" <> show (length checks) <> " structs)") + where + checks = + [ ("ArrowSchema", arrowSchemaLayout, c_layout_ArrowSchema) + , ("ArrowArray", arrowArrayLayout, c_layout_ArrowArray) + , ("ArrowArrayStream", arrowArrayStreamLayout, c_layout_ArrowArrayStream) + , ("DLPackVersion", dlPackVersionLayout, c_layout_DLPackVersion) + , ("DLDevice", dlDeviceLayout, c_layout_DLDevice) + , ("DLDataType", dlDataTypeLayout, c_layout_DLDataType) + , ("DLTensor", dlTensorLayout, c_layout_DLTensor) + , ("DLManagedTensorVersioned", dlManagedTensorVersionedLayout, c_layout_DLManagedTensorVersioned) + ] diff --git a/test/Managed.hs b/test/Managed.hs new file mode 100644 --- /dev/null +++ b/test/Managed.hs @@ -0,0 +1,186 @@ +-- | Unit tests for the managed layer ("Keel.Abi.Arrow") that need no +-- foreign producer: Haskell plays both exporter and consumer. +module Main (main) where + +import Control.Monad (forM_, unless) +import Data.IORef (modifyIORef', newIORef, readIORef) +import Foreign.ForeignPtr (finalizeForeignPtr, withForeignPtr) +import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats, getRTSStatsEnabled) +import System.Mem (performMajorGC) +import Foreign.Marshal.Alloc (callocBytes, free, mallocBytes) +import Foreign.Ptr (Ptr, castPtr, nullFunPtr, nullPtr) +import Foreign.Storable (peek, peekElemOff, poke, pokeElemOff, sizeOf) + +import Keel.Abi.Arrow +import Keel.Abi.Arrow.Raw +import Keel.Abi.DLPack +import Keel.Abi.DLPack.Raw + +expect :: Bool -> String -> IO () +expect ok msg = unless ok (fail msg) + +main :: IO () +main = do + -- 1. releasing a never-filled import target is a no-op (with* variant) + r <- withArrowArrayImport $ \p -> do + a <- peek p + expect (arrayRelease a == nullFunPtr) "import target not zeroed" + pure (42 :: Int) + expect (r == 42) "with-import did not return the action result" + + -- 2. export -> consumer-release lifecycle, array side + cnt <- newIORef (0 :: Int) + p <- callocBytes (sizeOf (undefined :: ArrowArray)) + exportArrowArray p emptyArrowArray { arrayLength = 3 } (modifyIORef' cnt (+ 1)) + a0 <- peek p + expect (arrayLength a0 == 3) "template field lost on export" + expect (arrayRelease a0 /= nullFunPtr) "export installed no release" + releaseArrowArray p + releaseArrowArray p -- double release must be a no-op + n <- readIORef cnt + expect (n == 1) ("cleanup ran " <> show n <> " times, want 1") + a1 <- peek p + expect (arrayRelease a1 == nullFunPtr) "release not nulled" + expect (arrayPrivateData a1 == nullPtr) "private_data not nulled" + free p + + -- 3. same lifecycle, schema side + scnt <- newIORef (0 :: Int) + sp <- callocBytes (sizeOf (undefined :: ArrowSchema)) + exportArrowSchema sp emptyArrowSchema (modifyIORef' scnt (+ 1)) + releaseArrowSchema sp + releaseArrowSchema sp + sn <- readIORef scnt + expect (sn == 1) ("schema cleanup ran " <> show sn <> " times, want 1") + free sp + + -- 4. with-import releases an exported struct on scope exit + wcnt <- newIORef (0 :: Int) + withArrowArrayImport $ \wp -> + exportArrowArray wp emptyArrowArray (modifyIORef' wcnt (+ 1)) + wn <- readIORef wcnt + expect (wn == 1) ("bracket cleanup ran " <> show wn <> " times, want 1") + + -- 5. malloc variant: finalizeForeignPtr runs the release deterministically + fcnt <- newIORef (0 :: Int) + fp <- mallocArrowArrayImport + withForeignPtr fp $ \rawp -> + exportArrowArray rawp emptyArrowArray (modifyIORef' fcnt (+ 1)) + finalizeForeignPtr fp + fn <- readIORef fcnt + expect (fn == 1) ("finalizer cleanup ran " <> show fn <> " times, want 1") + + -- 6. exported stream: drive it as a consumer, entirely in-process + streamCnt <- newIORef (0 :: Int) + gnCalls <- newIORef (0 :: Int) + stp <- callocBytes (sizeOf (undefined :: ArrowArrayStream)) + exportArrowArrayStream stp + ArrowStreamProducer + { producerGetSchema = \o -> do + exportArrowSchema o emptyArrowSchema (pure ()) + pure 0 + , producerGetNext = \o -> do + i <- readIORef gnCalls + modifyIORef' gnCalls (+ 1) + if i < 2 + then exportArrowArray o emptyArrowArray { arrayLength = fromIntegral (i + 1) } (pure ()) + else poke' o -- end-of-stream: zeroed struct, null release + pure 0 + , producerGetLastError = pure nullPtr + , producerCleanup = modifyIORef' streamCnt (+ 1) + } + st <- peek stp + -- schema call + withArrowSchemaImport $ \so -> do + rc <- callStreamGetSchema (streamGetSchema st) stp so + expect (rc == 0) "get_schema returned nonzero" + -- two batches then end-of-stream + lens <- mapM (const (nextLen st stp)) [1 :: Int, 2, 3] + expect (lens == [Just 1, Just 2, Nothing]) ("stream batches: " <> show lens) + releaseArrowArrayStream stp + releaseArrowArrayStream stp -- double release no-op + stn <- readIORef streamCnt + expect (stn == 1) ("stream cleanup ran " <> show stn <> " times, want 1") + free stp + + -- 7. DLPack: produce and consume entirely in-process + dcnt <- newIORef (0 :: Int) + dbuf <- mallocBytes (6 * 8) :: IO (Ptr Double) + mapM_ (uncurry (pokeElemOff dbuf)) (zip [0 ..] [1 .. 6]) + mt <- + newManagedTensor + (DLDataType kDLFloat 64 1) + [2, 3] + (castPtr dbuf) + dlpackFlagReadOnly + (free dbuf >> modifyIORef' dcnt (+ 1)) + dvals <- consumeManagedTensor mt $ \m -> do + expect (dlverMajor (mtvVersion m) == dlpackMajorVersion) "tensor version major" + expect (mtvFlags m == dlpackFlagReadOnly) "tensor flags lost" + let t = mtvTensor m + expect (dltNDim t == 2) ("tensor ndim: " <> show (dltNDim t)) + dsh <- mapM (peekElemOff (dltShape t)) [0, 1] + expect (dsh == [2, 3]) ("tensor shape: " <> show dsh) + expect (dltStrides t == nullPtr) "strides not null (compact expected)" + mapM (peekElemOff (castPtr (dltData t) :: Ptr Double)) [0 .. 5] + expect (dvals == [1 .. 6]) ("tensor values: " <> show dvals) + dn <- readIORef dcnt + expect (dn == 1) ("tensor cleanup ran " <> show dn <> " times, want 1") + + -- 8. a throwing cleanup must not escape into the (foreign) caller: + -- the struct is still marked released and the process survives + ecnt <- newIORef (0 :: Int) + ep <- callocBytes (sizeOf (undefined :: ArrowArray)) + exportArrowArray ep emptyArrowArray + (modifyIORef' ecnt (+ 1) >> ioError (userError "cleanup boom")) + releaseArrowArray ep + e1 <- peek ep + expect (arrayRelease e1 == nullFunPtr) "throwing cleanup: release not nulled" + releaseArrowArray ep -- and release stays a no-op afterwards + en <- readIORef ecnt + expect (en == 1) ("throwing cleanup ran " <> show en <> " times, want 1") + free ep + + -- 9. a throwing stream callback surfaces as errno EIO, not a crash + tp <- callocBytes (sizeOf (undefined :: ArrowArrayStream)) + exportArrowArrayStream tp + ArrowStreamProducer + { producerGetSchema = \_ -> ioError (userError "schema boom") + , producerGetNext = \_ -> ioError (userError "next boom") + , producerGetLastError = pure nullPtr + , producerCleanup = pure () + } + tst <- peek tp + trc <- withArrowSchemaImport (callStreamGetSchema (streamGetSchema tst) tp) + expect (trc == 5) ("throwing get_schema returned " <> show trc <> ", want 5 (EIO)") + releaseArrowArrayStream tp + free tp + + -- 10. leak gate: 10k full export/release cycles must not grow the live + -- Haskell heap (a leaked StablePtr or cleanup closure would). C-side + -- malloc leaks are invisible here — that is the publish-stage valgrind + -- lane's job. + statsOn <- getRTSStatsEnabled + expect statsOn "RTS stats disabled - test suite must be built with -with-rtsopts=-T" + performMajorGC + live0 <- gcdetails_live_bytes . gc <$> getRTSStats + forM_ [1 :: Int .. 10000] $ \_ -> + withArrowArrayImport $ \lp -> + exportArrowArray lp emptyArrowArray (pure ()) + performMajorGC + live1 <- gcdetails_live_bytes . gc <$> getRTSStats + let grownKiB = (fromIntegral live1 - fromIntegral live0) `div` 1024 :: Integer + expect (grownKiB < 1024) + ("live heap grew " <> show grownKiB <> " KiB over 10k cycles (leak)") + + putStrLn "keel-abi: managed-layer lifecycle tests passed (10 scenarios)" + where + poke' o = poke o emptyArrowArray + nextLen st stp = withArrowArrayImport $ \ao -> do + rc <- callStreamGetNext (streamGetNext st) stp ao + expect (rc == 0) "get_next returned nonzero" + a <- peek ao + pure $ + if arrayRelease a == nullFunPtr + then Nothing + else Just (arrayLength a) diff --git a/test/PyArrow.hs b/test/PyArrow.hs new file mode 100644 --- /dev/null +++ b/test/PyArrow.hs @@ -0,0 +1,310 @@ +-- | Conformance test: round-trip the Arrow C Data Interface AND the +-- C Stream Interface against pyarrow, in-process, in both directions. +-- +-- CPython is loaded at run time through keel-dyn (see "PyEmbed"). +-- Four directions: +-- +-- 1. array pyarrow -> Haskell: @[1, 2, 3, None, 5] :: int64@; verify +-- format\/length\/null bitmap\/values, release, check the callbacks +-- null themselves per spec (raw-layer conformance, managed alloc); +-- 2. array Haskell -> pyarrow: hand-built @[10, 20, 30]@ with raw +-- @\"wrapper\"@ release callbacks; Python verifies and drops it, +-- which must call back into Haskell and free our buffers; +-- 3. stream pyarrow -> Haskell: a two-batch RecordBatchReader; verify +-- struct schema (+s with int64 child \"x\"), both batches, and the +-- end-of-stream convention; +-- 4. stream Haskell -> pyarrow: an 'ArrowStreamProducer' serving two +-- batches through the managed export layer; Python @read_all()@s, +-- verifies, and drops it — our cleanup must run. +module Main (main) where + +import Control.Monad (forM) +import Data.Bits ((.&.)) +import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef) +import Data.Int (Int64) +import Data.Word (Word8) +import Foreign.C.String (newCString, peekCString) +import Foreign.Marshal.Alloc (callocBytes, free, mallocBytes) +import Foreign.Ptr (FunPtr, Ptr, castPtr, freeHaskellFunPtr, nullFunPtr, nullPtr) +import Foreign.Storable (peek, peekElemOff, poke, pokeElemOff, sizeOf) + +import Keel.Abi.Arrow +import Keel.Abi.Arrow.Raw +import PyEmbed + +foreign import ccall "wrapper" + mkSchemaRelease :: (Ptr ArrowSchema -> IO ()) -> IO (FunPtr (Ptr ArrowSchema -> IO ())) + +foreign import ccall "wrapper" + mkArrayRelease :: (Ptr ArrowArray -> IO ()) -> IO (FunPtr (Ptr ArrowArray -> IO ())) + +main :: IO () +main = withEmbeddedPython "pyarrow" "KEEL_ABI_REQUIRE_PYARROW" $ \runScript -> do + arrayImportDirection runScript + arrayExportDirection runScript + streamImportDirection runScript + streamExportDirection runScript + putStrLn "keel-abi-pyarrow: arrays and streams round-tripped both directions" + +-- --------------------------------------------------------------------- +-- Direction 1: array, pyarrow -> Haskell + +arrayImportDirection :: RunScript -> IO () +arrayImportDirection runScript = + withArrowArrayImport $ \impArr -> + withArrowSchemaImport $ \impSch -> do + runScript "array-import" $ + "import pyarrow as pa\n\ + \arr = pa.array([1, 2, 3, None, 5], type=pa.int64())\n\ + \arr._export_to_c(" <> addr impArr <> ", " <> addr impSch <> ")\n" + + sch <- peek impSch + fmt <- peekCString (schemaFormat sch) + expect (fmt == "l") ("schema format: expected \"l\", got " <> show fmt) + expect (schemaFlags sch .&. arrowFlagNullable /= 0) "schema not flagged nullable" + + arr <- peek impArr + expect (arrayLength arr == 5) ("length: " <> show (arrayLength arr)) + expect (arrayNullCount arr == 1) ("null_count: " <> show (arrayNullCount arr)) + expect (arrayNBuffers arr == 2) ("n_buffers: " <> show (arrayNBuffers arr)) + expect (arrayOffset arr == 0) ("offset: " <> show (arrayOffset arr)) + + validity <- peekElemOff (arrayBuffers arr) 0 + expect (validity /= nullPtr) "validity bitmap missing despite a null" + vbyte <- peek (castPtr validity :: Ptr Word8) + expect (vbyte .&. 0x1F == 0x17) -- rows 0,1,2,4 valid, row 3 null + ("validity bitmap byte: " <> show vbyte) + + dataBuf <- peekElemOff (arrayBuffers arr) 1 + vals <- mapM (peekElemOff (castPtr dataBuf :: Ptr Int64)) [0, 1, 2, 4] + expect (vals == [1, 2, 3, 5]) ("values: " <> show vals) + + -- release early (the with-brackets would anyway), then verify the + -- callbacks nulled themselves per spec + releaseArrowArray impArr + releaseArrowSchema impSch + arrAfter <- peek impArr + schAfter <- peek impSch + expect (arrayRelease arrAfter == nullFunPtr) "array release not nulled" + expect (schemaRelease schAfter == nullFunPtr) "schema release not nulled" + +-- --------------------------------------------------------------------- +-- Direction 2: array, Haskell -> pyarrow (raw-layer wrappers on purpose) + +arrayExportDirection :: RunScript -> IO () +arrayExportDirection runScript = do + expData <- mallocBytes (3 * 8) :: IO (Ptr Int64) + mapM_ (uncurry (pokeElemOff expData)) (zip [0 ..] [10, 20, 30]) + bufArr <- mallocBytes (2 * 8) :: IO (Ptr (Ptr ())) + pokeElemOff bufArr 0 nullPtr + pokeElemOff bufArr 1 (castPtr expData) + expFmt <- newCString "l" + + schemaReleased <- newIORef False + arrayReleased <- newIORef False + sRel <- mkSchemaRelease (schemaReleaseAction schemaReleased) + aRel <- mkArrayRelease (arrayReleaseAction arrayReleased) + + expArr <- callocBytes (sizeOf (undefined :: ArrowArray)) + expSch <- callocBytes (sizeOf (undefined :: ArrowSchema)) + poke expSch + emptyArrowSchema + { schemaFormat = expFmt + , schemaFlags = arrowFlagNullable + , schemaRelease = sRel + } + poke expArr + emptyArrowArray + { arrayLength = 3 + , arrayNBuffers = 2 + , arrayBuffers = bufArr + , arrayRelease = aRel + } + + runScript "array-export" $ + "import pyarrow as pa, gc\n\ + \imp = pa.Array._import_from_c(" <> addr expArr <> ", " <> addr expSch <> ")\n\ + \got = imp.to_pylist()\n\ + \if got != [10, 20, 30]:\n\ + \ raise RuntimeError('roundtrip mismatch: %r' % (got,))\n\ + \del imp\n\ + \gc.collect()\n" + + sDone <- readIORef schemaReleased + aDone <- readIORef arrayReleased + expect sDone "pyarrow never called our schema release callback" + expect aDone "pyarrow never called our array release callback" + free expArr + free expSch + freeHaskellFunPtr sRel + freeHaskellFunPtr aRel + +-- Raw exporter-side release callbacks: free what we allocated, then null +-- the release member per spec. +schemaReleaseAction :: IORef Bool -> Ptr ArrowSchema -> IO () +schemaReleaseAction flag p = do + s <- peek p + free (schemaFormat s) + poke p s { schemaRelease = nullFunPtr } + writeIORef flag True + +arrayReleaseAction :: IORef Bool -> Ptr ArrowArray -> IO () +arrayReleaseAction flag p = do + a <- peek p + d <- peekElemOff (arrayBuffers a) 1 + free d + free (arrayBuffers a) + poke p a { arrayRelease = nullFunPtr } + writeIORef flag True + +-- --------------------------------------------------------------------- +-- Direction 3: stream, pyarrow -> Haskell + +streamImportDirection :: RunScript -> IO () +streamImportDirection runScript = + withArrowArrayStreamImport $ \stp -> do + runScript "stream-import" $ + "import pyarrow as pa\n\ + \schema = pa.schema([('x', pa.int64())])\n\ + \batches = [pa.record_batch([pa.array([1, 2, 3])], schema=schema),\n\ + \ pa.record_batch([pa.array([4, 5])], schema=schema)]\n\ + \reader = pa.RecordBatchReader.from_batches(schema, batches)\n\ + \reader._export_to_c(" <> addr stp <> ")\n" + + st <- peek stp + withArrowSchemaImport $ \so -> do + rc <- callStreamGetSchema (streamGetSchema st) stp so + expect (rc == 0) "stream get_schema returned nonzero" + sch <- peek so + fmt <- peekCString (schemaFormat sch) + expect (fmt == "+s") ("stream schema format: " <> show fmt) + expect (schemaNChildren sch == 1) "stream schema child count" + childP <- peekElemOff (schemaChildren sch) 0 + child <- peek childP + cfmt <- peekCString (schemaFormat child) + cname <- peekCString (schemaName child) + expect (cfmt == "l") ("child format: " <> show cfmt) + expect (cname == "x") ("child name: " <> show cname) + + b1 <- readBatch st stp + b2 <- readBatch st stp + b3 <- readBatch st stp + expect (b1 == Just [1, 2, 3]) ("batch 1: " <> show b1) + expect (b2 == Just [4, 5]) ("batch 2: " <> show b2) + expect (b3 == Nothing) ("expected end-of-stream, got " <> show b3) + +-- Read one record batch (struct array, one int64 child) from an +-- imported stream; Nothing = end-of-stream per the null-release rule. +readBatch :: ArrowArrayStream -> Ptr ArrowArrayStream -> IO (Maybe [Int64]) +readBatch st stp = + withArrowArrayImport $ \ao -> do + rc <- callStreamGetNext (streamGetNext st) stp ao + expect (rc == 0) "stream get_next returned nonzero" + a <- peek ao + if arrayRelease a == nullFunPtr + then pure Nothing + else do + expect (arrayNChildren a == 1) "batch child count" + childP <- peekElemOff (arrayChildren a) 0 + child <- peek childP + dataBuf <- peekElemOff (arrayBuffers child) 1 + vals <- forM [0 .. fromIntegral (arrayLength child) - 1] $ + peekElemOff (castPtr dataBuf :: Ptr Int64) + pure (Just vals) + +-- --------------------------------------------------------------------- +-- Direction 4: stream, Haskell -> pyarrow (managed export layer) + +streamExportDirection :: RunScript -> IO () +streamExportDirection runScript = do + cleanupRan <- newIORef False + batchIx <- newIORef (0 :: Int) + stp <- callocBytes (sizeOf (undefined :: ArrowArrayStream)) + exportArrowArrayStream stp + ArrowStreamProducer + { producerGetSchema = \o -> buildStructSchema o >> pure 0 + , producerGetNext = \o -> do + i <- readIORef batchIx + modifyIORef' batchIx (+ 1) + case batchVals i of + Just vs -> buildStructBatch o vs >> pure 0 + Nothing -> poke o emptyArrowArray >> pure 0 + , producerGetLastError = pure nullPtr + , producerCleanup = writeIORef cleanupRan True + } + + runScript "stream-export" $ + "import pyarrow as pa, gc\n\ + \r = pa.RecordBatchReader._import_from_c(" <> addr stp <> ")\n\ + \tbl = r.read_all()\n\ + \got = tbl.column('x').to_pylist()\n\ + \if got != [10, 20, 30, 40, 50]:\n\ + \ raise RuntimeError('stream mismatch: %r' % (got,))\n\ + \del r, tbl\n\ + \gc.collect()\n" + + done <- readIORef cleanupRan + expect done "pyarrow never released our exported stream" + free stp + where + batchVals :: Int -> Maybe [Int64] + batchVals 0 = Just [10, 20, 30] + batchVals 1 = Just [40, 50] + batchVals _ = Nothing + +-- Build the stream schema: struct "+s" with a single nullable int64 +-- child named "x". Parent release cascades into the child per spec. +buildStructSchema :: Ptr ArrowSchema -> IO () +buildStructSchema out = do + childFmt <- newCString "l" + childName <- newCString "x" + child <- callocBytes (sizeOf (undefined :: ArrowSchema)) + exportArrowSchema child + emptyArrowSchema + { schemaFormat = childFmt + , schemaName = childName + , schemaFlags = arrowFlagNullable + } + (free childFmt >> free childName) + kids <- mallocBytes 8 + pokeElemOff kids 0 child + parentFmt <- newCString "+s" + exportArrowSchema out + emptyArrowSchema + { schemaFormat = parentFmt + , schemaNChildren = 1 + , schemaChildren = kids + } + (releaseArrowSchema child >> free child >> free kids >> free parentFmt) + +-- Build one record batch: a struct array (validity-only buffer list) +-- with a single int64 child. Parent release cascades into the child. +buildStructBatch :: Ptr ArrowArray -> [Int64] -> IO () +buildStructBatch out vs = do + let n = length vs + dataBuf <- mallocBytes (n * 8) :: IO (Ptr Int64) + mapM_ (uncurry (pokeElemOff dataBuf)) (zip [0 ..] vs) + childBufs <- mallocBytes (2 * 8) :: IO (Ptr (Ptr ())) + pokeElemOff childBufs 0 nullPtr + pokeElemOff childBufs 1 (castPtr dataBuf) + child <- callocBytes (sizeOf (undefined :: ArrowArray)) + exportArrowArray child + emptyArrowArray + { arrayLength = fromIntegral n + , arrayNBuffers = 2 + , arrayBuffers = childBufs + } + (free dataBuf >> free childBufs) + kids <- mallocBytes 8 + pokeElemOff kids 0 child + parentBufs <- mallocBytes 8 :: IO (Ptr (Ptr ())) + pokeElemOff parentBufs 0 nullPtr + exportArrowArray out + emptyArrowArray + { arrayLength = fromIntegral n + , arrayNBuffers = 1 + , arrayBuffers = parentBufs + , arrayNChildren = 1 + , arrayChildren = kids + } + (releaseArrowArray child >> free child >> free kids >> free parentBufs) diff --git a/test/PyEmbed.hs b/test/PyEmbed.hs new file mode 100644 --- /dev/null +++ b/test/PyEmbed.hs @@ -0,0 +1,113 @@ +-- | Shared harness for conformance tests that embed CPython in-process +-- through keel-dyn: discovery (out of process), load, init, script +-- runner, finalize, and the SKIP-vs-required policy. +module PyEmbed + ( RunScript + , withEmbeddedPython + , expect + , addr + ) where + +import Control.Exception (IOException, try) +import Control.Monad (unless, when) +import Foreign.C.String (CString, withCString) +import Foreign.C.Types (CInt (..)) +import Foreign.Ptr (FunPtr, Ptr, ptrToIntPtr) +import System.Environment (lookupEnv) +import System.Exit (ExitCode (..)) +import System.Process (readProcessWithExitCode) + +import Keel.Dyn + +-- | @runScript label script@ — run python source, fail the test with +-- @label@ if it raises (the traceback goes to stderr). +type RunScript = String -> String -> IO () + +expect :: Bool -> String -> IO () +expect ok msg = unless ok (fail msg) + +-- | A pointer, rendered for splicing into python source as an integer. +addr :: Ptr a -> String +addr = show . ptrToIntPtr + +foreign import ccall safe "dynamic" + mkPyInitEx :: FunPtr (CInt -> IO ()) -> CInt -> IO () + +foreign import ccall safe "dynamic" + mkPyRun :: FunPtr (CString -> IO CInt) -> CString -> IO CInt + +foreign import ccall safe "dynamic" + mkPyFinalize :: FunPtr (IO CInt) -> IO CInt + +runPy :: String -> [String] -> IO (Maybe String) +runPy exe args = do + r <- try (readProcessWithExitCode exe args "") + :: IO (Either IOException (ExitCode, String, String)) + pure $ case r of + Right (ExitSuccess, out, _) -> Just (filter (`notElem` "\r\n") out) + _ -> Nothing + +-- Find an interpreter that can import the given module, and the shared +-- library behind it. Out of process, so a broken python cannot kill us. +-- POSIX candidates cover plain builds (LIBDIR + INSTSONAME/LDLIBRARY) +-- and macOS framework builds, where INSTSONAME is framework-relative +-- and must be joined to PYTHONFRAMEWORKPREFIX instead. +findPython :: String -> IO (Maybe FilePath) +findPython pymodule = go ["python", "python3"] + where + dllScript = + "import sys, sysconfig, os\n\ + \if os.name == 'nt':\n\ + \ print(os.path.join(sys.base_prefix, 'python%d%d.dll' % sys.version_info[:2]))\n\ + \else:\n\ + \ g = sysconfig.get_config_var\n\ + \ libdir = g('LIBDIR') or ''\n\ + \ cands = [os.path.join(libdir, n) for n in (g('INSTSONAME') or '', g('LDLIBRARY') or '') if n]\n\ + \ fw = g('PYTHONFRAMEWORKPREFIX') or ''\n\ + \ if fw:\n\ + \ cands += [os.path.join(fw, n) for n in (g('INSTSONAME') or '', g('LDLIBRARY') or '') if n]\n\ + \ print(next((c for c in cands if os.path.exists(c)), ''))\n" + go [] = pure Nothing + go (exe : rest) = do + ok <- runPy exe ["-c", "import " <> pymodule] + case ok of + Nothing -> go rest + Just _ -> do + p <- runPy exe ["-c", dllScript] + pure $ case p of + Just path | not (null path) -> Just path + _ -> Nothing + +-- | Load CPython through keel-dyn and hand a 'RunScript' to the action; +-- initialize before, finalize after. When no interpreter that imports +-- @pymodule@ exists, SKIP (print why, exit 0) — unless the named env +-- var is set non-empty\/non-zero, which turns absence into failure +-- (publish-stage CI sets it). +withEmbeddedPython + :: String -- ^ python module the test needs, e.g. @\"pyarrow\"@ + -> String -- ^ require-env-var, e.g. @\"KEEL_ABI_REQUIRE_PYARROW\"@ + -> (RunScript -> IO ()) + -> IO () +withEmbeddedPython pymodule requireVar action = do + found <- findPython pymodule + required <- lookupEnv requireVar + case found of + Nothing -> case required of + Just v | v /= "" && v /= "0" -> + fail (requireVar <> " set but no usable python+" <> pymodule <> " found") + _ -> putStrLn ("SKIP - no usable python+" <> pymodule <> " on this machine") + Just dll -> do + -- RTLD_GLOBAL: C extension modules (manylinux policy) leave + -- Python's own symbols undefined and need them process-visible + lib <- either (\e -> fail ("load python: " <> show e)) pure =<< loadLibraryGlobal dll + pyInitEx <- requireSym lib "Py_InitializeEx" + pyRun <- requireSym lib "PyRun_SimpleString" + pyFin <- requireSym lib "Py_FinalizeEx" + mkPyInitEx pyInitEx 0 + let runScript label script = do + rc <- withCString script (mkPyRun pyRun) + expect (rc == 0) (label <> ": python script failed (traceback above)") + action runScript + fin <- mkPyFinalize pyFin + when (fin /= 0) (putStrLn "note: Py_FinalizeEx reported errors (ignored)") + closeLibrary lib diff --git a/test/cbits/layout_gate.c b/test/cbits/layout_gate.c new file mode 100644 --- /dev/null +++ b/test/cbits/layout_gate.c @@ -0,0 +1,247 @@ +/* Test-only layout gate for keel-abi. The shipped library has NO C + * sources; this file exists so the test suite fails to COMPILE (via + * _Static_assert) or fails at RUN TIME (via the keel_layout_* probes, + * compared against the Haskell-side layout tables) if the hand-written + * Storable offsets in Keel.Abi.*.Raw ever disagree with a real C + * compiler on the build platform. + * + * Struct definitions vendored from their frozen-ABI specifications: + * - Apache Arrow C Data / C Stream Interface (Apache-2.0), which the + * spec instructs consumers to copy verbatim: + * https://arrow.apache.org/docs/format/CDataInterface.html + * https://arrow.apache.org/docs/format/CStreamInterface.html + * - dlpack.h v1.1 (Apache-2.0), trimmed to the exchanged structs: + * https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h + */ +#include +#include + +/* ------------------------------------------------------------------ */ +/* Arrow C Data Interface */ + +struct ArrowSchema { + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + void (*release)(struct ArrowSchema*); + void* private_data; +}; + +struct ArrowArray { + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + void (*release)(struct ArrowArray*); + void* private_data; +}; + +struct ArrowArrayStream { + int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); + int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); + const char* (*get_last_error)(struct ArrowArrayStream*); + void (*release)(struct ArrowArrayStream*); + void* private_data; +}; + +/* ------------------------------------------------------------------ */ +/* DLPack v1.x (versioned exchange structs only) */ + +typedef struct { + uint32_t major; + uint32_t minor; +} DLPackVersion; + +typedef enum { + kDLCPU = 1, + kDLCUDA = 2, + kDLCUDAHost = 3, + kDLOpenCL = 4, + kDLVulkan = 7, + kDLMetal = 8, + kDLVPI = 9, + kDLROCM = 10 +} DLDeviceType; + +typedef struct { + DLDeviceType device_type; + int32_t device_id; +} DLDevice; + +typedef struct { + uint8_t code; + uint8_t bits; + uint16_t lanes; +} DLDataType; + +typedef struct { + void* data; + DLDevice device; + int32_t ndim; + DLDataType dtype; + int64_t* shape; + int64_t* strides; + uint64_t byte_offset; +} DLTensor; + +struct DLManagedTensorVersioned { + DLPackVersion version; + void* manager_ctx; + void (*deleter)(struct DLManagedTensorVersioned* self); + uint64_t flags; + DLTensor dl_tensor; +}; + +/* ------------------------------------------------------------------ */ +/* Compile-time gate: the literals below are the exact numbers the */ +/* Haskell Storable instances use. 64-bit only by design. */ + +_Static_assert(sizeof(void*) == 8, "keel-abi targets 64-bit platforms only"); + +_Static_assert(sizeof(struct ArrowSchema) == 72, "ArrowSchema size"); +_Static_assert(offsetof(struct ArrowSchema, format) == 0, "ArrowSchema.format"); +_Static_assert(offsetof(struct ArrowSchema, name) == 8, "ArrowSchema.name"); +_Static_assert(offsetof(struct ArrowSchema, metadata) == 16, "ArrowSchema.metadata"); +_Static_assert(offsetof(struct ArrowSchema, flags) == 24, "ArrowSchema.flags"); +_Static_assert(offsetof(struct ArrowSchema, n_children) == 32, "ArrowSchema.n_children"); +_Static_assert(offsetof(struct ArrowSchema, children) == 40, "ArrowSchema.children"); +_Static_assert(offsetof(struct ArrowSchema, dictionary) == 48, "ArrowSchema.dictionary"); +_Static_assert(offsetof(struct ArrowSchema, release) == 56, "ArrowSchema.release"); +_Static_assert(offsetof(struct ArrowSchema, private_data) == 64, "ArrowSchema.private_data"); + +_Static_assert(sizeof(struct ArrowArray) == 80, "ArrowArray size"); +_Static_assert(offsetof(struct ArrowArray, length) == 0, "ArrowArray.length"); +_Static_assert(offsetof(struct ArrowArray, null_count) == 8, "ArrowArray.null_count"); +_Static_assert(offsetof(struct ArrowArray, offset) == 16, "ArrowArray.offset"); +_Static_assert(offsetof(struct ArrowArray, n_buffers) == 24, "ArrowArray.n_buffers"); +_Static_assert(offsetof(struct ArrowArray, n_children) == 32, "ArrowArray.n_children"); +_Static_assert(offsetof(struct ArrowArray, buffers) == 40, "ArrowArray.buffers"); +_Static_assert(offsetof(struct ArrowArray, children) == 48, "ArrowArray.children"); +_Static_assert(offsetof(struct ArrowArray, dictionary) == 56, "ArrowArray.dictionary"); +_Static_assert(offsetof(struct ArrowArray, release) == 64, "ArrowArray.release"); +_Static_assert(offsetof(struct ArrowArray, private_data) == 72, "ArrowArray.private_data"); + +_Static_assert(sizeof(struct ArrowArrayStream) == 40, "ArrowArrayStream size"); +_Static_assert(offsetof(struct ArrowArrayStream, get_schema) == 0, "ArrowArrayStream.get_schema"); +_Static_assert(offsetof(struct ArrowArrayStream, get_next) == 8, "ArrowArrayStream.get_next"); +_Static_assert(offsetof(struct ArrowArrayStream, get_last_error) == 16, "ArrowArrayStream.get_last_error"); +_Static_assert(offsetof(struct ArrowArrayStream, release) == 24, "ArrowArrayStream.release"); +_Static_assert(offsetof(struct ArrowArrayStream, private_data) == 32, "ArrowArrayStream.private_data"); + +_Static_assert(sizeof(DLPackVersion) == 8, "DLPackVersion size"); +_Static_assert(offsetof(DLPackVersion, major) == 0, "DLPackVersion.major"); +_Static_assert(offsetof(DLPackVersion, minor) == 4, "DLPackVersion.minor"); + +_Static_assert(sizeof(DLDevice) == 8, "DLDevice size"); +_Static_assert(offsetof(DLDevice, device_type) == 0, "DLDevice.device_type"); +_Static_assert(offsetof(DLDevice, device_id) == 4, "DLDevice.device_id"); + +_Static_assert(sizeof(DLDataType) == 4, "DLDataType size"); +_Static_assert(offsetof(DLDataType, code) == 0, "DLDataType.code"); +_Static_assert(offsetof(DLDataType, bits) == 1, "DLDataType.bits"); +_Static_assert(offsetof(DLDataType, lanes) == 2, "DLDataType.lanes"); + +_Static_assert(sizeof(DLTensor) == 48, "DLTensor size"); +_Static_assert(offsetof(DLTensor, data) == 0, "DLTensor.data"); +_Static_assert(offsetof(DLTensor, device) == 8, "DLTensor.device"); +_Static_assert(offsetof(DLTensor, ndim) == 16, "DLTensor.ndim"); +_Static_assert(offsetof(DLTensor, dtype) == 20, "DLTensor.dtype"); +_Static_assert(offsetof(DLTensor, shape) == 24, "DLTensor.shape"); +_Static_assert(offsetof(DLTensor, strides) == 32, "DLTensor.strides"); +_Static_assert(offsetof(DLTensor, byte_offset) == 40, "DLTensor.byte_offset"); + +_Static_assert(sizeof(struct DLManagedTensorVersioned) == 80, "DLManagedTensorVersioned size"); +_Static_assert(offsetof(struct DLManagedTensorVersioned, version) == 0, "DLManagedTensorVersioned.version"); +_Static_assert(offsetof(struct DLManagedTensorVersioned, manager_ctx) == 8, "DLManagedTensorVersioned.manager_ctx"); +_Static_assert(offsetof(struct DLManagedTensorVersioned, deleter) == 16, "DLManagedTensorVersioned.deleter"); +_Static_assert(offsetof(struct DLManagedTensorVersioned, flags) == 24, "DLManagedTensorVersioned.flags"); +_Static_assert(offsetof(struct DLManagedTensorVersioned, dl_tensor) == 32, "DLManagedTensorVersioned.dl_tensor"); + +/* ------------------------------------------------------------------ */ +/* Run-time probes: fill `out` with offsetof per field in declaration */ +/* order, return sizeof. The Haskell test compares these against the */ +/* layout tables exported by the Raw modules, closing the loop between */ +/* the C compiler's view and the Storable instances. */ + +size_t keel_layout_ArrowSchema(size_t* out) { + out[0] = offsetof(struct ArrowSchema, format); + out[1] = offsetof(struct ArrowSchema, name); + out[2] = offsetof(struct ArrowSchema, metadata); + out[3] = offsetof(struct ArrowSchema, flags); + out[4] = offsetof(struct ArrowSchema, n_children); + out[5] = offsetof(struct ArrowSchema, children); + out[6] = offsetof(struct ArrowSchema, dictionary); + out[7] = offsetof(struct ArrowSchema, release); + out[8] = offsetof(struct ArrowSchema, private_data); + return sizeof(struct ArrowSchema); +} + +size_t keel_layout_ArrowArray(size_t* out) { + out[0] = offsetof(struct ArrowArray, length); + out[1] = offsetof(struct ArrowArray, null_count); + out[2] = offsetof(struct ArrowArray, offset); + out[3] = offsetof(struct ArrowArray, n_buffers); + out[4] = offsetof(struct ArrowArray, n_children); + out[5] = offsetof(struct ArrowArray, buffers); + out[6] = offsetof(struct ArrowArray, children); + out[7] = offsetof(struct ArrowArray, dictionary); + out[8] = offsetof(struct ArrowArray, release); + out[9] = offsetof(struct ArrowArray, private_data); + return sizeof(struct ArrowArray); +} + +size_t keel_layout_ArrowArrayStream(size_t* out) { + out[0] = offsetof(struct ArrowArrayStream, get_schema); + out[1] = offsetof(struct ArrowArrayStream, get_next); + out[2] = offsetof(struct ArrowArrayStream, get_last_error); + out[3] = offsetof(struct ArrowArrayStream, release); + out[4] = offsetof(struct ArrowArrayStream, private_data); + return sizeof(struct ArrowArrayStream); +} + +size_t keel_layout_DLPackVersion(size_t* out) { + out[0] = offsetof(DLPackVersion, major); + out[1] = offsetof(DLPackVersion, minor); + return sizeof(DLPackVersion); +} + +size_t keel_layout_DLDevice(size_t* out) { + out[0] = offsetof(DLDevice, device_type); + out[1] = offsetof(DLDevice, device_id); + return sizeof(DLDevice); +} + +size_t keel_layout_DLDataType(size_t* out) { + out[0] = offsetof(DLDataType, code); + out[1] = offsetof(DLDataType, bits); + out[2] = offsetof(DLDataType, lanes); + return sizeof(DLDataType); +} + +size_t keel_layout_DLTensor(size_t* out) { + out[0] = offsetof(DLTensor, data); + out[1] = offsetof(DLTensor, device); + out[2] = offsetof(DLTensor, ndim); + out[3] = offsetof(DLTensor, dtype); + out[4] = offsetof(DLTensor, shape); + out[5] = offsetof(DLTensor, strides); + out[6] = offsetof(DLTensor, byte_offset); + return sizeof(DLTensor); +} + +size_t keel_layout_DLManagedTensorVersioned(size_t* out) { + out[0] = offsetof(struct DLManagedTensorVersioned, version); + out[1] = offsetof(struct DLManagedTensorVersioned, manager_ctx); + out[2] = offsetof(struct DLManagedTensorVersioned, deleter); + out[3] = offsetof(struct DLManagedTensorVersioned, flags); + out[4] = offsetof(struct DLManagedTensorVersioned, dl_tensor); + return sizeof(struct DLManagedTensorVersioned); +}