PyTorch and Bazel

TimZaman: <p>Hello!</p>
<p>I’m creating PyTorch C++ extensions and building them with bazel.</p>
<p>The documentation (<a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="nofollow noopener">source</a>) explains how to build your extensions with either setuptools or JIT. However, Bazel likes to take the building into its own hands.</p>
<p>I’d like to get some advise on how to proceed. I currently have two working solutions:</p>
<ol>
<li>
<p>Hacky solution<br>
Per extension, create a bazel genrule, that just invokes a <code>python setup.py build</code>, and set the resulting .so file as the output artifact. This can then be loaded in the code. This leverages all the nice abstractions and build arguments that are set through the torch utilities (<code>torch.utils.cpp_extension.$</code>)</p>
</li>
<li>
<p>Proper solution<br>
Create the library through bazel’s <code>cc_library</code>. This is nice, but everything (arguments, flags, includes, directories, etc) need to be set manually.</p>
</li>
</ol>
<p>Is anyone using pytorch extensions with bazel already, or does anyone have any general advice here.</p>

TimZaman: <p>Just to help out some other people, here is the gist of it. The solution does prerequires having setup for bazel the (1) python headers and (2) pip requirements and (3) cuda.</p>
<p>Create a .bzl file containing something like</p>
<pre><code class="lang-auto">load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
load("@local_config_cuda//cuda:build_defs.bzl", "cuda_default_copts")

load("@pip_deps//:requirements.bzl", "requirement")

def pytorch_cpp_extension(name, srcs=[], gpu_srcs=[], deps=[], copts=[], defines=[],  
                          binary=True, linkopts=[]):
    """Create a pytorch cpp extension as a cpp and importable python library.
    
    All options defined below should stay close to the official torch cpp extension options as
    defined in https://github.com/pytorch/pytorch/blob/master/torch/utils/cpp_extension.py.
    """
    name_so = name + ".so"
    torch_deps = [
        requirement("torch", target = "cpp"),
    ]

    cuda_deps = [
        "@local_config_cuda//cuda:cudart_static",
        "@local_config_cuda//cuda:cuda_headers",
    ]

    copts = copts +[
        "-fPIC",
        "-D_GLIBCXX_USE_CXX11_ABI=0",
        "-DTORCH_API_INCLUDE_EXTENSION_H",
        "-fno-strict-aliasing",
        "-fopenmp",
        "-fstack-protector-strong",
        "-fwrapv",
        "-O2",
        "-std=c++14",
        "-DTORCH_EXTENSION_NAME=" + name
    ]

    if gpu_srcs:
        native.cc_library(
            name = name_so + "_gpu",
            srcs = gpu_srcs,
            deps = deps + torch_deps + if_cuda(cuda_deps),
            copts = copts + cuda_default_copts(),
            defines = defines,
            linkopts = linkopts,
        )
        cuda_deps.extend([":" + name_so + "_gpu"])

    if binary:
        native.cc_binary(
            name = name_so,
            srcs = srcs,
            deps = deps + torch_deps + if_cuda(cuda_deps),
            linkshared = 1,
            copts = copts,
            defines = defines,
            linkopts = linkopts,
        )
    else:
        native.cc_library(
            name = name_so,
            srcs = srcs,
            deps = deps + torch_deps + if_cuda(cuda_deps),
            copts = copts,
            defines = defines,
            linkopts = linkopts,
        )

    native.py_library(
        name = name,
        data = [":" + name_so],
    )
</code></pre>
<p>And be sure you can actually <code>require</code> torch as a cpp target library like like so;</p>
<pre><code class="lang-auto">genrule_directory(
    name = "include",
    srcs = [":extracted"],
    cmd = "mkdir -p $@ &amp;&amp; cp -a $&lt;/torch/lib/include/. $@",
)

# NOTE: Make sure this yields the same includes as `include_paths()`:
# See https://github.com/pytorch/pytorch/blob/master/torch/utils/cpp_extension.py#L494

cc_library(
    name = "cpp",
    hdrs = [":include"],
    visibility = ["//visibility:public"],
    includes = [
        "include",
        "include/torch/csrc/api/include",
        "include/TH",
        "include/THC",
    ],
    deps = [
        "@//third_party/python:headers",
    ]
)
</code></pre>

dmadisetti: <p>I don’t understand where <code>:extracted</code> comes from?</p>

shreyash: <p>Does anyone else have more context on this.</p>
<p><a class="mention" href="/u/timzaman">@TimZaman</a> I’d be extremely interested in understanding your solution to this and helping write a post about this to help other people who might face this issue.</p>
<p>Thanks!</p>

Po-Jen: <p><a class="mention" href="/u/timzaman">@TimZaman</a> Could you share a little bit more about how the hacky solution is done exactly?</p>
<p>For example, in the description:</p>
<p>Per extension, create a bazel genrule, that just invokes a <code>python setup.py build</code>, and set the resulting .so file as the output artifact. This can then be loaded in the code. This leverages all the nice abstractions and build arguments that are set through the torch utilities (<code>torch.utils.cpp_extension.$</code>)</p>
<p>How do these files looks like?</p>
<ol>
<li>BUILD</li>
<li>The .cpp file of custom C++ extension</li>
<li>The .py file calling the custom C++ extension</li>
</ol>
<p>I am also interesting in writing a post about this for future reference. Thank you.</p>

aikez: <p><a class="mention" href="/u/timzaman">@TimZaman</a> Thanks for your solution shared.<br>
But there’s a step that I don’t understand quite clear. Should we put the <code>genrule_directory</code> in each C++ extension <code>BUILD</code> file?  How does the <code>genrule_directory</code> work?</p>

ArsenK: <p>For anyone who ends up on this thread looking ways to build PyTorch C++ extensions with Bazel ~8.0.0</p>
<p>Here’s a setup that eventually worked in my case</p>
<p><code>BUILD</code> file</p>
<pre data-code-wrap="starlark"><code class="lang-starlark">load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
load("@rules_cuda//cuda:defs.bzl", "cuda_library")

cuda_library(
    name = "kernel_cc_lib",
    srcs = [ "kernel.cu", ],
    deps = [
        "@rules_cuda//cuda:runtime",
        "@libtorch_archive//:libtorch",
        "@rules_python//python/cc:current_py_cc_headers",
    ],
    copts = ["-D_GLIBCXX_USE_CXX11_ABI=0"],
)
pybind_extension(
    name = "kernel_wrapper",
    srcs = ["kernel.cpp"],
    deps = [":kernel_cc_lib"],
    copts = ["-DTORCH_EXTENSION_NAME= kernel_wrapper", "-D_GLIBCXX_USE_CXX11_ABI=0"],
)
py_library(
    name = "extension_name",
    srcs = [],
    deps = [],
    data = [":kernel_wrapper"],
    imports = ["."],
)
</code></pre>
<p><code>MODULE.bazel</code> file looks something like</p>
<pre data-code-wrap="starlark"><code class="lang-starlark">bazel_dep(name = "rules_python", version = "1.0.0")
bazel_dep(name = "pybind11_bazel", version = "2.13.6")

http_archive = use_repo_rule("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
http_archive(
  name = "libtorch_archive",
  strip_prefix = "libtorch",
  type="zip",
  build_file = "@//:libtorch.BUILD",
  urls = ["https://download.pytorch.org/libtorch/cu124/libtorch-shared-with-deps-2.5.0%2Bcu124.zip"],
  sha256 = "SHA_THAT_WORKS_FORY_YOU",
)

# https://github.com/bazel-contrib/rules_cuda
bazel_dep(name = "rules_cuda", version = "0.2.3")
cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
cuda.local_toolchain(
    name = "local_cuda",
    toolkit_path = "/usr/local/cuda",
)
use_repo(cuda, "local_cuda")
</code></pre>
<p><code>lib torch.BUILD</code> is something like</p>
<pre data-code-wrap="starlark"><code class="lang-starlark">package(default_visibility = ["//visibility:public"])

cc_library(
    name = "libtorch",
    srcs = glob(["lib/lib*.so*"]),
    deps = [],
    hdrs = glob(["include/**/*.h"]) + glob(["include/**/*.cuh"]),
    includes = [
        "include",
        "include/torch/csrc/api/include",
    ],
)
</code></pre>
<p>One can access the wrapper from python like so</p>
<pre data-code-wrap="python"><code class="lang-python">
import kernel_wrapper

...
</code></pre>
<p>Extension follows generic <a href="https://pytorch.org/tutorials/advanced/cpp_extension.html" rel="noopener nofollow ugc">PyTorch guidance</a></p>

jteuwen: <p><a class="mention" href="/u/arsenk">@ArsenK</a> Thanks for the explanation. This indeed works. However, it doesn’t work when you also installed torch through pip and do <code>import torch; import kernel_wrapper</code> as there seems to be an ABI mismatch.</p>
<p>Have you worked through that problem?</p>