[Mlir-commits] [mlir] [mlir][memref] Fix offset update in emulating narrow type for strided memref (PR #68043)

[Kunwar Grover]

https://github.com/Groverkss created https://github.com/llvm/llvm-project/pull/68043

The offset when converting type in emulating narrow types did not account for the offset in strided memrefs. This patch fixes this.

>From 48513591a49085ee15c017ff953a72edc6b0894b Mon Sep 17 00:00:00 2001
From: Groverkss <groverkss at gmail.com>
Date: Tue, 3 Oct 2023 03:38:58 +0530
Subject: [PATCH] [mlir][memref] Fix offset update in emulating narrow type for
 strided memref

---
 .../MemRef/Transforms/EmulateNarrowType.cpp     | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
index 2a524ceb9db887b..2a9b27debaece3f 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
@@ -271,9 +271,22 @@ void memref::populateMemRefNarrowTypeEmulationConversions(
           return std::nullopt;
 
         StridedLayoutAttr layoutAttr;
+        // If the offset is 0, we do not need a strided layout as the stride is
+        // 1, so we only use the strided layout if the offset is not 0.
         if (offset != 0) {
-          layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
-                                              ArrayRef<int64_t>{1});
+          if (offset == ShapedType::kDynamic) {
+            layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
+                                                ArrayRef<int64_t>{1});
+          } else {
+            // Check if the number of bytes are a multiple of the loadStoreWidth
+            // and if so, divide it by the loadStoreWidth to get the offset.
+            if ((offset * width) % loadStoreWidth != 0)
+              return std::nullopt;
+            offset = (offset * width) / loadStoreWidth;
+
+            layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
+                                                ArrayRef<int64_t>{1});
+          }
         }
 
         return MemRefType::get(getLinearizedShape(ty, width, loadStoreWidth),