diff --git a/tensorflow/compiler/mlir/tensorflow/tests/tf_executor_ops_location_roundtrip.mlir b/tensorflow/compiler/mlir/tensorflow/tests/tf_executor_ops_location_roundtrip.mlir index 82e4205440b..24808692481 100644 --- a/tensorflow/compiler/mlir/tensorflow/tests/tf_executor_ops_location_roundtrip.mlir +++ b/tensorflow/compiler/mlir/tensorflow/tests/tf_executor_ops_location_roundtrip.mlir @@ -17,8 +17,8 @@ // When parsing it back, we should recover all 3 locations (the // tf_executor.island, tf.Identity, and tf_executor.yield). -// CHECK-LABEL: func @island_one_op_all_locs_same(%{{.*}}: tensor) -> tensor { -// CHECK-NEXT: "tf_executor.graph"() ( { +// CHECK-LABEL: "func" +// CHECK: "tf_executor.graph"() ( { // CHECK-NEXT: "tf_executor.island"() ( { // CHECK-NEXT: "tf.Identity"(%{{.*}}) : (tensor) -> tensor loc("identity@some_function") // CHECK-NEXT: "tf_executor.yield"(%{{.*}}) : (tensor) -> () loc("identity@some_function") @@ -26,7 +26,7 @@ // CHECK-NEXT: "tf_executor.fetch"(%{{.*}}) : (tensor) -> () loc(unknown) // CHECK-NEXT: }) : () -> tensor loc(unknown) // CHECK-NEXT: "std.return"(%{{.*}}) : (tensor) -> () loc(unknown) -// CHECK-NEXT: } loc(unknown) +// CHECK-NEXT: sym_name = "island_one_op_all_locs_same" func @island_one_op_all_locs_same(%arg0: tensor) -> tensor { %0 = "tf_executor.graph"() ( { @@ -44,8 +44,8 @@ func @island_one_op_all_locs_same(%arg0: tensor) -> tensor { // it is incorrect to use that syntax if the island, wrapped op, and yield // don't have identical locations. -// CHECK-LABEL: func @island_one_op_all_locs_NOT_same(%{{.*}}: tensor) -> tensor { -// CHECK-NEXT: "tf_executor.graph"() ( { +// CHECK-LABEL: "func" +// CHECK: "tf_executor.graph"() ( { // CHECK-NEXT: "tf_executor.island"() ( { // CHECK-NEXT: "tf.Identity"(%{{.*}}) : (tensor) -> tensor loc("identity@some_function") // CHECK-NEXT: "tf_executor.yield"(%{{.*}}) : (tensor) -> () loc("identity@some_function") @@ -53,7 +53,7 @@ func @island_one_op_all_locs_same(%arg0: tensor) -> tensor { // CHECK-NEXT: "tf_executor.fetch"(%{{.*}}) : (tensor) -> () loc(unknown) // CHECK-NEXT: }) : () -> tensor loc(unknown) // CHECK-NEXT: "std.return"(%{{.*}}) : (tensor) -> () loc(unknown) -// CHECK-NEXT: } loc(unknown) +// CHECK-NEXT: sym_name = "island_one_op_all_locs_NOT_same" func @island_one_op_all_locs_NOT_same(%arg0: tensor) -> tensor { %0 = "tf_executor.graph"() ( { diff --git a/tensorflow/compiler/mlir/xla/BUILD b/tensorflow/compiler/mlir/xla/BUILD index 717c0d0535d..0fd57b4f6f3 100644 --- a/tensorflow/compiler/mlir/xla/BUILD +++ b/tensorflow/compiler/mlir/xla/BUILD @@ -155,8 +155,8 @@ cc_library( "@com_google_absl//absl/memory", "@llvm-project//llvm:support", "@llvm-project//mlir:IR", - "@llvm-project//mlir:Linalg", "@llvm-project//mlir:LinalgDialectRegistration", + "@llvm-project//mlir:LinalgOps", "@llvm-project//mlir:Pass", "@llvm-project//mlir:StandardOps", "@llvm-project//mlir:Transforms", @@ -174,7 +174,7 @@ cc_library( "@llvm-project//llvm:support", "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:IR", - "@llvm-project//mlir:Linalg", + "@llvm-project//mlir:LinalgOps", "@llvm-project//mlir:LoopOps", "@llvm-project//mlir:Pass", "@llvm-project//mlir:StandardOps", @@ -189,8 +189,10 @@ cc_library( deps = [ ":lhlo", "@com_google_absl//absl/memory", - "@llvm-project//mlir:Linalg", + "@llvm-project//mlir:EDSC", "@llvm-project//mlir:LinalgDialectRegistration", + "@llvm-project//mlir:LinalgOps", + "@llvm-project//mlir:LinalgTransforms", "@llvm-project//mlir:Pass", ], alwayslink = 1, diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc index 8ad6717a3f1..9514422569b 100644 --- a/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc +++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc @@ -19,6 +19,7 @@ limitations under the License. #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h" #include "absl/memory/memory.h" #include "mlir/Dialect/Linalg/Utils/Utils.h" // TF:llvm-project +#include "mlir/EDSC/Helpers.h" // TF:llvm-project #include "mlir/Pass/Pass.h" // TF:llvm-project namespace mlir { @@ -52,7 +53,7 @@ struct LhloFuseLinalg : public FunctionPass { const SmallVector tile_sizes( generic_op.getNumInputsAndOutputs(), 1); auto op = cast(generic_op.getOperation()); - for (const Value result : op.getOutputs()) { + for (const Value result : op.getOutputBuffers()) { if (!func_args.count(result)) continue; if (linalg::tileLinalgOp(b, op, tile_sizes, /*permutation=*/{}, &folder)) { diff --git a/tensorflow/compiler/xla/service/mlir_gpu/BUILD b/tensorflow/compiler/xla/service/mlir_gpu/BUILD index 20b448286d5..775385bc731 100644 --- a/tensorflow/compiler/xla/service/mlir_gpu/BUILD +++ b/tensorflow/compiler/xla/service/mlir_gpu/BUILD @@ -160,9 +160,10 @@ cc_library( "@llvm-project//mlir:IR", "@llvm-project//mlir:LLVMDialect", "@llvm-project//mlir:LLVMTransforms", - "@llvm-project//mlir:Linalg", "@llvm-project//mlir:LinalgDialectRegistration", + "@llvm-project//mlir:LinalgOps", "@llvm-project//mlir:LinalgToLLVM", + "@llvm-project//mlir:LinalgTransforms", "@llvm-project//mlir:LoopDialectRegistration", "@llvm-project//mlir:LoopOps", "@llvm-project//mlir:LoopsToGPUPass", diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc index c878c90ef2a..e7f99161e44 100644 --- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc +++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc @@ -284,7 +284,7 @@ Status LowerLHLOToGPU(mlir::ModuleOp module) { pm.addPass(::mlir::xla_lhlo::createLegalizeToGpuPass()); // Fuse linalg operations. This will yield a single tiled loop nest where // Go from linalg to normal loops. - pm.addPass(::mlir::linalg::createConvertLinalgToLoopsPass()); + pm.addPass(::mlir::createConvertLinalgToLoopsPass()); // Canonicalize the code to simplify index computations. pm.addNestedPass<::mlir::FuncOp>(::mlir::createCanonicalizerPass()); // The innermost loops will be single-trip. diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl index 4f5a75d7262..1c3ce5f40c8 100755 --- a/tensorflow/workspace.bzl +++ b/tensorflow/workspace.bzl @@ -570,8 +570,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""): ) # Check out LLVM and MLIR from llvm-project. - LLVM_COMMIT = "41b520188820a732e6de4865c08704f412013209" - LLVM_SHA256 = "4cdf03a17f3acc0b6e23f97291ab266933df40a8dc5851ca39cf0209466eb37c" + LLVM_COMMIT = "0133cc60e4e230ee2c176c23eff5aa2f4ee17a75" + LLVM_SHA256 = "b660732cc9c2075916cd29b1719c1328e9d994568c838352d8e267ecba7bfa0a" LLVM_URLS = [ "https://storage.googleapis.com/mirror.tensorflow.org/github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT), "https://github.com/llvm/llvm-project/archive/{commit}.tar.gz".format(commit = LLVM_COMMIT), diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD index bf177752f3c..45b32f9328f 100644 --- a/third_party/mlir/BUILD +++ b/third_party/mlir/BUILD @@ -433,8 +433,8 @@ cc_library( ]), includes = ["include"], deps = [ + ":AffineOps", ":DialectUtils", - ":EDSC", ":IR", ":StandardOps", ":Support", @@ -1345,7 +1345,7 @@ cc_library( ":AffineToStandardTransforms", ":GPUDialect", ":IR", - ":Linalg", + ":LinalgTransforms", ":LoopOps", ":StandardOps", ":Support", @@ -2195,7 +2195,8 @@ cc_library( ":IR", ":LLVMDialect", ":LLVMTransforms", - ":Linalg", + ":LinalgOps", + ":LinalgTransforms", ":Pass", ":StandardOps", ":Support", @@ -2207,14 +2208,37 @@ cc_library( alwayslink = 1, ) -# TODO(ntv): Update these to make mapping with cmake simpler. cc_library( - name = "Linalg", + name = "LinalgOps", + srcs = [ + "lib/Dialect/Linalg/IR/LinalgOps.cpp", + "lib/Dialect/Linalg/IR/LinalgTypes.cpp", + ], + hdrs = [ + "include/mlir/Dialect/Linalg/IR/LinalgOps.h", + "include/mlir/Dialect/Linalg/IR/LinalgTraits.h", + "include/mlir/Dialect/Linalg/IR/LinalgTypes.h", + ], + includes = ["include"], + deps = [ + ":DialectUtils", + ":IR", + ":LinalgOpsIncGen", + ":LinalgStructuredOpsIncGen", + ":LinalgTransformPatternsIncGen", + ":Parser", + ":Support", + "@llvm-project//llvm:core", + "@llvm-project//llvm:support", + ], + alwayslink = 1, +) + +cc_library( + name = "LinalgTransforms", srcs = [ "lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp", "lib/Dialect/Linalg/EDSC/Builders.cpp", - "lib/Dialect/Linalg/IR/LinalgOps.cpp", - "lib/Dialect/Linalg/IR/LinalgTypes.cpp", "lib/Dialect/Linalg/Transforms/Fusion.cpp", "lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp", "lib/Dialect/Linalg/Transforms/LinalgTransforms.cpp", @@ -2226,9 +2250,6 @@ cc_library( "include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h", "include/mlir/Dialect/Linalg/EDSC/Builders.h", "include/mlir/Dialect/Linalg/EDSC/Intrinsics.h", - "include/mlir/Dialect/Linalg/IR/LinalgOps.h", - "include/mlir/Dialect/Linalg/IR/LinalgTraits.h", - "include/mlir/Dialect/Linalg/IR/LinalgTypes.h", "include/mlir/Dialect/Linalg/Passes.h", "include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h", "include/mlir/Dialect/Linalg/Utils/Intrinsics.h", @@ -2245,6 +2266,7 @@ cc_library( ":IR", ":LLVMDialect", ":LLVMTransforms", + ":LinalgOps", ":LinalgOpsIncGen", ":LinalgStructuredOpsIncGen", ":LinalgTransformPatternsIncGen", @@ -2264,8 +2286,8 @@ cc_library( cc_library( name = "LinalgDialectRegistration", - srcs = ["lib/Dialect/Linalg/LinalgRegistration.cpp"], - deps = [":Linalg"], + srcs = ["lib/Dialect/Linalg/IR/LinalgRegistration.cpp"], + deps = [":LinalgOps"], alwayslink = 1, ) diff --git a/third_party/mlir/test.BUILD b/third_party/mlir/test.BUILD index a0a05aa1356..943ed51c8ab 100644 --- a/third_party/mlir/test.BUILD +++ b/third_party/mlir/test.BUILD @@ -156,7 +156,8 @@ cc_library( "@llvm-project//mlir:EDSC", "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:IR", - "@llvm-project//mlir:Linalg", + "@llvm-project//mlir:LinalgOps", + "@llvm-project//mlir:LinalgTransforms", "@llvm-project//mlir:LoopOps", "@llvm-project//mlir:Pass", "@llvm-project//mlir:StandardOps",