[mlir] Support complex numbers in Linalg promotion

[tpopp]

Summary

FillOp allows complex ops, and filling a properly sized buffer with
a default zero complex number is implemented.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

63

should we start being more deliberate about using the data layout here and in other places ?
@ftynse

[ftynse]

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

63

Yes, it would be more future-proof to query the size in bytes from the data layout, even if it's for the element type for now (there is no support for complex type in the data layout, yet). I would also suggest to construct the DataLayout object at the pass entry and forward it down here to avoid recomputation.

285

Nit: just cast to catch cases of other element types?

288

Nit: Please make braces symmetric

[tpopp]

Switch to using DataLayout.

[tpopp]

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

63

I switched to using DataLayout, but to cache the computation will require changing AllocBufferCallbackFn (https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h#L213). I don't know if that would be a good change?

The reason that change would be needed is because the alternative is to only use this in the default callback, but that can be constructed and used across multiple modules which leads to segfaults from using dangling pointers.

[tpopp]

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

55–58 ↗

(On Diff #336768)

Actually implemented here https://reviews.llvm.org/D100289

[ftynse]

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

55–58 ↗

(On Diff #336768)

You can add that diff as parent for this to create a "stack" of diffs only showing relevant changes. Click "edit related revisions" on the top right, and then "edit parent revisions"

[tpopp]

Rebase onto parent commit

[ftynse]

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

56

Use something like DataLayout::closest(size.getDefiningOp()), otherwise it's the default layout that ignores the specification. Incidentally, that is also why you may want to pass it in from above.

63

I expect other things will eventually want data layout as well, so it's okay to change the callback signature.

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

63

+1 re. ok to change signature.
Note however that this will break a bunch of things in IREE.

@ftynse is there a future in which the DataLayout could be auto-whatever from a builder insertion point (haven't thought this through, just basic idea that passed through my mind) ?

Also @maheshkhanwalkar for visibility if signature changes.

[ftynse]

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

63

You can already do DataLayout::closest(builder.getInsertionPoint()). Doing it the other way around feels like a layering violation, builder isn't supposed to know about individual interface.

[tpopp]

Forward DataLayout object to buffer allocation.

[tpopp]

I forwarded the DataLayout through several layers of functions. I started at the last point that the call sequence has the originating LinalgOp and passed it through as an additional argument rather than inside of any of the LinalgPromotionOption data structures. Please let me know what you think.

[tpopp]

Friendly ping.

[ftynse]

Sorry for the delay.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

def CopyOp : LinalgStructured_Op<"copy", [CopyOpInterface]> {
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
   let skipDefaultBuilders = 1;
 }
 
 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<AnyRankedTensor>:$result);
   let regions = (region AnyRegion:$region);
   let extraClassDeclaration = structuredOpsDecls # [{
     ValueRange inputs() { return {}; }
     ValueRange outputs() { return getOperands().take_front(); }
 
     // Rank-polymorphic.
     //   filling_value -> O(ivs) with parallel iterators.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 void interchange(PatternRewriter &rewriter, LinalgOp op,
                  ArrayRef<unsigned> interchangeVector);
 
 /// Callback function type used to perform the allocation for the promoted
 /// `subView`. In `boundingSubViewsize` a best attempt is made to find the
 /// 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<Optional<Value>(
-    OpBuilder &b, memref::SubViewOp subView,
-    ArrayRef<Value> boundingSubViewSize, OperationFolder *folder)>;
+using AllocBufferCallbackFn =
+    std::function<Optional<Value>(OpBuilder &b, memref::SubViewOp subView,
+                                  ArrayRef<Value> boundingSubViewSize,
+                                  DataLayout &layout, OperationFolder *folder)>;
 
 /// Callback function type used to deallocate the buffers used to hold the
 /// promoted subview.
 using DeallocBufferCallbackFn =
     std::function<LogicalResult(OpBuilder &b, Value buffer)>;
 
 /// Callback function type used to insert copy from original subview to subview
 /// of the promoted region for the read operands/subview of promoted region to

 /// of subview operation as `partialLocalView`.
 struct PromotionInfo {
   Value fullLocalView;
   Value partialLocalView;
 };
 Optional<PromotionInfo>
 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
 /// `b`. Promotion occurs in 3 steps:
 ///   1. Create a new buffer for a full tile (i.e. not clipped at the boundary).
 ///   2. Take a full view on the buffer.
 ///   3. Take a partial slice of the full view in step 2. and copy into it.
 /// Infers statically sized buffers from subViews unless `dynamicBuffers` is

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

 //===- Promotion.cpp - Implementation of linalg Promotion -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the linalg dialect Promotion pass.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/EDSC/Intrinsics.h"
 #include "mlir/Dialect/SCF/SCF.h"

 using folded_memref_view = FoldedValueBuilder<memref::ViewOp>;
 
 #define DEBUG_TYPE "linalg-promotion"
 
 /// Alloc a new buffer of `size`. If `dynamicBuffers` is true allocate exactly
 /// 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 =
         IntegerAttr::get(IntegerType::get(ctx, 64), alignment.getValue());
   if (!dynamicBuffers)
     if (auto cst = size.getDefiningOp<ConstantIndexOp>())
       return options.useAlloca
                  ? memref_alloca(MemRefType::get(width * cst.getValue(),
                                                  IntegerType::get(ctx, 8)),
                                  ValueRange{}, alignment_attr)
                        .value
                  : memref_alloc(MemRefType::get(width * cst.getValue(),
                                                 IntegerType::get(ctx, 8)),
                                 ValueRange{}, alignment_attr)
                        .value;

 /// no call back to do so is provided. The default is to allocate a
 /// memref<..xi8> and return a view to get a memref type of shape
 /// boundingSubViewSize.
 static Optional<Value>
 defaultAllocBufferCallBack(const LinalgPromotionOptions &options,
                            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;
   assert(rank > 0 && boundingSubViewSize.size() == static_cast<size_t>(rank));
   auto zero = folded_std_constant_index(folder, 0);
   auto one = folded_std_constant_index(folder, 1);
 
   Value allocSize = one;
   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(
       folder, MemRefType::get(dynSizes, viewType.getElementType()), buffer,
       zero, boundingSubViewSize);
   return view;
 }
 
if (options.operandsToPromote && !options.operandsToPromote->count(idx))
       continue;
     auto *op = linalgOp.getShapedOperand(idx).getDefiningOp();
     if (auto sv = dyn_cast_or_null<memref::SubViewOp>(op)) {
       subViews[idx] = sv;
       useFullTileBuffers[sv] = vUseFullTileBuffers[idx];
     }
   }
 
-  allocationFn = (options.allocationFn
-                      ? *(options.allocationFn)
-                      : [&](OpBuilder &builder, memref::SubViewOp subViewOp,
-                            ArrayRef<Value> boundingSubViewSize,
-                            OperationFolder *folder) -> Optional<Value> {
+  allocationFn =
+      (options.allocationFn
+                 ArrayRef<Value> boundingSubViewSize, DataLayout &layout,
         return defaultAllocBufferCallBack(options, builder, subViewOp,
                                           boundingSubViewSize, dynamicBuffers,
-                                      alignment, folder);
+                                          alignment, layout, folder);
       });
   deallocationFn =
       (options.deallocationFn
            ? *(options.deallocationFn)
            : [&](OpBuilder &b, Value buffer) {
                return defaultDeallocBufferCallBack(options, b, buffer);
              });
   auto defaultCopyCallBack = [&](OpBuilder &builder, Value src,
                                  Value dst) -> LogicalResult {

 // The point of the full tile buffer is that constant static tile sizes are
 // folded and result in a buffer type with statically known size and alignment
 // properties.
 // To account for general boundary effects, padding must be performed on the
 // boundary tiles. For now this is done with an unconditional `fill` op followed
 // 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<Value, 4> fullSizes;
   SmallVector<OpFoldResult> partialSizes;
   fullSizes.reserve(rank);
   partialSizes.reserve(rank);
   for (auto en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) {
     auto rangeValue = en.value();
     // Try to extract a tight constant.
     LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n");
     IntegerAttr sizeAttr = getSmallestBoundingIndex(rangeValue.size);
     Value size =
         (!sizeAttr) ? rangeValue.size : b.create<ConstantOp>(loc, sizeAttr);
     LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n");
     fullSizes.push_back(size);
     partialSizes.push_back(
         folded_memref_dim(folder, subView, en.index()).value);
   }
   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));
   SmallVector<OpFoldResult, 4> ones(fullSizes.size(), b.getIndexAttr(1));
   auto partialLocalView =
       folded_memref_subview(folder, *fullLocalView, zeros, partialSizes, ones);
   return PromotionInfo{*fullLocalView, partialLocalView};
 }
 
 static Optional<MapVector<unsigned, PromotionInfo>>
 promoteSubViews(OpBuilder &b, Location loc,
-                LinalgOpInstancePromotionOptions options,
+                LinalgOpInstancePromotionOptions options, DataLayout &layout,
                 OperationFolder *folder) {
   if (options.subViews.empty())
     return {};
 
   ScopedContext scope(b, loc);
   MapVector<unsigned, PromotionInfo> promotionInfoMap;
 
   for (auto v : options.subViews) {
     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;
 
     // Only fill the buffer if the full local view is used
     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);
   }
 
   // Copy data into the promoted buffers. Use callback if provided.
   for (auto v : options.subViews) {
     auto info = promotionInfoMap.find(v.first);
     if (info == promotionInfoMap.end())
       continue;
     if (failed(options.copyInFn(
             b, cast<memref::SubViewOp>(v.second.getDefiningOp()),
             info->second.partialLocalView)))
       return {};
   }
   return promotionInfoMap;
 }
 
 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");
 
   if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
     // TODO: add a level of indirection to linalg.generic.
     if (convOp.padding())
       return {};
   }
 
   // 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 {};
 
   // 2. Append all other operands as they appear, this enforces that such
   // operands are not views. This is to support cases such as FillOp taking
   // extra scalars etc.  Keep a reference to output buffers;
   SmallVector<Value, 8> opViews;
mlir::linalg::promoteSubviewsPrecondition(Operation *op,
   return failure();
 }
 
 Optional<LinalgOp> mlir::linalg::promoteSubViews(OpBuilder &b,
                                                  LinalgOp linalgOp,
                                                  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 {
 struct LinalgPromotionPass : public LinalgPromotionBase<LinalgPromotionPass> {
   LinalgPromotionPass() = default;
   LinalgPromotionPass(bool dynamicBuffers, bool useAlloca) {
     this->dynamicBuffers = dynamicBuffers;
     this->useAlloca = useAlloca;

mlir/test/Dialect/Linalg/transform-patterns.mlir

 // CHECK:         %[[s0:.*]] = memref.subview {{%.*}}[{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] [{{%.*}}, {{%.*}}] : memref<?x?xf32, #map{{.*}}> to memref<?x?xf32, #map{{.*}}>
 // CHECK:         %[[a0:.*]] = memref.alloc({{%.*}}) {alignment = 32 : i64} : memref<?xi8>
 // CHECK:         %[[v0:.*]] = memref.view %[[a0]][{{.*}}][{{%.*}}, {{%.*}}] : memref<?xi8> to memref<?x?xf32>
 // CHECK:         %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xf32> to memref<?x?xf32, #[[$STRIDED_2D_u_1]]>
 // CHECK:         linalg.fill(%[[v0]], {{%.*}}) : memref<?x?xf32>, f32
 // 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>) {
   linalg.matmul {__internal_linalg_transform__ = "par__with_perm__"}
     ins(%arg0, %arg1: memref<?x?xf32>, memref<?x?xf32>)
    outs(%arg2: memref<?x?xf32>)
   return
 }

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

 
 //===----------------------------------------------------------------------===//
 // Test promotion callbacks
 //===----------------------------------------------------------------------===//
 
 // Allocation call back
 static Optional<Value> allocCallBackFn(OpBuilder &b, memref::SubViewOp subView,
                                        ArrayRef<Value> boundingSubViewSize,
+                                       DataLayout &layout,
                                        OperationFolder *folder) {
   SmallVector<int64_t, 4> shape(boundingSubViewSize.size(), -1);
   return b
       .create<memref::AllocOp>(
           subView.getLoc(),
           MemRefType::get(shape, subView.getType().getElementType(),
                           /*affineMapComposition =*/{}, 3),
           boundingSubViewSize)