#include "PassDetail.h"
#include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Linalg/EDSC/FoldedIntrinsics.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
/// the size needed, otherwise try to allocate a static bounding box.
static Value allocBuffer(const LinalgPromotionOptions &options,
Type elementType, Value size, bool dynamicBuffers,
- OperationFolder *folder,
+ DataLayout &layout, OperationFolder *folder,
Optional<unsigned> alignment = None) {
auto *ctx = size.getContext();
- auto width = llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
+ auto width = layout.getTypeSize(elementType);
IntegerAttr alignment_attr;
if (alignment.hasValue())
alignment_attr =
OpBuilder &builder, memref::SubViewOp subView,
ArrayRef<Value> boundingSubViewSize,
bool dynamicBuffers, Optional<unsigned> alignment,
- OperationFolder *folder) {
+ DataLayout &layout, OperationFolder *folder) {
ShapedType viewType = subView.getType();
int64_t rank = viewType.getRank();
(void)rank;
for (auto size : llvm::enumerate(boundingSubViewSize))
allocSize = folded_std_muli(folder, allocSize, size.value());
Value buffer = allocBuffer(options, viewType.getElementType(), allocSize,
- dynamicBuffers, folder, alignment);
+ dynamicBuffers, layout, folder, alignment);
SmallVector<int64_t, 4> dynSizes(boundingSubViewSize.size(),
ShapedType::kDynamicSize);
Value view = folded_memref_view(
}
}
- allocationFn = (options.allocationFn
- ? *(options.allocationFn)
- : [&](OpBuilder &builder, memref::SubViewOp subViewOp,
- ArrayRef<Value> boundingSubViewSize,
- OperationFolder *folder) -> Optional<Value> {
- return defaultAllocBufferCallBack(options, builder, subViewOp,
- boundingSubViewSize, dynamicBuffers,
- alignment, folder);
- });
+ allocationFn =
+ (options.allocationFn
+ ? *(options.allocationFn)
+ : [&](OpBuilder &builder, memref::SubViewOp subViewOp,
+ ArrayRef<Value> boundingSubViewSize, DataLayout &layout,
+ OperationFolder *folder) -> Optional<Value> {
+ return defaultAllocBufferCallBack(options, builder, subViewOp,
+ boundingSubViewSize, dynamicBuffers,
+ alignment, layout, folder);
+ });
deallocationFn =
(options.deallocationFn
? *(options.deallocationFn)
// by a partial `copy` op.
Optional<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
OpBuilder &b, Location loc, memref::SubViewOp subView,
- AllocBufferCallbackFn allocationFn, OperationFolder *folder) {
+ AllocBufferCallbackFn allocationFn, DataLayout &layout,
+ OperationFolder *folder) {
ScopedContext scopedContext(b, loc);
auto viewType = subView.getType();
auto rank = viewType.getRank();
SmallVector<int64_t, 4> dynSizes(fullSizes.size(), -1);
// If a callback is not specified, then use the default implementation for
// allocating the promoted buffer.
- Optional<Value> fullLocalView = allocationFn(b, subView, fullSizes, folder);
+ Optional<Value> fullLocalView =
+ allocationFn(b, subView, fullSizes, layout, folder);
if (!fullLocalView)
return {};
SmallVector<OpFoldResult, 4> zeros(fullSizes.size(), b.getIndexAttr(0));
static Optional<MapVector<unsigned, PromotionInfo>>
promoteSubViews(OpBuilder &b, Location loc,
- LinalgOpInstancePromotionOptions options,
+ LinalgOpInstancePromotionOptions options, DataLayout &layout,
OperationFolder *folder) {
if (options.subViews.empty())
return {};
memref::SubViewOp subView =
cast<memref::SubViewOp>(v.second.getDefiningOp());
Optional<PromotionInfo> promotionInfo = promoteSubviewAsNewBuffer(
- b, loc, subView, options.allocationFn, folder);
+ b, loc, subView, options.allocationFn, layout, folder);
if (!promotionInfo)
return {};
promotionInfoMap[v.first] = *promotionInfo;
if (!options.useFullTileBuffers[v.second])
continue;
Value fillVal;
- if (auto t = subView.getType().getElementType().dyn_cast<FloatType>())
+ if (auto t = subView.getType().getElementType().dyn_cast<FloatType>()) {
fillVal = folded_std_constant(folder, FloatAttr::get(t, 0.0));
- else if (auto t =
- subView.getType().getElementType().dyn_cast<IntegerType>())
+ } else if (auto t =
+ subView.getType().getElementType().dyn_cast<IntegerType>()) {
fillVal = folded_std_constant_int(folder, 0, t);
+ } else if (auto t =
+ subView.getType().getElementType().dyn_cast<ComplexType>()) {
+ if (auto et = t.getElementType().dyn_cast<FloatType>())
+ fillVal = folded_std_constant(folder, FloatAttr::get(et, 0.0));
+ else if (auto et = t.getElementType().cast<IntegerType>())
+ fillVal = folded_std_constant_int(folder, 0, et);
+ fillVal = b.create<complex::CreateOp>(loc, t, fillVal, fillVal);
+ } else {
+ return {};
+ }
linalg_fill(promotionInfo->fullLocalView, fillVal);
}
static Optional<LinalgOp>
promoteSubViews(OpBuilder &b, LinalgOp op,
- LinalgOpInstancePromotionOptions options,
+ LinalgOpInstancePromotionOptions options, DataLayout &layout,
OperationFolder *folder) {
assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
// 1. Promote the specified views and use them in the new op.
auto loc = op.getLoc();
- auto promotedBuffersAndViews = promoteSubViews(b, loc, options, folder);
+ auto promotedBuffersAndViews =
+ promoteSubViews(b, loc, options, layout, folder);
if (!promotedBuffersAndViews ||
promotedBuffersAndViews->size() != options.subViews.size())
return {};
LinalgPromotionOptions options,
OperationFolder *folder) {
LinalgOpInstancePromotionOptions linalgOptions(linalgOp, options);
- return ::promoteSubViews(
- b, linalgOp, LinalgOpInstancePromotionOptions(linalgOp, options), folder);
+ auto layout = DataLayout::closest(linalgOp);
+ return ::promoteSubViews(b, linalgOp, linalgOptions, layout, folder);
}
namespace {
// CHECK: linalg.copy(%[[s0]], %[[l0]]) : memref<?x?xf32, #map{{.*}}>, memref<?x?xf32, #map{{.*}}>
// CHECK: linalg.fill(%[[v0]], %[[cf]]) : memref<?x?xf32>, f32
+func @aligned_promote_fill_complex(%arg0: memref<?x?xcomplex<f32>, offset: ?, strides: [?, 1]>) {
+ %c2000 = constant 2000 : index
+ %c4000 = constant 4000 : index
+ %c0 = constant 0 : index
+ %c1 = constant 1 : index
+ %cf = constant 1.0 : f32
+ %cc = complex.create %cf, %cf : complex<f32>
+ %3 = memref.subview %arg0[%c0, %c0][%c2000, %c4000][%c1, %c1] :
+ memref<?x?xcomplex<f32>, offset: ?, strides: [?, 1]> to memref<?x?xcomplex<f32>, offset: ?, strides: [?, ?]>
+ linalg.fill(%3, %cc) { __internal_linalg_transform__ = "_promote_views_aligned_"}
+ : memref<?x?xcomplex<f32>, offset: ?, strides: [?, ?]>, complex<f32>
+ return
+}
+// CHECK-LABEL: func @aligned_promote_fill_complex
+// CHECK: %[[cc:.*]] = complex.create {{.*}} : complex<f32>
+// CHECK: %[[s0:.*]] = memref.subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xcomplex<f32>, #map{{.*}}> to memref<?x?xcomplex<f32>, #map{{.*}}>
+// CHECK: %[[a0:.*]] = memref.alloc({{%.*}}) {alignment = 32 : i64} : memref<?xi8>
+// CHECK: %[[v0:.*]] = memref.view %[[a0]][{{.*}}][{{%.*}}, {{%.*}}] : memref<?xi8> to memref<?x?xcomplex<f32>>
+// CHECK: %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xcomplex<f32>> to memref<?x?xcomplex<f32>, #[[$STRIDED_2D_u_1]]>
+// CHECK: linalg.fill(%[[v0]], {{%.*}}) : memref<?x?xcomplex<f32>>, complex<f32>
+// CHECK: linalg.copy(%[[s0]], %[[l0]]) : memref<?x?xcomplex<f32>, #map{{.*}}>, memref<?x?xcomplex<f32>, #map{{.*}}>
+// CHECK: linalg.fill(%[[v0]], %[[cc]]) : memref<?x?xcomplex<f32>>, complex<f32>
+
func @tile_permute_parallel_loop(%arg0: memref<?x?xf32>,
%arg1: memref<?x?xf32>,
%arg2: memref<?x?xf32>) {