241de80d38
This commit defines interfaces for operations that perform certain kinds of indexed access on a memref. These interfaces are defined so that passes like fold-memref-alias-ops and the memref flattener can be made generic over operations that, informally, have the forms `op ... %m[%i0, %i1, ...] ...` (an IndexedAccessOpInterface) or the form `op %src[%s0, %s1, ...], %dst[%d0, %d1, ...] size ...` (an IndexedMemCopyOpInterface). These interfaces have been designed such that all the passes under MemRef/Transforms that currently have a big switch-case on memref.load, vector.load, nvgpu.ldmatrix, etc. can be migrated to use them. (This'll also let us get rid of the awkward fact that we have memref transforms depending on the GPU and NVGPU dialects) While the interface doesn't currently contemplate changing element types (enabling, for example, writing a bf16 => u16 update to be done in place), future extensions to allow such transformations could be worth exploring. This commit only defines the interfaces so that it'll be easier to review the design - the implementation is in a future PR. --------- Co-authored-by: Jakub Kuderski <jakub@nod-labs.com>
66 lines
2.4 KiB
C++
66 lines
2.4 KiB
C++
//===- MemoryAccessOpInterfaces.cpp ---------------------------------------===//
|
|
//
|
|
// 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/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/IR/Operation.h"
|
|
#include "mlir/IR/Value.h"
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// IndexedAccessOpInterface and IndexedMemCpyOpInterface
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace mlir::memref {
|
|
#include "mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.cpp.inc"
|
|
|
|
LogicalResult detail::verifyIndexedAccessOpInterface(Operation *op) {
|
|
auto iface = dyn_cast<IndexedAccessOpInterface>(op);
|
|
if (!iface)
|
|
return failure();
|
|
|
|
TypedValue<MemRefType> memref = iface.getAccessedMemref();
|
|
if (!memref) {
|
|
// Some operations can carry tensors, this is fine.
|
|
return success();
|
|
}
|
|
if (memref.getType().getRank() !=
|
|
static_cast<int64_t>(iface.getIndices().size()))
|
|
return op->emitOpError(
|
|
"invalid number of indices for accessed memref, expected ")
|
|
<< memref.getType().getRank() << " but got "
|
|
<< iface.getIndices().size();
|
|
return success();
|
|
}
|
|
|
|
LogicalResult detail::verifyIndexedMemCopyOpInterface(Operation *op) {
|
|
auto iface = dyn_cast<IndexedMemCopyOpInterface>(op);
|
|
if (!iface)
|
|
return failure();
|
|
|
|
TypedValue<MemRefType> src = iface.getSrc();
|
|
TypedValue<MemRefType> dst = iface.getDst();
|
|
if (!src || !dst) {
|
|
// Allow operations to not always have memref arguments.
|
|
return success();
|
|
}
|
|
if (src.getType().getRank() !=
|
|
static_cast<int64_t>(iface.getSrcIndices().size()))
|
|
return op->emitOpError(
|
|
"invalid number of indices for source memref, expected " +
|
|
Twine(src.getType().getRank()) + ", got " +
|
|
Twine(iface.getSrcIndices().size()));
|
|
if (dst.getType().getRank() !=
|
|
static_cast<int64_t>(iface.getDstIndices().size()))
|
|
return op->emitOpError(
|
|
"invalid number of indices for destination memref, expected ")
|
|
<< dst.getType().getRank() << ", got "
|
|
<< iface.getDstIndices().size();
|
|
return success();
|
|
}
|
|
} // namespace mlir::memref
|