[Encoding] Implement isIdentityLayout interface method. (iree-org#20286)

The revision introduces `isIdentityLayout` interface method to
SerializableEncodingAttrInterface, and implements the method for
`EncodingAttr` and `PadEncodingLayoutAttr` attributes.

It also fixes a bug in `EncodingAttr` builders, where we should use NULL
when maps are not present.

---------

Signed-off-by: hanhanW <[email protected]>

Author: hanhanW

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

@@ -69,13 +69,14 @@ EncodingAttr EncodingAttr::get(MLIRContext *ctx, int64_t operandIndex,
                                ArrayRef<Attribute> layouts) {
   Builder b(ctx);
   auto opTypeAttr = EncodingOpTypeAttr::get(ctx, opType);
+  auto mapsAttr = maps.empty() ? ArrayAttr() : b.getAffineMapArrayAttr(maps);
   auto roundDimsToAttr = roundDimsTo.empty()
                              ? DenseI64ArrayAttr()
                              : b.getDenseI64ArrayAttr(roundDimsTo);
   auto layoutsAttr = layouts.empty() ? ArrayAttr() : b.getArrayAttr(layouts);
   return get(ctx, b.getIndexAttr(operandIndex), opTypeAttr,
-             b.getTypeArrayAttr(elemTypes), b.getAffineMapArrayAttr(maps),
-             roundDimsToAttr, layoutsAttr);
+             b.getTypeArrayAttr(elemTypes), mapsAttr, roundDimsToAttr,
+             layoutsAttr);
 }
 
 LogicalResult
@@ -193,6 +194,25 @@ EncodingAttr::cloneWithNewOperandIndexingMap(AffineMap newIndexingMap) {
 
 bool EncodingAttr::isSerialized() const { return getLayouts() ? true : false; }
 
+bool EncodingAttr::isIdentityLayout() const {
+  if (!isSerialized()) {
+    return false;
+  }
+  ArrayAttr layoutsAttr = getLayouts();
+  if (!llvm::all_of(layoutsAttr.getValue(),
+                    llvm::IsaPred<SerializableEncodingAttrInterface>)) {
+    return false;
+  }
+  return llvm::all_of(layoutsAttr.getValue(), [](Attribute attr) {
+    auto serializableAttr =
+        llvm::dyn_cast<SerializableEncodingAttrInterface>(attr);
+    if (!serializableAttr) {
+      return false;
+    }
+    return serializableAttr.isIdentityLayout();
+  });
+}
+
 Attribute EncodingAttr::cloneWithLayouts(ArrayRef<Attribute> layouts) const {
   MLIRContext *ctx = getContext();
   return get(ctx, getOperandIndex(), getOpType(), getElementTypes(),
@@ -338,6 +358,11 @@ PadEncodingLayoutAttr PadEncodingLayoutAttr::getIdentityAttr(MLIRContext *ctx,
   return get(ctx, zeros);
 }
 
+bool PadEncodingLayoutAttr::isIdentityLayout() const {
+  ArrayRef<int32_t> padding = getPadding().asArrayRef();
+  return llvm::all_of(padding, [](int32_t val) { return val == 0; });
+}
+
 Value PadEncodingLayoutAttr::calculateStorageSizeInBytes(
     Location loc, OpBuilder &builder, RankedTensorType type,
     ValueRange dynamicDims) const {

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

@@ -56,6 +56,7 @@ def EncodingAttr :
     IREEEncoding_Attr<"Encoding", [
       DeclareAttrInterfaceMethods<IREEEncoding_SerializableEncodingAttrInterface, [
         "isSerialized",
+        "isIdentityLayout",
         "cloneWithLayouts",
         "calculateStorageSizeInBytes",
       ]>
@@ -149,8 +150,10 @@ def EncodingAttr :
 //===---------------------------------------------------------------------===//
 
 def PadEncodingLayoutAttr : IREEEncoding_Attr<"PadEncodingLayout", [
-      DeclareAttrInterfaceMethods<IREEEncoding_SerializableEncodingAttrInterface,
-        ["calculateStorageSizeInBytes"]>
+      DeclareAttrInterfaceMethods<IREEEncoding_SerializableEncodingAttrInterface, [
+         "calculateStorageSizeInBytes",
+         "isIdentityLayout",
+        ]>
     ]> {
   let mnemonic = "pad_encoding_layout";
   let assemblyFormat = "`<` $padding `>`";

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

@@ -121,6 +121,19 @@ def IREEEncoding_SerializableEncodingAttrInterface :
         return true;
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Returns true if the encoding is known as an identity layout.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isIdentityLayout",
+      /*args=*/(ins
+      ),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return false;
+      }]
+    >,
     InterfaceMethod<
       /*desc=*/[{
         Creates an encoding with a new layout list. It is valid to drop any

compiler/src/iree/compiler/Dialect/Encoding/IR/unittests/BUILD.bazel

@@ -0,0 +1,25 @@
+# Copyright 2025 The IREE Authors
+#
+# Licensed 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
+
+load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_cc_test")
+
+package(
+    features = ["layering_check"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+iree_compiler_cc_test(
+    name = "EncodingAttrTest",
+    testonly = True,
+    srcs = ["EncodingAttrTest.cpp"],
+    deps = [
+        "//compiler/src/iree/compiler/Dialect/Encoding/IR",
+        "//compiler/src/iree/testing:gtest_main",
+        "@com_google_googletest//:gtest",
+        "@llvm-project//mlir:DialectUtils",
+        "@llvm-project//mlir:IR",
+    ],
+)

compiler/src/iree/compiler/Dialect/Encoding/IR/unittests/CMakeLists.txt

@@ -0,0 +1,26 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from           #
+# compiler/src/iree/compiler/Dialect/Encoding/IR/unittests/BUILD.bazel         #
+#                                                                              #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary   #
+# CMake-only content.                                                          #
+#                                                                              #
+# To disable autogeneration for this file entirely, delete this header.        #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_test(
+  NAME
+    EncodingAttrTest
+  SRCS
+    "EncodingAttrTest.cpp"
+  DEPS
+    MLIRIR
+    gmock
+    gtest
+    iree::compiler::Dialect::Encoding::IR
+    iree::testing::gtest_main
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###

compiler/src/iree/compiler/Dialect/Encoding/IR/unittests/EncodingAttrTest.cpp

@@ -0,0 +1,59 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed 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
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include "iree/compiler/Dialect/Encoding/IR/EncodingDialect.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
+
+namespace mlir::iree_compiler::IREE::Encoding {
+namespace {
+
+class EncodingAttrsTest : public ::testing::Test {
+protected:
+  EncodingAttrsTest() {
+    reg.insert<IREEEncodingDialect>();
+    ctx.appendDialectRegistry(reg);
+    ctx.loadAllAvailableDialects();
+  }
+  ~EncodingAttrsTest() override {}
+
+  MLIRContext *getContext() { return &ctx; }
+
+private:
+  MLIRContext ctx;
+  DialectRegistry reg;
+};
+
+TEST_F(EncodingAttrsTest, EncodingAttr) {
+  MLIRContext *ctx = getContext();
+  Builder builder(ctx);
+  SmallVector<Type> elemTypes(3, builder.getF32Type());
+  auto attr = cast<SerializableEncodingAttrInterface>(EncodingAttr::get(
+      ctx, /*operandIndex=*/0, EncodingOpType::matmul, elemTypes));
+  EXPECT_FALSE(attr.isIdentityLayout());
+
+  attr = cast<SerializableEncodingAttrInterface>(attr.cloneWithLayouts(
+      PadEncodingLayoutAttr::getIdentityAttr(ctx, /*rank=*/2)));
+  EXPECT_TRUE(attr.isIdentityLayout());
+}
+
+TEST_F(EncodingAttrsTest, PadEncodingLayoutAttr) {
+  MLIRContext *ctx = getContext();
+  auto zeroPaddingAttr =
+      PadEncodingLayoutAttr::getIdentityAttr(ctx, /*rank=*/2);
+  EXPECT_TRUE(cast<SerializableEncodingAttrInterface>(zeroPaddingAttr)
+                  .isIdentityLayout());
+
+  SmallVector<int32_t> paddings = {4, 2};
+  auto nonZeroPaddingAttr = PadEncodingLayoutAttr::get(ctx, paddings);
+  EXPECT_FALSE(cast<SerializableEncodingAttrInterface>(nonZeroPaddingAttr)
+                   .isIdentityLayout());
+}
+
+} // namespace
+} // namespace mlir::iree_compiler::IREE::Encoding