From 45d00a1700515bcf322e327c37323d059f4aca05 Mon Sep 17 00:00:00 2001
From: lorenzo chelini <l.chelini@icloud.com>
Date: Mon, 27 Nov 2023 09:28:45 -0500
Subject: [PATCH] Improve readability for `isInVnniLayout` (NFC) (#808)

---
 lib/TPP/Transforms/Utils/VNNIUtils.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/TPP/Transforms/Utils/VNNIUtils.cpp b/lib/TPP/Transforms/Utils/VNNIUtils.cpp
index 634196349..434a443ad 100644
--- a/lib/TPP/Transforms/Utils/VNNIUtils.cpp
+++ b/lib/TPP/Transforms/Utils/VNNIUtils.cpp
@@ -27,8 +27,7 @@ std::optional<int64_t> getVnniBlockingFactor(Type type) {
 bool isInVnniLayout(MemRefType memref) {
   if (memref.getRank() < 3 || !memref.getElementType().isBF16())
     return false;
-  return memref.getShape()[memref.getRank() - 1] ==
-         vnni::utils::getVnniBlockingFactor(memref);
+  return memref.getShape().back() == vnni::utils::getVnniBlockingFactor(memref);
 }
 
 } // namespace utils