diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc index 7a1b3f4f416..1ce93b977c6 100644 --- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc +++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc @@ -541,10 +541,6 @@ struct HloLegalizeToLhlo return std::all_of(op.operand_type_begin(), op.operand_type_end(), isMemRefType); }); - target.addDynamicallyLegalOp([&](shape::AssumingOp op) { - return std::all_of(op.result_type_begin(), op.result_type_end(), - isMemRefType); - }); auto kind = results_escape_function ? BufferizeTypeConverter::KeepAsFunctionResult @@ -557,7 +553,8 @@ struct HloLegalizeToLhlo populateWithBufferizeOpConversionPatterns( &context, converter, patterns); - populateShapeTypeConversionPatterns(&context, converter, patterns); + populateShapeStructuralTypeConversionsAndLegality(&context, converter, + patterns, target); if (failed(applyPartialConversion(getOperation(), target, patterns))) signalPassFailure(); } diff --git a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/lhlo_fuse_linalg.cc b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/lhlo_fuse_linalg.cc index 8f50ad0667f..232948775c0 100644 --- a/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/lhlo_fuse_linalg.cc +++ b/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/lhlo_fuse_linalg.cc @@ -121,7 +121,7 @@ class LhloFuseLinalgPass for (unsigned id = 0, e = LinalgOp(op).getNumInputs(); id < e; ++id) { linalg::Aliases aliases; linalg::LinalgDependenceGraph graph(aliases, linalg_ops); - if (auto info = fuseProducerOf(b, op, id, graph, &folder)) { + if (auto info = fuseProducerOfBuffer(b, op, id, graph, &folder)) { auto originalOp = info->originalProducer.getOperation(); erase_set.insert(originalOp); auto originalOpInLinalgOpsVector = std::find_if( diff --git a/tensorflow/compiler/mlir/runlit.cfg.py b/tensorflow/compiler/mlir/runlit.cfg.py index 17410b4e5b2..8b81f8a7a4c 100644 --- a/tensorflow/compiler/mlir/runlit.cfg.py +++ b/tensorflow/compiler/mlir/runlit.cfg.py @@ -37,7 +37,7 @@ config.name = 'MLIR ' + os.path.basename(config.mlir_test_dir) config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell) # suffixes: A list of file extensions to treat as test files. -config.suffixes = ['.cc', '.hlo', '.hlotxt', '.mlir', '.pbtxt', '.py'] +config.suffixes = ['.cc', '.hlo', '.hlotxt', '.json', '.mlir', '.pbtxt', '.py'] # test_source_root: The root path where tests are located. config.test_source_root = config.mlir_test_dir diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize_pass.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize_pass.cc index 9a531515012..bba32fc56dc 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize_pass.cc +++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize_pass.cc @@ -97,8 +97,7 @@ struct BufferizePass : public BufferizePassBase { return converter.isLegal(inputs) && converter.isLegal(results) && converter.isLegal(&op.getBody()); }); - target.addDynamicallyLegalOp( - typesAreLegal); + target.addDynamicallyLegalOp(typesAreLegal); OwningRewritePatternList patterns; mhlo::populateHLOToLHLOConversionPattern(&context, &converter, &patterns); @@ -106,7 +105,8 @@ struct BufferizePass : public BufferizePassBase { lmhlo::CopyOp>( &context, converter, patterns); populateStandardBufferizePattern(&context, &converter, &patterns); - populateShapeTypeConversionPatterns(&context, converter, patterns); + populateShapeStructuralTypeConversionsAndLegality(&context, converter, + patterns, target); patterns.insert(&context); auto module = getOperation(); diff --git a/tensorflow/compiler/xla/service/mlir_gpu/tests/add_in_gpu_dialect.hlo b/tensorflow/compiler/xla/service/mlir_gpu/tests/add_in_gpu_dialect.hlo index c640130d245..db39919a803 100644 --- a/tensorflow/compiler/xla/service/mlir_gpu/tests/add_in_gpu_dialect.hlo +++ b/tensorflow/compiler/xla/service/mlir_gpu/tests/add_in_gpu_dialect.hlo @@ -8,7 +8,9 @@ ENTRY %Add (x: f32[2,2], y: f32[2,2]) -> f32[2,2] { } // CHECK: func @add(%[[ARG0:.*]]: [[TYPE:.*]], %[[ARG1:.*]]: [[TYPE]], %[[ARG2:.*]]: [[TYPE]]) { -// CHECK: "gpu.launch_func"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[ARG0]], %[[ARG1]], %[[ARG2]] +// CHECK: gpu.launch_func +// CHECK-SAME: blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) args +// CHECK-SAME: (%[[ARG0]] : [[TYPE]], %[[ARG1]] : [[TYPE]], %[[ARG2]] : [[TYPE]]) // CHECK: } // CHECK: func @add_kernel(%[[ARG0]]: [[TYPE]], %[[ARG1]]: [[TYPE]], %[[ARG2]]: [[TYPE]] // CHECK-DAG: subview %[[ARG0]]{{\[}}[[INDEX:.*]]] diff --git a/tensorflow/compiler/xla/service/mlir_gpu/tests/passes/rewrite_kernel_signatures.mlir b/tensorflow/compiler/xla/service/mlir_gpu/tests/passes/rewrite_kernel_signatures.mlir index cff1989f05b..befbae7382e 100644 --- a/tensorflow/compiler/xla/service/mlir_gpu/tests/passes/rewrite_kernel_signatures.mlir +++ b/tensorflow/compiler/xla/service/mlir_gpu/tests/passes/rewrite_kernel_signatures.mlir @@ -19,11 +19,12 @@ func @caller(%arg0: memref<32xf32>, %arg1: memref<16xf32>) -> memref<8xf32> { %res = alloc() : memref<8xf32> // CHECK: gpu.launch_func - // CHECK-SAME: index, memref<32xf32>, memref<16xf32>, memref<8xf32>) - "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %res, %arg1, %arg0) - { kernel = @kernel_module::@kernel } - : (index, index, index, index, index, index, - memref<8xf32>, memref<16xf32>, memref<32xf32>) -> () + // CHECK-SAME: memref<32xf32>, + // CHECK-SAME: memref<16xf32>, + // CHECK-SAME: memref<8xf32>) + gpu.launch_func @kernel_module::@kernel + blocks in (%cst, %cst, %cst) threads in(%cst, %cst, %cst) + args(%res : memref<8xf32>, %arg1 : memref<16xf32>, %arg0 : memref<32xf32>) return %res : memref<8xf32> } @@ -45,10 +46,9 @@ func @caller(%arg0: memref<32xf32>, %arg1: memref<16xf32>) -> memref<8xf32> { %cst = constant 8 : index %res = alloc() : memref<8xf32> - "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %arg1, %arg0) - { kernel = @kernel_module::@kernel } - : (index, index, index, index, index, index, - memref<16xf32>, memref<32xf32>) -> () + gpu.launch_func @kernel_module::@kernel + blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) + args(%arg1 : memref<16xf32>, %arg0 : memref<32xf32>) return %res : memref<8xf32> } @@ -72,10 +72,9 @@ func @caller(%arg0: memref<32xf32>, %arg1: memref<16xf32>) -> memref<8xf32> { %res = alloc() : memref<8xf32> %fake = alloc() : memref<8xf32> - "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %arg1, %arg0, %fake) - { kernel = @kernel_module::@kernel } - : (index, index, index, index, index, index, - memref<16xf32>, memref<32xf32>, memref<8xf32>) -> () + gpu.launch_func @kernel_module::@kernel + blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) + args(%arg1 : memref<16xf32>, %arg0 : memref<32xf32>, %fake : memref<8xf32>) return %res : memref<8xf32> } @@ -99,10 +98,9 @@ func @caller(%arg0: memref<32xf32>, %arg1: memref<16xf32>) -> memref<8xf32> { %res = alloc() : memref<8xf32> %fake = alloc() : memref<16xf32> - "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %fake, %arg0, %res) - { kernel = @kernel_module::@kernel } - : (index, index, index, index, index, index, - memref<16xf32>, memref<32xf32>, memref<8xf32>) -> () + gpu.launch_func @kernel_module::@kernel + blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) + args(%fake : memref<16xf32>, %arg0 : memref<32xf32>, %res : memref<8xf32>) return %res : memref<8xf32> } @@ -127,10 +125,9 @@ func @caller(%arg0: memref<32xf32>, %arg1: memref<16xf32>) -> memref<8xf32> { br ^bb1 ^bb1: - "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %res, %arg1, %arg0) - { kernel = @kernel_module::@kernel } - : (index, index, index, index, index, index, - memref<8xf32>, memref<16xf32>, memref<32xf32>) -> () + gpu.launch_func @kernel_module::@kernel + blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) + args(%res : memref<8xf32>, %arg1 : memref<16xf32>, %arg0 : memref<32xf32>) return %res : memref<8xf32> } diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl index 80580ac4930..952e786704a 100755 --- a/tensorflow/workspace.bzl +++ b/tensorflow/workspace.bzl @@ -680,8 +680,8 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""): ) # Check out LLVM and MLIR from llvm-project. - LLVM_COMMIT = "c89447b65984c97145f63be21e42cfa98da60dd2" - LLVM_SHA256 = "b35dd27eace459897c07faa333b0cb9ddc0ef260b20582dd04b6910d548a7e08" + LLVM_COMMIT = "26750a1264b3df114a1efae7cde6f0784206b2ce" + LLVM_SHA256 = "eb360bbcd3e4b505689a21756c89fdf087064882899bab2f98258f8cf0546218" 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), @@ -696,7 +696,6 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""): "//third_party/mlir:BUILD": "mlir/BUILD", "//third_party/mlir:test.BUILD": "mlir/test/BUILD", }, - patch_file = clean_dep("//third_party:llvm_fix_windows.patch"), ) # Intel openMP that is part of LLVM sources. diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD index d129c475a0d..04533ae5f24 100644 --- a/third_party/mlir/BUILD +++ b/third_party/mlir/BUILD @@ -120,6 +120,7 @@ cc_library( cc_library( name = "CAPIIR", srcs = [ + "lib/CAPI/IR/AffineExpr.cpp", "lib/CAPI/IR/AffineMap.cpp", "lib/CAPI/IR/Diagnostics.cpp", "lib/CAPI/IR/IR.cpp", @@ -129,6 +130,7 @@ cc_library( "lib/CAPI/Standard/StandardDialect.cpp", ], hdrs = [ + "include/mlir-c/AffineExpr.h", "include/mlir-c/AffineMap.h", "include/mlir-c/Diagnostics.h", "include/mlir-c/IR.h", @@ -136,6 +138,7 @@ cc_library( "include/mlir-c/StandardDialect.h", "include/mlir-c/StandardTypes.h", "include/mlir-c/Support.h", + "include/mlir/CAPI/AffineExpr.h", "include/mlir/CAPI/AffineMap.h", "include/mlir/CAPI/Diagnostics.h", "include/mlir/CAPI/IR.h", @@ -582,6 +585,28 @@ cc_library( ], ) +cc_library( + name = "AsyncTransforms", + srcs = glob([ + "lib/Dialect/Async/Transforms/*.cpp", + "lib/Dialect/Async/Transforms/*.h", + ]), + includes = ["include"], + deps = [ + ":Async", + ":IR", + ":Pass", + ":SCFDialect", + ":StandardOps", + ":Support", + ":TransformUtils", + ":Transforms", + ":TransformsPassIncGen", + "@llvm-project//llvm:Core", + "@llvm-project//llvm:Support", + ], +) + cc_library( name = "AffineUtils", srcs = glob( @@ -665,6 +690,7 @@ cc_library( deps = [ ":AVX512ToLLVM", ":AffineToStandard", + ":AsyncToLLVM", ":ConversionPassIncGen", ":GPUToGPURuntimeTransforms", ":GPUToNVVMTransforms", @@ -688,6 +714,27 @@ cc_library( ], ) +cc_library( + name = "AsyncToLLVM", + srcs = glob([ + "lib/Conversion/AsyncToLLVM/*.cpp", + "lib/Conversion/AsyncToLLVM/*.h", + ]) + ["lib/Conversion/PassDetail.h"], + hdrs = glob(["include/mlir/Conversion/AsyncToLLVM/*.h"]), + includes = ["include"], + deps = [ + ":Async", + ":ConversionPassIncGen", + ":IR", + ":LLVMDialect", + ":Pass", + ":StandardOps", + ":Support", + ":Transforms", + "@llvm-project//llvm:Support", + ], +) + cc_library( name = "AffineToStandard", srcs = glob([ @@ -977,6 +1024,7 @@ cc_library( ":ControlFlowInterfaces", ":IR", ":Pass", + ":SCFDialect", ":StandardOps", ":StandardOpsTransformsPassIncGen", ":Support", @@ -1659,6 +1707,7 @@ cc_library( includes = ["include"], deps = [ ":ConversionPassIncGen", + ":GPUDialect", ":IR", ":LLVMDialect", ":Pass", @@ -2805,11 +2854,13 @@ cc_library( deps = [ ":Affine", ":CallOpInterfaces", + ":ControlFlowInterfaces", ":IR", ":LinalgOps", ":SCFDialect", ":StandardOps", ":Support", + ":ViewLikeInterface", "@llvm-project//llvm:Support", ], ) @@ -2972,6 +3023,7 @@ cc_library( deps = [ "@llvm-project//llvm:Analysis", "@llvm-project//llvm:Core", + "@llvm-project//llvm:Coroutines", "@llvm-project//llvm:IPO", "@llvm-project//llvm:Support", "@llvm-project//llvm:Target", @@ -3073,6 +3125,8 @@ cc_library( ":AffineToStandard", ":AffineTransforms", ":Async", + ":AsyncToLLVM", + ":AsyncTransforms", ":ConversionPasses", ":GPUDialect", ":GPUPassIncGen", @@ -3201,6 +3255,17 @@ cc_library( includes = ["include"], ) +cc_library( + name = "mlir_async_runtime", + srcs = [ + "lib/ExecutionEngine/AsyncRuntime.cpp", + ], + hdrs = [ + "include/mlir/ExecutionEngine/AsyncRuntime.h", + ], + includes = ["include"], +) + cc_library( name = "mlir_runner_utils", srcs = [