[Encoding] Add new identity encoding attribute (iree-org#21258)

Currently, the padding encoding attribute with sizes all zero is used to
represent the identity encoding. This PR replaces that with an explicit
identity encoding attribute which avoids the confusion of padding
attributes showing up in non-padding related data-tiling pipelines.

Signed-off-by: Jorn Tuyls <[email protected]>
Signed-off-by: keshavvinayak01 <[email protected]>

Author: jtuyls

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp

@@ -496,6 +496,39 @@ LogicalResult PaddingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
+//===---------------------------------------------------------------------===//
+// iree_encoding.identity
+//===---------------------------------------------------------------------===//
+
+Value IdentityAttr::calculateStorageSizeInBytes(Location loc,
+                                                OpBuilder &builder,
+                                                RankedTensorType type,
+                                                ValueRange dynamicDims) const {
+  const int64_t elementSize = getRoundedElementByteWidth(type.getElementType());
+  int64_t staticProduct = elementSize;
+  Value dynamicProduct = builder.create<arith::ConstantIndexOp>(loc, 1);
+
+  size_t dynamicDimIdx = 0;
+  for (int64_t dimSize : type.getShape()) {
+    if (!ShapedType::isDynamic(dimSize)) {
+      staticProduct *= dimSize;
+      continue;
+    }
+
+    Value dynamicDimSize = dynamicDims[dynamicDimIdx];
+    ++dynamicDimIdx;
+    dynamicProduct = builder.createOrFold<arith::MulIOp>(
+        loc, dynamicProduct, dynamicDimSize, arith::IntegerOverflowFlags::nsw);
+  }
+  return builder.createOrFold<arith::MulIOp>(
+      loc, builder.create<arith::ConstantIndexOp>(loc, staticProduct),
+      dynamicProduct, arith::IntegerOverflowFlags::nsw);
+}
+
+bool IdentityAttr::isIdentityLayout() const { return true; }
+
+bool IdentityAttr::isSerialized() const { return true; }
+
 //===---------------------------------------------------------------------===//
 // iree_encoding.identity_resolver
 //===---------------------------------------------------------------------===//
@@ -506,9 +539,7 @@ IdentityResolverAttr::cloneWithSimplifiedConfig(DictionaryAttr) const {
 }
 
 Attribute IdentityResolverAttr::getLayout(RankedTensorType type) const {
-  MLIRContext *ctx = getContext();
-  SmallVector<int64_t> zeros(type.getRank(), 0);
-  return Encoding::PaddingAttr::get(ctx, DenseI64ArrayAttr::get(ctx, zeros));
+  return Encoding::IdentityAttr::get(getContext());
 }
 
 //===---------------------------------------------------------------------===//

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td

@@ -297,6 +297,28 @@ def PaddingAttr : IREEEncoding_Attr<"Padding", [
   let genVerifyDecl = 1;
 }
 
+//===---------------------------------------------------------------------===//
+// iree_encoding.identity
+//===---------------------------------------------------------------------===//
+
+def IdentityAttr : IREEEncoding_Attr<"Identity", [
+      DeclareAttrInterfaceMethods<IREEEncoding_SerializableAttr, [
+        "calculateStorageSizeInBytes",
+        "isIdentityLayout",
+        "isSerialized"
+        ]>
+    ]> {
+  let mnemonic = "identity";
+
+  let summary = [{The identity encoding.}];
+  let description = [{
+    An encoding attribute that represents the identity function on a type, i.e.
+    it represents the same type as if there was no encoding.
+  }];
+
+  let genVerifyDecl = 0;
+}
+
 //===---------------------------------------------------------------------===//
 // iree_encoding.identity_resolver
 //===---------------------------------------------------------------------===//

compiler/src/iree/compiler/Dialect/Encoding/IR/test/roundtrip.mlir

@@ -262,3 +262,11 @@ func.func @dynamic_layout_encoding(%arg0: tensor<?x?x?xf32, #encoding>) -> tenso
 }
 //      CHECK: #[[ENCODING:.+]] = #iree_encoding.layout<[#iree_encoding.padding<[0, ?, 64]>]>
 //      CHECK: func.func @dynamic_layout_encoding(%[[ARG0:.+]]: tensor<?x?x?xf32, #[[ENCODING]]>)
+
+// -----
+
+#encoding = #iree_encoding.identity
+func.func @identity_encoding(%arg0: tensor<?x?xf32, #encoding>) -> tensor<?x?xf32, #encoding> {
+  return %arg0 : tensor<?x?xf32, #encoding>
+}
+//      CHECK: func.func @identity_encoding(%[[ARG0:.+]]: tensor<?x?xf32, #iree_encoding.identity>

compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir

@@ -459,7 +459,7 @@ util.func public @drop_encoding(%arg0: index, %arg1: index, %scalar_f32 : f32) {
   %0 = stream.tensor.empty on(#hal.device.affinity<@device_a>) : tensor<?x0xf32, #encoding>{%arg0} in !stream.resource<*>{%arg1}
   util.return
 }
-// CHECK-DAG:   #[[$IDENTITY_ENCODING:.+]] = #iree_encoding.testing<[#iree_encoding.padding<[0, 0]>]>
+// CHECK-DAG:   #[[$IDENTITY_ENCODING:.+]] = #iree_encoding.testing<[#iree_encoding.identity]>
 // CHECK-LABEL: util.func public @drop_encoding
 // CHECK:         stream.tensor.empty {{.+}} : tensor<?x0xf32, #[[$IDENTITY_ENCODING]]>
 
@@ -476,7 +476,7 @@ util.func public @ignore_encoding_by_identity_resolver(%arg0: index, %arg1: inde
   %0 = stream.tensor.empty on(#hal.device.affinity<@device_a>) : tensor<?x0xf32, #encoding>{%arg0} in !stream.resource<*>{%arg1}
   util.return
 }
-// CHECK-DAG:   #[[$IDENTITY_ENCODING:.+]] = #iree_encoding.testing<[#iree_encoding.padding<[0, 0]>]>
+// CHECK-DAG:   #[[$IDENTITY_ENCODING:.+]] = #iree_encoding.testing<[#iree_encoding.identity]>
 // CHECK-LABEL: util.func public @ignore_encoding_by_identity_resolver
 // CHECK:         stream.tensor.empty {{.+}} : tensor<?x0xf32, #[[$IDENTITY_ENCODING]]>