[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.

283

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

286

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

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

[ftynse]

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

55–58

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–57

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/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"

 
 /// 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,
                          Optional<unsigned> alignment = None) {
   auto *ctx = size.getContext();
-  auto width = llvm::divideCeil(elementType.getIntOrFloatBitWidth(), 8);
+  DataLayout layout;
+  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)
for (auto v : options.subViews) {
     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;

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

mlir::detail::getDefaultTypeSize(Type type, const DataLayout &dataLayout,
   return llvm::divideCeil(bits, 8);
 }
 
 unsigned mlir::detail::getDefaultTypeSizeInBits(Type type,
                                                 const DataLayout &dataLayout,
                                                 DataLayoutEntryListRef params) {
   if (type.isa<IntegerType, FloatType>())
     return type.getIntOrFloatBitWidth();
 
+  if (auto ctype = type.dyn_cast<ComplexType>())
+    return 2 *
+           getDefaultTypeSizeInBits(ctype.getElementType(), dataLayout, params);
+
   // Index is an integer of some bitwidth.
   if (type.isa<IndexType>())
     return dataLayout.getTypeSizeInBits(
         IntegerType::get(type.getContext(), getIndexBitwidth(params)));
 
   // Sizes of vector types are rounded up to those of types with closest
   // power-of-two number of elements in the innermost dimension. We also assume
   // there is no bit-packing at the moment element sizes are taken in bytes and

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
 }