From 1a6f45475b3b8458071ec93340a60a25ff320b7c Mon Sep 17 00:00:00 2001
From: Alan Li <me@alanli.org>
Date: Fri, 17 Jan 2025 04:17:30 +0000
Subject: [PATCH] Another commit

---
 .../Codegen/Common/test/type_propagation_packing.mlir      | 2 +-
 .../iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp    | 7 ++-----
 .../src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h  | 3 ++-
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/compiler/src/iree/compiler/Codegen/Common/test/type_propagation_packing.mlir b/compiler/src/iree/compiler/Codegen/Common/test/type_propagation_packing.mlir
index 8595bba90703..74b0284f665c 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/type_propagation_packing.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/type_propagation_packing.mlir
@@ -73,4 +73,4 @@ func.func @generic_op_i1_packed() {
 //   CHECK-DAG:       %[[ADD:.+]] = arith.addi %[[ARG0]], %[[ARG0]] : i1
 //       CHECK:       linalg.yield %[[ADD]]
 //       CHECK:   -> tensor<?xi1, #iree_encoding.packed_storage>
-//       CHECK:   flow.dispatch.tensor.store %[[GENERIC]], %[[OUT]]
\ No newline at end of file
+//       CHECK:   flow.dispatch.tensor.store %[[GENERIC]], %[[OUT]]
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
index 7185f2e85c00..84ea520e2bb1 100644
--- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
+++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
@@ -291,11 +291,8 @@ EncodingAttr getEncodingAttr(RankedTensorType type) {
 
 bool hasPackedStorageAttr(Type type) {
   if (auto tensorType = dyn_cast<RankedTensorType>(type)) {
-    auto encoding = tensorType.getEncoding();
-    if (!encoding) {
-      return false;
-    }
-    return dyn_cast_or_null<PackedStorageAttr>(encoding) != nullptr;
+    return dyn_cast_or_null<PackedStorageAttr>(tensorType.getEncoding()) !=
+           nullptr;
   }
   return false;
 }
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h
index 3f357df16a09..74aeb29a8686 100644
--- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h
+++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h
@@ -38,7 +38,8 @@ namespace mlir::iree_compiler::IREE::Encoding {
 /// Otherwise, returns null.
 EncodingAttr getEncodingAttr(RankedTensorType type);
 
-/// Returns true if the type contains packed_storage attribute.
+/// Returns true if the type is a RankedTensorType type and it contains
+/// packed_storage attribute.
 bool hasPackedStorageAttr(Type type);
 
 /// Returns the ContractionDimensions for the encoding user_indexing_maps.