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/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/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; 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/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",