• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Description:
2#   MLIR-GPU-specific components in XLA service implementation.
3
4load("//third_party/mlir:tblgen.bzl", "gentbl")
5
6# buildifier: disable=same-origin-load
7load("//tensorflow:tensorflow.bzl", "filegroup")
8
9# buildifier: disable=same-origin-load
10load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")
11load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
12load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
13
14package(
15    default_visibility = [":friends"],
16    licenses = ["notice"],  # Apache 2.0
17)
18
19package_group(
20    name = "friends",
21    includes = ["//tensorflow/compiler/xla:friends"],
22)
23
24# Filegroup used to collect source files for dependency checking.
25filegroup(
26    name = "c_srcs",
27    data = glob([
28        "**/*.cc",
29        "**/*.h",
30    ]),
31)
32
33gentbl(
34    name = "passes_inc_gen",
35    compatible_with = get_compatible_with_cloud(),
36    tbl_outs = [("-gen-pass-decls -name XlaMlirGpu", "passes.h.inc")],
37    tblgen = "@llvm-project//mlir:mlir-tblgen",
38    td_file = "passes.td",
39    td_srcs = ["@llvm-project//mlir:PassBaseTdFiles"],
40)
41
42cc_library(
43    name = "passes",
44    srcs = ["passes.cc"],
45    hdrs = ["passes.h"],
46    deps = [
47        ":passes_inc_gen",
48        "//tensorflow/compiler/mlir/hlo:lhlo",
49        "@com_google_absl//absl/memory",
50        "@llvm-project//llvm:Support",
51        "@llvm-project//mlir:GPUDialect",
52        "@llvm-project//mlir:GPUTransforms",
53        "@llvm-project//mlir:IR",
54        "@llvm-project//mlir:Pass",
55        "@llvm-project//mlir:SCFDialect",
56        "@llvm-project//mlir:SCFTransforms",
57        "@llvm-project//mlir:SideEffects",
58        "@llvm-project//mlir:StandardOps",
59        "@llvm-project//mlir:Transforms",
60    ],
61)
62
63cc_library(
64    name = "kernel_lowering",
65    srcs = ["kernel_lowering.cc"],
66    hdrs = ["kernel_lowering.h"],
67    deps = [
68        ":passes",
69        "//tensorflow/compiler/mlir/hlo",
70        "//tensorflow/compiler/mlir/hlo:hlo_legalize_to_lhlo",
71        "//tensorflow/compiler/mlir/hlo:legalize_to_linalg",
72        "//tensorflow/compiler/mlir/hlo:legalize_trigonometric_to_approximation",
73        "//tensorflow/compiler/mlir/hlo:lhlo",
74        "//tensorflow/compiler/mlir/hlo:lhlo_fuse_linalg",
75        "//tensorflow/compiler/mlir/hlo:lhlo_legalize_to_affine",
76        "//tensorflow/compiler/mlir/hlo:lhlo_legalize_to_gpu",
77        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
78        "//tensorflow/compiler/xla:status",
79        "//tensorflow/compiler/xla:statusor",
80        "//tensorflow/compiler/xla:util",
81        "@com_google_absl//absl/memory",
82        "@llvm-project//mlir:AffineToStandardTransforms",
83        "@llvm-project//mlir:CFGTransforms",
84        "@llvm-project//mlir:GPUDialect",
85        "@llvm-project//mlir:GPUToNVVMTransforms",
86        "@llvm-project//mlir:GPUToROCDLTransforms",
87        "@llvm-project//mlir:GPUTransforms",
88        "@llvm-project//mlir:IR",
89        "@llvm-project//mlir:LLVMDialect",
90        "@llvm-project//mlir:LLVMTransforms",
91        "@llvm-project//mlir:LinalgToLLVM",
92        "@llvm-project//mlir:LinalgTransforms",
93        "@llvm-project//mlir:NVVMDialect",
94        "@llvm-project//mlir:Pass",
95        "@llvm-project//mlir:ROCDLDialect",
96        "@llvm-project//mlir:SCFDialect",
97        "@llvm-project//mlir:SCFToGPUPass",
98        "@llvm-project//mlir:SCFTransforms",
99        "@llvm-project//mlir:StandardOps",
100        "@llvm-project//mlir:Transforms",
101    ],
102)
103
104tf_cc_binary(
105    name = "xla-mlir-gpu-opt",
106    srcs = ["xla_mlir_gpu_opt.cc"],
107    visibility = ["//tensorflow/compiler/xla/service/mlir_gpu/tests:__subpackages__"],
108    deps = [
109        ":passes",
110        "//tensorflow/compiler/mlir/hlo:all_passes",
111        "//tensorflow/compiler/mlir/hlo:hlo_dialect_registration",
112        "@llvm-project//llvm:Support",
113        "@llvm-project//mlir:AllPassesAndDialectsNoRegistration",
114        "@llvm-project//mlir:IR",
115        "@llvm-project//mlir:MlirOptLib",
116        "@llvm-project//mlir:Pass",
117        "@llvm-project//mlir:SideEffects",
118        "@llvm-project//mlir:Support",
119    ],
120)
121