[Mlir-commits] [mlir] 5c6c424 - [mlir/memref] handle rank-0 contiguous check (#198541)

[llvmlistbot at llvm.org]

Author: Alan Li
Date: 2026-05-20T09:58:35-04:00
New Revision: 5c6c424a50b840a39a6410a490af668a26d3a97a

URL: https://github.com/llvm/llvm-project/commit/5c6c424a50b840a39a6410a490af668a26d3a97a
DIFF: https://github.com/llvm/llvm-project/commit/5c6c424a50b840a39a6410a490af668a26d3a97a.diff

LOG: [mlir/memref] handle rank-0 contiguous check (#198541)

Added: 
    mlir/unittests/Dialect/MemRef/MemRefUtilsTest.cpp

Modified: 
    mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
    mlir/unittests/Dialect/MemRef/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
index bc5a327df9140..2aea597dd5b90 100644
--- a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
+++ b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
@@ -24,6 +24,10 @@ bool isStaticShapeAndContiguousRowMajor(MemRefType type) {
   if (!type.hasStaticShape())
     return false;
 
+  int64_t rank = type.getRank();
+  if (rank == 0)
+    return true;
+
   SmallVector<int64_t> strides;
   int64_t offset;
   if (failed(type.getStridesAndOffset(strides, offset)))
@@ -32,7 +36,7 @@ bool isStaticShapeAndContiguousRowMajor(MemRefType type) {
   // MemRef is contiguous if outer dimensions are size-1 and inner
   // dimensions have unit strides.
   int64_t runningStride = 1;
-  int64_t curDim = strides.size() - 1;
+  int64_t curDim = rank - 1;
   // Finds all inner dimensions with unit strides.
   while (curDim >= 0 && strides[curDim] == runningStride) {
     runningStride *= type.getDimSize(curDim);
@@ -123,7 +127,8 @@ getLinearizedMemRefOffsetAndSize(OpBuilder &builder, Location loc, int srcBits,
     strides.back() = builder.getIndexAttr(1);
     AffineExpr s0, s1;
     bindSymbols(builder.getContext(), s0, s1);
-    for (int index = sizes.size() - 1; index > 0; --index) {
+    for (int64_t index = static_cast<int64_t>(sizes.size()) - 1; index > 0;
+         --index) {
       strides[index - 1] = affine::makeComposedFoldedAffineApply(
           builder, loc, s0 * s1,
           ArrayRef<OpFoldResult>{strides[index], sizes[index]});
@@ -185,7 +190,7 @@ computeSuffixProductIRBlockImpl(Location loc, OpBuilder &builder,
   AffineExpr s0, s1;
   bindSymbols(builder.getContext(), s0, s1);
 
-  for (int64_t r = strides.size() - 1; r > 0; --r) {
+  for (int64_t r = static_cast<int64_t>(strides.size()) - 1; r > 0; --r) {
     strides[r - 1] = affine::makeComposedFoldedAffineApply(
         builder, loc, s0 * s1, {strides[r], sizes[r]});
   }

diff  --git a/mlir/unittests/Dialect/MemRef/CMakeLists.txt b/mlir/unittests/Dialect/MemRef/CMakeLists.txt
index dede3ba0a885c..face68c587d39 100644
--- a/mlir/unittests/Dialect/MemRef/CMakeLists.txt
+++ b/mlir/unittests/Dialect/MemRef/CMakeLists.txt
@@ -1,7 +1,9 @@
 add_mlir_unittest(MLIRMemRefTests
   InferShapeTest.cpp
+  MemRefUtilsTest.cpp
 )
 mlir_target_link_libraries(MLIRMemRefTests
   PRIVATE
   MLIRMemRefDialect
+  MLIRMemRefUtils
   )

diff  --git a/mlir/unittests/Dialect/MemRef/MemRefUtilsTest.cpp b/mlir/unittests/Dialect/MemRef/MemRefUtilsTest.cpp
new file mode 100644
index 0000000000000..92f447be969c2
--- /dev/null
+++ b/mlir/unittests/Dialect/MemRef/MemRefUtilsTest.cpp
@@ -0,0 +1,33 @@
+//===- MemRefUtilsTest.cpp - MemRef utils unit tests ---------------------===//
+//
+// Part of the LLVM Project, 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 "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+TEST(MemRefUtilsTest, IsStaticShapeAndContiguousRowMajor) {
+  MLIRContext ctx;
+  OpBuilder builder(&ctx);
+
+  // Rank-0 memrefs are scalar values and should return before stride analysis.
+  EXPECT_TRUE(memref::isStaticShapeAndContiguousRowMajor(
+      MemRefType::get({}, builder.getI32Type())));
+}
+
+TEST(MemRefUtilsTest, ComputeSuffixProductIRBlockEmptySizes) {
+  MLIRContext ctx;
+  OpBuilder builder(&ctx);
+
+  // Empty size lists should not underflow the reverse loop's initial index.
+  EXPECT_TRUE(
+      memref::computeSuffixProductIRBlock(builder.getUnknownLoc(), builder, {})
+          .empty());
+}