1a44a1f9dd
This leaves some passes and the kernel_lowering in place for now. Those are pending removal,, too. PiperOrigin-RevId: 348621555 Change-Id: I4ff2ed11fb646bf76ceaa780956c6257d89f54ae
121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
# Description:
|
|
# MLIR-GPU-specific components in XLA service implementation.
|
|
|
|
load("//third_party/mlir:tblgen.bzl", "gentbl")
|
|
|
|
# buildifier: disable=same-origin-load
|
|
load("//tensorflow:tensorflow.bzl", "filegroup")
|
|
|
|
# buildifier: disable=same-origin-load
|
|
load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")
|
|
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
|
|
load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
|
|
|
|
package(
|
|
default_visibility = [":friends"],
|
|
licenses = ["notice"], # Apache 2.0
|
|
)
|
|
|
|
package_group(
|
|
name = "friends",
|
|
includes = ["//tensorflow/compiler/xla:friends"],
|
|
)
|
|
|
|
# Filegroup used to collect source files for dependency checking.
|
|
filegroup(
|
|
name = "c_srcs",
|
|
data = glob([
|
|
"**/*.cc",
|
|
"**/*.h",
|
|
]),
|
|
)
|
|
|
|
gentbl(
|
|
name = "passes_inc_gen",
|
|
compatible_with = get_compatible_with_cloud(),
|
|
tbl_outs = [("-gen-pass-decls -name XlaMlirGpu", "passes.h.inc")],
|
|
tblgen = "@llvm-project//mlir:mlir-tblgen",
|
|
td_file = "passes.td",
|
|
td_srcs = ["@llvm-project//mlir:PassBaseTdFiles"],
|
|
)
|
|
|
|
cc_library(
|
|
name = "passes",
|
|
srcs = ["passes.cc"],
|
|
hdrs = ["passes.h"],
|
|
deps = [
|
|
":passes_inc_gen",
|
|
"//tensorflow/compiler/mlir/hlo:lhlo",
|
|
"@com_google_absl//absl/memory",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:GPUDialect",
|
|
"@llvm-project//mlir:GPUTransforms",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:Pass",
|
|
"@llvm-project//mlir:SCFDialect",
|
|
"@llvm-project//mlir:SCFTransforms",
|
|
"@llvm-project//mlir:SideEffects",
|
|
"@llvm-project//mlir:StandardOps",
|
|
"@llvm-project//mlir:Transforms",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "kernel_lowering",
|
|
srcs = ["kernel_lowering.cc"],
|
|
hdrs = ["kernel_lowering.h"],
|
|
deps = [
|
|
":passes",
|
|
"//tensorflow/compiler/mlir/hlo",
|
|
"//tensorflow/compiler/mlir/hlo:hlo_legalize_to_lhlo",
|
|
"//tensorflow/compiler/mlir/hlo:legalize_to_linalg",
|
|
"//tensorflow/compiler/mlir/hlo:legalize_trigonometric_to_approximation",
|
|
"//tensorflow/compiler/mlir/hlo:lhlo",
|
|
"//tensorflow/compiler/mlir/hlo:lhlo_fuse_linalg",
|
|
"//tensorflow/compiler/mlir/hlo:lhlo_legalize_to_affine",
|
|
"//tensorflow/compiler/mlir/hlo:lhlo_legalize_to_gpu",
|
|
"//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
|
|
"//tensorflow/compiler/xla:status",
|
|
"//tensorflow/compiler/xla:statusor",
|
|
"//tensorflow/compiler/xla:util",
|
|
"@com_google_absl//absl/memory",
|
|
"@llvm-project//mlir:AffineToStandardTransforms",
|
|
"@llvm-project//mlir:CFGTransforms",
|
|
"@llvm-project//mlir:GPUDialect",
|
|
"@llvm-project//mlir:GPUToNVVMTransforms",
|
|
"@llvm-project//mlir:GPUToROCDLTransforms",
|
|
"@llvm-project//mlir:GPUTransforms",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:LLVMDialect",
|
|
"@llvm-project//mlir:LLVMTransforms",
|
|
"@llvm-project//mlir:LinalgToLLVM",
|
|
"@llvm-project//mlir:LinalgTransforms",
|
|
"@llvm-project//mlir:NVVMDialect",
|
|
"@llvm-project//mlir:Pass",
|
|
"@llvm-project//mlir:ROCDLDialect",
|
|
"@llvm-project//mlir:SCFDialect",
|
|
"@llvm-project//mlir:SCFToGPUPass",
|
|
"@llvm-project//mlir:SCFTransforms",
|
|
"@llvm-project//mlir:StandardOps",
|
|
"@llvm-project//mlir:Transforms",
|
|
],
|
|
)
|
|
|
|
tf_cc_binary(
|
|
name = "xla-mlir-gpu-opt",
|
|
srcs = ["xla_mlir_gpu_opt.cc"],
|
|
visibility = ["//tensorflow/compiler/xla/service/mlir_gpu/tests:__subpackages__"],
|
|
deps = [
|
|
":passes",
|
|
"//tensorflow/compiler/mlir/hlo:all_passes",
|
|
"//tensorflow/compiler/mlir/hlo:hlo_dialect_registration",
|
|
"@llvm-project//llvm:Support",
|
|
"@llvm-project//mlir:AllPassesAndDialectsNoRegistration",
|
|
"@llvm-project//mlir:IR",
|
|
"@llvm-project//mlir:MlirOptLib",
|
|
"@llvm-project//mlir:Pass",
|
|
"@llvm-project//mlir:SideEffects",
|
|
"@llvm-project//mlir:Support",
|
|
],
|
|
)
|