[mlir][spirv][complex] Support convert complex.constant

[kuhar]

mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp

37–38

If we expect this to be a shaped type later on, we can do convertType<ShapedType>(constOp.getType());

mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir

44

Should we also have test for vector of complex?

[antiagainst]

Address comments

mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp

37–38

Ah, interesting. Didn't notice this overwrite!

mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir

44

complext.constant does not support creating vectors. https://mlir.llvm.org/docs/Dialects/ComplexOps/#complexconstant-mlircomplexconstantop. You mean using arith.constant or something? Also right now we don't support vector of complex type yet---see callers of https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp#L338.

[kuhar]

LGTM

mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir

44

ah, TIL

mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp

 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Operation conversion
 //===----------------------------------------------------------------------===//
 
 namespace {
 
+struct ConstantOpPattern final : OpConversionPattern<complex::ConstantOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::ConstantOp constOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto spirvType =
+        getTypeConverter()->convertType<ShapedType>(constOp.getType());
+    if (!spirvType)
+      return rewriter.notifyMatchFailure(constOp,
+                                         "unable to convert result type");
+
+    rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
+        constOp, spirvType,
+        DenseElementsAttr::get(spirvType, constOp.getValue().getValue()));
+    return success();
+  }
+};
+
 struct CreateOpPattern final : OpConversionPattern<complex::CreateOp> {
   using OpConversionPattern::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(complex::CreateOp createOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type spirvType = getTypeConverter()->convertType(createOp.getType());
     if (!spirvType)

 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 void mlir::populateComplexToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                           RewritePatternSet &patterns) {
   MLIRContext *context = patterns.getContext();
 
-  patterns.add<CreateOpPattern, ReOpPattern, ImOpPattern>(typeConverter,
-                                                          context);
+  patterns.add<ConstantOpPattern, CreateOpPattern, ReOpPattern, ImOpPattern>(
+      typeConverter, context);
 }

mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir

 }
 
 // CHECK-LABEL: func.func @imaginary_number
 //  CHECK-SAME: %[[ARG:.+]]: complex<f32>
 //       CHECK:   %[[CAST:.+]] = builtin.unrealized_conversion_cast %[[ARG]] : complex<f32> to vector<2xf32>
 //       CHECK:   %[[IM:.+]] = spirv.CompositeExtract %[[CAST]][1 : i32] : vector<2xf32>
 //       CHECK:   return %[[IM]] : f32
 
+// -----
+
+func.func @complex_const() -> complex<f32> {
+  %cst = complex.constant [0x7FC00000 : f32, 0.000000e+00 : f32] : complex<f32>
+  return %cst : complex<f32>
+}
+
+// CHECK-LABEL: func.func @complex_const()
+//       CHECK:   spirv.Constant dense<[0x7FC00000, 0.000000e+00]> : vector<2xf32>