[mlir][spirv][complex] Convert complex ops to SPIR-V ops

[antiagainst]

Summary

This commit adds conversion from complex construction and
extraction ops to SPIR-V. Other arithemtic ops can be done
via ComplexToStandard patterns.

[kuhar]

mlir/include/mlir/Conversion/Passes.td

231

I think this constructor should be auto generated if you omit this line

mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp

32

nit: public is redundant for structs. Also elsewhere below.

mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRVPass.cpp

30–34

ubernit: why not define it inline?

[antiagainst]

Address comments

[antiagainst]

Drop unnecessary headers for Bazel

mlir/include/mlir/Conversion/ComplexToSPIRV/ComplexToSPIRV.h

+//===- ComplexToSPIRV.h - Complex to SPIR-V Patterns ------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Provides patterns to convert Complex dialect to SPIR-V dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_COMPLEXTOSPIRV_COMPLEXTOSPIRV_H
+#define MLIR_CONVERSION_COMPLEXTOSPIRV_COMPLEXTOSPIRV_H
+
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+class SPIRVTypeConverter;
+
+/// Appends to a pattern list additional patterns for translating Complex ops
+/// to SPIR-V ops.
+void populateComplexToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
+                                    RewritePatternSet &patterns);
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_COMPLEXTOSPIRV_COMPLEXTOSPIRV_H

mlir/include/mlir/Conversion/ComplexToSPIRV/ComplexToSPIRVPass.h

+//===- ComplexToSPIRVPass.h - Complex to SPIR-V Passes ----------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Provides passes to convert Complex dialect to SPIR-V dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_COMPLEXTOSPIRV_COMPLEXTOSPIRVPASS_H
+#define MLIR_CONVERSION_COMPLEXTOSPIRV_COMPLEXTOSPIRVPASS_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+class ModuleOp;
+
+#define GEN_PASS_DECL_CONVERTCOMPLEXTOSPIRVPass
+#include "mlir/Conversion/Passes.h.inc"
+
+/// Creates a pass to convert Complex ops to SPIR-V ops.
+std::unique_ptr<OperationPass<>> createConvertComplexToSPIRVPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_COMPLEXTOSPIRV_COMPLEXTOSPIRVPASS_H

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h"
 #include "mlir/Conversion/ArmNeon2dToIntr/ArmNeon2dToIntr.h"
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/BufferizationToMemRef/BufferizationToMemRef.h"
 #include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
 #include "mlir/Conversion/ComplexToLibm/ComplexToLibm.h"
+#include "mlir/Conversion/ComplexToSPIRV/ComplexToSPIRVPass.h"
 #include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRV.h"
 #include "mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRVPass.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
 #include "mlir/Conversion/FuncToSPIRV/FuncToSPIRVPass.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"

mlir/include/mlir/Conversion/Passes.td

def ConvertComplexToLibm : Pass<"convert-complex-to-libm", "ModuleOp"> {
   }];
   let constructor = "mlir::createConvertComplexToLibmPass()";
   let dependentDialects = [
     "func::FuncDialect",
   ];
 }
 
 //===----------------------------------------------------------------------===//
+// ComplexToSPIRV
+//===----------------------------------------------------------------------===//
+
+def ConvertComplexToSPIRVPass : Pass<"convert-complex-to-spirv"> {
+  let summary = "Convert Complex dialect to SPIRV dialect";
+  let constructor = "mlir::createConvertComplexToSPIRVPass()";
+  let dependentDialects = ["spirv::SPIRVDialect"];
+}
+
+//===----------------------------------------------------------------------===//
 // ComplexToStandard
 //===----------------------------------------------------------------------===//
 
 def ConvertComplexToStandard : Pass<"convert-complex-to-standard"> {
   let summary = "Convert Complex dialect to standard dialect";
   let constructor = "mlir::createConvertComplexToStandardPass()";
   let dependentDialects = ["math::MathDialect"];
 }

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AMDGPUToROCDL)
 add_subdirectory(ArithCommon)
 add_subdirectory(ArithToLLVM)
 add_subdirectory(ArithToSPIRV)
 add_subdirectory(ArmNeon2dToIntr)
 add_subdirectory(AsyncToLLVM)
 add_subdirectory(BufferizationToMemRef)
 add_subdirectory(ComplexToLLVM)
 add_subdirectory(ComplexToLibm)
+add_subdirectory(ComplexToSPIRV)
 add_subdirectory(ComplexToStandard)
 add_subdirectory(ControlFlowToLLVM)
 add_subdirectory(ControlFlowToSPIRV)
 add_subdirectory(FuncToLLVM)
 add_subdirectory(FuncToSPIRV)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)

mlir/lib/Conversion/ComplexToSPIRV/CMakeLists.txt

+add_mlir_conversion_library(MLIRComplexToSPIRV
+  ComplexToSPIRV.cpp
+  ComplexToSPIRVPass.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SPIRV
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRComplexDialect
+  MLIRIR
+  MLIRPass
+  MLIRSPIRVDialect
+  MLIRSPIRVConversion
+  MLIRSupport
+  MLIRTransformUtils
+  )
+

mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRV.cpp

+//===- ComplexToSPIRV.cpp - Complex to SPIR-V Patterns --------------------===//
+//
+// 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 patterns to convert Complex dialect to SPIR-V dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/ComplexToSPIRV/ComplexToSPIRV.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
+#include "mlir/Dialect/SPIRV/Utils/LayoutUtils.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "complex-to-spirv-pattern"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Operation conversion
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+struct CreateOpPattern final : public 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)
+      return rewriter.notifyMatchFailure(createOp,
+                                         "unable to convert result type");
+
+    rewriter.replaceOpWithNewOp<spirv::CompositeConstructOp>(
+        createOp, spirvType, adaptor.getOperands());
+    return success();
+  }
+};
+
+struct ReOpPattern final : public OpConversionPattern<complex::ReOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::ReOp reOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type spirvType = getTypeConverter()->convertType(reOp.getType());
+    if (!spirvType)
+      return rewriter.notifyMatchFailure(reOp, "unable to convert result type");
+
+    rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
+        reOp, adaptor.getComplex(), llvm::ArrayRef(0));
+    return success();
+  }
+};
+
+struct ImOpPattern final : public OpConversionPattern<complex::ImOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(complex::ImOp imOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type spirvType = getTypeConverter()->convertType(imOp.getType());
+    if (!spirvType)
+      return rewriter.notifyMatchFailure(imOp, "unable to convert result type");
+
+    rewriter.replaceOpWithNewOp<spirv::CompositeExtractOp>(
+        imOp, adaptor.getComplex(), llvm::ArrayRef(1));
+    return success();
+  }
+};
+
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Pattern population
+//===----------------------------------------------------------------------===//
+
+void mlir::populateComplexToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
+                                          RewritePatternSet &patterns) {
+  MLIRContext *context = patterns.getContext();
+
+  patterns.add<CreateOpPattern, ReOpPattern, ImOpPattern>(typeConverter,
+                                                          context);
+}

mlir/lib/Conversion/ComplexToSPIRV/ComplexToSPIRVPass.cpp

+//===- ComplexToSPIRVPass.cpp - Complex to SPIR-V Passes ------------------===//
+//
+// 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 a pass to convert Complex dialect to SPIR-V dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/ComplexToSPIRV/ComplexToSPIRVPass.h"
+
+#include "mlir/Conversion/ComplexToSPIRV/ComplexToSPIRV.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTCOMPLEXTOSPIRVPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+/// A pass converting MLIR Complex operations into the SPIR-V dialect.
+class ConvertComplexToSPIRVPass
+    : public impl::ConvertComplexToSPIRVPassBase<ConvertComplexToSPIRVPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void ConvertComplexToSPIRVPass::runOnOperation() {
+  MLIRContext *context = &getContext();
+  Operation *op = getOperation();
+
+  auto targetAttr = spirv::lookupTargetEnvOrDefault(op);
+  std::unique_ptr<ConversionTarget> target =
+      SPIRVConversionTarget::get(targetAttr);
+
+  SPIRVConversionOptions options;
+  SPIRVTypeConverter typeConverter(targetAttr, options);
+
+  // Use UnrealizedConversionCast as the bridge so that we don't need to pull in
+  // patterns for other dialects.
+  auto addUnrealizedCast = [](OpBuilder &builder, Type type, ValueRange inputs,
+                              Location loc) {
+    auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+    return std::optional<Value>(cast.getResult(0));
+  };
+  typeConverter.addSourceMaterialization(addUnrealizedCast);
+  typeConverter.addTargetMaterialization(addUnrealizedCast);
+  target->addLegalOp<UnrealizedConversionCastOp>();
+
+  RewritePatternSet patterns(context);
+  populateComplexToSPIRVPatterns(typeConverter, patterns);
+
+  if (failed(applyPartialConversion(op, *target, std::move(patterns))))
+    return signalPassFailure();
+}
+
+std::unique_ptr<OperationPass<>> mlir::createConvertComplexToSPIRVPass() {
+  return std::make_unique<ConvertComplexToSPIRVPass>();
+}

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

+// RUN: mlir-opt -split-input-file -convert-complex-to-spirv %s | FileCheck %s
+
+func.func @create_complex(%real: f32, %imag: f32) -> complex<f32> {
+  %0 = complex.create %real, %imag : complex<f32>
+  return %0 : complex<f32>
+}
+
+// CHECK-LABEL: func.func @create_complex
+//  CHECK-SAME: (%[[RE:.+]]: f32, %[[IM:.+]]: f32)
+//       CHECK:   %[[CC:.+]] = spirv.CompositeConstruct %[[RE]], %[[IM]] : (f32, f32) -> vector<2xf32>
+//       CHECK:   %[[CAST:.+]] = builtin.unrealized_conversion_cast %[[CC]] : vector<2xf32> to complex<f32>
+//       CHECK:   return %[[CAST]] : complex<f32>
+
+
+// -----
+
+func.func @real_number(%arg: complex<f32>) -> f32 {
+  %real = complex.re %arg : complex<f32>
+  return %real : f32
+}
+
+// CHECK-LABEL: func.func @real_number
+//  CHECK-SAME: %[[ARG:.+]]: complex<f32>
+//       CHECK:   %[[CAST:.+]] = builtin.unrealized_conversion_cast %[[ARG]] : complex<f32> to vector<2xf32>
+//       CHECK:   %[[RE:.+]] = spirv.CompositeExtract %[[CAST]][0 : i32] : vector<2xf32>
+//       CHECK:   return %[[RE]] : f32
+
+// -----
+
+func.func @imaginary_number(%arg: complex<f32>) -> f32 {
+  %imaginary = complex.im %arg : complex<f32>
+  return %imaginary: f32
+}
+
+// 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
+

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":Support",
         ":Transforms",
         "//llvm:Core",
         "//llvm:Support",
     ],
 )
 
 cc_library(
+    name = "ComplexToSPIRV",
+    srcs = glob([
+        "lib/Conversion/ComplexToSPIRV/*.cpp",
+        "lib/Conversion/ComplexToSPIRV/*.h",
+    ]),
+    hdrs = glob([
+        "include/mlir/Conversion/ComplexToSPIRV/*.h",
+    ]),
+    includes = ["include"],
+    deps = [
+        ":ComplexDialect",
+        ":ConversionPassIncGen",
+        ":IR",
+        ":Pass",
+        ":SPIRVCommonConversion",
+        ":SPIRVConversion",
+        ":SPIRVDialect",
+        ":Support",
+        ":Transforms",
+        "//llvm:Core",
+        "//llvm:Support",
+    ],
+)
+
+cc_library(
     name = "ComplexToStandard",
     srcs = glob([
         "lib/Conversion/ComplexToStandard/*.cpp",
         "lib/Conversion/ComplexToStandard/*.h",
     ]),
     hdrs = glob([
         "include/mlir/Conversion/ComplexToStandard/*.h",
     ]),