Fix crash on ComplexType in PointwiseToLinalgMapConverter (#2754)

A recent change added code to greedily materialize splat constants, but
the code would crash when used with `complex<..>` types.

Author: christopherbate

stablehlo/conversions/linalg/tests/pointwise.mlir

@@ -23,6 +23,29 @@ func.func @float_add(%lhs: tensor<2x2xf32>,
 
 // -----
 
+// CHECK: #map = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-LABEL: func @complex_add_const
+// CHECK-PRIMITIVE-LABEL: func @complex_add_const
+func.func @complex_add_const(%lhs: tensor<2x2xcomplex<f32>>,
+                             %rhs: tensor<2x2xcomplex<f32>>)
+                             -> tensor<2x2xcomplex<f32>> {
+
+  // CHECK: %[[CST:.+]] = complex.constant [1.000000e-01 : f32, 2.000000e-01 : f32] : complex<f32>
+  // CHECK: linalg.generic
+  // CHECK: ^bb0(%[[IN:.+]]: complex<f32>, %[[OUT:.+]]: complex<f32>)
+  // CHECK: %[[RESULT:[a-zA-Z0-9_]*]] = complex.add %[[IN]], %[[CST]]
+  // CHECK: linalg.yield %[[RESULT]]
+
+  // CHECK-PRIMITIVE: linalg.map
+  // CHECK-PRIMITIVE: complex.add
+  %cst = stablehlo.constant dense<(0.1, 0.2)> : tensor<2x2xcomplex<f32>>
+  %0 = "stablehlo.add"(%lhs, %cst) {someattr}
+      : (tensor<2x2xcomplex<f32>>, tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
+  func.return %0 : tensor<2x2xcomplex<f32>>
+}
+
+// -----
+
 // CHECK-LABEL: func @float_add_dynamic_encoding
 // CHECK-PRIMITIVE-LABEL: func @float_add_dynamic_encoding
 func.func @float_add_dynamic_encoding(

stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp

@@ -114,6 +114,28 @@ FailureOr<PointwiseConversionInfo> checkOperandsAndResults(
   return PointwiseConversionInfo{maxRank, resultTy};
 }
 
+/// If `input` is a splat constant value, materialize the scalar splat
+/// value. Otherwise, return nullopt.
+std::optional<Value> materializeSplatScalarConstant(RewriterBase &rewriter,
+                                                    Location loc, Value input) {
+  SplatElementsAttr attr;
+  Type elementType = mlir::getElementTypeOrSelf(input.getType());
+  if (!matchPattern(input, m_Constant(&attr))) return {};
+  if (isa<IntegerType, FloatType, IndexType>(elementType)) {
+    return rewriter
+        .create<arith::ConstantOp>(loc, elementType,
+                                   attr.getSplatValue<TypedAttr>())
+        .getResult();
+  }
+  if (isa<ComplexType>(elementType)) {
+    return rewriter
+        .create<complex::ConstantOp>(loc, elementType,
+                                     attr.getSplatValue<ArrayAttr>())
+        .getResult();
+  }
+  return {};
+}
+
 /// Converts a HLO operation to a linalg.map op that contains the corresponding
 /// scalar operations.
 template <typename OpTy>
@@ -160,11 +182,9 @@ struct PointwiseToLinalgMapConverter : OpConversionPattern<OpTy> {
     SmallVector<Value> mappedInputs;
     SmallVector<Value> scalarInputs;
     for (Value input : adaptor.getOperands()) {
-      DenseElementsAttr attr;
-      if (matchPattern(input, m_Constant(&attr)) && attr.isSplat()) {
-        scalarInputs.push_back(rewriter.create<arith::ConstantOp>(
-            loc, cast<ShapedType>(input.getType()).getElementType(),
-            attr.getSplatValue<TypedAttr>()));
+      if (std::optional<Value> splatVal =
+              materializeSplatScalarConstant(rewriter, loc, input)) {
+        scalarInputs.push_back(*splatVal);
       } else if (getRank(input) == maxRank) {
         mappedInputs.push_back(coerceTensorShape(
             rewriter, loc, cast<TypedValue<ShapedType>>(input),