From 8fc976574e66186e3b4c4b94a6477eb090618cab Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Thu, 21 May 2020 18:03:31 -0700 Subject: [PATCH] Integrate LLVM at https://github.com/llvm/llvm-project/commit/1108f5c737db PiperOrigin-RevId: 312775865 Change-Id: Iee2170660e6b2cd0a81695e8843bebfb311c480b --- tensorflow/compiler/mlir/xla/ir/lhlo_ops.td | 4 ++ .../xla/transforms/hlo_legalize_to_lhlo.cc | 4 +- .../mlir/xla/transforms/lhlo_fuse_linalg.cc | 26 ++++++---- third_party/mlir/BUILD | 49 ++++++++++++++++++- 4 files changed, 70 insertions(+), 13 deletions(-) diff --git a/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td b/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td index 020859aa0bf..9a2168d3088 100644 --- a/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td +++ b/tensorflow/compiler/mlir/xla/ir/lhlo_ops.td @@ -444,6 +444,10 @@ def TerminatorOp : let description = [{ Terminator operation for the LHLO dialect. }]; + let builders = [OpBuilder< + "OpBuilder &b, OperationState &result, ValueRange operands", + [{ build(b, result, llvm::None, operands, llvm::None); }] + >]; } #endif // LHLO_OPS diff --git a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc index 11b2ae65d8e..5851bad4565 100644 --- a/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc +++ b/tensorflow/compiler/mlir/xla/transforms/hlo_legalize_to_lhlo.cc @@ -43,8 +43,8 @@ constexpr StringRef kTempBufferAttr = "temp"; template using BaseOpConversion = BufferAssignmentOpConversionPattern; using StdReturnOpConverter = - NonVoidToVoidReturnOpConverter; + NoBufferOperandsReturnOpConverter; Value InsertDynamicAllocAndDealloc(Location loc, Value result, Value shape_operand, diff --git a/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc b/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc index 43c0911a4a6..ddbb672c70a 100644 --- a/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc +++ b/tensorflow/compiler/mlir/xla/transforms/lhlo_fuse_linalg.cc @@ -57,8 +57,9 @@ class LhloFuseLinalg : public PassWrapper { for (auto func_arg : func.getArguments()) { func_args.insert(func_arg); } + MLIRContext* ctx = func.getContext(); OpBuilder b(func); - OperationFolder folder(func.getContext()); + OperationFolder folder(ctx); func.walk([&](linalg::GenericOp generic_op) { SmallVector tile_sizes(tile_sizes_.begin(), tile_sizes_.end()); @@ -68,12 +69,14 @@ class LhloFuseLinalg : public PassWrapper { auto op = cast(generic_op.getOperation()); for (const Value result : op.getOutputBuffers()) { if (!func_args.count(result)) continue; - if (tileGenericOp(op, tile_sizes, &b, &folder)) { + if (tileGenericOp(op, tile_sizes, &b)) { generic_op.erase(); return; } } }); + auto patterns = linalg::getLinalgTilingCanonicalizationPatterns(ctx); + applyPatternsAndFoldGreedily(func, patterns); // Fuse producers of tiled linalg ops. llvm::SmallDenseSet erase_set; @@ -92,19 +95,22 @@ class LhloFuseLinalg : public PassWrapper { *originalOpInLinalgOpsVector = info->fusedProducer.getOperation(); } } + + auto patterns = linalg::getLinalgTilingCanonicalizationPatterns(ctx); + applyPatternsAndFoldGreedily(func, patterns); } for (auto* e : erase_set) e->erase(); } private: - bool tileGenericOp(LinalgOp op, ArrayRef tile_sizes, OpBuilder* b, - OperationFolder* folder) { - auto tiled_generic_op = - use_parallel_loops_ - ? linalg::tileLinalgOpToParallelLoops(*b, op, tile_sizes, - /*permutation=*/{}, folder) - : linalg::tileLinalgOp(*b, op, tile_sizes, - /*permutation=*/{}, folder); + bool tileGenericOp(LinalgOp op, ArrayRef tile_sizes, OpBuilder* b) { + auto loopType = use_parallel_loops_ + ? linalg::LinalgTilingLoopType::ParallelLoops + : linalg::LinalgTilingLoopType::Loops; + auto tiled_generic_op = linalg::tileLinalgOp(*b, op, + linalg::LinalgTilingOptions() + .setTileSizes(tile_sizes) + .setLoopType(loopType)); return tiled_generic_op.hasValue(); } diff --git a/third_party/mlir/BUILD b/third_party/mlir/BUILD index 5ebcbb6e3d2..a57088432e2 100644 --- a/third_party/mlir/BUILD +++ b/third_party/mlir/BUILD @@ -175,6 +175,7 @@ filegroup( filegroup( name = "AffineOpsTdFiles", srcs = [ + "include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td", "include/mlir/Dialect/Affine/IR/AffineOps.td", "include/mlir/Dialect/Affine/IR/AffineOpsBase.td", "include/mlir/Interfaces/LoopLikeInterface.td", @@ -207,6 +208,26 @@ gentbl( ], ) +gentbl( + name = "AffineMemoryOpInterfacesIncGen", + strip_include_prefix = "include", + tbl_outs = [ + ( + "-gen-op-interface-decls", + "include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.h.inc", + ), + ( + "-gen-op-interface-defs", + "include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.cpp.inc", + ), + ], + tblgen = ":mlir-tblgen", + td_file = "include/mlir/Dialect/Affine/IR/AffineMemoryOpInterfaces.td", + td_srcs = [ + ":AffineOpsTdFiles", + ], +) + ##---------------------------------------------------------------------------## # AVX512 dialect. ##---------------------------------------------------------------------------## @@ -462,6 +483,7 @@ cc_library( ]), includes = ["include"], deps = [ + ":AffineMemoryOpInterfacesIncGen", ":AffineOpsIncGen", ":EDSC", ":IR", @@ -677,6 +699,7 @@ cc_library( deps = [ ":CallOpInterfaces", ":CommonFolders", + ":ControlFlowInterfaces", ":Dialect", ":IR", ":InferTypeOpInterface", @@ -1153,6 +1176,28 @@ cc_library( ], ) +cc_library( + name = "GPURuntimeTransforms", + srcs = [ + "lib/Conversion/GPUCommon/ConvertLaunchFuncToRuntimeCalls.cpp", + "lib/Conversion/PassDetail.h", + ], + hdrs = [ + "include/mlir/Conversion/GPUCommon/GPUCommonPass.h", + ], + includes = ["include"], + deps = [ + ":ConversionPassIncGen", + ":GPUDialect", + ":IR", + ":LLVMDialect", + ":Pass", + ":Support", + "@llvm-project//llvm:core", + "@llvm-project//llvm:support", + ], +) + gentbl( name = "GPUToNVVMGen", strip_include_prefix = "lib/Conversion/GPUToNVVM", @@ -1265,7 +1310,6 @@ cc_library( name = "GPUToCUDATransforms", srcs = [ "lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp", - "lib/Conversion/GPUToCUDA/ConvertLaunchFuncToCudaCalls.cpp", "lib/Conversion/PassDetail.h", ], hdrs = ["include/mlir/Conversion/GPUToCUDA/GPUToCUDAPass.h"], @@ -2446,6 +2490,7 @@ cc_library( includes = ["include"], deps = [ ":Analysis", + ":GPURuntimeTransforms", ":GPUToNVVMTransforms", ":GPUToROCDLTransforms", ":GPUToSPIRVTransforms", @@ -2525,6 +2570,7 @@ cc_library( ":ConversionPassIncGen", ":GPUDialect", ":GPUPassIncGen", + ":GPURuntimeTransforms", ":GPUToCUDATransforms", ":GPUToNVVMTransforms", ":GPUToROCDLTransforms", @@ -2730,6 +2776,7 @@ cc_binary( ":AllPassesAndDialectsNoRegistration", ":ExecutionEngineUtils", ":GPUDialect", + ":GPURuntimeTransforms", ":GPUToNVVMTransforms", ":GPUToROCDLTransforms", ":GPUTransforms",