"Fossies" - the Fresh Open Source Software Archive

Member "pytorch-1.8.2/aten.bzl" (23 Jul 2021, 1353 Bytes) of package /linux/misc/pytorch-1.8.2.tar.gz:


As a special service "Fossies" has tried to format the requested source page into HTML format using (guessed) Python source code syntax highlighting (style: standard) with prefixed line numbers. Alternatively you can here view or download the uninterpreted source code file. See also the last Fossies "Diffs" side-by-side code changes report for "aten.bzl": 1.12.1_vs_1.13.0.

    1 load("@rules_cc//cc:defs.bzl", "cc_library")
    2 
    3 CPU_CAPABILITY_NAMES = ["DEFAULT", "AVX", "AVX2"]
    4 CAPABILITY_COMPILER_FLAGS = {
    5     "AVX2": ["-mavx2", "-mfma"],
    6     "AVX": ["-mavx"],
    7     "DEFAULT": [],
    8 }
    9 
   10 PREFIX = "aten/src/ATen/native/"
   11 
   12 def intern_build_aten_ops(copts, deps):
   13     for cpu_capability in CPU_CAPABILITY_NAMES:
   14         srcs = []
   15         for impl in native.glob(
   16             [
   17                 PREFIX + "cpu/*.cpp",
   18                 PREFIX + "quantized/cpu/kernels/*.cpp",
   19             ]):
   20             name = impl.replace(PREFIX, "")
   21             out = PREFIX + name + "." + cpu_capability + ".cpp"
   22             native.genrule(
   23                 name = name + "_" + cpu_capability + "_cp",
   24                 srcs = [impl],
   25                 outs = [out],
   26                 cmd = "cp $< $@",
   27             )
   28             srcs.append(out)
   29 
   30         cc_library(
   31             name = "ATen_CPU_" + cpu_capability,
   32             srcs = srcs,
   33             copts = copts + [
   34                 "-DCPU_CAPABILITY=" + cpu_capability,
   35                 "-DCPU_CAPABILITY_" + cpu_capability,
   36             ] + CAPABILITY_COMPILER_FLAGS[cpu_capability],
   37             deps = deps,
   38             linkstatic = 1,
   39         )
   40     cc_library(
   41         name = "ATen_CPU",
   42         deps = [":ATen_CPU_" + cpu_capability for cpu_capability in CPU_CAPABILITY_NAMES],
   43         linkstatic = 1,
   44     )