load("@bazel_skylib//rules:build_test.bzl", "build_test") load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library") load("//tensorflow:tensorflow.bzl", "tf_cc_test") load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable") # MLIR passes for DTensor support. load("//tensorflow/core/platform:rules_cc.bzl", "cc_library") package( # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], default_visibility = [ "//tensorflow/dtensor:dtensor-internal", # Allow visibility from the mlir language server. "//learning/brain/mlir/mlir_lsp_server:__pkg__", "//third_party/py/tensorflow_addons/custom_ops/dtensor:__subpackages__", ], licenses = ["notice"], ) gentbl_cc_library( name = "tensorflow_dtensor_ops_inc_gen", compatible_with = get_compatible_with_portable(), tbl_outs = { "ir/tf_dtensor.h.inc": ["-gen-op-decls"], "ir/tf_dtensor.cc.inc": ["-gen-op-defs"], }, tblgen = "@llvm-project//mlir:mlir-tblgen", td_file = "ir/tf_dtensor.td", td_srcs = [ "//tensorflow/compiler/mlir/tensorflow:ir/tf_op_base.td", "//tensorflow/compiler/mlir/tensorflow:ir/tf_op_interfaces.td", ], deps = [ "@llvm-project//mlir:CallInterfacesTdFiles", "@llvm-project//mlir:FuncTdFiles", "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", "@llvm-project//mlir:OpBaseTdFiles", "@llvm-project//mlir:SideEffectInterfacesTdFiles", ], ) gentbl_cc_library( name = "dtensor_passes_inc_gen", compatible_with = get_compatible_with_portable(), tbl_outs = {"dtensor_passes.h.inc": [ "-gen-pass-decls", "-name=DTensor", ]}, tblgen = "@llvm-project//mlir:mlir-tblgen", td_file = "Passes.td", deps = ["@llvm-project//mlir:PassBaseTdFiles"], ) cc_library( name = "tf_dtensor_dialect", srcs = ["ir/tf_dtensor.cc"], hdrs = ["ir/tf_dtensor.h"], includes = ["include"], deps = [ ":tensorflow_dtensor_ops_inc_gen", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes", "//tensorflow/compiler/mlir/tensorflow:tensorflow_op_interfaces", "//tensorflow/compiler/mlir/tensorflow:tensorflow_side_effects", "//tensorflow/compiler/mlir/tensorflow:tensorflow_traits", "//tensorflow/compiler/mlir/tensorflow:tensorflow_types", "//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes", "@llvm-project//llvm:Support", "@llvm-project//mlir:DerivedAttributeOpInterface", "@llvm-project//mlir:Dialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:InferTypeOpInterface", "@llvm-project//mlir:SideEffectInterfaces", "@llvm-project//mlir:Support", ], alwayslink = 1, ) cc_library( name = "collectives", srcs = ["collectives.cc"], hdrs = ["collectives.h"], deps = [ ":collectives_common", ":dtensor_location", ":layout_parsing", ":shape_utils", ":sparse_expander_common", ":spmd_expander_common", ":tf_dtensor_dialect", ":value_utils", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes", "//tensorflow/core:lib", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:dtensor_utils", "//tensorflow/dtensor/cc:tensor_layout", "//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], ) cc_library( name = "collectives_common", srcs = ["collectives_common.cc"], hdrs = ["collectives_common.h"], deps = [ "//tensorflow/core:portable_gif_internal", "//tensorflow/core/platform:errors", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:tensor_layout", "@com_google_absl//absl/container:flat_hash_set", ], ) cc_library( name = "device_utils", srcs = ["device_utils.cc"], hdrs = ["device_utils.h"], deps = [ "//tensorflow/core/platform:errors", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:tensor_layout", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], alwayslink = True, ) cc_library( name = "dtensor_location", srcs = ["dtensor_location.cc"], hdrs = ["dtensor_location.h"], deps = [ "//tensorflow/compiler/mlir/utils:name_utils", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], ) cc_library( name = "create_dtensor_mlir_passes", hdrs = [ "create_dtensor_mlir_passes.h", ], deps = [ ":device_utils", ":dtensor_passes_inc_gen", ":dtensor_send_recv", ":layout_parsing", ":op_utils", ":shape_utils", ":sparse_expander", ":spmd_expander", ":spmd_expander_common", ":tf_dtensor_dialect", ":value_utils", "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes", "//tensorflow/core:core_cpu_base", "//tensorflow/core:lib", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", ], alwayslink = 1, ) cc_library( name = "dtensor_mlir_passes", srcs = [ "annotate_global_shape.cc", "cluster_function_conversion.cc", "constant_folding.cc", "dce.cc", "designate_resource_handle_mesh.cc", "device_mesh_cluster_coarsening.cc", "dtensor_allreduce_combine_optimization.cc", "dtensor_allreduce_scatter_optimization.cc", "dtensor_allreduce_sum_optimization.cc", "dtensor_collective_type_lowering.cc", "dtensor_layout_to_xla_sharding_op.cc", "dtensor_mixed_precision_reduce.cc", "dtensor_mlir_passes.cc", "dtensor_multi_device_expansion.cc", "dtensor_remove_dtensorlayout.cc", "dtensor_replace_auxiliary_layout_op.cc", "dtensor_replace_relayout_with_identity.cc", "dtensor_set_hlo_sharding.cc", "function_renaming.cc", "handle_cross_cluster_dependencies.cc", "handle_sparsetensors.cc", "layout_propagation_v2.cc", "lower_send_recv.cc", "merge_clusters.cc", "mesh_propagation.cc", "move_compilation_to_host.cc", "op_to_device_cluster.cc", "propagate_default_layout.cc", "propagate_device_id_to_function_args.cc", "restore_shape_inference.cc", "set_default_sharding.cc", "sparse_expansion.cc", "spmd_expansion.cc", "tpu_add_resource_device_attribute.cc", "tpu_integration.cc", "undo_merge_const_across_mesh.cc", ], hdrs = ["dtensor_mlir_passes.h"], deps = [ ":collectives_common", ":create_dtensor_mlir_passes", ":device_utils", ":dtensor_location", ":dtensor_passes_inc_gen", ":dtensor_send_recv", ":layout_parsing", ":op_utils", ":shape_utils", ":sparse_expander", ":spmd_expander", ":spmd_expander_common", ":tf_dtensor_dialect", ":topological_iterator", ":value_utils", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:attribute_utils", "//tensorflow/compiler/mlir/tensorflow:bridge_logger", "//tensorflow/compiler/mlir/tensorflow:convert_tensor", "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util", "//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis", "//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes", "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops", "//tensorflow/compiler/mlir/tensorflow:tensorflow_types", "//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util", "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes", "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:runtime_passes", "//tensorflow/compiler/mlir/utils:name_utils", "//tensorflow/core:core_cpu_base", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/dtensor/cc:constants", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:dtensor_utils", "//tensorflow/dtensor/cc:layout_to_xla_sharding", "//tensorflow/dtensor/cc:tensor_layout", "//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes", "//tensorflow/dtensor/mlir/utils:dtensor_mlir_passes_internal", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", "@llvm-project//llvm:Support", "@llvm-project//mlir:Analysis", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:FunctionInterfaces", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@llvm-project//mlir:SideEffectInterfaces", "@llvm-project//mlir:SparseTensorDialect", "@llvm-project//mlir:Support", "@llvm-project//mlir:TransformUtils", "@llvm-project//mlir:Transforms", "@xla//xla:xla_data_proto_cc", "@xla//xla/hlo/builder:sharding_builder", "@xla//xla/tsl/platform:status", ], alwayslink = 1, ) cc_library( name = "dtensor_send_recv", srcs = ["dtensor_send_recv.cc"], hdrs = ["dtensor_send_recv.h"], deps = [ ":collectives", ":device_utils", ":layout_parsing", ":op_utils", ":shape_utils", ":spmd_expander_common", ":tf_dtensor_dialect", ":value_utils", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:convert_tensor", "//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes", "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes", "//tensorflow/core/platform:errors", "//tensorflow/dtensor/cc:constants", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:tensor_layout", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/status", "@com_google_absl//absl/types:span", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], alwayslink = True, ) cc_library( name = "layout_parsing", srcs = [ "layout_parsing.cc", ], hdrs = ["layout_parsing.h"], deps = [ ":tf_dtensor_dialect", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/core:lib", "//tensorflow/dtensor/cc:constants", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:tensor_layout", "//tensorflow/dtensor/proto:layout_proto_cc", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], alwayslink = 1, ) cc_library( name = "op_utils", srcs = ["op_utils.cc"], hdrs = ["op_utils.h"], deps = [ ":tf_dtensor_dialect", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/dtensor/cc:constants", "@llvm-project//llvm:Support", "@llvm-project//mlir:CallOpInterfaces", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], alwayslink = True, ) cc_library( name = "shape_utils", srcs = ["shape_utils.cc"], hdrs = ["shape_utils.h"], deps = [ ":tf_dtensor_dialect", ":value_utils", "//tensorflow/compiler/mlir/tensorflow:shape_inference_utils", "//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes", "//tensorflow/compiler/mlir/tensorflow:tensorflow_types", "//tensorflow/core:framework", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:status", "//tensorflow/dtensor/cc:constants", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:tensor_layout", "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:DerivedAttributeOpInterface", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:InferTypeOpInterface", "@llvm-project//mlir:Support", ], alwayslink = True, ) cc_library( name = "sparse_expander", srcs = [ "sparse_expanders.cc", ] + glob([ "*sparse_expander.cc", "sparse_expansions/*sparse_expander.cc", ]), hdrs = glob([ "*sparse_expander.h", "sparse_expansions/*sparse_expander.h", ]), deps = [ ":op_utils", ":sparse_expander_common", ":tf_dtensor_dialect", ":value_utils", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops", "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes", "//tensorflow/core:framework", "//tensorflow/core/platform:errors", "//tensorflow/core/platform:status", "//tensorflow/dtensor/cc:dstatus", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], alwayslink = 1, ) cc_library( name = "sparse_expander_common", srcs = ["sparse_expander_common.cc"], hdrs = ["sparse_expander_common.h"], deps = [ ":tf_dtensor_dialect", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops", "//tensorflow/core/platform:errors", "//tensorflow/dtensor/cc:dstatus", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:optional", "@llvm-project//llvm:Support", "@llvm-project//mlir:BytecodeOpInterface", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", ], alwayslink = True, ) cc_library( name = "spmd_expander_common", srcs = ["spmd_expander_common.cc"], hdrs = ["spmd_expander_common.h"], deps = [ ":device_utils", ":layout_parsing", ":op_utils", ":shape_utils", ":tf_dtensor_dialect", ":value_utils", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops", "//tensorflow/compiler/mlir/tensorflow:tensorflow_types", "//tensorflow/core:lib", "//tensorflow/dtensor/cc:constants", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:tensor_layout", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@llvm-project//llvm:Support", "@llvm-project//mlir:Dialect", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], alwayslink = True, ) cc_library( name = "spmd_expander", srcs = [ "spmd_expanders.cc", ] + glob([ "*spmd_expander.cc", "expansions/*spmd_expander.cc", "expansions/*spmd_expander_v2.cc", ]), hdrs = glob([ "*spmd_expander.h", "expansions/*spmd_expander.h", "expansions/*spmd_expander_v2.h", ]), deps = [ ":collectives", ":device_utils", ":dtensor_location", ":dtensor_send_recv", ":layout_parsing", ":op_utils", ":shape_utils", ":spmd_expander_common", ":tf_dtensor_dialect", ":value_utils", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes", "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops", "//tensorflow/compiler/mlir/tensorflow:tensorflow_types", "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/dtensor/cc:constants", "//tensorflow/dtensor/cc:dstatus", "//tensorflow/dtensor/cc:dtensor_utils", "//tensorflow/dtensor/cc:save_restore_util", "//tensorflow/dtensor/cc:slice_util", "//tensorflow/dtensor/cc:tensor_layout", "//tensorflow/dtensor/mlir/dtensor_dialect:ir/dtensor_attributes", "//tensorflow/dtensor/proto:layout_proto_cc", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/numeric:bits", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", "@llvm-project//llvm:Support", "@llvm-project//mlir:BytecodeOpInterface", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", "@xla//xla/mlir_hlo:convert_op_folder", ], alwayslink = 1, ) cc_library( name = "value_utils", srcs = ["value_utils.cc"], hdrs = ["value_utils.h"], deps = [ ":op_utils", ":tf_dtensor_dialect", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils", "//tensorflow/compiler/mlir/tensorflow:tensorflow_types", "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes", "//tensorflow/core:lib", "//tensorflow/dtensor/cc:dstatus", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc", ], alwayslink = True, ) cc_library( name = "topological_iterator", srcs = ["topological_iterator.cc"], hdrs = ["topological_iterator.h"], deps = [ ":device_utils", ":layout_parsing", ":op_utils", ":shape_utils", ":tf_dtensor_dialect", ":value_utils", "//tensorflow/compiler/mlir/tensorflow", "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops", "@llvm-project//llvm:Support", "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:IR", ], ) tf_cc_test( name = "dtensor_location_test", srcs = ["dtensor_location_test.cc"], deps = [ ":dtensor_location", "//tensorflow/compiler/mlir/utils:name_utils", "//tensorflow/core:test", "//tensorflow/core:test_main", "@llvm-project//mlir:IR", "@llvm-project//mlir:Support", ], ) build_test( name = "mlir_build_test", targets = [ ":tf_dtensor_dialect", ":tensorflow_dtensor_ops_inc_gen", ], )