From 1a36ba72a657cff84a7681e8507c3d009a72b6e5 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 23 Apr 2020 04:03:38 -0700 Subject: [PATCH] Integrate LLVM at https://github.com/llvm/llvm-project/commit/35cf2f42dda4 PiperOrigin-RevId: 308019661 Change-Id: Ie6f60e5ff4afab31dcfeabe9543c348ab1313fb6 --- .../xla/service/mlir_gpu/kernel_lowering.cc | 6 +-- .../service/mlir_gpu/mlir_compiler_impl.cc | 3 +- third_party/mlir/BUILD | 28 ++++++++++++- third_party/mlir/linalggen.bzl | 39 +++++++++++++++++++ third_party/mlir/tblgen.bzl | 1 + 5 files changed, 70 insertions(+), 7 deletions(-) create mode 100644 third_party/mlir/linalggen.bzl diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc index 2cfd5110f13..33d3690d4ab 100644 --- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc +++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc @@ -312,11 +312,8 @@ struct FixKernelFunctionSignatures mlir::FuncOp func = getFunction(); mlir::ModuleOp module = func.getParentOfType(); getFunction().walk([&](mlir::gpu::LaunchFuncOp launchOp) { - mlir::gpu::GPUModuleOp gpu_module = - module.lookupSymbol( - launchOp.getKernelModuleName()); mlir::gpu::GPUFuncOp kernel = - gpu_module.lookupSymbol(launchOp.kernel()); + module.lookupSymbol(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::OpBuilder kernel_builder(gpu_module.body()); auto new_kernel = kernel_builder.create( kernel.getLoc(), kernel.getName(), func.getType()); diff --git a/tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc b/tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc index c258d532f8e..35ac3b2bf63 100644 --- a/tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc +++ b/tensorflow/compiler/xla/service/mlir_gpu/mlir_compiler_impl.cc @@ -404,7 +404,8 @@ StatusOr> TransformKernelToXlaThunk( return Unimplemented("No kernel was generated."); } - auto kernel = kernel_module.lookupSymbol(launchOp.kernel()); + auto kernel = + kernel_module.lookupSymbol(launchOp.getKernelName()); // Store the assignment of operands to block arguments. Note that an operand // might be used in multiple argument positions, hence the vector. diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD index eedac459509..ffb811f5cc3 100644 --- a/third_party/mlir/BUILD +++ b/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", diff --git a/third_party/mlir/linalggen.bzl b/third_party/mlir/linalggen.bzl new file mode 100644 index 00000000000..5162911720f --- /dev/null +++ b/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, + ) diff --git a/third_party/mlir/tblgen.bzl b/third_party/mlir/tblgen.bzl index 6434bba762c..623d085ccc3 100644 --- a/third_party/mlir/tblgen.bzl +++ b/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