maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir] Fold memref.dim of OffsetSizeAndStrideOpInterface outputs 

This previously handled memref::SubviewOp, but this can be extended to
all ops implementing the interface.

Differential Revision: https://reviews.llvm.org/D103076
This commit is contained in:
Tres Popp 2021-05-25 11:35:14 +02:00
parent 536447eb20
commit 9ccdc2e23b
2 changed files with 20 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
View File

@ -18,6 +18,7 @@
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "llvm/ADT/STLExtras.h"
using namespace mlir;
@ -679,10 +680,11 @@ OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
return *(view.getDynamicSizes().begin() +
memrefType.getDynamicDimIndex(unsignedIndex));
if (auto subview = dyn_cast_or_null<SubViewOp>(definingOp)) {
assert(subview.isDynamicSize(unsignedIndex) &&
if (auto sizeInterface =
dyn_cast_or_null<OffsetSizeAndStrideOpInterface>(definingOp)) {
assert(sizeInterface.isDynamicSize(unsignedIndex) &&
"Expected dynamic subview size");
return subview.getDynamicSize(unsignedIndex);
return sizeInterface.getDynamicSize(unsignedIndex);
}
// dim(memrefcast) -> dim

15
mlir/test/Dialect/MemRef/canonicalize.mlir
View File

@ -192,3 +192,18 @@ func @alias_is_freed(%arg0 : memref<?xf32>) {
memref.dealloc %1 : memref<32xf32>
return
}
// -----
// CHECK-LABEL: func @dim_of_sized_view
// CHECK-SAME: %{{[a-z0-9A-Z_]+}}: memref<?xi8>
// CHECK-SAME: %[[SIZE:.[a-z0-9A-Z_]+]]: index
// CHECK: return %[[SIZE]] : index
func @dim_of_sized_view(%arg : memref<?xi8>, %size: index) -> index {
%c0 = constant 0 : index
%0 = memref.reinterpret_cast %arg to offset: [0], sizes: [%size], strides: [0] : memref<?xi8> to memref<?xi8>
%1 = memref.dim %0, %c0 : memref<?xi8>
return %1 : index
}