[NFC] Create a AllocLikeOpInterface and make memref::AllocOp, memref::AllocaOp and gpu::AllocOp implement it.

[arnab-oss]

Summary

• Create AllocLikeOpInterface and put the common methods used by memref.alloc, mermen.alloca and gpu.alloc there.

• Instead of checking explicitly for alloc or alloca ops, use hasSingleEffect().

[arnab-oss]

Make gpu.alloc op implement AllocLikeOpInterface

[arnab-oss]

Refactoring

[bondhugula]

This is looking good to me. Please pull out changes unrelated to AllocLikeOpInterface into a separate revision (the ones just switching to hasSingleEffect).

mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp

1967

This change isn't related to the AllocLikeOpInterface; this and other similar changes can go into another revision.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

1071

The change to include GetGlobalOp is unrelated to this PR (and not proper as well) - it should be dropped.

[ftynse]

I would expect a new top-level interface (in lib/Interfaces) to go through the RFC process.

I have concerns about the interface design. Specifically, it appears to be promoting details specific to operations from one dialect (memref.alloc(a)) such as the fact that the result of an allocation is a memref (it's not, we have allocations in LLVM and SPIR-V that produce different types), that it has alignment defined as integer, and that it has some symbol operands presumably related to memref's affine layouts. If this were to live in lib/Interfaces, it should make either cover wider range of allocation-like operations or make a clear case as to why this is not desired accompanied by a name change.

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

937–954

I suppose these are only necessary because the dialect doesn't emit accessor using the prefixed form? Make sure to rebase this commit because the GPU dialect must have been switched to the prefixed form - https://discourse.llvm.org/t/psa-raw-accessors-are-being-removed/65629.

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

103–120

Ditto.

mlir/include/mlir/Interfaces/AllocLikeOpInterface.h

1

Please use the correct header line.

mlir/include/mlir/Interfaces/AllocLikeOpInterface.td

1

Please follow the (implicit) convention from other files for this line: the file name is left aligned, the file type is right aligned, and the entire line fits into 80 cols.

9

This looks wrong.

20

Type is not mandatory for allocations, can be just bytes.

26

Memref is not a mandatory concept for allocating memory, I would rather not enshrine it in the top-level interface.

27

Does it have to be a vector (copy) of values? Can't it be an OperandRange, a ValueRange or something similar?

mlir/lib/Interfaces/AllocLikeOpInterface.cpp

1

Nit: 80 cols.

clang/docs/tools/clang-formatted-files.txt

 mlir/include/mlir/Dialect/X86Vector/X86VectorDialect.h
 mlir/include/mlir/ExecutionEngine/AsyncRuntime.h
 mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
 mlir/include/mlir/ExecutionEngine/JitRunner.h
 mlir/include/mlir/ExecutionEngine/MemRefUtils.h
 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
 mlir/include/mlir/Interfaces/CopyOpInterface.h
 mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
 mlir/include/mlir/Interfaces/DecodeAttributesInterfaces.h
 mlir/include/mlir/Interfaces/DerivedAttributeOpInterface.h
 mlir/include/mlir/Interfaces/FoldInterfaces.h

 mlir/lib/ExecutionEngine/CRunnerUtils.cpp
 mlir/lib/ExecutionEngine/CudaRuntimeWrappers.cpp
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp
 mlir/lib/ExecutionEngine/JitRunner.cpp
 mlir/lib/ExecutionEngine/OptUtils.cpp
 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
 mlir/lib/Interfaces/CopyOpInterface.cpp
 mlir/lib/Interfaces/DataLayoutInterfaces.cpp
 mlir/lib/Interfaces/DerivedAttributeOpInterface.cpp
 mlir/lib/Interfaces/InferTypeOpInterface.cpp
 mlir/lib/Interfaces/LoopLikeInterface.cpp

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

 #include "mlir/Dialect/DLTI/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #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"
 
 namespace mlir {
 namespace gpu {
 
 /// Utility class for the GPU dialect to represent triples of `Value`s

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

 #define GPU_OPS
 
 include "mlir/Dialect/DLTI/DLTIBase.td"
 include "mlir/Dialect/GPU/IR/GPUBase.td"
 include "mlir/Dialect/GPU/IR/ParallelLoopMapperAttr.td"
 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"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 //===----------------------------------------------------------------------===//
 // GPU Dialect operations.
 //===----------------------------------------------------------------------===//
def GPU_WaitOp : GPU_Op<"wait", [GPU_AsyncOpInterface]> {
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) attr-dict
   }];
 
   let hasCanonicalizer = 1;
 }
 
 def GPU_AllocOp : GPU_Op<"alloc", [
+    AllocLikeOpInterface,
     GPU_AsyncOpInterface,
     AttrSizedOperandSegments
   ]> {
 
   let summary = "GPU memory allocation operation.";
   let description = [{
     The `gpu.alloc` operation allocates a region of memory on the GPU. It is
     similar to the `memref.alloc` op, but supports asynchronous GPU execution.
def GPU_AllocOp : GPU_Op<"alloc", [
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    Variadic<Index>:$dynamicSizes, Variadic<Index>:$symbolOperands);
   let results = (outs Res<AnyMemRef, "", [MemAlloc]>:$memref,
                  Optional<GPU_AsyncToken>:$asyncToken);
 
   let extraClassDeclaration = [{
     MemRefType getType() { return memref().getType().cast<MemRefType>(); }
+
+    /// 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 = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) ` `
     `(` $dynamicSizes `)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
   }];
 
   let hasVerifier = 1;

mlir/include/mlir/Dialect/MemRef/IR/MemRef.h

 //===- MemRef.h - MemRef dialect --------------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MEMREF_IR_MEMREF_H_
 #define MLIR_DIALECT_MEMREF_IR_MEMREF_H_
 
 #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"
 #include "mlir/Interfaces/CopyOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 //===- MemRefOps.td - MemRef op definitions ----------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MEMREF_OPS
 #define MEMREF_OPS
 
 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"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 
 /// A TypeAttr for memref types.

 //
 //   %0 = alloclike(%m)[%s] : memref<8x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>>
 //
 class AllocLikeOp<string mnemonic,
                   Resource resource,
                   list<Trait> traits = []> :
     MemRef_Op<mnemonic,
     !listconcat([
+      AllocLikeOpInterface,
       AttrSizedOperandSegments
     ], traits)> {
 
   let arguments = (ins Variadic<Index>:$dynamicSizes,
                        // The symbolic operands (the ones in square brackets)
                        // bind to the symbols of the memref's layout map.
                        Variadic<Index>:$symbolOperands,
                        ConfinedAttr<OptionalAttr<I64Attr>,
OpBuilder<(ins "MemRefType":$memrefType, "ValueRange":$dynamicSizes,
       if (alignment)
         $_state.addAttribute(getAlignmentAttrStrName(), alignment);
     }]>];
 
   let extraClassDeclaration = [{
     static StringRef getAlignmentAttrStrName() { return "alignment"; }
 
     MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
+
+    /// 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 = [{
     `(`$dynamicSizes`)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
   }];
 
   let hasCanonicalizer = 1;
   let hasVerifier = 1;

mlir/include/mlir/Interfaces/AllocLikeOpInterface.h

+//===- 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_

mlir/include/mlir/Interfaces/AllocLikeOpInterface.td

+//===-------------- 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

mlir/include/mlir/Interfaces/CMakeLists.txt

+add_mlir_interface(AllocLikeOpInterface)
 add_mlir_interface(CallInterfaces)
 add_mlir_interface(CastInterfaces)
 add_mlir_interface(ControlFlowInterfaces)
 add_mlir_interface(CopyOpInterface)
 add_mlir_interface(DerivedAttributeOpInterface)
 add_mlir_interface(InferIntRangeInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)

mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp

for (auto &pair : mdg->memrefEdgeCount) {
       if (pair.second > 0)
         continue;
       auto memref = pair.first;
       // Skip if there exist other uses (return operation or function calls).
       if (!memref.use_empty())
         continue;
       // Use list expected to match the dep graph info.
       auto *op = memref.getDefiningOp();
-      if (isa_and_nonnull<memref::AllocOp>(op))
+      if (hasSingleEffect<MemoryEffects::Allocate>(op))
         op->erase();
     }
   }
 };
 
 } // namespace
 
 void LoopFusion::runOnOperation() {

mlir/lib/Dialect/Affine/Utils/Utils.cpp

void mlir::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
 
   // Check if the store fwd'ed memrefs are now left with only stores and
   // deallocs and can thus be completely deleted. Note: the canonicalize pass
   // should be able to do this as well, but we'll do it here since we collected
   // these anyway.
   for (auto memref : memrefsToErase) {
     // If the memref hasn't been locally alloc'ed, skip.
     Operation *defOp = memref.getDefiningOp();
-    if (!defOp || !hasSingleEffect<MemoryEffects::Allocate>(defOp, memref))
+    if (!defOp || (!hasSingleEffect<MemoryEffects::Allocate>(defOp) &&
+                   !isa<memref::GetGlobalOp>(defOp)))
       // TODO: if the memref was returned by a 'call' operation, we
       // could still erase it if the call had no side-effects.
       continue;
     if (llvm::any_of(memref.getUsers(), [&](Operation *ownerOp) {
           return !isa<AffineWriteOpInterface>(ownerOp) &&
                  !hasSingleEffect<MemoryEffects::Free>(ownerOp, memref);
         }))
       continue;

 /// 4. Create AffineApplyOp to apply the new maps. The output of AffineApplyOp
 /// is used in dynamicSizes of new AllocOp.
 ///   %0 = dim %arg0, %c1 : memref<4x?xf32>
 ///   %c4 = arith.constant 4 : index
 ///   %1 = affine.apply #map1(%c4, %0)
 ///   %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<Value> &newDynamicSizes) {
   // Create new input for AffineApplyOp.
   SmallVector<Value, 4> inAffineApply;
   ArrayRef<int64_t> oldMemRefShape = oldMemRefType.getShape();
   unsigned dynIdx = 0;
   for (unsigned d = 0; d < oldMemRefType.getRank(); ++d) {
     if (oldMemRefShape[d] < 0) {
       // Use dynamicSizes of allocOp for dynamic dimension.
if (newMemRefShape[newDimIdx] < 0) {
           b.create<AffineApplyOp>(allocOp->getLoc(), newMap, inAffineApply);
       newDynamicSizes.emplace_back(affineApp);
     }
     newDimIdx++;
   }
 }
 
 // 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);
 
   // Fetch a new memref type after normalizing the old memref to have an
   // identity map layout.
   MemRefType newMemRefType =
       normalizeMemRefType(memrefType, b, allocOp->getSymbolOperands().size());
   if (newMemRefType == memrefType)

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

 // the resulting code will be invalid with respect to program semantics.
 // However, structured control-flow loops are fully supported. Furthermore, it
 // doesn't accept functions which return buffers already.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "llvm/ADT/SetOperations.h"
 
 namespace mlir {
 namespace bufferization {

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

 LogicalResult LinalgCopyVTRForwardingPattern::matchAndRewrite(
     vector::TransferReadOp xferOp, PatternRewriter &rewriter) const {
 
   // TODO: support mask.
   if (xferOp.getMask())
     return failure();
 
   // Transfer into `view`.
-  Value viewOrAlloc = xferOp.getSource();
-  if (!viewOrAlloc.getDefiningOp<memref::ViewOp>() &&
-      !viewOrAlloc.getDefiningOp<memref::AllocOp>())
+  Value viewOrAlloc = xferOp.source();
+  Operation *viewOrAllocDefOp = viewOrAlloc.getDefiningOp();
+  if (!viewOrAllocDefOp)
+    return failure();
+  if (!isa<memref::ViewOp>(viewOrAllocDefOp) &&
+      !hasSingleEffect<MemoryEffects::Allocate>(viewOrAllocDefOp))
     return failure();
 
   LDBG(viewOrAlloc);
 
   // Ensure there is exactly one subview of `viewOrAlloc` defining `subView`.
   memref::SubViewOp subViewOp = getSubViewUseIfUnique(viewOrAlloc);
   if (!subViewOp)
     return failure();

 /// when available.
 LogicalResult LinalgCopyVTWForwardingPattern::matchAndRewrite(
     vector::TransferWriteOp xferOp, PatternRewriter &rewriter) const {
   // TODO: support mask.
   if (xferOp.getMask())
     return failure();
 
   // Transfer into `viewOrAlloc`.
-  Value viewOrAlloc = xferOp.getSource();
-  if (!viewOrAlloc.getDefiningOp<memref::ViewOp>() &&
-      !viewOrAlloc.getDefiningOp<memref::AllocOp>())
+  Value viewOrAlloc = xferOp.source();
+  auto viewOrAllocDefOp = viewOrAlloc.getDefiningOp();
+  if (!viewOrAllocDefOp)
+    return failure();
+  if (!isa<memref::ViewOp>(viewOrAllocDefOp) &&
+      !hasSingleEffect<MemoryEffects::Allocate>(viewOrAllocDefOp))
     return failure();
 
   // Ensure there is exactly one subview of `viewOrAlloc` defining `subView`.
   memref::SubViewOp subViewOp = getSubViewUseIfUnique(viewOrAlloc);
   if (!subViewOp)
     return failure();
   Value subView = subViewOp.getResult();
 

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

 /// still label the entire function as normalizable. We assume external
 /// functions to be normalizable.
 bool NormalizeMemRefs::areMemRefsNormalizable(func::FuncOp funcOp) {
   // We assume external functions to be normalizable.
   if (funcOp.isExternal())
     return true;
 
   if (funcOp
-          .walk([&](memref::AllocOp allocOp) -> WalkResult {
+          .walk([&](AllocLikeOpInterface) -> WalkResult {
             Value oldMemRef = allocOp.getResult();
             if (!allocOp.getType().getLayout().isIdentity() &&
                 !isMemRefNormalizable(oldMemRef.getUsers()))
               return WalkResult::interrupt();
             return WalkResult::advance();
           })
           .wasInterrupted())
     return false;

 /// Normalizes the memrefs within a function which includes those arising as a
 /// result of AllocOps, CallOps and function's argument. The ModuleOp argument
 /// is used to help update function's signature after normalization.
 void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
                                               ModuleOp moduleOp) {
   // 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<memref::AllocOp, 4> allocOps;
-  funcOp.walk([&](memref::AllocOp op) { allocOps.push_back(op); });
-  for (memref::AllocOp allocOp : allocOps)
+  SmallVector<AllocLikeOpInterface, 4> 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.
   OpBuilder b(funcOp);
 
   FunctionType functionType = funcOp.getFunctionType();
   SmallVector<Location> functionArgLocs(llvm::map_range(
       funcOp.getArguments(), [](BlockArgument arg) { return arg.getLoc(); }));

mlir/lib/Interfaces/AllocLikeOpInterface.cpp

+//===- 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"

mlir/lib/Interfaces/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
+  AllocLikeOpInterface.cpp
   CallInterfaces.cpp
   CastInterfaces.cpp
   ControlFlowInterfaces.cpp
   CopyOpInterface.cpp
   DataLayoutInterfaces.cpp
   DerivedAttributeOpInterface.cpp
   InferIntRangeInterface.cpp
   InferTypeOpInterface.cpp
add_mlir_library(MLIR${name}
     MLIR${name}IncGen
 
     LINK_LIBS PUBLIC
     MLIRIR
     )
 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)
 add_mlir_interface_library(CopyOpInterface)
 add_mlir_interface_library(DataLayoutInterfaces)
 add_mlir_interface_library(DerivedAttributeOpInterface)
 add_mlir_interface_library(InferIntRangeInterface)
 add_mlir_interface_library(InferTypeOpInterface)
 add_mlir_interface_library(ParallelCombiningOpInterface)
 add_mlir_interface_library(SideEffectInterfaces)
 add_mlir_interface_library(TilingInterface)
 add_mlir_interface_library(VectorInterfaces)
 add_mlir_interface_library(ViewLikeInterface)
 add_mlir_interface_library(LoopLikeInterface)

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

 //===- TestVectorTransforms.cpp - Test Vector transforms and lowerings ----===//
 //
 // 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 <type_traits>
 
+#include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

td_library(
     ],
     includes = ["include"],
     deps = [
         ":AttrTdFiles",
     ],
 )
 
 td_library(
+    name = "AllocLikeOpInterfaceTdFiles",
+    srcs = ["include/mlir/Interfaces/AllocLikeOpInterface.td"],
+    includes = ["include"],
+    deps = [":OpBaseTdFiles"],
+)
+
+td_library(
     name = "CallInterfacesTdFiles",
     srcs = ["include/mlir/Interfaces/CallInterfaces.td"],
     includes = ["include"],
     deps = [":OpBaseTdFiles"],
 )
 
 td_library(
     name = "CastInterfacesTdFiles",
deps = [
         ":SCFDialect",
         ":Support",
         ":Transforms",
         "//llvm:Support",
     ],
 )
 
 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"],
     hdrs = ["include/mlir/Interfaces/DataLayoutInterfaces.h"],
     includes = ["include"],
     deps = [
         ":DataLayoutInterfacesIncGen",
         ":IR",
         "//llvm:Support",
tbl_outs = [
         ),
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Interfaces/ViewLikeInterface.td",
     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",
     tbl_outs = [
         (
             ["-gen-op-interface-decls"],
             "include/mlir/Interfaces/CopyOpInterface.h.inc",
         ),
         (
deps = [
         ":CopyOpInterfaceTdFiles",
         ":OpBaseTdFiles",
         ":SideEffectInterfacesTdFiles",
         ":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 = "MemRefBaseIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             [
                 "-gen-dialect-decls",
                 "-dialect=memref",
             ],
srcs = glob(
         ],
     ),
     hdrs = [
         "include/mlir/Dialect/MemRef/IR/MemRef.h",
         "include/mlir/Dialect/MemRef/Utils/MemRefUtils.h",
     ],
     includes = ["include"],
     deps = [
+        ":AllocLikeOpInterface",
         ":ArithmeticDialect",
         ":ArithmeticUtils",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefBaseIncGen",