diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td @@ -182,7 +182,8 @@ def FillOp : LinalgStructured_Op<"fill", []> { let arguments = (ins AnyShaped:$output, - AnyTypeOf<[AnyFloat, AnySignlessInteger, AnyVector]>:$value); + AnyTypeOf<[AnyComplex, AnyFloat, AnySignlessInteger, + AnyVector]>:$value); let results = (outs Optional:$result); let regions = (region AnyRegion:$region); let extraClassDeclaration = structuredOpsDecls # [{ diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h --- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h @@ -215,9 +215,10 @@ /// smallest constant value for the size of the buffer needed for each /// dimension. If that is not possible, contains the dynamic size of the /// subview. The call back should return the buffer to use. -using AllocBufferCallbackFn = std::function( - OpBuilder &b, memref::SubViewOp subView, - ArrayRef boundingSubViewSize, OperationFolder *folder)>; +using AllocBufferCallbackFn = + std::function(OpBuilder &b, memref::SubViewOp subView, + ArrayRef boundingSubViewSize, + DataLayout &layout, OperationFolder *folder)>; /// Callback function type used to deallocate the buffers used to hold the /// promoted subview. @@ -315,6 +316,7 @@ Optional promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView, AllocBufferCallbackFn allocationFn, + DataLayout &layout, OperationFolder *folder = nullptr); /// Promotes the `subViews` into a new buffer allocated at the insertion point diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp @@ -12,6 +12,7 @@ #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" @@ -49,10 +50,10 @@ /// 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 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 = @@ -88,7 +89,7 @@ OpBuilder &builder, memref::SubViewOp subView, ArrayRef boundingSubViewSize, bool dynamicBuffers, Optional alignment, - OperationFolder *folder) { + DataLayout &layout, OperationFolder *folder) { ShapedType viewType = subView.getType(); int64_t rank = viewType.getRank(); (void)rank; @@ -100,7 +101,7 @@ 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 dynSizes(boundingSubViewSize.size(), ShapedType::kDynamicSize); Value view = folded_memref_view( @@ -170,15 +171,16 @@ } } - allocationFn = (options.allocationFn - ? *(options.allocationFn) - : [&](OpBuilder &builder, memref::SubViewOp subViewOp, - ArrayRef boundingSubViewSize, - OperationFolder *folder) -> Optional { - return defaultAllocBufferCallBack(options, builder, subViewOp, - boundingSubViewSize, dynamicBuffers, - alignment, folder); - }); + allocationFn = + (options.allocationFn + ? *(options.allocationFn) + : [&](OpBuilder &builder, memref::SubViewOp subViewOp, + ArrayRef boundingSubViewSize, DataLayout &layout, + OperationFolder *folder) -> Optional { + return defaultAllocBufferCallBack(options, builder, subViewOp, + boundingSubViewSize, dynamicBuffers, + alignment, layout, folder); + }); deallocationFn = (options.deallocationFn ? *(options.deallocationFn) @@ -213,7 +215,8 @@ // by a partial `copy` op. Optional 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(); @@ -236,7 +239,8 @@ SmallVector dynSizes(fullSizes.size(), -1); // If a callback is not specified, then use the default implementation for // allocating the promoted buffer. - Optional fullLocalView = allocationFn(b, subView, fullSizes, folder); + Optional fullLocalView = + allocationFn(b, subView, fullSizes, layout, folder); if (!fullLocalView) return {}; SmallVector zeros(fullSizes.size(), b.getIndexAttr(0)); @@ -248,7 +252,7 @@ static Optional> promoteSubViews(OpBuilder &b, Location loc, - LinalgOpInstancePromotionOptions options, + LinalgOpInstancePromotionOptions options, DataLayout &layout, OperationFolder *folder) { if (options.subViews.empty()) return {}; @@ -260,7 +264,7 @@ memref::SubViewOp subView = cast(v.second.getDefiningOp()); Optional promotionInfo = promoteSubviewAsNewBuffer( - b, loc, subView, options.allocationFn, folder); + b, loc, subView, options.allocationFn, layout, folder); if (!promotionInfo) return {}; promotionInfoMap[v.first] = *promotionInfo; @@ -269,11 +273,21 @@ if (!options.useFullTileBuffers[v.second]) continue; Value fillVal; - if (auto t = subView.getType().getElementType().dyn_cast()) + if (auto t = subView.getType().getElementType().dyn_cast()) { fillVal = folded_std_constant(folder, FloatAttr::get(t, 0.0)); - else if (auto t = - subView.getType().getElementType().dyn_cast()) + } else if (auto t = + subView.getType().getElementType().dyn_cast()) { fillVal = folded_std_constant_int(folder, 0, t); + } else if (auto t = + subView.getType().getElementType().dyn_cast()) { + if (auto et = t.getElementType().dyn_cast()) + fillVal = folded_std_constant(folder, FloatAttr::get(et, 0.0)); + else if (auto et = t.getElementType().cast()) + fillVal = folded_std_constant_int(folder, 0, et); + fillVal = b.create(loc, t, fillVal, fillVal); + } else { + return {}; + } linalg_fill(promotionInfo->fullLocalView, fillVal); } @@ -292,7 +306,7 @@ static Optional promoteSubViews(OpBuilder &b, LinalgOp op, - LinalgOpInstancePromotionOptions options, + LinalgOpInstancePromotionOptions options, DataLayout &layout, OperationFolder *folder) { assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics"); @@ -304,7 +318,8 @@ // 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 {}; @@ -376,8 +391,8 @@ 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 { diff --git a/mlir/test/Dialect/Linalg/transform-patterns.mlir b/mlir/test/Dialect/Linalg/transform-patterns.mlir --- a/mlir/test/Dialect/Linalg/transform-patterns.mlir +++ b/mlir/test/Dialect/Linalg/transform-patterns.mlir @@ -345,6 +345,29 @@ // CHECK: linalg.copy(%[[s0]], %[[l0]]) : memref, memref // CHECK: linalg.fill(%[[v0]], %[[cf]]) : memref, f32 +func @aligned_promote_fill_complex(%arg0: memref, 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 + %3 = memref.subview %arg0[%c0, %c0][%c2000, %c4000][%c1, %c1] : + memref, offset: ?, strides: [?, 1]> to memref, offset: ?, strides: [?, ?]> + linalg.fill(%3, %cc) { __internal_linalg_transform__ = "_promote_views_aligned_"} + : memref, offset: ?, strides: [?, ?]>, complex + return +} +// CHECK-LABEL: func @aligned_promote_fill_complex +// CHECK: %[[cc:.*]] = complex.create {{.*}} : complex +// CHECK: %[[s0:.*]] = memref.subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref, #map{{.*}}> to memref, #map{{.*}}> +// CHECK: %[[a0:.*]] = memref.alloc({{%.*}}) {alignment = 32 : i64} : memref +// CHECK: %[[v0:.*]] = memref.view %[[a0]][{{.*}}][{{%.*}}, {{%.*}}] : memref to memref> +// CHECK: %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref> to memref, #[[$STRIDED_2D_u_1]]> +// CHECK: linalg.fill(%[[v0]], {{%.*}}) : memref>, complex +// CHECK: linalg.copy(%[[s0]], %[[l0]]) : memref, #map{{.*}}>, memref, #map{{.*}}> +// CHECK: linalg.fill(%[[v0]], %[[cc]]) : memref>, complex + func @tile_permute_parallel_loop(%arg0: memref, %arg1: memref, %arg2: memref) { diff --git a/mlir/test/lib/Transforms/TestLinalgTransforms.cpp b/mlir/test/lib/Transforms/TestLinalgTransforms.cpp --- a/mlir/test/lib/Transforms/TestLinalgTransforms.cpp +++ b/mlir/test/lib/Transforms/TestLinalgTransforms.cpp @@ -276,6 +276,7 @@ // Allocation call back static Optional allocCallBackFn(OpBuilder &b, memref::SubViewOp subView, ArrayRef boundingSubViewSize, + DataLayout &layout, OperationFolder *folder) { SmallVector shape(boundingSubViewSize.size(), -1); return b