[mlir][spirv] Support convert complex types

[antiagainst]

Summary

Complex types are converted to a two-element vector type to contain
the real and imaginary numbers.

[antiagainst]

Update

[kuhar]

Do we have to also support vector of complex, e.g. ,vector<2xcomplex<f32>>. This doesn't seem supported at the moment.

[antiagainst]

In D147188#4233702, @kuhar wrote:

Do we have to also support vector of complex, e.g. ,vector<2xcomplex<f32>>. This doesn't seem supported at the moment.

vector type does not allow to have complex element types (at least for now).

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
+#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 
 #include <functional>
+#include <optional>
 
 #define DEBUG_TYPE "mlir-spirv-conversion"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Utility functions
 //===----------------------------------------------------------------------===//
if (type.isa<spirv::ScalarType>()) {
     // be used with logical addressing operations, not physical, and only with
     // non-externally visible shader Storage Classes: Workgroup, CrossWorkgroup,
     // Private, Function, Input, and Output."
     if (bitWidth == 1)
       return std::nullopt;
     return bitWidth / 8;
   }
 
+  if (auto complexType = type.dyn_cast<ComplexType>()) {
+    auto elementSize = getTypeNumBytes(options, complexType.getElementType());
+    if (!elementSize)
+      return std::nullopt;
+    return 2 * *elementSize;
+  }
+
   if (auto vecType = type.dyn_cast<VectorType>()) {
     auto elementSize = getTypeNumBytes(options, vecType.getElementType());
     if (!elementSize)
       return std::nullopt;
     return vecType.getNumElements() * *elementSize;
   }
 
   if (auto memRefType = type.dyn_cast<MemRefType>()) {
convertVectorType(const spirv::TargetEnv &targetEnv,
 
   auto elementType =
       convertScalarType(targetEnv, options, scalarType, storageClass);
   if (elementType)
     return VectorType::get(type.getShape(), elementType);
   return nullptr;
 }
 
+static Type
+convertComplexType(const spirv::TargetEnv &targetEnv,
+                   const SPIRVConversionOptions &options, ComplexType type,
+                   std::optional<spirv::StorageClass> storageClass = {}) {
+  auto scalarType = dyn_cast_or_null<spirv::ScalarType>(type.getElementType());
+  if (!scalarType) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: cannot convert non-scalar element type\n");
+    return nullptr;
+  }
+
+  auto elementType =
+      convertScalarType(targetEnv, options, scalarType, storageClass);
+  if (!elementType)
+    return nullptr;
+  if (elementType != type.getElementType()) {
+    LLVM_DEBUG(llvm::dbgs()
+               << type << " illegal: complex type emulation unsupported\n");
+    return nullptr;
+  }
+
+  return VectorType::get(2, elementType);
+}
+
 /// Converts a tensor `type` to a suitable type under the given `targetEnv`.
 ///
 /// Note that this is mainly for lowering constant tensors. In SPIR-V one can
 /// create composite constants with OpConstantComposite to embed relative large
 /// constant values and use OpCompositeExtract and OpCompositeInsert to
 /// manipulate, like what we do for vectors.
 static Type convertTensorType(const spirv::TargetEnv &targetEnv,
                               const SPIRVConversionOptions &options,
static Type convertBoolMemrefType(const spirv::TargetEnv &targetEnv,
   std::optional<int64_t> arrayElemSize =
       getTypeNumBytes(options, arrayElemType);
   if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return nullptr;
   }
 
-
   if (!type.hasStaticShape()) {
     // For OpenCL Kernel, dynamic shaped memrefs convert into a pointer pointing
     // to the element.
     if (targetEnv.allows(spirv::Capability::Kernel))
       return spirv::PointerType::get(arrayElemType, storageClass);
     int64_t stride = needsExplicitLayout(storageClass) ? *arrayElemSize : 0;
     auto arrayType = spirv::RuntimeArrayType::get(arrayElemType, stride);
     // For Vulkan we need extra wrapping struct and array to satisfy interface
if (type.getElementType().isa<IntegerType>() &&
     return convertBoolMemrefType(targetEnv, options, type, storageClass);
   }
 
   Type arrayElemType;
   Type elementType = type.getElementType();
   if (auto vecType = elementType.dyn_cast<VectorType>()) {
     arrayElemType =
         convertVectorType(targetEnv, options, vecType, storageClass);
+  } else if (auto complexType = elementType.dyn_cast<ComplexType>()) {
+    arrayElemType =
+        convertComplexType(targetEnv, options, complexType, storageClass);
   } else if (auto scalarType = elementType.dyn_cast<spirv::ScalarType>()) {
     arrayElemType =
         convertScalarType(targetEnv, options, scalarType, storageClass);
   } else if (auto indexType = elementType.dyn_cast<IndexType>()) {
     type = convertIndexElementType(type, options).cast<MemRefType>();
     arrayElemType = type.getElementType();
   } else {
     LLVM_DEBUG(
         llvm::dbgs()
         << type
         << " unhandled: can only convert scalar or vector element type\n");
     return nullptr;
   }
   if (!arrayElemType)
     return nullptr;
 
   std::optional<int64_t> arrayElemSize =
       getTypeNumBytes(options, arrayElemType);
   if (!arrayElemSize) {
     LLVM_DEBUG(llvm::dbgs()
                << type << " illegal: cannot deduce converted element size\n");
     return nullptr;
   }
 
-
   if (!type.hasStaticShape()) {
     // For OpenCL Kernel, dynamic shaped memrefs convert into a pointer pointing
     // to the element.
     if (targetEnv.allows(spirv::Capability::Kernel))
       return spirv::PointerType::get(arrayElemType, storageClass);
     int64_t stride = needsExplicitLayout(storageClass) ? *arrayElemSize : 0;
     auto arrayType = spirv::RuntimeArrayType::get(arrayElemType, stride);
     // For Vulkan we need extra wrapping struct and array to satisfy interface
SPIRVTypeConverter::SPIRVTypeConverter(spirv::TargetEnvAttr targetAttr,
   });
 
   addConversion([this](FloatType floatType) -> std::optional<Type> {
     if (auto scalarType = floatType.dyn_cast<spirv::ScalarType>())
       return convertScalarType(this->targetEnv, this->options, scalarType);
     return Type();
   });
 
+  addConversion([this](ComplexType complexType) {
+    return convertComplexType(this->targetEnv, this->options, complexType);
+  });
+
   addConversion([this](VectorType vectorType) {
     return convertVectorType(this->targetEnv, this->options, vectorType);
   });
 
   addConversion([this](TensorType tensorType) {
     return convertTensorType(this->targetEnv, this->options, tensorType);
   });
 

mlir/test/Conversion/FuncToSPIRV/types-to-spirv.mlir

 // CHECK-NOT: spirv.func @bf16_type
 func.func @bf16_type(%arg0: bf16) { return }
 
 } // end module
 
 // -----
 
 //===----------------------------------------------------------------------===//
+// Complex types
+//===----------------------------------------------------------------------===//
+
+// Check that capabilities for scalar types affects complex types too: having
+// special capabilities means keep vector types untouched.
+module attributes {
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0,
+    [Float64, StorageUniform16, StorageBuffer16BitAccess],
+    [SPV_KHR_16bit_storage, SPV_KHR_8bit_storage]>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: func @complex_types
+// CHECK-SAME: vector<2xf32>
+// CHECK-SAME: vector<2xf64>
+func.func @complex_types(
+    %arg0: complex<f32>,
+    %arg2: complex<f64>
+) { return }
+
+// CHECK-LABEL: func @memref_complex_types_with_cap
+// CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<4 x vector<2xf16>, stride=4> [0])>, StorageBuffer>
+// CHECK-SAME: !spirv.ptr<!spirv.struct<(!spirv.array<16 x vector<2xf16>, stride=4> [0])>, Uniform>
+func.func @memref_complex_types_with_cap(
+    %arg0: memref<4xcomplex<f16>, #spirv.storage_class<StorageBuffer>>,
+    %arg1: memref<2x8xcomplex<f16>, #spirv.storage_class<Uniform>>
+) { return }
+
+} // end module
+
+// -----
+
+// Check that capabilities for scalar types affects complex types too: no special
+// capabilities available means widening element types to 32-bit.
+
+module attributes {
+  spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, #spirv.resource_limits<>>
+} {
+
+// Emulation is unimplemented right now.
+// CHECK-LABEL: func @memref_complex_types_no_cap
+// CHECK-SAME: memref<4xcomplex<f16>, #spirv.storage_class<StorageBuffer>>
+// CHECK-SAME: memref<2x8xcomplex<f16>, #spirv.storage_class<Uniform>>
+// NOEMU-LABEL: func @memref_complex_types_no_cap
+// NOEMU-SAME: memref<4xcomplex<f16>, #spirv.storage_class<StorageBuffer>>
+// NOEMU-SAME: memref<2x8xcomplex<f16>, #spirv.storage_class<Uniform>>
+func.func @memref_complex_types_no_cap(
+    %arg0: memref<4xcomplex<f16>, #spirv.storage_class<StorageBuffer>>,
+    %arg1: memref<2x8xcomplex<f16>, #spirv.storage_class<Uniform>>
+) { return }
+
+} // end module
+
+// -----
+
+//===----------------------------------------------------------------------===//
 // Vector types
 //===----------------------------------------------------------------------===//
 
 // Check that capabilities for scalar types affects vector types too: no special
 // capabilities available means widening element types to 32-bit.
 module attributes {
   spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [], []>, #spirv.resource_limits<>>
 } {