diff --git a/compiler/lib/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.cpp b/compiler/lib/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.cpp index 279d7ff11..f5680051b 100644 --- a/compiler/lib/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/compiler/lib/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.cpp @@ -36,12 +36,31 @@ namespace {} // namespace namespace { +// Returns a map with a symbolic offset for each dimension, i.e., for N +// dimensions, it returns +// +// [d1, d2, ..., dN](s1, s2, ..., sN) -> (d1 + s1, d2 + s2, ..., dN + sN) +// +AffineMap getMultiDimSymbolicOffsetMap(mlir::RewriterBase &rewriter, + unsigned rank) { + SmallVector dimExprs; + dimExprs.reserve(rank); + + for (unsigned i = 0; i < rank; ++i) + dimExprs.push_back(rewriter.getAffineDimExpr(i) + + rewriter.getAffineSymbolExpr(i)); + + return AffineMap::get(/*dimCount=*/rank, /*symbolCount=*/rank, dimExprs, + rewriter.getContext()); +} + mlir::Type getDynamicMemrefWithUnknownOffset(mlir::RewriterBase &rewriter, size_t rank) { mlir::MLIRContext *ctx = rewriter.getContext(); std::vector shape(rank, -1); + return mlir::MemRefType::get(shape, rewriter.getI64Type(), - rewriter.getMultiDimIdentityMap(rank)); + getMultiDimSymbolicOffsetMap(rewriter, rank)); } // Returns `memref.cast %0 : memref<...xAxT> to memref<...x?xT>`