Integrate LLVM at llvm/llvm-project@26750a1264

Updates LLVM usage to match
[26750a1264b3](https://github.com/llvm/llvm-project/commit/26750a1264b3)

PiperOrigin-RevId: 339127601
Change-Id: I006d610bd1bcc5054d94533ab93dce32db75600a

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>([&](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<mlir::ReturnOp, mlir::ReturnOp,
                                               lmhlo::CopyOp>(
         &context, converter, patterns);
-    populateShapeTypeConversionPatterns(&context, converter, patterns);
+    populateShapeStructuralTypeConversionsAndLegality(&context, converter,
+                                                      patterns, target);
     if (failed(applyPartialConversion(getOperation(), target, patterns)))
       signalPassFailure();
   }

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(

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

tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize_pass.cc

@@ -97,8 +97,7 @@ struct BufferizePass : public BufferizePassBase<BufferizePass> {
       return converter.isLegal(inputs) && converter.isLegal(results) &&
              converter.isLegal(&op.getBody());
     });
-    target.addDynamicallyLegalOp<CallOp, ReturnOp, SelectOp, shape::AssumingOp>(
+    target.addDynamicallyLegalOp<CallOp, ReturnOp, SelectOp>(typesAreLegal);
-        typesAreLegal);
 
     OwningRewritePatternList patterns;
     mhlo::populateHLOToLHLOConversionPattern(&context, &converter, &patterns);
@@ -106,7 +105,8 @@ struct BufferizePass : public BufferizePassBase<BufferizePass> {
                                               lmhlo::CopyOp>(
         &context, converter, patterns);
     populateStandardBufferizePattern(&context, &converter, &patterns);
-    populateShapeTypeConversionPatterns(&context, converter, patterns);
+    populateShapeStructuralTypeConversionsAndLegality(&context, converter,
+                                                      patterns, target);
     patterns.insert<UnrankedTensorStoreTestOnlyPattern>(&context);
 
     auto module = getOperation();

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:.*]]]

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>)
+  // CHECK-SAME: memref<32xf32>,
-  "gpu.launch_func"(%cst, %cst, %cst, %cst, %cst, %cst, %res, %arg1, %arg0)
+  // CHECK-SAME: memref<16xf32>,
-      { kernel = @kernel_module::@kernel }
+  // CHECK-SAME: memref<8xf32>)
-      : (index, index, index, index, index, index,
+  gpu.launch_func @kernel_module::@kernel
-         memref<8xf32>, memref<16xf32>, memref<32xf32>) -> ()
+      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)
+  gpu.launch_func @kernel_module::@kernel
-      { kernel = @kernel_module::@kernel }
+      blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
-      : (index, index, index, index, index, index,
+      args(%arg1 : memref<16xf32>, %arg0 : memref<32xf32>)
-         memref<16xf32>, 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)
+  gpu.launch_func @kernel_module::@kernel
-      { kernel = @kernel_module::@kernel }
+      blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
-      : (index, index, index, index, index, index,
+      args(%arg1 : memref<16xf32>, %arg0 : memref<32xf32>, %fake : memref<8xf32>)
-         memref<16xf32>, memref<32xf32>, 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)
+  gpu.launch_func @kernel_module::@kernel
-      { kernel = @kernel_module::@kernel }
+      blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
-      : (index, index, index, index, index, index,
+      args(%fake : memref<16xf32>, %arg0 : memref<32xf32>, %res : memref<8xf32>)
-         memref<16xf32>, memref<32xf32>, 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)
+  gpu.launch_func @kernel_module::@kernel
-      { kernel = @kernel_module::@kernel }
+      blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
-      : (index, index, index, index, index, index,
+      args(%res : memref<8xf32>, %arg1 : memref<16xf32>, %arg0 : memref<32xf32>)
-         memref<8xf32>, memref<16xf32>, memref<32xf32>) -> ()
 
   return %res : memref<8xf32>
 }

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_COMMIT = "26750a1264b3df114a1efae7cde6f0784206b2ce"
-    LLVM_SHA256 = "b35dd27eace459897c07faa333b0cb9ddc0ef260b20582dd04b6910d548a7e08"
+    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.

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 = [