[mlir] Add `ComplexType` conversion support for `convertScalarToDtype`

[rsuderman]

Summary

Linalg operations can include complex types in the src/target types.
This should include conversion between arith and complex types when
constructing linalg operations.

[kuhar]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

126

nit: this function seems to be getting long, I wonder if a type switch and dispatching to helper functions would help here

149

nit: use free cast functions: https://mlir.llvm.org/deprecation/

also below

[rsuderman]

Factored into helper functions and performed general var name cleanup.

[jpienaar]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

127

Should operand be returned if width is equal or is it expected that canonicalized or fold would have handled that?

[rsuderman]

Add operand bailout

[rsuderman]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

127

Added the return operand;. It was technically unnecessary due to the parent function checking quality, but best to handle defense in depth.

[jpienaar]

mlir/lib/Dialect/Arith/Utils/CMakeLists.txt

9

This feels off.

[rsuderman]

mlir/lib/Dialect/Arith/Utils/CMakeLists.txt

9

I agree. The question is where would be the right place to move the conversion?

[kuhar]

Do you plan to add some tests, or is this covered by some existing one?

[rsuderman]

In D154740#4494857, @kuhar wrote:

Do you plan to add some tests, or is this covered by some existing one?

Happy to add some. I am less familiar with this area in the codebase so thoughts on how / where to test is appreciate.

[mravishankar]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

97

Nit: I think the static methods need some comments to say what each one of these are for? Also maybe call this convertScalarToIntDtype and so on to make it clearer that these are implementations of parts of convertScalarToDtype?

[kuhar]

In D154740#4495452, @rsuderman wrote:

In D154740#4494857, @kuhar wrote:

Do you plan to add some tests, or is this covered by some existing one?

Happy to add some. I am less familiar with this area in the codebase so thoughts on how / where to test is appreciate.

I checked and the code from this patch is not exercised by any lit test. The helper is used with the other types in these:

Dialect/Linalg/IR/LinalgOps.cpp
Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
Dialect/SparseTensor/Transforms/CodegenUtils.cpp

So I'd think we either add some linalg/sparse tensor tests, or a new test pass.

[rsuderman]

Added testing and added complex type support for ConvertConv2DtoImg2Col

[rsuderman]

In D154740#4495836, @kuhar wrote:

In D154740#4495452, @rsuderman wrote:

In D154740#4494857, @kuhar wrote:

Do you plan to add some tests, or is this covered by some existing one?

Happy to add some. I am less familiar with this area in the codebase so thoughts on how / where to test is appreciate.

I checked and the code from this patch is not exercised by any lit test. The helper is used with the other types in these:

Dialect/Linalg/IR/LinalgOps.cpp
Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
Dialect/Linalg/Transforms/ConvertConv2DToImg2Col.cpp
Dialect/SparseTensor/Transforms/CodegenUtils.cpp

So I'd think we either add some linalg/sparse tensor tests, or a new test pass.

Added tests to the Img2Col case. This included fixing the pass to work for complex types and expanding support for mixed types.

[kuhar]

LGTM

mlir/lib/Dialect/Arith/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRArithUtils
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith/Utils
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
+  MLIRComplexDialect
   MLIRDialect
   MLIRIR
   )

mlir/lib/Dialect/Arith/Utils/Utils.cpp

 //===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
 //
 // 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 utilities for the Linalg dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 
 /// Matches a ConstantIndexOp.
 /// TODO: This should probably just be a general matcher that uses matchConstant
 /// and checks the operation for an index type.
 detail::op_matcher<arith::ConstantIndexOp> mlir::matchConstantIndex() {
Value mlir::convertScalarToDtype(OpBuilder &b, Location loc, Value operand,
                                  Type toType, bool isUnsignedCast) {
   if (operand.getType() == toType)
     return operand;
   if (auto toIntType = dyn_cast<IntegerType>(toType)) {
     // If operand is floating point, cast directly to the int type.
     if (isa<FloatType>(operand.getType())) {
       if (isUnsignedCast)
         return b.create<arith::FPToUIOp>(loc, toType, operand);
       return b.create<arith::FPToSIOp>(loc, toType, operand);
     }
     // Cast index operands directly to the int type.
     if (operand.getType().isIndex())
       return b.create<arith::IndexCastOp>(loc, toType, operand);
     if (auto fromIntType = dyn_cast<IntegerType>(operand.getType())) {
       // Either extend or truncate.
       if (toIntType.getWidth() > fromIntType.getWidth()) {
         if (isUnsignedCast)
if (isa<IntegerType>(operand.getType())) {
       return b.create<arith::SIToFPOp>(loc, toFloatType, operand);
     }
     if (auto fromFloatType = dyn_cast<FloatType>(operand.getType())) {
       if (toFloatType.getWidth() > fromFloatType.getWidth())
         return b.create<arith::ExtFOp>(loc, toFloatType, operand);
       if (toFloatType.getWidth() < fromFloatType.getWidth())
         return b.create<arith::TruncFOp>(loc, toFloatType, operand);
     }
+  } else if (auto toComplexType = dyn_cast<ComplexType>(toType)) {
+    if (auto fromComplexType = dyn_cast<ComplexType>(operand.getType())) {
+      if (isa<FloatType>(toComplexType.getElementType()) &&
+          isa<FloatType>(fromComplexType.getElementType())) {
+        Value real = b.create<complex::ReOp>(loc, operand);
+        Value imag = b.create<complex::ImOp>(loc, operand);
+        if (toComplexType.getElementType().getIntOrFloatBitWidth() <
+            fromComplexType.getElementType().getIntOrFloatBitWidth()) {
+          real = b.create<arith::TruncFOp>(loc, toComplexType.getElementType(),
+                                           real);
+          imag = b.create<arith::TruncFOp>(loc, toComplexType.getElementType(),
+                                           imag);
+        } else {
+          real = b.create<arith::ExtFOp>(loc, toComplexType.getElementType(),
+                                         real);
+          imag = b.create<arith::ExtFOp>(loc, toComplexType.getElementType(),
+                                         imag);
+        }
+        return b.create<complex::CreateOp>(loc, toComplexType, real, imag);
       }
+    }
+
+    if (auto fromFpType = dyn_cast<FloatType>(operand.getType())) {
+      FloatType toFloatTy = toComplexType.getElementType().cast<FloatType>();
+      int32_t toBitwidth = toFloatTy.getIntOrFloatBitWidth();
+      Value from = operand;
+      if (from.getType().getIntOrFloatBitWidth() < toBitwidth) {
+        from = b.create<arith::ExtFOp>(loc, toFloatTy, from);
+      }
+      if (from.getType().getIntOrFloatBitWidth() > toBitwidth) {
+        from = b.create<arith::TruncFOp>(loc, toFloatTy, from);
+      }
+      Value zero = b.create<mlir::arith::ConstantFloatOp>(
+          loc, mlir::APFloat(toFloatTy.getFloatSemantics(), 0), toFloatTy);
+      return b.create<complex::CreateOp>(loc, toComplexType, from, zero);
+    }
+
+    if (auto fromIntType = dyn_cast<IntegerType>(operand.getType())) {
+      FloatType toFloatTy = toComplexType.getElementType().cast<FloatType>();
+      Value from = operand;
+      if (isUnsignedCast) {
+        from = b.create<arith::UIToFPOp>(loc, toFloatTy, from);
+      } else {
+        from = b.create<arith::SIToFPOp>(loc, toFloatTy, from);
+      }
+      Value zero = b.create<mlir::arith::ConstantFloatOp>(
+          loc, mlir::APFloat(toFloatTy.getFloatSemantics(), 0), toFloatTy);
+      return b.create<complex::CreateOp>(loc, toComplexType, from, zero);
+    }
+  }
+
   emitWarning(loc) << "could not cast operand of type " << operand.getType()
                    << " to " << toType;
   return operand;
 }
 
 SmallVector<Value>
 mlir::getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
                                       ArrayRef<OpFoldResult> valueOrAttrVec) {