[Mlir-commits] [mlir] 52aaac6 - [mlir][spirv][complex] Support convert complex.constant

[Lei Zhang]

Author: Lei Zhang
Date: 2023-05-29T17:54:28-07:00
New Revision: 52aaac635a0f084fbbe4a4209afba3447b2af068

URL: https://github.com/llvm/llvm-project/commit/52aaac635a0f084fbbe4a4209afba3447b2af068
DIFF: https://github.com/llvm/llvm-project/commit/52aaac635a0f084fbbe4a4209afba3447b2af068.diff

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

Reviewed By: kuhar

Differential Revision: https://reviews.llvm.org/D151622

Added: 
    

Modified: 
    mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
    mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
index cb1362dc7f120..d531659e0623a 100644
--- a/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
+++ b/mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp
@@ -28,6 +28,25 @@ using namespace mlir;
 
 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;
 
@@ -87,6 +106,6 @@ 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);
 }

diff  --git a/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir b/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
index 6b56a4488b882..45f38d435c50b 100644
--- a/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
+++ b/mlir/test/Conversion/ComplexToSPIRV/complex-to-spirv.mlir
@@ -38,3 +38,12 @@ func.func @imaginary_number(%arg: complex<f32>) -> f32 {
 //       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>