diff --git a/clang/docs/tools/clang-formatted-files.txt b/clang/docs/tools/clang-formatted-files.txt --- a/clang/docs/tools/clang-formatted-files.txt +++ b/clang/docs/tools/clang-formatted-files.txt @@ -7777,6 +7777,7 @@ mlir/include/mlir/ExecutionEngine/OptUtils.h mlir/include/mlir/ExecutionEngine/RunnerUtils.h mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h +mlir/include/mlir/Interfaces/AllocLikeOpInterface.h mlir/include/mlir/Interfaces/CallInterfaces.h mlir/include/mlir/Interfaces/CastInterfaces.h mlir/include/mlir/Interfaces/ControlFlowInterfaces.h @@ -8299,6 +8300,7 @@ mlir/lib/ExecutionEngine/RocmRuntimeWrappers.cpp mlir/lib/ExecutionEngine/RunnerUtils.cpp mlir/lib/ExecutionEngine/SparseTensorUtils.cpp +mlir/lib/Interfaces/AllocLikeOpInterface.cpp mlir/lib/Interfaces/CallInterfaces.cpp mlir/lib/Interfaces/CastInterfaces.cpp mlir/lib/Interfaces/ControlFlowInterfaces.cpp diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h --- a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h @@ -22,6 +22,7 @@ #include "mlir/IR/OpDefinition.h" #include "mlir/IR/OpImplementation.h" #include "mlir/IR/SymbolTable.h" +#include "mlir/Interfaces/AllocLikeOpInterface.h" #include "mlir/Interfaces/InferIntRangeInterface.h" #include "mlir/Interfaces/InferTypeOpInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td @@ -19,6 +19,7 @@ include "mlir/IR/EnumAttr.td" include "mlir/IR/FunctionInterfaces.td" include "mlir/IR/SymbolInterfaces.td" +include "mlir/Interfaces/AllocLikeOpInterface.td" include "mlir/Interfaces/DataLayoutInterfaces.td" include "mlir/Interfaces/InferIntRangeInterface.td" include "mlir/Interfaces/InferTypeOpInterface.td" @@ -902,6 +903,7 @@ } def GPU_AllocOp : GPU_Op<"alloc", [ + AllocLikeOpInterface, GPU_AsyncOpInterface, AttrSizedOperandSegments ]> { @@ -932,6 +934,24 @@ let extraClassDeclaration = [{ MemRefType getType() { return memref().getType().cast(); } + + /// Returns the dynamic sizes for this alloc operation if specified. + operand_range getDynamicSizes() { return dynamicSizes(); } + + /// Returns the symbol operands for this alloc operation if specified. + operand_range getSymbolOperands() { return symbolOperands(); } + + /// Returns the allocated memref. + Value getMemref() { return memref(); } + + /// Returns the alignment of the allocated memref. For gpu.alloc op + /// this is guaranteed to be 256. + uint64_t getAlignmentAttr() {return 256; } + + /// We cannot set alignment for gpu.alloc op, so we do nothing. + void setAlignmentAttr(unsigned requiredAlignment) { + assert(requiredAlignment == 256 && "wrong alignment value for gpu.alloc op."); + } }]; let assemblyFormat = [{ diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRef.h @@ -11,6 +11,7 @@ #include "mlir/Dialect/Utils/ReshapeOpsUtils.h" #include "mlir/IR/Dialect.h" +#include "mlir/Interfaces/AllocLikeOpInterface.h" #include "mlir/Interfaces/CallInterfaces.h" #include "mlir/Interfaces/CastInterfaces.h" #include "mlir/Interfaces/ControlFlowInterfaces.h" diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td @@ -11,6 +11,8 @@ include "mlir/Dialect/Arithmetic/IR/ArithmeticBase.td" include "mlir/Dialect/MemRef/IR/MemRefBase.td" +include "mlir/IR/OpBase.td" +include "mlir/Interfaces/AllocLikeOpInterface.td" include "mlir/Interfaces/CastInterfaces.td" include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/CopyOpInterface.td" @@ -59,6 +61,7 @@ list traits = []> : MemRef_Op { @@ -97,6 +100,24 @@ static StringRef getAlignmentAttrStrName() { return "alignment"; } MemRefType getType() { return getResult().getType().cast(); } + + /// Returns the dynamic sizes for this alloc operation if specified. + operand_range getDynamicSizes() { return dynamicSizes(); } + + /// Returns the symbol operands for this alloc operation if specified. + operand_range getSymbolOperands() { return symbolOperands(); } + + /// Returns the allocated memref. + Value getMemref() { return memref(); } + + /// Returns the alignment of the allocated memref. + uint64_t getAlignmentAttr() { return alignment().getValue(); } + + /// Sets alignment value for the allocated memref. + void setAlignmentAttr(unsigned requiredAlignment) { + OpBuilder b(memref().getContext()); + alignmentAttr(b.getI64IntegerAttr(requiredAlignment)); + } }]; let assemblyFormat = [{ diff --git a/mlir/include/mlir/Interfaces/AllocLikeOpInterface.h b/mlir/include/mlir/Interfaces/AllocLikeOpInterface.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Interfaces/AllocLikeOpInterface.h @@ -0,0 +1,22 @@ +//===- AllocLikeOpInterface.h - alloc like operations interface-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file implements the operation interface for alloc-like operations. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_INTERFACES_ALLOCLIKEOPINTERFACE_H_ +#define MLIR_INTERFACES_ALLOCLIKEOPINTERFACE_H_ + +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" + +// Include the generated interface declarations. +#include "mlir/Interfaces/AllocLikeOpInterface.h.inc" + +#endif // MLIR_INTERFACES_ALLOCLIKEOPINTERFACE_H_ diff --git a/mlir/include/mlir/Interfaces/AllocLikeOpInterface.td b/mlir/include/mlir/Interfaces/AllocLikeOpInterface.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Interfaces/AllocLikeOpInterface.td @@ -0,0 +1,59 @@ +//===-------------- AllocLikeOpInterface.td -*- tablegen -*--------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Defines the interface for copy-like operations. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_INTERFACES_ALLOCLIKEOPINTERFACE +#define MLIR_INTERFACES_ALLOCLIKEOPINTERFACE + +include "mlir/IR/OpBase.td" + +def AllocLikeOpInterface : OpInterface<"AllocLikeOpInterface"> { + let description = [{ + A alloc-like operation is one that allocates memory of a given type. + }]; + let cppNamespace = "::mlir"; + + let methods = [ + InterfaceMethod< + /*desc=*/"Returns the dynamic dimension sizes of the resultant memref.", + /*retTy=*/"::llvm::SmallVector<::mlir::Value>", + /*methodName=*/"getDynamicSizes" + >, + InterfaceMethod< + /*desc=*/"Returns the symbolic operands which bind to the symbols of the memref's layout map.", + /*retTy=*/"::llvm::SmallVector<::mlir::Value>", + /*methodName=*/"getSymbolOperands" + >, + InterfaceMethod< + /*desc=*/"Returns the allocated memref.", + /*retTy=*/"::mlir::Value", + /*methodName=*/"getMemref" + >, + InterfaceMethod< + /*desc=*/"Returns the type of the allocated memref.", + /*retTy=*/"MemRefType", + /*methodName=*/"getType" + >, + InterfaceMethod< + /*desc=*/"Returns the alignment attribute of the allocated memref.", + /*retTy=*/"uint64_t", + /*methodName=*/"getAlignmentAttr" + >, + InterfaceMethod< + /*desc=*/"Sets the alignment attribute of the allocated memref.", + /*retTy=*/"void", + /*methodName=*/"setAlignmentAttr", + /*args=*/(ins "unsigned":$requiredAlignment) + > + ]; +} + +#endif // MLIR_INTERFACES_ALLOCLIKEOPINTERFACE diff --git a/mlir/include/mlir/Interfaces/CMakeLists.txt b/mlir/include/mlir/Interfaces/CMakeLists.txt --- a/mlir/include/mlir/Interfaces/CMakeLists.txt +++ b/mlir/include/mlir/Interfaces/CMakeLists.txt @@ -1,3 +1,4 @@ +add_mlir_interface(AllocLikeOpInterface) add_mlir_interface(CallInterfaces) add_mlir_interface(CastInterfaces) add_mlir_interface(ControlFlowInterfaces) diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp --- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp @@ -1964,7 +1964,7 @@ continue; // Use list expected to match the dep graph info. auto *op = memref.getDefiningOp(); - if (isa_and_nonnull(op)) + if (hasSingleEffect(op)) op->erase(); } } diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp --- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp @@ -1067,7 +1067,8 @@ for (auto memref : memrefsToErase) { // If the memref hasn't been locally alloc'ed, skip. Operation *defOp = memref.getDefiningOp(); - if (!defOp || !hasSingleEffect(defOp, memref)) + if (!defOp || (!hasSingleEffect(defOp) && + !isa(defOp))) // TODO: if the memref was returned by a 'call' operation, we // could still erase it if the call had no side-effects. continue; @@ -1639,7 +1640,7 @@ /// %2 = affine.apply #map2(%c4, %0) static void createNewDynamicSizes(MemRefType oldMemRefType, MemRefType newMemRefType, AffineMap map, - memref::AllocOp *allocOp, OpBuilder b, + AllocLikeOpInterface *allocOp, OpBuilder b, SmallVectorImpl &newDynamicSizes) { // Create new input for AffineApplyOp. SmallVector inAffineApply; @@ -1687,7 +1688,7 @@ } // TODO: Currently works for static memrefs with a single layout map. -LogicalResult mlir::normalizeMemRef(memref::AllocOp *allocOp) { +LogicalResult mlir::normalizeMemRef(AllocLikeOpInterface *allocOp) { MemRefType memrefType = allocOp->getType(); OpBuilder b(*allocOp); diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp @@ -52,7 +52,6 @@ #include "mlir/Dialect/Bufferization/Transforms/Passes.h" -#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h" #include "mlir/Dialect/Func/IR/FuncOps.h" diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -1157,9 +1157,12 @@ return failure(); // Transfer into `view`. - Value viewOrAlloc = xferOp.getSource(); - if (!viewOrAlloc.getDefiningOp() && - !viewOrAlloc.getDefiningOp()) + Value viewOrAlloc = xferOp.source(); + Operation *viewOrAllocDefOp = viewOrAlloc.getDefiningOp(); + if (!viewOrAllocDefOp) + return failure(); + if (!isa(viewOrAllocDefOp) && + !hasSingleEffect(viewOrAllocDefOp)) return failure(); LDBG(viewOrAlloc); @@ -1240,9 +1243,12 @@ return failure(); // Transfer into `viewOrAlloc`. - Value viewOrAlloc = xferOp.getSource(); - if (!viewOrAlloc.getDefiningOp() && - !viewOrAlloc.getDefiningOp()) + Value viewOrAlloc = xferOp.source(); + auto viewOrAllocDefOp = viewOrAlloc.getDefiningOp(); + if (!viewOrAllocDefOp) + return failure(); + if (!isa(viewOrAllocDefOp) && + !hasSingleEffect(viewOrAllocDefOp)) return failure(); // Ensure there is exactly one subview of `viewOrAlloc` defining `subView`. diff --git a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp --- a/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp +++ b/mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp @@ -162,7 +162,7 @@ return true; if (funcOp - .walk([&](memref::AllocOp allocOp) -> WalkResult { + .walk([&](AllocLikeOpInterface) -> WalkResult { Value oldMemRef = allocOp.getResult(); if (!allocOp.getType().getLayout().isIdentity() && !isMemRefNormalizable(oldMemRef.getUsers())) @@ -339,9 +339,9 @@ // Turn memrefs' non-identity layouts maps into ones with identity. Collect // alloc ops first and then process since normalizeMemRef replaces/erases ops // during memref rewriting. - SmallVector allocOps; - funcOp.walk([&](memref::AllocOp op) { allocOps.push_back(op); }); - for (memref::AllocOp allocOp : allocOps) + SmallVector allocOps; + funcOp.walk([&](AllocLikeOpInterface op) { allocOps.push_back(op); }); + for (AllocLikeOpInterface allocOp : allocOps) (void)normalizeMemRef(&allocOp); // We use this OpBuilder to create new memref layout later. diff --git a/mlir/lib/Interfaces/AllocLikeOpInterface.cpp b/mlir/lib/Interfaces/AllocLikeOpInterface.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Interfaces/AllocLikeOpInterface.cpp @@ -0,0 +1,14 @@ +//===- AllocLikeOpInterface.cpp - Alloc like operations interface in MLIR-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Interfaces/AllocLikeOpInterface.h" + +using namespace mlir; + +/// Include the definitions of the copy operation interface. +#include "mlir/Interfaces/AllocLikeOpInterface.cpp.inc" diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt --- a/mlir/lib/Interfaces/CMakeLists.txt +++ b/mlir/lib/Interfaces/CMakeLists.txt @@ -1,4 +1,5 @@ set(LLVM_OPTIONAL_SOURCES + AllocLikeOpInterface.cpp CallInterfaces.cpp CastInterfaces.cpp ControlFlowInterfaces.cpp @@ -31,6 +32,7 @@ endfunction(add_mlir_interface_library) +add_mlir_interface_library(AllocLikeOpInterface) add_mlir_interface_library(CallInterfaces) add_mlir_interface_library(CastInterfaces) add_mlir_interface_library(ControlFlowInterfaces) diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp --- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp +++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp @@ -8,6 +8,7 @@ #include +#include "mlir/Analysis/AffineAnalysis.h" #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -922,6 +922,13 @@ ], ) +td_library( + name = "AllocLikeOpInterfaceTdFiles", + srcs = ["include/mlir/Interfaces/AllocLikeOpInterface.td"], + includes = ["include"], + deps = [":OpBaseTdFiles"], +) + td_library( name = "CallInterfacesTdFiles", srcs = ["include/mlir/Interfaces/CallInterfaces.td"], @@ -2726,6 +2733,17 @@ ], ) +cc_library( + name = "AllocLikeOpInterface", + srcs = ["lib/Interfaces/AllocLikeOpInterface.cpp"], + hdrs = ["include/mlir/Interfaces/AllocLikeOpInterface.h"], + includes = ["include"], + deps = [ + ":AllocLikeOpInterfaceIncGen", + ":IR", + ], +) + cc_library( name = "DataLayoutInterfaces", srcs = ["lib/Interfaces/DataLayoutInterfaces.cpp"], @@ -5301,6 +5319,24 @@ deps = [":ViewLikeInterfaceTdFiles"], ) +gentbl_cc_library( + name = "AllocLikeOpInterfaceIncGen", + strip_include_prefix = "include", + tbl_outs = [ + ( + ["-gen-op-interface-decls"], + "include/mlir/Interfaces/AllocLikeOpInterface.h.inc", + ), + ( + ["-gen-op-interface-defs"], + "include/mlir/Interfaces/AllocLikeOpInterface.cpp.inc", + ), + ], + tblgen = ":mlir-tblgen", + td_file = "include/mlir/Interfaces/AllocLikeOpInterface.td", + deps = [":AllocLikeOpInterfaceTdFiles"], +) + gentbl_cc_library( name = "CopyOpInterfaceIncGen", strip_include_prefix = "include", @@ -8766,6 +8802,24 @@ ], ) +gentbl_cc_library( + name = "AllocLikeOpInterfaceIncGen", + strip_include_prefix = "include", + tbl_outs = [ + ( + ["-gen-op-interface-decls"], + "include/mlir/Interfaces/AllocLikeOpInterface.h.inc", + ), + ( + ["-gen-op-interface-defs"], + "include/mlir/Interfaces/AllocLikeOpInterface.cpp.inc", + ), + ], + tblgen = ":mlir-tblgen", + td_file = "include/mlir/Interfaces/AllocLikeOpInterface.td", + deps = [":AllocLikeOpInterfaceTdFiles"], +) + gentbl_cc_library( name = "MemRefBaseIncGen", strip_include_prefix = "include", @@ -8823,6 +8877,7 @@ ], includes = ["include"], deps = [ + ":AllocLikeOpInterface", ":ArithmeticDialect", ":ArithmeticUtils", ":ControlFlowInterfaces",