[MLIR] Add getSizeInBits() for tensor of complex

[rriddle]

mlir/lib/IR/StandardTypes.cpp

235

nit: Add a short message here.

[jurahul]

mlir/lib/IR/StandardTypes.cpp

235

Do Types get verification similar to Ops? Wondering if the assert belong in ComplexType verification if it exists.

[rriddle]

mlir/lib/IR/StandardTypes.cpp

235

Yes, types are verified on construction, so this assert could be removed. For some reason I was thinking that complex was relaxed recently to allow non int/float. That must have been just a request.

[timshen]

Remove the assert.

mlir/lib/IR/StandardTypes.cpp

 int64_t ShapedType::getSizeInBits() const {
   assert(hasStaticShape() &&
          "cannot get the bit size of an aggregate with a dynamic shape");
 
   auto elementType = getElementType();
   if (elementType.isIntOrFloat())
     return elementType.getIntOrFloatBitWidth() * getNumElements();
 
+  if (auto complexType = elementType.dyn_cast<ComplexType>()) {
+    elementType = complexType.getElementType();
+    return elementType.getIntOrFloatBitWidth() * getNumElements() * 2;
+  }
+
   // Tensors can have vectors and other tensors as elements, other shaped types
   // cannot.
   assert(isa<TensorType>() && "unsupported element type");
   assert((elementType.isa<VectorType, TensorType>()) &&
          "unsupported tensor element type");
   return getNumElements() * elementType.cast<ShapedType>().getSizeInBits();
 }
 

mlir/test/mlir-tblgen/op-derived-attribute.mlir

 // RUN: mlir-opt -test-derived-attr -verify-diagnostics %s | FileCheck %s
 
 // CHECK-LABEL: verifyDerivedAttributes
 func @verifyDerivedAttributes() {
   // expected-remark @+2 {{element_dtype = f32}}
   // expected-remark @+1 {{size = 320}}
   %0 = "test.derived_type_attr"() : () -> tensor<10xf32>
+
   // expected-remark @+2 {{element_dtype = i79}}
   // expected-remark @+1 {{size = 948}}
   %1 = "test.derived_type_attr"() : () -> tensor<12xi79>
 
+  // expected-remark @+2 {{element_dtype = complex<f32>}}
+  // expected-remark @+1 {{size = 768}}
+  %2 = "test.derived_type_attr"() : () -> tensor<12xcomplex<f32>>
+
   return
 }