Cross-compilation support

[ahans]

First up, thanks for this great project, we've been using it successfully for quite a while now!

I'd like to add support for cross compilation for Jetson/aarch64 using redist.json as a source. I have hacked already a bit and have a version that can at least compile examples/basic. The libraries still need some work, but before I put that in, I'd like to hear if there are already thoughts on how to accomplish that. I guess I could finish my version and it'd be fine for my needs, but ideally I'd like to contribute a PR so that others benefit from this as well.

Here's a rough outline of what I did:

• Extend redist_json_helper so that it does not filter for x86_64, but also includes aarch64. I then added a field arch to the struct that gets appended to specs for each component.
• Then in _redist_json_impl in cuda/extensions.bzl, I modify the specs list a bit so that for each component I have an entry where there are dicts urls, shas256, and strip_prefixes that each have two keys, x86_64 and aarch64.
• The result gets fed into a new repository rule cuda_multi_platform_component which iterates over the keys, effectively calling download_and_extract twice.
• Finally, the templates needed a bit of work to have targets for each architecture that are used by a top-level target (such as cuda_runtime), which has a bunch of selects, very similar to what is currently present to support both Windows and Linux.

As I said, this seems to be working, however, I'm not certain which usecases I broke along the way. Also, some cleanup and refactoring would be necessary to avoid redundancy.

One thing that complicates things a bit is that for most things I need the target architecture archive, but nvcc I need for the host. Just like clang, you only need to care about the host architecture here. In the repository rule I extract archives into x86_64 and aarch64 subdirectories. For nvcc I added some special handling, there I also create some symlinks to the directories in x86_64, so that all files are available where current main expects them.

Another thing I thought about was having subdirectories only in lib, but then I'd have to modify the original archive structure, which is something I'd like to avoid.

Initially I also thought about having a repository per architecture, but then I'd need a new place for the select. Ultimately, I want to be able to depend on @cuda//:cublas in my top-level target and it picks the correct libraries automatically depending on the architecture I'm building for.

I only need support for Linux x86_64 as the host and x86_64 and aarch64 as targets. However, AFAIU, support for a Linux aarch64 host would be possible with the redist.json archives, but would complicate the work required for rules_cuda even further.

If the overall approach outlined above sounds reasonable, I would suggest to start with the Linux x86_64 -> aarch64 support without breaking existing Linux x86_64 -> Linux x86_64 and Windows -> Windows support.

Another thing; I only care about bzlmod. Not sure how much of a priority new features for workspace are here.

Anyway, let me know what you think! Thanks!

[ahans]

One thing I forgot: My current approach does download and extract both x86_64 and aarch64 archives, regardless of which ones are actually required. I guess to avoid that we would need some supported_archs attribute for cuda.redist_json.

[cloudhan]

The overall direction sounds great.

I don't think you need to change component repo structure or merge the archives from different archs. Instead, you just spawn all repos for each arch with different name. IIRC, we currently encode it as <component_name>_<version>, you just make it as <component_name>_<version>_<arch>, and then change targets in the @cuda repo with more selects and arch match. For eaxample,

# some where
config_setting(name="version", ...)  # ctk version to use
config_setting(name="arch", ...)     # target arch

# @cuda//:BUILD.bazel
cc_library(  # we need a special rule "cc_alias", rules_python has something that can be borrowed here
    name="cublas",  # this is our standard entry point for a cuda component
    deps=select(
      # if version == 11.8 and arch == x86_64 then cublas_11.8_x86_64
      # if version == 12.9 and arch == x86_64 then cublas_12.9_x86_64 
    )
)

# @cublas_11.8_x86_64 is generated via cuda_component
# @cublas_12.9_x86_64 is generated via cuda_component

This will not complicate the overall design at all, you just setup a bunch more repos, you don't need arch specific subdirectories, also make the archive downloading on demand. The nvcc needs not be changed very much, but you need setup host and target arch repo, where host repo for compiling and target repo for linking. A select should correctly resolve everything for us :)

For workspace, I think we can de-prioritize it.

[ahans]

Thanks for the feedback! That sounds great! I wasn't fully aware that there is already some redirection. A target such as @cuda//:cublas appeared to me as being directly from BUILD.cublas. I now found the place where the alias is created. Looking at the final BUILD file confirmed that. So I guess I would use a meta target with the select instead of the alias, right? If things work out, a first PR is hopefully not too far away!

[ahans]

With the rules_python equivalent for cc_alias you mean this? https://rules-python.readthedocs.io/en/latest/api/rules_python/python/private/pypi/pkg_aliases.html#python-private-pypi-pkg-aliases-bzl
However, I don't quite see why we can't just use cc_library?

[cloudhan]

With the rules_python equivalent for cc_alias you mean this? https://rules-python.readthedocs.io/en/latest/api/rules_python/python/private/pypi/pkg_aliases.html#python-private-pypi-pkg-aliases-bzl However, I don't quite see why we can't just use cc_library?

Yes, that's it.

You can use cc_library for now. But ultimately it needs to be swap away. For example, cufft callback needs to be device linked, a cc_library cannot achieve it (which we don't support now). I am also suspicious that the alwayslink will not work here due to indirect deps.

[ahans]

So this overall works. Turns out you can use select in an alias, so the overall approach with the alias in @cuda BUILD can remain the same. However, it is a bit unclear how to best change the mapping from component to repo. I opted for <comp>__<arch> now, e.g., cublas__x86_64. The double underscore makes it possible to split on that later. Next, the host compiler needs --target (when using clang via toolchains_llvm). I didn't manage to pass that through cleanly yet and would appreciate some guidance. I will try to get this to a somewhat working state and open a draft PR, so we can continue discussing there.

[cloudhan]

There repos where they may use manually configured cuda.component and mapping them into cuda.toolkit, so you shouldn't assume naming convention.

The component_mappings needs to be extended. It currently supports a single set of mappings, we need multiple sets of mappings, maybe reading it from a json file.

[ZhenshengLee]

actually, I use rules_cuda 0.2.5 in my multi-platform crosscompile bazel development environment and it works well with something workaround, that is use a custom cuda repo definition.

the app BUILD may look like this

gaos_cuda_library(
    name="module_cuda",
    srcs = ["src/module_demo/module_cuda.cu"],
    hdrs = ["include/module_demo/module_cuda.cuh"],
    includes = ["include"],
    visibility = ["//main:__pkg__"],
)

gaos_cc_binary(
    name = "module_cuda_main",
    srcs = ["tool/module_cuda_main.cpp"],
    includes = ["include"],
    tags = ["tool"],
    visibility = ["//main:__pkg__"],
    deps = [
        ":module_cuda"
    ] + select({
        "@platforms//cpu:aarch64": [
            "@cuda_aarch64_linux//:cudart",
            "@cuda_aarch64_linux//:cuda",
        ],
        "@platforms//cpu:x86_64": [
            "@cuda_x86_64_linux//:cudart",
            "@cuda_x86_64_linux//:cuda",
        ]
    }),
)

the cuda_x86_64_linux repo may look like this

native.new_local_repository(
        name = "cuda_x86_64_linux",
        build_file = clean_dep("//bazel/3rdparty/cuda:cuda.BUILD"),
        path = "/usr/local/cuda/targets/x86_64-linux/",
    )
  native.new_local_repository(
        name = "cuda_aarch64_linux",
        build_file = clean_dep("//bazel/3rdparty/cuda:cuda.BUILD"),
        path = "/usr/local/cuda/targets/aarch64-linux/",
  )

Hope it helps.