{-# LANGUAGE ViewPatterns        #-}
--------------------------------------------------------------------------------
-- |
-- Module      : ArrayFire.Device
-- Copyright   : David Johnson (c) 2019-2026
-- License     : BSD3
-- Maintainer  : David Johnson <code@dmj.io>
-- Stability   : Experimental
-- Portability : GHC
--
-- Information about ArrayFire API and devices
--
-- >>> info
--  ArrayFire v3.6.4 (OpenCL, 64-bit Mac OSX, build 1b8030c5)
-- [0] APPLE: AMD Radeon Pro 555X Compute Engine, 4096 MB
-- -1- APPLE: Intel(R) UHD Graphics 630, 1536 MB
--
--------------------------------------------------------------------------------
module ArrayFire.Device where

import Control.Exception (finally)
import Foreign.C.String
import Foreign.Ptr (castPtr)
import ArrayFire.Internal.Device
import ArrayFire.FFI

foreign import ccall unsafe "af_notify_shutdown"
  afNotifyShutdown :: IO ()

-- | Bracket for ArrayFire usage.  Wrap your @main@ (or top-level IO action)
-- with this to ensure the safe-finalizer shutdown flag is set before GHC's
-- finalizer thread runs, preventing a "double free or corruption" abort when
-- GC-managed array handles outlive ArrayFire's C++ allocator teardown.
--
-- @
-- main :: IO ()
-- main = withArrayFire $ do
--   ...
-- @
withArrayFire :: IO a -> IO a
withArrayFire :: forall a. IO a -> IO a
withArrayFire IO a
action = IO a
action IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`finally` IO ()
afNotifyShutdown

-- | Retrieve info from ArrayFire API
--
-- @
-- ArrayFire v3.6.4 (OpenCL, 64-bit Mac OSX, build 1b8030c5)
-- [0] APPLE: AMD Radeon Pro 555X Compute Engine, 4096 MB
-- -1- APPLE: Intel(R) UHD Graphics 630, 1536 MB
-- @
info :: IO ()
info :: IO ()
info = IO AFErr -> IO ()
afCall IO AFErr
af_info

-- | Calls /af_init/ C function from ArrayFire API
--
-- >>> afInit
-- ()
afInit :: IO ()
afInit :: IO ()
afInit = IO AFErr -> IO ()
afCall IO AFErr
af_init

-- | Retrieves ArrayFire device information as 'String', same as 'info'.
--
-- >>> getInfoString
-- "ArrayFire v3.6.4 (OpenCL, 64-bit Mac OSX, build 1b8030c5)\n[0] APPLE: AMD Radeon Pro 555X Compute Engine, 4096 MB\n-1- APPLE: Intel(R) UHD Graphics 630, 1536 MB\n"
getInfoString :: IO String
getInfoString :: IO String
getInfoString = do
  strPtr <- (Ptr (Ptr CChar) -> IO AFErr) -> IO (Ptr CChar)
forall a. Storable a => (Ptr a -> IO AFErr) -> IO a
afCall1 ((Ptr (Ptr CChar) -> CBool -> IO AFErr)
-> CBool -> Ptr (Ptr CChar) -> IO AFErr
forall a b c. (a -> b -> c) -> b -> a -> c
flip Ptr (Ptr CChar) -> CBool -> IO AFErr
af_info_string CBool
1)
  str <- peekCString strPtr
  -- allocated by ArrayFire with af_alloc_host; free to avoid leaking
  _ <- af_free_host (castPtr strPtr)
  pure str

-- | Retrieves count of devices
--
-- >>> getDeviceCount
-- 2
getDeviceCount :: IO Int
getDeviceCount :: IO Int
getDeviceCount = CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ptr CInt -> IO AFErr) -> IO CInt
forall a. Storable a => (Ptr a -> IO AFErr) -> IO a
afCall1 Ptr CInt -> IO AFErr
af_get_device_count

-- | Sets a device by 'Int'
--
-- >>> setDevice 0
-- ()
setDevice :: Int -> IO ()
setDevice :: Int -> IO ()
setDevice (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
x) = IO AFErr -> IO ()
afCall (CInt -> IO AFErr
af_set_device CInt
x)

-- | Retrieves device identifier
--
-- >>> getDevice
-- 0
getDevice :: IO Int
getDevice :: IO Int
getDevice = CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Ptr CInt -> IO AFErr) -> IO CInt
forall a. Storable a => (Ptr a -> IO AFErr) -> IO a
afCall1 Ptr CInt -> IO AFErr
af_get_device

-- | Runs the device garbage collector, freeing any cached memory buffers.
deviceGC :: IO ()
deviceGC :: IO ()
deviceGC = IO AFErr -> IO ()
afCall IO AFErr
af_device_gc