Add missing includes for MLIR dialects
These were transitively included before
https://github.com/llvm/llvm-project/commit/eda6f907d2581a956aec511ae1d13adc3f97d3ad
PiperOrigin-RevId: 444224180
diff --git a/tensorflow/compiler/mlir/hlo/BUILD b/tensorflow/compiler/mlir/hlo/BUILD
index 07baa4e..c01ad75 100644
--- a/tensorflow/compiler/mlir/hlo/BUILD
+++ b/tensorflow/compiler/mlir/hlo/BUILD
@@ -944,6 +944,7 @@
":type_conversion",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithmeticDialect",
+ "@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Shape",
@@ -987,6 +988,7 @@
":hlo",
":mhlo_pass_details",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:InferTypeOpInterface",
"@llvm-project//mlir:Pass",
@@ -1007,8 +1009,10 @@
":hlo",
":mhlo_pass_details",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
+ "@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:Transforms",
],
)
@@ -1021,8 +1025,11 @@
":hlo",
":mhlo_pass_details",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:DialectUtils",
+ "@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
+ "@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:Transforms",
],
)
@@ -1104,6 +1111,8 @@
":hlo",
":mhlo_pass_details",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:ArithmeticDialect",
+ "@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:TensorDialect",
],
)
@@ -1391,6 +1400,7 @@
deps = [
":hlo",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:TensorDialect",
@@ -1643,6 +1653,7 @@
deps = [
":gml_st_test_passes_inc_gen",
":gml_st_transforms",
+ "@llvm-project//mlir:Affine",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Transforms",
@@ -1715,6 +1726,7 @@
":hlo",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:ArithmeticDialect",
+ "@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Shape",
"@llvm-project//mlir:TensorDialect",
@@ -1798,10 +1810,13 @@
":shape_component_analysis",
":transforms_pass_inc_gen",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:ArithmeticDialect",
+ "@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:LinalgOps",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:Shape",
+ "@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:Transforms",
],
)
@@ -1818,6 +1833,7 @@
deps = [
":hlo",
":transforms_pass_inc_gen",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
diff --git a/tensorflow/compiler/mlir/hlo/lib/Analysis/shape_component_analysis.cc b/tensorflow/compiler/mlir/hlo/lib/Analysis/shape_component_analysis.cc
index ae68d30..628814d 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Analysis/shape_component_analysis.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Analysis/shape_component_analysis.cc
@@ -21,6 +21,7 @@
#include "llvm/ADT/STLExtras.h"
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExpr.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Analysis/test_shape_component_analysis.cc b/tensorflow/compiler/mlir/hlo/lib/Analysis/test_shape_component_analysis.cc
index d4ace69..73aa34b 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Analysis/test_shape_component_analysis.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Analysis/test_shape_component_analysis.cc
@@ -17,6 +17,7 @@
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "mlir-hlo/Transforms/PassDetail.h"
#include "mlir-hlo/Transforms/passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/gml_st/transforms/test_passes.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/gml_st/transforms/test_passes.cc
index c9a9fad..e97f317 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/gml_st/transforms/test_passes.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/gml_st/transforms/test_passes.cc
@@ -19,6 +19,7 @@
#include <utility>
#include "mlir-hlo/Dialect/gml_st/transforms/transforms.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/broadcast_propagation.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/broadcast_propagation.cc
index 90bc20d..48c2345 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/broadcast_propagation.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/broadcast_propagation.cc
@@ -26,6 +26,7 @@
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/MLIRContext.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/collapse_elementwise_map.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/collapse_elementwise_map.cc
index a6d98ed..139714a 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/collapse_elementwise_map.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/collapse_elementwise_map.cc
@@ -19,6 +19,7 @@
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/group_reduction_dimensions.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/group_reduction_dimensions.cc
index 12c564c..df25806 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/group_reduction_dimensions.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/group_reduction_dimensions.cc
@@ -24,6 +24,8 @@
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Pass/Pass.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_shape_ops_to_standard.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_shape_ops_to_standard.cc
index 9cde04a..0b4b959 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_shape_ops_to_standard.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_shape_ops_to_standard.cc
@@ -25,6 +25,7 @@
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
#include "mlir-hlo/Dialect/mhlo/transforms/type_conversion.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExpr.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_einsum_to_dot_general.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_einsum_to_dot_general.cc
index a5d5a13..a335f2a 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_einsum_to_dot_general.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_einsum_to_dot_general.cc
@@ -20,6 +20,7 @@
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_gather_to_torch_index_select.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_gather_to_torch_index_select.cc
index 16a73b1..378e344 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_gather_to_torch_index_select.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_gather_to_torch_index_select.cc
@@ -17,6 +17,7 @@
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
index 2ee510d..eb2bbef 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
@@ -35,6 +35,7 @@
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Arithmetic/Utils/Utils.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/merge_assuming_ops.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/merge_assuming_ops.cc
index d982936..f671462 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/merge_assuming_ops.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/merge_assuming_ops.cc
@@ -24,6 +24,7 @@
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BlockAndValueMapping.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/mhlo_canonicalize_reduction.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/mhlo_canonicalize_reduction.cc
index 9612f90..8296d5f 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/mhlo_canonicalize_reduction.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/mhlo_canonicalize_reduction.cc
@@ -18,6 +18,8 @@
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
namespace mlir {
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/mhlo_flatten_tuple.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/mhlo_flatten_tuple.cc
index 7034e50..e1fd9b2 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/mhlo_flatten_tuple.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/mhlo_flatten_tuple.cc
@@ -27,6 +27,7 @@
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Location.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/shape_reification_pass.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/shape_reification_pass.cc
index 12bf9d0..5b49313 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/shape_reification_pass.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/shape_reification_pass.cc
@@ -22,6 +22,7 @@
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
#include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
#include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc
index 507532d..f902db9 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/test_infer_shaped_type_pass.cc
@@ -16,6 +16,7 @@
#include <utility>
#include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/MLIRContext.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/unfuse_batch_norm.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/unfuse_batch_norm.cc
index 3c95d2a..dc1dd36 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/unfuse_batch_norm.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/unfuse_batch_norm.cc
@@ -15,6 +15,7 @@
#include "llvm/ADT/SmallVector.h"
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Attributes.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Transforms/lower_index_cast_pass.cc b/tensorflow/compiler/mlir/hlo/lib/Transforms/lower_index_cast_pass.cc
index 545d305..e4e244e 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Transforms/lower_index_cast_pass.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Transforms/lower_index_cast_pass.cc
@@ -20,6 +20,7 @@
#include "mlir-hlo/Transforms/PassDetail.h"
#include "mlir-hlo/Transforms/passes.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/PatternMatch.h"
diff --git a/tensorflow/compiler/mlir/hlo/lib/Transforms/symbolic_shape_optimization.cc b/tensorflow/compiler/mlir/hlo/lib/Transforms/symbolic_shape_optimization.cc
index c11f870..56730b7 100644
--- a/tensorflow/compiler/mlir/hlo/lib/Transforms/symbolic_shape_optimization.cc
+++ b/tensorflow/compiler/mlir/hlo/lib/Transforms/symbolic_shape_optimization.cc
@@ -27,8 +27,11 @@
#include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "mlir-hlo/Transforms/PassDetail.h"
#include "mlir-hlo/Transforms/passes.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index e9884a5..95e0f7a 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -386,6 +386,7 @@
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:status",
"//tensorflow/stream_executor/lib",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
],
diff --git a/tensorflow/compiler/mlir/lite/utils/constant_utils.cc b/tensorflow/compiler/mlir/lite/utils/constant_utils.cc
index 3ad345d..635c04c 100644
--- a/tensorflow/compiler/mlir/lite/utils/constant_utils.cc
+++ b/tensorflow/compiler/mlir/lite/utils/constant_utils.cc
@@ -15,6 +15,7 @@
#include "tensorflow/compiler/mlir/lite/utils/constant_utils.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/mangling_util.h"
#include "tensorflow/core/framework/tensor.pb.h"
diff --git a/tensorflow/compiler/mlir/lite/utils/constant_utils.h b/tensorflow/compiler/mlir/lite/utils/constant_utils.h
index deb1e24..12f9822 100644
--- a/tensorflow/compiler/mlir/lite/utils/constant_utils.h
+++ b/tensorflow/compiler/mlir/lite/utils/constant_utils.h
@@ -16,6 +16,7 @@
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONSTANT_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONSTANT_UTILS_H_
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Location.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tensorflow/BUILD b/tensorflow/compiler/mlir/tensorflow/BUILD
index 7b27f3b..79e1d11 100644
--- a/tensorflow/compiler/mlir/tensorflow/BUILD
+++ b/tensorflow/compiler/mlir/tensorflow/BUILD
@@ -364,6 +364,7 @@
"//tensorflow/core:lib",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Analysis",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
@@ -708,6 +709,7 @@
"@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:CallOpInterfacesIncGen",
"@llvm-project//mlir:ControlFlowInterfaces",
+ "@llvm-project//mlir:ControlFlowOps",
"@llvm-project//mlir:DerivedAttributeOpInterface",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
@@ -965,6 +967,7 @@
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework_internal",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
@@ -1293,6 +1296,7 @@
"@llvm-project//mlir:AffineAnalysis",
"@llvm-project//mlir:AffineUtils",
"@llvm-project//mlir:Analysis",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:ControlFlowOps",
"@llvm-project//mlir:Dialect",
"@llvm-project//mlir:FuncDialect",
@@ -1399,6 +1403,7 @@
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:threadpool_options",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:Pass",
diff --git a/tensorflow/compiler/mlir/tensorflow/dialect_registration.h b/tensorflow/compiler/mlir/tensorflow/dialect_registration.h
index 52579d7..f59b8e7 100644
--- a/tensorflow/compiler/mlir/tensorflow/dialect_registration.h
+++ b/tensorflow/compiler/mlir/tensorflow/dialect_registration.h
@@ -17,6 +17,7 @@
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_DIALECT_REGISTRATION_H_
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Dialect.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc b/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
index 04dfc3d..e9220c6 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/device_index_selector.cc
@@ -15,6 +15,7 @@
// Converts DeviceIndex to constant device.
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/einsum.cc b/tensorflow/compiler/mlir/tensorflow/transforms/einsum.cc
index 12eb368..9bf6e32 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/einsum.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/einsum.cc
@@ -35,6 +35,7 @@
#include "llvm/Support/Debug.h"
#include "llvm/Support/Regex.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h" // from @llvm-project
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Traits.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc b/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc
index 4bb8245..9b10982 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc
@@ -18,6 +18,7 @@
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/MemoryBuffer.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import_test_pass.cc b/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import_test_pass.cc
index 374e267..7132368 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import_test_pass.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import_test_pass.cc
@@ -18,6 +18,7 @@
#include "llvm/Support/Casting.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/ToolOutputFile.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/OperationSupport.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/initialize_variables_in_session_init.cc b/tensorflow/compiler/mlir/tensorflow/transforms/initialize_variables_in_session_init.cc
index a40c484..0252a05 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/initialize_variables_in_session_init.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/initialize_variables_in_session_init.cc
@@ -17,6 +17,7 @@
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSet.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/IR/SymbolTable.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/legalize_hlo.cc b/tensorflow/compiler/mlir/tensorflow/transforms/legalize_hlo.cc
index 17ac3a2..490d538 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/legalize_hlo.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/legalize_hlo.cc
@@ -34,6 +34,7 @@
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/BlockAndValueMapping.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/optimize.cc b/tensorflow/compiler/mlir/tensorflow/transforms/optimize.cc
index 54bd98a..e87852e 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/optimize.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/optimize.cc
@@ -14,6 +14,7 @@
==============================================================================*/
#include <iostream>
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tfr/passes/canonicalize.cc b/tensorflow/compiler/mlir/tfr/passes/canonicalize.cc
index e7cd160..3ea59d3 100644
--- a/tensorflow/compiler/mlir/tfr/passes/canonicalize.cc
+++ b/tensorflow/compiler/mlir/tfr/passes/canonicalize.cc
@@ -20,6 +20,7 @@
#include "llvm/Support/raw_ostream.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h" // from @llvm-project
#include "mlir/Dialect/Affine/LoopUtils.h" // from @llvm-project
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/SCF/SCF.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tfrt/jit/transforms/BUILD b/tensorflow/compiler/mlir/tfrt/jit/transforms/BUILD
index 393dac0..13cc833 100644
--- a/tensorflow/compiler/mlir/tfrt/jit/transforms/BUILD
+++ b/tensorflow/compiler/mlir/tfrt/jit/transforms/BUILD
@@ -75,6 +75,7 @@
"//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
"//tensorflow/compiler/mlir/xla:xla_legalize_tf",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:Affine",
"@llvm-project//mlir:AffineAnalysis",
"@llvm-project//mlir:Analysis",
"@llvm-project//mlir:ArithmeticDialect",
@@ -85,6 +86,7 @@
"@llvm-project//mlir:LinalgOps",
"@llvm-project//mlir:LinalgTransforms",
"@llvm-project//mlir:MathDialect",
+ "@llvm-project//mlir:MemRefDialect",
"@llvm-project//mlir:MemRefTransforms",
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SCFDialect",
diff --git a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_fuse_fill_into_tiled_reduction.cc b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_fuse_fill_into_tiled_reduction.cc
index 68c9bf3..a2f0539 100644
--- a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_fuse_fill_into_tiled_reduction.cc
+++ b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_fuse_fill_into_tiled_reduction.cc
@@ -23,6 +23,7 @@
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h"
diff --git a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_fusion.cc b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_fusion.cc
index fb3ace51..feb6bef 100644
--- a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_fusion.cc
+++ b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_fusion.cc
@@ -20,6 +20,7 @@
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h"
namespace tensorflow {
diff --git a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h
index 61a6576..b956213 100644
--- a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h
+++ b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h
@@ -22,6 +22,7 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Pass/Pass.h"
diff --git a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_symbolic_shape_optimization.cc b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_symbolic_shape_optimization.cc
index fe4e171..7aa6e56 100644
--- a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_symbolic_shape_optimization.cc
+++ b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_symbolic_shape_optimization.cc
@@ -30,6 +30,8 @@
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Analysis/shape_component_analysis.h"
#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h"
diff --git a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_tile_reduction.cc b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_tile_reduction.cc
index 637da75..135e3f4 100644
--- a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_tile_reduction.cc
+++ b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_tile_reduction.cc
@@ -25,6 +25,7 @@
#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/IR/Diagnostics.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h"
diff --git a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_tile_transpose.cc b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_tile_transpose.cc
index c0b6cd8..355903d 100644
--- a/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_tile_transpose.cc
+++ b/tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_tile_transpose.cc
@@ -23,6 +23,7 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/Transforms/CodegenStrategy.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Linalg/Transforms/Transforms.h" // from @llvm-project
#include "mlir/Dialect/Utils/StructuredOpsUtils.h" // from @llvm-project
#include "tensorflow/compiler/mlir/tfrt/jit/transforms/tf_jitrt_passes.h"
diff --git a/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/kernel_ops_pattern.cc b/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/kernel_ops_pattern.cc
index 0326a48..d4694a8 100644
--- a/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/kernel_ops_pattern.cc
+++ b/tensorflow/compiler/mlir/tfrt/transforms/lmhlo_to_gpu/kernel_ops_pattern.cc
@@ -22,6 +22,7 @@
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/GPU/GPUDialect.h" // from @llvm-project
#include "mlir/Dialect/MemRef/IR/MemRef.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/BUILD b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/BUILD
index f0ea306..cc8a52e 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/BUILD
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/BUILD
@@ -54,6 +54,7 @@
deps = [
"//tensorflow/compiler/mlir/hlo",
"//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:BufferizationTransforms",
"@llvm-project//mlir:ComplexDialect",
"@llvm-project//mlir:FuncDialect",
diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc
index 8cf95a6..54f611f 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc
@@ -17,6 +17,7 @@
#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h" // from @llvm-project
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Complex/IR/Complex.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/MemRef/IR/MemRef.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/embed_memref_prints.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/embed_memref_prints.cc
index 4f8c84e..ccd0d17 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/embed_memref_prints.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/embed_memref_prints.cc
@@ -16,6 +16,7 @@
#include <memory>
#include <string>
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Linalg/IR/Linalg.h" // from @llvm-project
#include "mlir/Dialect/MemRef/IR/MemRef.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_to_jit_invocations.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_to_jit_invocations.cc
index e8aab38..ca1e76c 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_to_jit_invocations.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_to_jit_invocations.cc
@@ -22,6 +22,7 @@
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/Shape/IR/Shape.h" // from @llvm-project
#include "mlir/IR/Block.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tosa/transforms/fuse_bias_tf.cc b/tensorflow/compiler/mlir/tosa/transforms/fuse_bias_tf.cc
index 1f7abe8..3ac97d9 100644
--- a/tensorflow/compiler/mlir/tosa/transforms/fuse_bias_tf.cc
+++ b/tensorflow/compiler/mlir/tosa/transforms/fuse_bias_tf.cc
@@ -21,6 +21,7 @@
#include <iterator>
#include <numeric>
+#include "mlir/Dialect/Quant/QuantOps.h" // from @llvm-project
#include "mlir/Dialect/Tosa/IR/TosaOps.h" // from @llvm-project
#include "mlir/IR/MLIRContext.h" // from @llvm-project
#include "mlir/Pass/Pass.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tosa/transforms/legalize_tf.cc b/tensorflow/compiler/mlir/tosa/transforms/legalize_tf.cc
index df2c45e..180eeed 100644
--- a/tensorflow/compiler/mlir/tosa/transforms/legalize_tf.cc
+++ b/tensorflow/compiler/mlir/tosa/transforms/legalize_tf.cc
@@ -21,6 +21,7 @@
#include <iterator>
#include <numeric>
+#include "mlir/Dialect/Quant/QuantOps.h" // from @llvm-project
#include "mlir/Dialect/Tosa/IR/TosaOps.h" // from @llvm-project
#include "mlir/Support/LLVM.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/tosa/transforms/legalize_utils.cc b/tensorflow/compiler/mlir/tosa/transforms/legalize_utils.cc
index e77fa22..10e35b5 100644
--- a/tensorflow/compiler/mlir/tosa/transforms/legalize_utils.cc
+++ b/tensorflow/compiler/mlir/tosa/transforms/legalize_utils.cc
@@ -15,6 +15,7 @@
#include "tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
#include "mlir/Dialect/Tosa/IR/TosaOps.h" // from @llvm-project
#include "mlir/Dialect/Tosa/Utils/QuantUtils.h" // from @llvm-project
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/xla/BUILD b/tensorflow/compiler/mlir/xla/BUILD
index 851bbbf..87a902e 100644
--- a/tensorflow/compiler/mlir/xla/BUILD
+++ b/tensorflow/compiler/mlir/xla/BUILD
@@ -538,6 +538,7 @@
"//tensorflow/stream_executor/lib",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:Analysis",
+ "@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:MemRefDialect",
@@ -641,6 +642,7 @@
"@llvm-project//mlir:ArithmeticDialect",
"@llvm-project//mlir:FuncDialect",
"@llvm-project//mlir:IR",
+ "@llvm-project//mlir:TensorDialect",
"@llvm-project//mlir:Translation",
],
alwayslink = 1,
diff --git a/tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc b/tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc
index 3790326..3051faa 100644
--- a/tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc
+++ b/tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.cc
@@ -28,6 +28,7 @@
#include "llvm/Support/SMLoc.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/Dialect/MemRef/IR/MemRef.h" // from @llvm-project
#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/xla/transforms/adjust_layout.cc b/tensorflow/compiler/mlir/xla/transforms/adjust_layout.cc
index 8e29b56..32432f7 100644
--- a/tensorflow/compiler/mlir/xla/transforms/adjust_layout.cc
+++ b/tensorflow/compiler/mlir/xla/transforms/adjust_layout.cc
@@ -21,6 +21,7 @@
#include "absl/types/span.h"
#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
#include "mlir/IR/Attributes.h" // from @llvm-project
#include "mlir/IR/Builders.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
diff --git a/tensorflow/compiler/mlir/xla/xla_mlir_translate_registration.cc b/tensorflow/compiler/mlir/xla/xla_mlir_translate_registration.cc
index b6eba0d..a4b87c2 100644
--- a/tensorflow/compiler/mlir/xla/xla_mlir_translate_registration.cc
+++ b/tensorflow/compiler/mlir/xla/xla_mlir_translate_registration.cc
@@ -16,6 +16,7 @@
#include "llvm/Support/raw_os_ostream.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h" // from @llvm-project
#include "mlir/IR/BuiltinOps.h" // from @llvm-project
#include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project
#include "tensorflow/compiler/mlir/xla/mlir_hlo_to_hlo.h"
diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD
index 6ea8e9a..efecb4b 100644
--- a/tensorflow/compiler/xla/service/gpu/BUILD
+++ b/tensorflow/compiler/xla/service/gpu/BUILD
@@ -841,6 +841,7 @@
"//tensorflow/compiler/xla/service/llvm_ir:llvm_type_conversion_util",
"//tensorflow/core/platform:stream_executor_no_cuda",
"@llvm-project//llvm:Core",
+ "@llvm-project//mlir:ArithmeticDialect",
],
)
diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
index 97c4a55..c78ce0d 100644
--- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
+++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc
@@ -20,6 +20,7 @@
#include <vector>
#include "llvm/IR/IntrinsicsNVPTX.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" // from @llvm-project
#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
#include "tensorflow/compiler/mlir/xla/hlo_utils.h"
#include "tensorflow/compiler/mlir/xla/type_to_shape.h"
diff --git a/tensorflow/dtensor/mlir/BUILD b/tensorflow/dtensor/mlir/BUILD
index e86683b..04f7932 100644
--- a/tensorflow/dtensor/mlir/BUILD
+++ b/tensorflow/dtensor/mlir/BUILD
@@ -171,6 +171,7 @@
"@llvm-project//mlir:Pass",
"@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:Shape",
+ "@llvm-project//mlir:TensorDialect",
],
alwayslink = 1,
)
diff --git a/tensorflow/dtensor/mlir/dtensor_mlir_passes_classes.h b/tensorflow/dtensor/mlir/dtensor_mlir_passes_classes.h
index 7feccb3..3440c77 100644
--- a/tensorflow/dtensor/mlir/dtensor_mlir_passes_classes.h
+++ b/tensorflow/dtensor/mlir/dtensor_mlir_passes_classes.h
@@ -16,11 +16,12 @@
#ifndef TENSORFLOW_DTENSOR_MLIR_DTENSOR_MLIR_PASSES_CLASSES_H_
#define TENSORFLOW_DTENSOR_MLIR_DTENSOR_MLIR_PASSES_CLASSES_H_
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Dialect.h"
#include "mlir/Pass/Pass.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h" // from @llvm-project
#include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/chlo_ops.h"
namespace tensorflow {
