1  {-# LANGUAGE ForeignFunctionInterface #-}
  2  {-# LANGUAGE OverloadedStrings #-}
  3  
  4  {- | WASM plugin for typed package definitions.
  5  
  6  This module exports package definitions that can be called from Nix
  7  via @builtins.wasm@.
  8  
  9  = Usage
 10  
 11  @
 12  # From Nix (using straylight-nix with builtins.wasm)
 13  let
 14    wasm = builtins.wasm ./straylight-packages.wasm;
 15  in {
 16    zlib-ng = wasm "zlib_ng" {};
 17    fmt = wasm "fmt" {};
 18  }
 19  @
 20  
 21  = Building
 22  
 23  @
 24  wasm32-wasi-ghc -optl-mexec-model=reactor \
 25    -package text -package containers \
 26    Main.hs -o straylight-packages.wasm
 27  @
 28  
 29  NOTE: We use 'Main' module name because GHC WASM reactor modules need
 30  proper RTS initialization, which only happens when there's a Main module.
 31  The 'main' function is a dummy that's never called (reactor modules don't
 32  have a _start entry point).
 33  -}
 34  module Main where
 35  
 36  import Aleph.Nix
 37  import Aleph.Nix.Derivation
 38  import Aleph.Nix.Packages.AbseilCpp (abseilCpp)
 39  import Aleph.Nix.Packages.Catch2 (catch2)
 40  import Aleph.Nix.Packages.Cutlass (cutlass)
 41  import Aleph.Nix.Packages.Fmt (fmt)
 42  import Aleph.Nix.Packages.HelloWrapped (helloWrapped)
 43  import Aleph.Nix.Packages.Jq (jq)
 44  import Aleph.Nix.Packages.Mdspan (mdspan)
 45  import Aleph.Nix.Packages.NlohmannJson (nlohmannJson)
 46  import Aleph.Nix.Packages.Nvidia (cudnn, cusparselt, cutensor, nccl, tensorrt)
 47  import qualified Aleph.Nix.Packages.Nvidia as Nvidia (cutlass)
 48  import Aleph.Nix.Packages.Rapidjson (rapidjson)
 49  import Aleph.Nix.Packages.Spdlog (spdlog)
 50  import Aleph.Nix.Packages.ZlibNg (zlibNg)
 51  
 52  {- | Dummy main for GHC to link RTS properly.
 53  This is never called in reactor mode.
 54  -}
 55  main :: IO ()
 56  main = pure ()
 57  
 58  -- | Required initialization function.
 59  foreign export ccall "nix_wasm_init_v1" initPlugin :: IO ()
 60  
 61  initPlugin :: IO ()
 62  initPlugin = nixWasmInit
 63  
 64  {- | Export zlib-ng package specification.
 65  
 66  Called from Nix as: @wasm "zlib_ng" {}@
 67  
 68  Returns an attrset with the full derivation specification.
 69  -}
 70  foreign export ccall "zlib_ng" zlibNgExport :: Value -> IO Value
 71  
 72  zlibNgExport :: Value -> IO Value
 73  zlibNgExport _args = drvToNixAttrs zlibNg
 74  
 75  {- | Export fmt package specification.
 76  
 77  Called from Nix as: @wasm "fmt" {}@
 78  
 79  Returns an attrset with the full derivation specification.
 80  -}
 81  foreign export ccall "fmt" fmtExport :: Value -> IO Value
 82  
 83  fmtExport :: Value -> IO Value
 84  fmtExport _args = drvToNixAttrs fmt
 85  
 86  {- | Export mdspan package specification.
 87  
 88  Called from Nix as: @wasm "mdspan" {}@
 89  
 90  Returns an attrset with the full derivation specification.
 91  Includes typed postInstall phase to write C++23 shim header.
 92  -}
 93  foreign export ccall "mdspan" mdspanExport :: Value -> IO Value
 94  
 95  mdspanExport :: Value -> IO Value
 96  mdspanExport _args = drvToNixAttrs mdspan
 97  
 98  {- | Export CUTLASS package specification.
 99  
100  Called from Nix as: @wasm "cutlass" {}@
101  
102  Returns an attrset with the full derivation specification.
103  Header-only NVIDIA CUDA template library.
104  -}
105  foreign export ccall "cutlass" cutlassExport :: Value -> IO Value
106  
107  cutlassExport :: Value -> IO Value
108  cutlassExport _args = drvToNixAttrs cutlass
109  
110  {- | Export RapidJSON package specification.
111  
112  Called from Nix as: @wasm "rapidjson" {}@
113  -}
114  foreign export ccall "rapidjson" rapidjsonExport :: Value -> IO Value
115  
116  rapidjsonExport :: Value -> IO Value
117  rapidjsonExport _args = drvToNixAttrs rapidjson
118  
119  {- | Export nlohmann/json package specification.
120  
121  Called from Nix as: @wasm "nlohmann_json" {}@
122  
123  Header-only JSON library with pkg-config fixup.
124  -}
125  foreign export ccall "nlohmann_json" nlohmannJsonExport :: Value -> IO Value
126  
127  nlohmannJsonExport :: Value -> IO Value
128  nlohmannJsonExport _args = drvToNixAttrs nlohmannJson
129  
130  {- | Export spdlog package specification.
131  
132  Called from Nix as: @wasm "spdlog" {}@
133  
134  Super fast C++ logging library, uses external fmt.
135  -}
136  foreign export ccall "spdlog" spdlogExport :: Value -> IO Value
137  
138  spdlogExport :: Value -> IO Value
139  spdlogExport _args = drvToNixAttrs spdlog
140  
141  {- | Export Catch2 package specification.
142  
143  Called from Nix as: @wasm "catch2" {}@
144  
145  Modern C++ test framework.
146  -}
147  foreign export ccall "catch2" catch2Export :: Value -> IO Value
148  
149  catch2Export :: Value -> IO Value
150  catch2Export _args = drvToNixAttrs catch2
151  
152  {- | Export Abseil C++ package specification.
153  
154  Called from Nix as: @wasm "abseil_cpp" {}@
155  
156  Google's common C++ libraries.
157  -}
158  foreign export ccall "abseil_cpp" abseilCppExport :: Value -> IO Value
159  
160  abseilCppExport :: Value -> IO Value
161  abseilCppExport _args = drvToNixAttrs abseilCpp
162  
163  -- ============================================================================
164  -- NVIDIA SDK Packages
165  -- ============================================================================
166  
167  {- | Export NVIDIA NCCL package specification.
168  
169  Called from Nix as: @wasm "nvidia_nccl" {}@
170  
171  Multi-GPU communication library from PyPI wheel.
172  -}
173  foreign export ccall "nvidia_nccl" ncclExport :: Value -> IO Value
174  
175  ncclExport :: Value -> IO Value
176  ncclExport _args = drvToNixAttrs nccl
177  
178  {- | Export NVIDIA cuDNN package specification.
179  
180  Called from Nix as: @wasm "nvidia_cudnn" {}@
181  
182  Deep learning primitives from PyPI wheel.
183  -}
184  foreign export ccall "nvidia_cudnn" cudnnExport :: Value -> IO Value
185  
186  cudnnExport :: Value -> IO Value
187  cudnnExport _args = drvToNixAttrs cudnn
188  
189  {- | Export NVIDIA TensorRT package specification.
190  
191  Called from Nix as: @wasm "nvidia_tensorrt" {}@
192  
193  Inference optimization from PyPI wheel.
194  -}
195  foreign export ccall "nvidia_tensorrt" tensorrtExport :: Value -> IO Value
196  
197  tensorrtExport :: Value -> IO Value
198  tensorrtExport _args = drvToNixAttrs tensorrt
199  
200  {- | Export NVIDIA cuTensor package specification.
201  
202  Called from Nix as: @wasm "nvidia_cutensor" {}@
203  
204  Tensor operations from PyPI wheel.
205  -}
206  foreign export ccall "nvidia_cutensor" cutensorExport :: Value -> IO Value
207  
208  cutensorExport :: Value -> IO Value
209  cutensorExport _args = drvToNixAttrs cutensor
210  
211  {- | Export NVIDIA cuSPARSELt package specification.
212  
213  Called from Nix as: @wasm "nvidia_cusparselt" {}@
214  
215  Sparse matrix operations from PyPI wheel.
216  -}
217  foreign export ccall "nvidia_cusparselt" cusparseltExport :: Value -> IO Value
218  
219  cusparseltExport :: Value -> IO Value
220  cusparseltExport _args = drvToNixAttrs cusparselt
221  
222  {- | Export NVIDIA CUTLASS package specification (typed version).
223  
224  Called from Nix as: @wasm "nvidia_cutlass" {}@
225  
226  Header-only CUDA templates library.
227  This shadows the C++ cutlass package with proper nvidia- naming.
228  -}
229  foreign export ccall "nvidia_cutlass" nvidiaCutlassExport :: Value -> IO Value
230  
231  nvidiaCutlassExport :: Value -> IO Value
232  nvidiaCutlassExport _args = drvToNixAttrs Nvidia.cutlass
233  
234  -- ============================================================================
235  -- Test packages for typed actions
236  -- ============================================================================
237  
238  {- | Export jq package specification.
239  
240  Called from Nix as: @wasm "jq" {}@
241  
242  Demonstrates typed substitute action.
243  -}
244  foreign export ccall "jq" jqExport :: Value -> IO Value
245  
246  jqExport :: Value -> IO Value
247  jqExport _args = drvToNixAttrs jq
248  
249  {- | Export hello-wrapped package specification.
250  
251  Called from Nix as: @wasm "hello_wrapped" {}@
252  
253  Demonstrates typed wrap action.
254  -}
255  foreign export ccall "hello_wrapped" helloWrappedExport :: Value -> IO Value
256  
257  helloWrappedExport :: Value -> IO Value
258  helloWrappedExport _args = drvToNixAttrs helloWrapped