Integrate LLVM at https://github.com/llvm/llvm-project/commit/35cf2f42dda4

PiperOrigin-RevId: 308019661
Change-Id: Ie6f60e5ff4afab31dcfeabe9543c348ab1313fb6

tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc

@@ -312,11 +312,8 @@ struct FixKernelFunctionSignatures
     mlir::FuncOp func = getFunction();
     mlir::ModuleOp module = func.getParentOfType<mlir::ModuleOp>();
     getFunction().walk([&](mlir::gpu::LaunchFuncOp launchOp) {
-      mlir::gpu::GPUModuleOp gpu_module =
-          module.lookupSymbol<mlir::gpu::GPUModuleOp>(
-              launchOp.getKernelModuleName());
       mlir::gpu::GPUFuncOp kernel =
-          gpu_module.lookupSymbol<mlir::gpu::GPUFuncOp>(launchOp.kernel());
+          module.lookupSymbol<mlir::gpu::GPUFuncOp>(launchOp.kernel());
       // Compute a map from function arguments to kernel function operands.
       mlir::BlockAndValueMapping func_to_kernel;
       for (mlir::BlockArgument arg : func.getArguments()) {
@@ -331,6 +328,7 @@ struct FixKernelFunctionSignatures
       // Create a new kernel function with modified signature. We know that it
       // will have the same signature as the original function, so just reuse it
       // here.
+      auto gpu_module = kernel.getParentOfType<mlir::gpu::GPUModuleOp>();
       mlir::OpBuilder kernel_builder(gpu_module.body());
       auto new_kernel = kernel_builder.create<mlir::gpu::GPUFuncOp>(
           kernel.getLoc(), kernel.getName(), func.getType());

tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc

@@ -404,7 +404,8 @@ StatusOr<std::unique_ptr<gpu::KernelThunk>> TransformKernelToXlaThunk(
     return Unimplemented("No kernel was generated.");
   }
 
-  auto kernel = kernel_module.lookupSymbol<LLVMFuncOp>(launchOp.kernel());
+  auto kernel =
+      kernel_module.lookupSymbol<LLVMFuncOp>(launchOp.getKernelName());
 
   // Store the assignment of operands to block arguments. Note that an operand
   // might be used in multiple argument positions, hence the vector.

third_party/mlir/BUILD

@@ -2,6 +2,7 @@
 #   The MLIR "Multi-Level Intermediate Representation" Compiler Infrastructure
 
 load("@org_tensorflow//third_party/mlir:tblgen.bzl", "gentbl")
+load("@org_tensorflow//third_party/mlir:linalggen.bzl", "genlinalg")
 
 licenses(["notice"])
 
@@ -278,6 +279,7 @@ filegroup(
     name = "LoopOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/LoopOps/LoopOps.td",
+        "include/mlir/Interfaces/ControlFlowInterfaces.td",
         "include/mlir/Interfaces/LoopLikeInterface.td",
         "include/mlir/Interfaces/SideEffects.td",
         ":OpBaseTdFiles",
@@ -578,6 +580,7 @@ cc_library(
     ]),
     includes = ["include"],
     deps = [
+        ":ControlFlowInterfaces",
         ":EDSC",
         ":IR",
         ":LoopLikeInterface",
@@ -1898,6 +1901,7 @@ cc_library(
     deps = [
         ":Affine",
         ":Analysis",
+        ":ControlFlowInterfaces",
         ":IR",
         ":LoopLikeInterface",
         ":LoopOps",
@@ -1907,7 +1911,6 @@ cc_library(
         ":Support",
         ":TransformUtils",
         ":TransformsPassIncGen",
-        ":VectorOps",
         "@llvm-project//llvm:support",
     ],
 )
@@ -2869,9 +2872,26 @@ gentbl(
     ],
 )
 
+genlinalg(
+    name = "LinalgNamedStructuredOpsIncGen",
+    src = "include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc",
+    linalg_outs = [
+        (
+            "-gen-impl",
+            "include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.cpp.inc",
+        ),
+        (
+            "-gen-ods-decl",
+            "include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.td",
+        ),
+    ],
+    linalggen = ":mlir-linalg-ods-gen",
+)
+
 filegroup(
     name = "LinalgStructuredOpsTdFiles",
     srcs = [
+        "include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.td",
         "include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td",
         "include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td",
         ":AffineOpsTdFiles",
@@ -2905,6 +2925,7 @@ gentbl(
     td_file = "include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td",
     td_srcs = [
         ":LinalgStructuredOpsTdFiles",
+        "include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.td",
     ],
 )
 
@@ -3020,14 +3041,18 @@ cc_library(
         "lib/Dialect/Linalg/IR/LinalgTypes.cpp",
     ],
     hdrs = [
+        "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",
     ],
     includes = ["include"],
     deps = [
+        ":Affine",
         ":DialectUtils",
+        ":EDSC",
         ":IR",
+        ":LinalgNamedStructuredOpsIncGen",
         ":LinalgOpsIncGen",
         ":LinalgStructuredOpsIncGen",
         ":Parser",
@@ -3068,7 +3093,6 @@ cc_library(
         "include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h",
         "include/mlir/Dialect/Linalg/EDSC/Builders.h",
         "include/mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h",
-        "include/mlir/Dialect/Linalg/EDSC/Intrinsics.h",
         "include/mlir/Dialect/Linalg/Passes.h",
         "include/mlir/Dialect/Linalg/Transforms/LinalgTransforms.h",
         "include/mlir/Dialect/Linalg/Utils/Utils.h",

third_party/mlir/linalggen.bzl

@@ -0,0 +1,39 @@
+"""BUILD extensions for MLIR linalg generation."""
+
+def genlinalg(name, linalggen, src, linalg_outs):
+    """genlinalg() generates code from a tc spec file.
+
+    Args:
+      name: The name of the build rule for use in dependencies.
+      linalggen: The binary used to produce the output.
+      src: The tc spec file.
+      linalg_outs: A list of tuples (opts, out), where each opts is a string of
+        options passed to linalggen, and the out is the corresponding output file
+        produced.
+    """
+
+    for (opts, out) in linalg_outs:
+        # All arguments to generate the output except output destination.
+        base_args = [
+            "$(location %s)" % linalggen,
+            "%s" % opts,
+            "$(location %s)" % src,
+        ]
+        rule_suffix = "_".join(opts.replace("-", "_").replace("=", "_").split(" "))
+
+        # Rule to generate code using generated shell script.
+        native.genrule(
+            name = "%s_%s_genrule" % (name, rule_suffix),
+            srcs = [src],
+            outs = [out],
+            tools = [linalggen],
+            cmd = ("echo " + " ".join(base_args) + " -o $@; " + " ".join(base_args) + " -o $@"),
+        )
+
+    # List of opts that do not generate cc files.
+    hdrs = [f for (opts, f) in linalg_outs]
+    native.cc_library(
+        name = name,
+        hdrs = hdrs,
+        textual_hdrs = hdrs,
+    )

third_party/mlir/tblgen.bzl

@@ -21,6 +21,7 @@ def gentbl(name, tblgen, td_file, tbl_outs, td_srcs = [], td_includes = [], stri
         srcs += [td_file]
 
     td_includes_cmd = ["-I external/llvm-project/mlir/include -I external/org_tensorflow"]
+    td_includes_cmd += ["-I $(GENDIR)/third_party/llvm/llvm-project/mlir/include"]
     for td_include in td_includes:
         td_includes_cmd += ["-I%s" % td_include]
     local_inc = "-I $$(dirname $(location %s))" % td_file