hhlo 0.9.0.0 → 0.10.0.0

raw patch · 23 files changed

+859/−9 lines, 23 filesnew-component:exe:example-custom-callPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +46/−0
• cbits/pjrt_shim.c +77/−0
• cbits/pjrt_shim.h +7/−0
• examples/CustomCallPlugin.hs +57/−0
• hhlo.cabal +15/−1
• src/HHLO/Autograd/Rules.hs +31/−6
• src/HHLO/EDSL/Ops.hs +61/−1
• src/HHLO/IR/Builder.hs +25/−0
• src/HHLO/IR/Pretty.hs +11/−0
• src/HHLO/Runtime/CustomCall.hs +95/−0
• src/HHLO/Runtime/PJRT/FFI.hs +11/−0
• src/HHLO/Session.hs +1/−1
• test/Test/IR/Pretty.hs +28/−0
• test/Test/Runtime/EndToEndAutograd.hs +97/−0
• test/Test/Runtime/EndToEndAutogradGPU.hs +99/−0
• test/Test/Runtime/EndToEndDataMovement.hs +27/−0
• test/Test/Runtime/EndToEndDataMovementGPU.hs +29/−0
• test/Test/Runtime/EndToEndMatmul.hs +21/−0
• test/Test/Runtime/EndToEndMatmulGPU.hs +20/−0
• test/Test/Runtime/EndToEndNN.hs +34/−0
• test/Test/Runtime/EndToEndNNGPU.hs +31/−0
• test/Test/Runtime/EndToEndShape.hs +19/−0
• test/Test/Runtime/EndToEndShapeGPU.hs +17/−0

@@ -191,6 +191,52 @@ * New dependency: `vector-sized >= 1.5 && < 1.6`. * Fix `transposeConvolution` lhs_dilation bug — passing a 2-element spatial   dilation list no longer drops the second element.+  +## 0.10.0.0 -- 2026-05-02+* Fix `vjpConcatenate` slice offset bug — `splitAndAccumulate` now uses the+  cumulative `offset` in `start_indices` for all operands after the first.+  Previously only `limit_indices` used the offset, causing PJRT to reject+  gradient modules for any `concatenate2` with 2+ operands.+* Fix three autograd padding bugs:+  1. `transposeConvBackwardInput` — now applies `reversePad` so backward+     regular conv uses reversed padding instead of forward padding directly.+  2. `convBackwardInput` with `stride > 1` — now computes correct+     `targetPadTotal = input - (bar-1)*stride + kernel - 2` and adjusts+     reverse-pad to account for XLA floor division, restoring the original+     input spatial size in the transpose-conv backward pass.+  3. `vjpSlice` with `stride > 1` — `high'` padding now uses+     `n = ceil((limit-start)/stride)` instead of raw `(limit-start)`,+     preventing over-padded gradients. * `gather` and `scatter` kept as `[Int64]` for now. Their config vector lengths   depend on complex relationships between operand / indices / result ranks, so   a clean type-safe design requires a separate future phase.+* New E2E tests exercising the vector-sized configs:+  * `grad concatenate2`, `grad concatenate3` (regression tests for the concat bug)+  * `slice 2D`, `pad 2D symmetric`+  * `dotGeneral batched`+  * `transpose 3D`+  * GPU counterparts for all of the above.+* New regression tests for the padding bugs:+  * `grad conv2d stride` — strided conv `[1,4,4,1]` with stride=2+  * `grad transposeConvolution asymmetric pad` — transpose conv with+    `p2 (1,0) (1,0)` and `v2 2 2`+  * `grad pad interior` — pad with interior=1 on `[2]`+  * `transposeConvolution forward` — basic transpose conv smoke test+  * `conv2dWithPadding forward` — strided conv with explicit padding+* **Generic custom-call infrastructure** — first-class support for+  `stablehlo.custom_call`, enabling external packages to register their own+  XLA custom-call kernels (CUDA, CPU, or otherwise) without modifying HHLO.+  * `HHLO.IR.Builder.emitCustomCall` — emit `stablehlo.custom_call` with+    standard attributes (`call_target_name`, `has_side_effect`,+    `backend_config`, `api_version`).+  * `HHLO.EDSL.Ops.customCall1`, `customCall2`, `customCallRaw` — typed+    frontend wrappers. `customCallRaw` is the escape hatch for heterogeneous+    input types and arbitrary result counts.+  * `HHLO.Runtime.CustomCall.loadCustomCallLibrary` — `dlopen` wrapper with+    `RTLD_GLOBAL`, required for XLA's internal `dlsym` resolution.+  * `HHLO.IR.Pretty` — special case for `stablehlo.custom_call` emits the+    `@symbol` prefix before operands (e.g.+    `stablehlo.custom_call @foo(%arg0) {...} : ...`).+  * `examples/CustomCallPlugin.hs` + `examples/cbits/vector_add.cu` —+    minimal working example showing the full plugin contract.+* Test count: 205 CPU tests + 82 GPU tests = 287 total.

@@ -650,3 +650,80 @@     }     return err; }++// ---------------------------------------------------------------------------+// Custom calls (GPU)+// ---------------------------------------------------------------------------++// Local definitions for PJRT GPU Custom Call extension+// (matches xla/pjrt/c/pjrt_c_api_gpu_extension.h)++typedef struct {+    size_t struct_size;+    const char* function_name;+    size_t function_name_size;+    int api_version;+    void* handler_instantiate;+    void* handler_prepare;+    void* handler_initialize;+    void* handler_execute;+} PJRT_Gpu_Register_Custom_Call_Args;++typedef PJRT_Error* PJRT_Gpu_Register_Custom_Call(+    PJRT_Gpu_Register_Custom_Call_Args* args);++typedef struct {+    PJRT_Extension_Base base;+    PJRT_Gpu_Register_Custom_Call* custom_call;+} PJRT_Gpu_Custom_Call;++// Load a shared library, look up 'function_name', and register it with the+// PJRT GPU plugin via the PJRT_Gpu_Custom_Call extension.+//+// Returns 0 on success, or a negative code on failure.+// On failure, *out_error_msg is set to a static/dynamic error string.+int hhlo_pjrt_register_gpu_custom_call(PJRT_Api* api, const char* lib_path,+                                        const char* function_name,+                                        const char** out_error_msg) {+    *out_error_msg = NULL;++    // 1. Load the library (keep it open for the process lifetime)+    void* handle = dlopen(lib_path, RTLD_NOW | RTLD_LOCAL);+    if (!handle) {+        *out_error_msg = dlerror();+        return -1;+    }++    // 2. Look up the target symbol+    void* symbol = dlsym(handle, function_name);+    if (!symbol) {+        *out_error_msg = dlerror();+        return -2;+    }++    // 3. Walk the PJRT extension chain to find the GPU custom call extension+    PJRT_Extension_Base* ext = api->extension_start;+    while (ext != NULL) {+        if (ext->type == PJRT_Extension_Type_Gpu_Custom_Call) {+            PJRT_Gpu_Custom_Call* gpu_ext = (PJRT_Gpu_Custom_Call*)ext;++            PJRT_Gpu_Register_Custom_Call_Args args = {0};+            args.struct_size = sizeof(PJRT_Gpu_Register_Custom_Call_Args);+            args.function_name = function_name;+            args.function_name_size = strlen(function_name);+            args.api_version = 0;        // 0 = untyped / original ABI+            args.handler_execute = symbol;++            PJRT_Error* err = gpu_ext->custom_call(&args);+            if (err != NULL) {+                *out_error_msg = "PJRT_Gpu_Register_Custom_Call returned an error";+                return -3;+            }+            return 0;+        }+        ext = ext->next;+    }++    *out_error_msg = "PJRT GPU custom call extension not found";+    return -4;+}

@@ -130,6 +130,13 @@ PJRT_Error* hhlo_pjrt_error_destroy(PJRT_Api* api, PJRT_Error* error);  /* ---------------------------------------------------------------------------+ * Custom calls+ * --------------------------------------------------------------------------- */+int hhlo_pjrt_register_gpu_custom_call(PJRT_Api* api, const char* lib_path,+                                        const char* function_name,+                                        const char** out_error_msg);++/* ---------------------------------------------------------------------------  * Buffer type constants  * --------------------------------------------------------------------------- */ int hhlo_buffer_type_invalid(void);

@@ -0,0 +1,57 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example: using a custom-call plugin from HHLO.+--+-- Before running this example you must compile the CUDA kernel:+--+-- @+--   cd examples/cbits && bash build.sh+-- @+--+-- Then run with the GPU plugin:+--+-- @+--   cabal run example-custom-call --flag=examples+-- @+module Main where++import qualified Data.Text as T+import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.CustomCall (registerGpuCustomCall)+import HHLO.Session++main :: IO ()+main = withGPU $ \sess -> do+    -- 1. Register the GPU custom-call target with the PJRT CUDA plugin.+    --    This MUST happen before 'compile'.+    registerGpuCustomCall (sessionApi sess)+        "examples/cbits/libvector_add.so" "vector_add"++    -- 2. Build a StableHLO module that uses the custom call.+    let modu = moduleFromBuilder @'[4] @'F32 "main"+            [ FuncArg "a" (TensorType [4] F32)+            , FuncArg "b" (TensorType [4] F32)+            ]+            $ do+                a <- arg @'[4] @'F32+                b <- arg @'[4] @'F32+                c <- customCall1 "vector_add" [a, b] "" False+                return c++    putStrLn "=== Emitted MLIR ==="+    putStrLn (T.unpack (render modu))++    -- 3. Compile and execute via PJRT.+    compiled <- compile sess modu+    let aVals = hostFromList @'[4] @'F32 [1.0, 2.0, 3.0, 4.0]+        bVals = hostFromList @'[4] @'F32 [10.0, 20.0, 30.0, 40.0]+    result <- run sess compiled (aVals, bVals) :: IO (HostTensor '[4] 'F32)++    putStrLn "=== Result ==="+    print (hostToList result)

@@ -1,6 +1,6 @@ cabal-version:      3.0 name:               hhlo-version:            0.9.0.0+version:            0.10.0.0 synopsis:           Haskell Frontend for StableHLO — type-safe ML training/inference on CPU and GPU description:     HHLO is a Haskell library and runtime for building, compiling, and executing@@ -77,6 +77,7 @@         HHLO.Runtime.Execute         HHLO.Runtime.Async         HHLO.Runtime.Buffer+        HHLO.Runtime.CustomCall     build-depends:         base          >= 4.18.2 && < 5,         text          >= 2.0   && < 2.2,@@ -617,6 +618,19 @@ executable example-autograd-composite     import:           warnings     main-is:          36-autograd-composite.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-custom-call+    import:           warnings+    main-is:          CustomCallPlugin.hs     build-depends:         base          >= 4.18.2 && < 5,         hhlo,

@@ -7,7 +7,7 @@     ) where  import Data.Int (Int64)-import Data.List (sortOn)+import Data.List (sortOn, zipWith4) import Data.Maybe (isNothing) import Data.Text (Text) import qualified Data.Text as T@@ -412,7 +412,12 @@                 let start' = map fromIntegral start :: [Integer]                     limit' = map fromIntegral limit :: [Integer]                     stride' = map fromIntegral stride :: [Integer]-                    high' = map fromIntegral (zipWith (-) xShape (zipWith (+) start' (zipWith (*) (zipWith (-) limit' start') stride'))) :: [Int64]+                    -- Number of elements in the sliced bar per dimension:+                    -- n = ceil((limit - start) / stride)+                    n = zipWith3 (\l s st -> (l - s + st - 1) `div` st) limit' start' stride' :: [Integer]+                    -- Padded size = start + n * stride - stride + 1 + high = xShape+                    -- => high = xShape - start - n * stride + stride - 1+                    high' = zipWith3 (\x (s, st) n_ -> fromIntegral (x - s - n_ * st + st - 1)) xShape (zip start' stride') n :: [Int64]                     interior = map (\s -> max 0 (s - 1)) stride :: [Int64]                     zeroType = TensorType [] (ttDType xType)                 zero <- bconstant zeroType 0.0@@ -470,7 +475,7 @@     splitAndAccumulate _ _ _ [] [] [] acc = return acc     splitAndAccumulate bar dim offset (vid:vids) (itype:itypes) (sz:szs) acc = do         let shape = ttShape itype-            start = replicate (length shape) (0 :: Integer)+            start = zipWith (\i _ -> if i == dim then offset else 0) [0..] shape             limit = zipWith (\i s -> if i == dim then offset + s else s) [0..] shape             stride = replicate (length shape) (1 :: Integer)         piece <- bslice bar (map fromIntegral start) (map fromIntegral limit) (map fromIntegral stride) itype@@ -490,7 +495,7 @@         (low, _high, interior) <- parsePadAttrs (opAttributes op)         let xShape = ttShape xType             stride = map (+ 1) interior :: [Int64]-            limit = zipWith3 (\l s sz -> l + sz * s) low stride (map fromIntegral xShape) :: [Int64]+            limit = zipWith3 (\l s sz -> l + (sz - 1) * s + 1) low stride (map fromIntegral xShape) :: [Int64]         dx <- bslice bar low limit stride xType         accumulate cmap xVid dx     Nothing -> return cmap@@ -784,7 +789,24 @@     -- Window attributes only apply to spatial dims (first 2 of kernel shape).     let spatialKernelShape = take 2 (ttShape kernelType)         spatialReversePad = reversePad pad stride spatialKernelShape-        windowStr = buildWindowString [1, 1] spatialReversePad stride [1, 1]+        -- When stride == 1, the backward is a regular conv (lhs_dilate=1).+        -- When stride > 1, the backward is a transpose conv; we must adjust+        -- padding because XLA forward conv uses floor division.+        adjustedPad = if all (== 1) stride+            then spatialReversePad+            else+                let barShape = ttShape (btType bar)+                    inputShape = ttShape inputType+                    spatialBar = take 2 (tail barShape)+                    spatialInput = take 2 (tail inputShape)+                    strideInt = map fromIntegral stride :: [Integer]+                    targetPadTotal = zipWith4 (\inp bar_ stride_ k -> inp - (bar_ - 1) * stride_ + k - 2) (map fromIntegral spatialInput :: [Integer]) (map fromIntegral spatialBar :: [Integer]) strideInt (map fromIntegral spatialKernelShape :: [Integer])+                    actualPadTotal = map (fromIntegral . sum) spatialReversePad :: [Integer]+                    padDiffs = map fromIntegral (zipWith (-) targetPadTotal actualPadTotal) :: [Int64]+                in zipWith (\[l, h] d ->+                    let h' = h + d+                    in if h' >= 0 then [l, h'] else [l + h', 0]) spatialReversePad padDiffs+        windowStr = buildWindowString [1, 1] adjustedPad stride [1, 1]     bconvolution bar flippedKernel "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] inputType  convBackwardKernel :: BTensor -> TensorType -> BTensor -> [Int64] -> [[Int64]] -> Builder BTensor@@ -812,7 +834,10 @@ transposeConvBackwardInput :: BTensor -> BTensor -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor transposeConvBackwardInput bar kernel inputType lhsDilate pad = do     -- Backward input: conv(dy, kernel) with stride = lhs_dilate.-    let windowStr = buildWindowString lhsDilate pad [1, 1] [1, 1]+    -- Padding must be reversed for the backward regular conv.+    let spatialKernelShape = take 2 (ttShape (btType kernel))+        spatialReversePad = reversePad pad lhsDilate spatialKernelShape+        windowStr = buildWindowString lhsDilate spatialReversePad [1, 1] [1, 1]     bconvolution bar kernel "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] inputType  transposeConvBackwardKernel :: BTensor -> TensorType -> BTensor -> [Int64] -> [[Int64]] -> Builder BTensor

@@ -147,12 +147,16 @@     , pack2     , pack3     , slice1+    -- * Custom calls+    , customCall1+    , customCall2+    , customCallRaw     ) where  import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh, sqrt, sin, cos, tan, floor, ceiling)  import Control.Monad (when)-import Data.Int (Int64)+import Data.Int (Int64, Int32) import Data.List (elemIndex) import Data.Maybe (fromJust) import Data.Proxy@@ -2116,3 +2120,59 @@                     (left, ',':right) -> (left, right, outRest)                     _ -> error "einsum: expected two operands separated by comma"             _ -> error "einsum: expected -> in subscript string"+++-- ---------------------------------------------------------------------------+-- Custom calls+-- ---------------------------------------------------------------------------++-- | Single-result custom call. All inputs share the same shape / dtype.+--+-- The target name is the C symbol that XLA will resolve via @dlsym@.+-- Use 'customCallRaw' for heterogeneous input types or multiple results.+customCall1 :: forall s d. (KnownShape s, KnownDType d)+            => Text              -- ^ target symbol name+            -> [Tensor s d]      -- ^ inputs+            -> Text              -- ^ backend_config opaque payload+            -> Bool              -- ^ has_side_effect+            -> Builder (Tensor s d)+customCall1 target inputs backendConfig hasSideEffect = do+    let vids    = tensorValue <$> inputs+        inType  = tensorType (Proxy @s) (Proxy @d)+        outType = tensorType (Proxy @s) (Proxy @d)+    vidRes <- emitCustomCall target vids (replicate (length inputs) inType) backendConfig hasSideEffect 3 [outType]+    case vidRes of+        [vid] -> return (Tensor vid)+        _     -> error "customCall1: expected exactly one result"++-- | Two-result custom call.+customCall2 :: forall s1 d1 s2 d2. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+            => Text+            -> [Tensor s1 d1]    -- ^ inputs (uniform type for convenience)+            -> Text              -- ^ backend_config+            -> Bool              -- ^ has_side_effect+            -> Builder (Tensor s1 d1, Tensor s2 d2)+customCall2 target inputs backendConfig hasSideEffect = do+    let vids     = tensorValue <$> inputs+        inType   = tensorType (Proxy @s1) (Proxy @d1)+        outType1 = tensorType (Proxy @s1) (Proxy @d1)+        outType2 = tensorType (Proxy @s2) (Proxy @d2)+    vidsRes <- emitCustomCall target vids (replicate (length inputs) inType) backendConfig hasSideEffect 3 [outType1, outType2]+    case vidsRes of+        [v1, v2] -> return (Tensor v1, Tensor v2)+        _        -> error "customCall2: expected exactly two results"++-- | Low-level custom call for plugin authors.+--+-- Accepts raw 'ValueId's and 'TensorType's so that callers can mix shapes+-- and dtypes freely (e.g. sparse indices as 'I32' and values as 'F32').+-- The caller is responsible for ensuring the C kernel signature matches.+customCallRaw :: Text+              -> [ValueId]         -- ^ operand value ids+              -> [TensorType]      -- ^ operand types+              -> Text              -- ^ backend_config+              -> Bool              -- ^ has_side_effect+              -> Int32             -- ^ api_version+              -> [TensorType]      -- ^ result types+              -> Builder [ValueId]+customCallRaw = emitCustomCall

@@ -26,6 +26,7 @@     , emitOpN     , emitOpRegions     , emitOpRegionsN+    , emitCustomCall     , emitReduce     , emitReturn     , runBlockBuilder@@ -45,7 +46,9 @@     ) where  import Control.Monad.State+import Data.Int (Int32) import Data.Proxy+import qualified Data.Text as T import Data.Text (Text) import qualified Data.Text as T import GHC.TypeLits@@ -308,6 +311,28 @@ moduleFromBuilderT name args' action =     let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'     in Module [runBuilderT name renamed action]++-- | Emit a 'stablehlo.custom_call' operation.+--+-- The target name is the C symbol that XLA will look up via @dlsym@.+-- 'api_version' selects the ABI: @1@ for GPU (CUstream, void** buffers,+-- opaque, len); other values for CPU ABI.+emitCustomCall :: Text          -- ^ call_target_name+               -> [ValueId]     -- ^ operands+               -> [TensorType]  -- ^ operand types+               -> Text          -- ^ backend_config (opaque payload)+               -> Bool          -- ^ has_side_effect+               -> Int32         -- ^ api_version+               -> [TensorType]  -- ^ result types+               -> Builder [ValueId]+emitCustomCall target operands operandTypes backendConfig hasSideEffect apiVersion resultTypes =+    emitOpN "stablehlo.custom_call" operands operandTypes+        [ AttrString "call_target_name" target+        , AttrBool   "has_side_effect"  hasSideEffect+        , AttrString "backend_config"   backendConfig+        , AttrRaw    $ "api_version = " <> T.pack (show apiVersion) <> " : i32"+        ]+        resultTypes  -- | Emit a generic single-result operation into the builder. -- The caller must provide the operand types so that the pretty-printer

@@ -216,6 +216,17 @@         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> (if null attrs then mempty else " " <> prettyAttrs attrs)         <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.custom_call" operands operandTypes attrs _regions results resultTypes) =+        -- Custom form: @symbol prefix before operands, attribute dict after.+        -- Example:+        --   %0 = stablehlo.custom_call @foo(%arg0, %arg1) {call_target_name = "foo", ...}+        --          : (tensor<2xf32>) -> tensor<2xf32>+        let target = lookupAttrString "call_target_name" attrs+        in prettyResultVids results <> " = stablehlo.custom_call"+           <> (if T.null target then mempty else " @" <> fromText target)+           <> "(" <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+           <> (if null attrs then mempty else " " <> prettyAttrs attrs)+           <> " : " <> prettyResultType operandTypes resultTypes     pretty (Operation name operands operandTypes attrs regions results resultTypes) =         if null regions         then

@@ -0,0 +1,95 @@+{-# LANGUAGE ForeignFunctionInterface #-}++-- | Runtime support for XLA custom-call symbol loading.+--+-- Custom-call kernels live in separate shared libraries (e.g. @libfoo.so@).+-- Before compiling or executing any HHLO module that references a custom+-- target, the library must be loaded and its target registered with the+-- runtime.+--+-- For CPU plugins, 'loadCustomCallLibrary' promotes symbols to the global+-- namespace so that XLA can resolve them via @dlsym(RTLD_DEFAULT, ...)@.+--+-- For GPU plugins, 'registerGpuCustomCall' uses the PJRT GPU custom-call+-- extension to register the target directly with the PJRT CUDA plugin.+--+-- Typical usage in application code:+--+-- @+-- main = withGPU $ \\sess -> do+--     registerGpuCustomCall (sessionApi sess) "lib/libmyplugin.so" "my_kernel"+--     let modu = buildMyModule   -- uses 'customCall1' inside+--     compiled <- compile sess modu+--     ...+-- @+module HHLO.Runtime.CustomCall+    ( loadCustomCallLibrary+    , registerGpuCustomCall+    ) where++import Data.Bits ((.|.))+import Foreign.C.String (CString, withCString, peekCString)+import Foreign.C.Types (CInt(..))+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr (Ptr, nullPtr)+import Foreign.Storable (peek)+import System.IO.Error (ioeSetLocation, mkIOError, userErrorType)++import HHLO.Runtime.PJRT.Types (PJRTApi(..))+import HHLO.Runtime.PJRT.FFI (c_pjrtRegisterGpuCustomCall)++-- RTLD_NOW   = 0x00002+-- RTLD_GLOBAL = 0x00100+-- Combined   = 0x00102 = 258+flagRTLD_NOW, flagRTLD_GLOBAL, flagCombined :: CInt+flagRTLD_NOW    = 2+flagRTLD_GLOBAL = 256+flagCombined    = flagRTLD_NOW .|. flagRTLD_GLOBAL++-- | Open a dynamic library and promote its symbols to the global namespace.+--+-- This is a thin wrapper around @dlopen(path, RTLD_NOW | RTLD_GLOBAL)@.+-- It is sufficient for CPU custom calls, where XLA resolves symbols via+-- @dlsym(RTLD_DEFAULT, ...)@.+--+-- For GPU custom calls, use 'registerGpuCustomCall' instead.+loadCustomCallLibrary :: FilePath -> IO ()+loadCustomCallLibrary path = do+    handle <- withCString path $ \cstr ->+        c_dlopen cstr flagCombined+    if handle /= nullPtr+        then return ()+        else ioError $ ioeSetLocation+             (mkIOError userErrorType+                 ("cannot load custom-call library: " ++ path)+                 Nothing Nothing)+             "HHLO.Runtime.CustomCall.loadCustomCallLibrary"++foreign import ccall "dlopen"+    c_dlopen :: CString -> CInt -> IO (Ptr ())++-- | Register a GPU custom-call target with the PJRT CUDA plugin.+--+-- This function:+-- 1. Opens the shared library at @libPath@.+-- 2. Looks up the symbol @targetName@.+-- 3. Registers it with the PJRT GPU custom-call extension.+--+-- The library handle is intentionally kept open for the process lifetime.+--+-- Must be called before 'compile'.+registerGpuCustomCall :: PJRTApi -> FilePath -> String -> IO ()+registerGpuCustomCall (PJRTApi apiPtr) libPath targetName =+    withCString libPath $ \cLibPath ->+    withCString targetName $ \cTargetName ->+    alloca $ \errMsgPtr -> do+        rc <- c_pjrtRegisterGpuCustomCall apiPtr cLibPath cTargetName errMsgPtr+        if rc == 0+            then return ()+            else do+                errMsg <- peek errMsgPtr >>= peekCString+                ioError $ ioeSetLocation+                    (mkIOError userErrorType+                        ("registerGpuCustomCall failed: " ++ errMsg)+                        Nothing Nothing)+                    "HHLO.Runtime.CustomCall.registerGpuCustomCall"

@@ -189,3 +189,14 @@  foreign import ccall "pjrt_shim.h hhlo_pjrt_error_destroy"     c_pjrtErrorDestroy :: Ptr PJRTApi -> Ptr PJRTError -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Custom calls+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_register_gpu_custom_call"+    c_pjrtRegisterGpuCustomCall :: Ptr PJRTApi+                                -> CString            -- lib_path+                                -> CString            -- function_name+                                -> Ptr CString        -- out_error_msg+                                -> IO CInt

@@ -19,7 +19,7 @@ -- >     print (hostToList result) module HHLO.Session     ( -- * Session lifecycle-      Session+      Session(..)     , withCPU     , withGPU     , withGPUDevice

@@ -62,6 +62,34 @@             let rendered = render op             assertBool "should be generic form" $                 "\"stablehlo.return\"" `T.isInfixOf` rendered+        , testCase "custom_call with @symbol" $ do+            let op = Operation "stablehlo.custom_call" [ValueId 0, ValueId 1]+                    [TensorType [4] F32, TensorType [4] F32]+                    [ AttrString "call_target_name" "vector_add"+                    , AttrBool   "has_side_effect"  False+                    , AttrString "backend_config"   ""+                    , AttrRaw    "api_version = 3 : i32"+                    ] [] [ValueId 2] [TensorType [4] F32]+            let rendered = render op+            assertBool "should contain @symbol prefix" $+                "stablehlo.custom_call @vector_add(" `T.isInfixOf` rendered+            assertBool "should contain call_target_name attr" $+                "call_target_name = \"vector_add\"" `T.isInfixOf` rendered+            assertBool "should contain has_side_effect attr" $+                "has_side_effect = false" `T.isInfixOf` rendered+            assertBool "should contain api_version attr" $+                "api_version = 3 : i32" `T.isInfixOf` rendered+            assertBool "should end with function type" $+                ": (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>" `T.isSuffixOf` rendered+        , testCase "custom_call without operands" $ do+            let op = Operation "stablehlo.custom_call" []+                    []+                    [ AttrString "call_target_name" "rng_seed"+                    , AttrBool   "has_side_effect"  False+                    ] [] [ValueId 0] [TensorType [] I64]+            let rendered = render op+            assertBool "should contain @symbol with empty parens" $+                "stablehlo.custom_call @rng_seed()" `T.isInfixOf` rendered         ]     , testGroup "Constants"         [ testCase "scalar constant" $ do

@@ -114,6 +114,45 @@         let expected = V.fromList [1, 2, 1, 2, 4, 2, 1, 2, 1]         assertBool "conv2d grad close" $             V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad slice stride" $ withPJRTCPU $ \api client -> do+        let f x = do+                y <- slice @'[5] @'[3] x (v1 0) (v1 5) (v1 2)+                sumAll y+            gradModu = gradModule @'[5] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0]+        bufIn <- toDeviceF32 api client inp [5]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 5+        -- grad of sum(slice(x, [0:5:2])) = [1, 0, 1, 0, 1]+        let expected = V.fromList [1.0, 0.0, 1.0, 0.0, 1.0]+        result @?= expected+    , testCase "grad conv2d stride" $ withPJRTCPU $ \api client -> do+        let f x = do+                k <- constant @'[3, 3, 1, 1] @'F32 1.0+                y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @2 @2 (v2 2 2) (p2 (1,1) (1,1)) x k+                sumAll y+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 16+        -- grad should be non-zero and finite (shape validation is the main goal)+        assertBool "grad non-zero" $ V.sum result > 0+    , testCase "grad transposeConvolution asymmetric pad" $ withPJRTCPU $ \api client -> do+        let f x = do+                k <- constant @'[3, 3, 1, 1] @'F32 1.0+                y <- transposeConvolution @1 @2 @2 @1 @1 @3 @3 @2 @2 (v2 2 2) (p2 (1,0) (1,0)) x k+                sumAll y+            gradModu = gradModule @'[1, 2, 2, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32 api client inp [1, 2, 2, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- grad should be non-zero and finite (shape validation is the main goal)+        assertBool "grad non-zero" $ V.sum result > 0     , testCase "grad maxPool" $ withPJRTCPU $ \api client -> do         let f x = do                 let kernel = v2 2 2@@ -135,6 +174,64 @@         let expected = V.fromList [0,0,0,0, 0,1,0,1, 0,0,0,0, 0,1,0,1]         assertBool "maxPool grad close" $             V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad pad interior" $ withPJRTCPU $ \api client -> do+        let f x = do+                padVal <- constant @'[] @'F32 0.0+                y <- pad @'[2] @'[3] x padVal (v1 0) (v1 0) (v1 1)+                sumAll y+            gradModu = gradModule @'[2] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32 api client inp [2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        -- grad of sumAll(pad(x, interior=1)) = [1, 1]+        let expected = V.fromList [1.0, 1.0]+        result @?= expected+    , testCase "grad concatenate2" $ withPJRTCPU $ \api client -> do+        let f a b = do+                c <- concatenate2 @'[2, 4] @'[2, 4] @'[2, 8] @'F32 1 a b+                sumAll c+            modu = moduleFromBuilder @'[4, 4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                ] $ do+                    a <- arg @'[2, 4] @'F32+                    b <- arg @'[2, 4] @'F32+                    (da, db) <- grad2 f a b+                    concatenate 0 [da, db]+        exec <- compile api client (render modu)+        let inp1 = V.fromList [1..8]+            inp2 = V.fromList [1..8]+        bufIn1 <- toDeviceF32 api client inp1 [2, 4]+        bufIn2 <- toDeviceF32 api client inp2 [2, 4]+        [bufOut] <- execute api exec [bufIn1, bufIn2]+        result <- fromDeviceF32 api bufOut 16+        let expected = V.fromList (replicate 16 1.0)+        result @?= expected+    , testCase "grad concatenate3" $ withPJRTCPU $ \api client -> do+        let f a b c = do+                x <- concatenate @'[2, 4] @'[2, 12] @'F32 1 [a, b, c]+                sumAll x+            modu = moduleFromBuilder @'[6, 4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                ] $ do+                    a <- arg @'[2, 4] @'F32+                    b <- arg @'[2, 4] @'F32+                    c <- arg @'[2, 4] @'F32+                    (da, db, dc) <- grad3 f a b c+                    concatenate 0 [da, db, dc]+        exec <- compile api client (render modu)+        let inp = V.fromList [1..8]+        bufIn1 <- toDeviceF32 api client inp [2, 4]+        bufIn2 <- toDeviceF32 api client inp [2, 4]+        bufIn3 <- toDeviceF32 api client inp [2, 4]+        [bufOut] <- execute api exec [bufIn1, bufIn2, bufIn3]+        result <- fromDeviceF32 api bufOut 24+        let expected = V.fromList (replicate 24 1.0)+        result @?= expected     , testCase "grad2 multiply" $ withPJRTCPU $ \api client -> do         let f x y = do z <- multiply x y; sumAll z             modu = moduleFromBuilder @'[4] @'F32 "main"

@@ -113,6 +113,45 @@         let expected = V.fromList [1, 2, 1, 2, 4, 2, 1, 2, 1]         assertBool "conv2d grad close" $             V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad slice stride" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                y <- slice @'[5] @'[3] x (v1 0) (v1 5) (v1 2)+                sumAll y+            gradModu = gradModule @'[5] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0]+        bufIn <- toDeviceF32On api client dev inp [5]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 5+        let expected = V.fromList [1.0, 0.0, 1.0, 0.0, 1.0]+        result @?= expected+    , testCase "grad conv2d stride" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                k <- constant @'[3, 3, 1, 1] @'F32 1.0+                y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @2 @2 (v2 2 2) (p2 (1,1) (1,1)) x k+                sumAll y+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 16+        assertBool "grad non-zero" $ V.sum result > 0+    , testCase "grad transposeConvolution asymmetric pad" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                k <- constant @'[3, 3, 1, 1] @'F32 1.0+                y <- transposeConvolution @1 @2 @2 @1 @1 @3 @3 @2 @2 (v2 2 2) (p2 (1,0) (1,0)) x k+                sumAll y+            gradModu = gradModule @'[1, 2, 2, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [1, 2, 2, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        assertBool "grad non-zero" $ V.sum result > 0     , testCase "grad maxPool" $ do         GPUResource api client dev <- getGPU         let f x = do@@ -130,6 +169,66 @@         let expected = V.fromList [0,0,0,0, 0,1,0,1, 0,0,0,0, 0,1,0,1]         assertBool "maxPool grad close" $             V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad pad interior" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                padVal <- constant @'[] @'F32 0.0+                y <- pad @'[2] @'[3] x padVal (v1 0) (v1 0) (v1 1)+                sumAll y+            gradModu = gradModule @'[2] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32On api client dev inp [2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        let expected = V.fromList [1.0, 1.0]+        result @?= expected+    , testCase "grad concatenate2" $ do+        GPUResource api client dev <- getGPU+        let f a b = do+                c <- concatenate2 @'[2, 4] @'[2, 4] @'[2, 8] @'F32 1 a b+                sumAll c+            modu = moduleFromBuilder @'[4, 4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                ] $ do+                    a <- arg @'[2, 4] @'F32+                    b <- arg @'[2, 4] @'F32+                    (da, db) <- grad2 f a b+                    concatenate 0 [da, db]+        exec <- compile api client (render modu)+        let inp1 = V.fromList [1..8]+            inp2 = V.fromList [1..8]+        bufIn1 <- toDeviceF32On api client dev inp1 [2, 4]+        bufIn2 <- toDeviceF32On api client dev inp2 [2, 4]+        [bufOut] <- executeOn api exec dev [bufIn1, bufIn2]+        result <- fromDeviceF32 api bufOut 16+        let expected = V.fromList (replicate 16 1.0)+        result @?= expected+    , testCase "grad concatenate3" $ do+        GPUResource api client dev <- getGPU+        let f a b c = do+                x <- concatenate @'[2, 4] @'[2, 12] @'F32 1 [a, b, c]+                sumAll x+            modu = moduleFromBuilder @'[6, 4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                ] $ do+                    a <- arg @'[2, 4] @'F32+                    b <- arg @'[2, 4] @'F32+                    c <- arg @'[2, 4] @'F32+                    (da, db, dc) <- grad3 f a b c+                    concatenate 0 [da, db, dc]+        exec <- compile api client (render modu)+        let inp = V.fromList [1..8]+        bufIn1 <- toDeviceF32On api client dev inp [2, 4]+        bufIn2 <- toDeviceF32On api client dev inp [2, 4]+        bufIn3 <- toDeviceF32On api client dev inp [2, 4]+        [bufOut] <- executeOn api exec dev [bufIn1, bufIn2, bufIn3]+        result <- fromDeviceF32 api bufOut 24+        let expected = V.fromList (replicate 24 1.0)+        result @?= expected     , testCase "grad2 multiply" $ do         GPUResource api client dev <- getGPU         let f x y = do z <- multiply x y; sumAll z

@@ -36,6 +36,19 @@         [bufOut] <- execute api exec [bufIn]         result <- fromDeviceF32 api bufOut 3         result @?= V.fromList [1.0, 2.0, 3.0]+    , testCase "slice 2D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [4, 4] F32) ]+                $ do+                    x <- arg @'[4, 4] @'F32+                    y <- slice @'[4, 4] @'[2, 2] x (v2 1 1) (v2 3 3) (v2 1 1)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..16]+        bufIn <- toDeviceF32 api client inp [4, 4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [6, 7, 10, 11]     , testCase "slice with stride" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[2] @'F32 "main"                 [ FuncArg "arg0" (TensorType [5] F32) ]@@ -49,6 +62,20 @@         [bufOut] <- execute api exec [bufIn]         result <- fromDeviceF32 api bufOut 2         result @?= V.fromList [0.0, 2.0]+    , testCase "pad 2D symmetric" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[4, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg @'[2, 2] @'F32+                    padVal <- constant @'[] @'F32 0.0+                    y <- pad @'[2, 2] @'[4, 4] x padVal (v2 1 1) (v2 1 1) (v2 0 0)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1, 2, 3, 4]+        bufIn <- toDeviceF32 api client inp [2, 2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 16+        result @?= V.fromList [0,0,0,0, 0,1,2,0, 0,3,4,0, 0,0,0,0]     , testCase "pad edge" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[4] @'F32 "main"                 [ FuncArg "arg0" (TensorType [2] F32) ]

@@ -38,6 +38,20 @@         [bufOut] <- executeOn api exec dev [bufIn]         result <- fromDeviceF32 api bufOut 3         result @?= V.fromList [1.0, 2.0, 3.0]+    , testCase "slice 2D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [4, 4] F32) ]+                $ do+                    x <- arg @'[4, 4] @'F32+                    y <- slice @'[4, 4] @'[2, 2] x (v2 1 1) (v2 3 3) (v2 1 1)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..16]+        bufIn <- toDeviceF32On api client dev inp [4, 4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [6, 7, 10, 11]     , testCase "slice with stride" $ do         GPUResource api client dev <- getGPU         let modu = moduleFromBuilder @'[2] @'F32 "main"@@ -52,6 +66,21 @@         [bufOut] <- executeOn api exec dev [bufIn]         result <- fromDeviceF32 api bufOut 2         result @?= V.fromList [0.0, 2.0]+    , testCase "pad 2D symmetric" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg @'[2, 2] @'F32+                    padVal <- constant @'[] @'F32 0.0+                    y <- pad @'[2, 2] @'[4, 4] x padVal (v2 1 1) (v2 1 1) (v2 0 0)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1, 2, 3, 4]+        bufIn <- toDeviceF32On api client dev inp [2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 16+        result @?= V.fromList [0,0,0,0, 0,1,2,0, 0,3,4,0, 0,0,0,0]     , testCase "pad edge" $ do         GPUResource api client dev <- getGPU         let modu = moduleFromBuilder @'[4] @'F32 "main"

@@ -76,6 +76,27 @@         -- Row 0: [3.0+0.1, 3.0+0.1] = [3.1, 3.1]         -- Row 1: [7.5+0.1, 7.5+0.1] = [7.6, 7.6]         result @?= V.fromList [3.1, 3.1, 7.6, 7.6]+    , testCase "dotGeneral batched" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2, 3] F32)+                , FuncArg "arg1" (TensorType [2, 3, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 2, 3] @'F32+                    y <- arg @'[2, 3, 2] @'F32+                    z <- dotGeneral @'[2, 2, 3] @'[2, 3, 2] @'[2, 2, 2] @'F32 (v1 0) (v1 0) (v1 2) (v1 1) x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1,2,3, 4,5,6, 1,2,3, 4,5,6] :: V.Vector Float+            b = V.fromList [1,2, 3,4, 5,6, 1,2, 3,4, 5,6] :: V.Vector Float+        bufA <- toDeviceF32 api client a [2, 2, 3]+        bufB <- toDeviceF32 api client b [2, 3, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 8+        -- Batch 0: [[1,2,3],[4,5,6]] . [[1,2],[3,4],[5,6]] = [[22,28],[49,64]]+        -- Batch 1: same+        let expected = V.fromList [22,28,49,64, 22,28,49,64]+        result @?= expected     , testCase "dotGeneral 3D x 2D" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[1, 2, 2] @'F32 "main"                 [ FuncArg "arg0" (TensorType [1, 2, 3] F32)

@@ -76,6 +76,26 @@         [bufOut] <- executeOn api exec dev [bufIn]         result <- fromDeviceF32 api bufOut 4         result @?= V.fromList [3.1, 3.1, 7.6, 7.6]+    , testCase "dotGeneral batched" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2, 3] F32)+                , FuncArg "arg1" (TensorType [2, 3, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 2, 3] @'F32+                    y <- arg @'[2, 3, 2] @'F32+                    z <- dotGeneral @'[2, 2, 3] @'[2, 3, 2] @'[2, 2, 2] @'F32 (v1 0) (v1 0) (v1 2) (v1 1) x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1,2,3, 4,5,6, 1,2,3, 4,5,6] :: V.Vector Float+            b = V.fromList [1,2, 3,4, 5,6, 1,2, 3,4, 5,6] :: V.Vector Float+        bufA <- toDeviceF32On api client dev a [2, 2, 3]+        bufB <- toDeviceF32On api client dev b [2, 3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 8+        let expected = V.fromList [22,28,49,64, 22,28,49,64]+        result @?= expected     , testCase "dotGeneral 3D x 2D" $ do         GPUResource api client dev <- getGPU         let modu = moduleFromBuilder @'[1, 2, 2] @'F32 "main"

@@ -119,6 +119,40 @@         let row1 = V.slice 4 4 result         assertBool "layerNorm of uniform row has near-zero mean" $             abs (V.sum row1 / 4) < 0.1+    , testCase "conv2dWithPadding forward" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 3, 3, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 2, 1] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    y <- conv2dWithPadding @1 @2 @2 @1 @1 @2 @2 @3 @3 (v2 1 1) (p2 (1,1) (1,1)) x k+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32 api client inp [1, 2, 2, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 9+        -- With zero padding and 2x2 kernel all 1s:+        -- [0,0,0]    [1,3,2]+        -- [0,1,2] -> [4,10,6]+        -- [0,3,4]    [3,7,4]+        let expected = V.fromList [1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0]+        result @?= expected+    , testCase "transposeConvolution forward" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 2, 1] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x k+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32 api client inp [1, 2, 2, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- Output shape is [1,2,2,1]; verify non-zero and finite+        assertBool "transposeConv output non-zero" $ V.sum result > 0     , testCase "globalAvgPool" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[1, 2] @'F32 "main"                 [ FuncArg "arg0" (TensorType [1, 4, 4, 2] F32) ]

@@ -122,6 +122,37 @@         let row1 = V.slice 4 4 result         assertBool "layerNorm of uniform row has near-zero mean" $             abs (V.sum row1 / 4) < 0.1+    , testCase "conv2dWithPadding forward" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 3, 3, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 2, 1] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    y <- conv2dWithPadding @1 @2 @2 @1 @1 @2 @2 @3 @3 (v2 1 1) (p2 (1,1) (1,1)) x k+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [1, 2, 2, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 9+        let expected = V.fromList [1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0]+        result @?= expected+    , testCase "transposeConvolution forward" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 2, 1] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x k+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [1, 2, 2, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        assertBool "transposeConv output non-zero" $ V.sum result > 0     , testCase "globalAvgPool" $ do         GPUResource api client dev <- getGPU         let modu = moduleFromBuilder @'[1, 2] @'F32 "main"

@@ -47,6 +47,25 @@         [bufOut] <- execute api exec [bufIn]         result <- fromDeviceF32 api bufOut 4         result @?= V.fromList [1.0, 3.0, 2.0, 4.0]+    , testCase "transpose 3D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 4, 3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3, 4] F32) ]+                $ do+                    x <- arg @'[2, 3, 4] @'F32+                    y <- transpose @'[2, 3, 4] @'[2, 4, 3] (v3 0 2 1) x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..24]+        bufIn <- toDeviceF32 api client inp [2, 3, 4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 24+        -- Batch 0: transpose last two dims of [[1..4],[5..8],[9..12]]+        -- = [[1,5,9],[2,6,10],[3,7,11],[4,8,12]]+        -- Batch 1: same with 13..24+        let expected = V.fromList+                [1,5,9, 2,6,10, 3,7,11, 4,8,12,+                 13,17,21, 14,18,22, 15,19,23, 16,20,24]+        result @?= expected     , testCase "transpose identity" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[2, 2] @'F32 "main"                 [ FuncArg "arg0" (TensorType [2, 2] F32) ]

@@ -51,6 +51,23 @@         [bufOut] <- executeOn api exec dev [bufIn]         result <- fromDeviceF32 api bufOut 4         result @?= V.fromList [1.0, 3.0, 2.0, 4.0]+    , testCase "transpose 3D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 4, 3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3, 4] F32) ]+                $ do+                    x <- arg @'[2, 3, 4] @'F32+                    y <- transpose @'[2, 3, 4] @'[2, 4, 3] (v3 0 2 1) x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..24]+        bufIn <- toDeviceF32On api client dev inp [2, 3, 4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 24+        let expected = V.fromList+                [1,5,9, 2,6,10, 3,7,11, 4,8,12,+                 13,17,21, 14,18,22, 15,19,23, 16,20,24]+        result @?= expected     , testCase "transpose identity" $ do         GPUResource api client dev <- getGPU         let modu = moduleFromBuilder @'[2, 2] @'F32 "main"