Add support for serializing complex resource attributes (#15960)

Complex resource attributes are just a composite version of int and float types and can be serialized the exact same way.
iree-org · Dec 17, 2023 · e2dcbe1 · e2dcbe1
1 parent 18ed8d5
commit e2dcbe1
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 10 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Util/IR/UtilAttrs.cpp b/compiler/src/iree/compiler/Dialect/Util/IR/UtilAttrs.cpp
@@ -442,14 +442,18 @@ static LogicalResult serializeGenericResourceElementData(
     return emitError(loc) << "the endian of the "
                              "DenseResourceElementsAttr is not supported";
   }
-  if (llvm::isa<IntegerType>(resourceElementsAttr.getType().getElementType())) {
-    // Don't hoist bitWidth given `getElementTypeBitWidth()` asserts if the
-    // element type is not integer or floating-point.
+  // For complex resource types, we can just serialize based on the bit width of
+  // the underlying integer or floating point type.
+  Type elementType = resourceElementsAttr.getType().getElementType();
+  if (auto complexType = llvm::dyn_cast<ComplexType>(elementType)) {
+    elementType = complexType.getElementType();
+  }
+  if (auto integerType = llvm::dyn_cast<IntegerType>(elementType)) {
     // At the time of writing, DenseResourceElementsAttr byte aligned physical
     // element types only with the exception of i1, which is stored as a full
     // byte. This is in contrast to DenseElementsAttr which has an exception for
     // i1 where it is bit-packed.
-    unsigned bitWidth = resourceElementsAttr.getType().getElementTypeBitWidth();
+    unsigned bitWidth = integerType.getIntOrFloatBitWidth();
     switch (bitWidth) {
     case 1:
       return serializeResourceRawData(loc, resourceElementsAttr, os);
@@ -466,12 +470,9 @@ static LogicalResult serializeGenericResourceElementData(
              << "unhandled integer element bit width " << bitWidth
              << " for type " << resourceElementsAttr.getType();
     }
-  } else if (llvm::isa<FloatType>(
-                 resourceElementsAttr.getType().getElementType())) {
-    // Don't hoist bitWidth given `getElementTypeBitWidth()` asserts if the
-    // element type is not integer or floating-point.
-    // TODO(saienduri): implement float64 support (not neccesary now)
-    unsigned bitWidth = resourceElementsAttr.getType().getElementTypeBitWidth();
+  } else if (auto floatType = llvm::dyn_cast<FloatType>(elementType)) {
+    // TODO(saienduri): implement float64 support (not necessary now)
+    unsigned bitWidth = floatType.getIntOrFloatBitWidth();
     switch (bitWidth) {
     case 16:
       return serializeResourceRawData(loc, resourceElementsAttr, os);

diff --git a/compiler/src/iree/compiler/Dialect/VM/Target/Bytecode/test/constant_encoding.mlir b/compiler/src/iree/compiler/Dialect/VM/Target/Bytecode/test/constant_encoding.mlir
@@ -108,6 +108,30 @@ vm.module @constants {
   // CHECK-NEXT: ]
   vm.rodata private @dense_f32 dense<[1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<3xf32>
 
+
+  //      CHECK: "embedded_data": [
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   128,
+  // CHECK-NEXT:   63,
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   64,
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   64,
+  // CHECK-NEXT:   64,
+
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   0,
+  // CHECK-NEXT:   128,
+  // CHECK-NEXT:   64
+  // CHECK-NEXT: ]
+  vm.rodata private @dense_resource_complex_f32 dense<
+                      "0x0000803F000000400000404000008040"
+                    > : tensor<2xcomplex<f32>>
+
   //      CHECK: "embedded_data": [
   // CHECK-NEXT:   0,
   // CHECK-NEXT:   0,