[mlir][linalg] Add a padding case for `ComplexType`

[rsuderman]

Summary

If the paddingAttr is an ArrayAttr with two values we know that
the element type is a ComplexType and we should pad the value
accordingly.

[rsuderman]

Rebased on head

[mravishankar]

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

131

Is it a given that for complex types it is going to be an ArrayAttr? Seems like a fairly strong assumption.

[rsuderman]

Updated to check ComplexType directly.

[rsuderman]

git-clang-format

[rsuderman]

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

131

The type is an ArrayAttr for ComplexTy but you are right, instead we should directly check the return type. Updated appropriately.

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

 //===- Padding.cpp - Padding of Linalg ops --------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 #define DEBUG_TYPE "linalg-padding"
 
 using namespace mlir;
 using namespace mlir::linalg;
static FailureOr<Value> padOperandToSmallestStaticBoundingBox(
   if (!nofold && alreadyHasRequestedShape)
     return opOperand->get();
 
   // Fail if `paddingValues` specifies no padding value.
   if (opOperand->getOperandNumber() >= options.paddingValues.size()) {
     return rewriter.notifyMatchFailure(opToPad, "--no padding value specified");
   }
   Attribute paddingAttr = options.paddingValues[opOperand->getOperandNumber()];
-  Value paddingValue = rewriter.create<arith::ConstantOp>(
+
+  Value paddingValue;
+  if (auto complexTy = dyn_cast<ComplexType>(
+          getElementTypeOrSelf(opOperand->get().getType()))) {
+    auto complexAttr = cast<ArrayAttr>(paddingAttr);
+    paddingValue = rewriter.create<complex::ConstantOp>(opToPad.getLoc(),
+                                                        complexTy, complexAttr);
+  } else {
+    paddingValue = rewriter.create<arith::ConstantOp>(
         opToPad.getLoc(), cast<TypedAttr>(paddingAttr));
+  }
 
   // Pad the operand to the bounding box defined by `paddedShape`.
   auto paddedTensorType = RankedTensorType::get(
       paddedShape, getElementTypeOrSelf(opOperand->get()));
   LLVM_DEBUG(DBGS() << "--SUCCESS, makeComposedPadHighOp with type: "
                     << paddedTensorType);
   return makeComposedPadHighOp(rewriter, opToPad->getLoc(), paddedTensorType,
                                opOperand->get(), paddingValue, nofold);