XLA backend for Lc0 (LeelaChessZero#1949)

official-pikafish · Feb 24, 2024 · 04f73fc · 04f73fc
1 parent 3b9a87e
commit 04f73fc
Show file tree

Hide file tree

Showing 16 changed files with 5,375 additions and 0 deletions.
diff --git a/meson.build b/meson.build
@@ -620,6 +620,22 @@ if get_option('build_backends')
   endif
 
 
+  ## ~~~~~~~~
+  ## XLA
+  ## ~~~~~~~~
+  if get_option('xla')
+      files += [
+        'src/neural/xla/hlo_builder.cc',
+        'src/neural/xla/network_xla.cc',
+        'src/neural/xla/onnx2hlo.cc',
+        'src/neural/xla/print_hlo.cc',
+        'src/neural/xla/pjrt.cc',
+        'src/neural/xla/xla_runner.cc',
+      ]
+      files += gen_proto_src.process('src/neural/xla/hlo.proto',
+          preserve_path_from : meson.current_source_dir() + '/src/')
+      deps += cc.find_library('dl', required: false)
+  endif
 
 endif # if get_option('build_backends')
 

diff --git a/meson_options.txt b/meson_options.txt
@@ -187,3 +187,8 @@ option('onnx_include',
        type: 'string',
        value: '',
        description: 'Paths to ONNX runtime includes')
+
+option('xla',
+       type: 'boolean',
+       value: false,
+       description: 'Enable XLA backend')
diff --git a/scripts/compile_proto.py b/scripts/compile_proto.py
@@ -64,6 +64,7 @@
     'package',
     'message',
     'optional',
+    'required',
     'repeated',
     'enum',
 ] + list(TYPES.keys())