[mlir][memref] Fix layout map computation in inferRankReducedResultType

[springerm]

This fixes https://github.com/google/iree/issues/9556.

We have no tests for inferRankReducedResultType. Should I add a C++ unit test for this?

[jpienaar]

Thanks for quick fix.

In D128160#3595423, @springerm wrote:

This fixes https://github.com/google/iree/issues/9556.

We have no tests for inferRankReducedResultType. Should I add a C++ unit test for this?

Aren't these used during verification? E.g., could we test simply by having a couple of invalid ops that get flagged? If not, my preference would be a tester pass (that just runs infer and prints out) that way we can easily add tests and different configs.

[springerm]

Aren't these used during verification? E.g., could we test simply by having a couple of invalid ops that get flagged? If not, my preference would be a tester pass (that just runs infer and prints out) that way we can easily add tests and different configs.

There are two inferRankReducedResultType overloads. The verifier executes only one of them. The other one is only used when building a new op with the C++ API.

[springerm]

add test cases

[springerm]

fix cmake

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

if (rankDiff > 0) {
     auto shape = inferredType.getShape();
     llvm::SmallBitVector dimsToProject =
         getPositionsOfShapeOne(rankDiff, shape);
     SmallVector<int64_t> projectedShape;
     for (unsigned pos = 0, e = shape.size(); pos < e; ++pos)
       if (!dimsToProject.test(pos))
         projectedShape.push_back(shape[pos]);
 
-    AffineMap map = inferredType.getLayout().getAffineMap();
-    if (!map.isIdentity())
-      map = getProjectedMap(map, dimsToProject);
+    AffineMap map =
+        getProjectedMap(inferredType.getLayout().getAffineMap(), dimsToProject);
     inferredType =
         MemRefType::get(projectedShape, inferredType.getElementType(), map,
                         inferredType.getMemorySpace());
   }
   return inferredType;
 }
 
 Type SubViewOp::inferRankReducedResultType(unsigned resultRank,

mlir/unittests/Dialect/CMakeLists.txt

 add_mlir_unittest(MLIRDialectTests
   BroadcastShapeTest.cpp
 )
 target_link_libraries(MLIRDialectTests
   PRIVATE
   MLIRIR
   MLIRDialect)
 
 add_subdirectory(Affine)
+add_subdirectory(MemRef)
+
 add_subdirectory(Quant)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(Utils)

mlir/unittests/Dialect/MemRef/CMakeLists.txt

+add_mlir_unittest(MLIRMemRefTests
+  InferShapeTest.cpp
+)
+target_link_libraries(MLIRMemRefTests
+  PRIVATE
+  MLIRMemRefDialect
+  )

mlir/unittests/Dialect/MemRef/InferShapeTest.cpp

+//===- InferShapeTest.cpp - unit tests for shape inference ----------------===//
+//
+// 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/IR/MemRef.h"
+#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+using namespace mlir::memref;
+
+// Source memref has identity layout.
+TEST(InferShapeTest, inferRankReducedShapeIdentity) {
+  MLIRContext ctx;
+  OpBuilder b(&ctx);
+  auto sourceMemref = MemRefType::get({10, 5}, b.getIndexType());
+  auto reducedType = SubViewOp::inferRankReducedResultType(
+      /*resultRank=*/1, sourceMemref, {2, 3}, {1, 2}, {1, 1});
+  AffineExpr dim0;
+  bindDims(&ctx, dim0);
+  auto expectedType =
+      MemRefType::get({2}, b.getIndexType(), AffineMap::get(1, 0, dim0 + 13));
+  EXPECT_EQ(reducedType, expectedType);
+}
+
+// Source memref has non-identity layout.
+TEST(InferShapeTest, inferRankReducedShapeNonIdentity) {
+  MLIRContext ctx;
+  OpBuilder b(&ctx);
+  AffineExpr dim0, dim1;
+  bindDims(&ctx, dim0, dim1);
+  auto sourceMemref = MemRefType::get({10, 5}, b.getIndexType(),
+                                      AffineMap::get(2, 0, 1000 * dim0 + dim1));
+  auto reducedType = SubViewOp::inferRankReducedResultType(
+      /*resultRank=*/1, sourceMemref, {2, 3}, {1, 2}, {1, 1});
+  auto expectedType =
+      MemRefType::get({2}, b.getIndexType(), AffineMap::get(1, 0, dim0 + 2003));
+  EXPECT_EQ(reducedType, expectedType);
+}
+
+TEST(InferShapeTest, inferRankReducedShapeToScalar) {
+  MLIRContext ctx;
+  OpBuilder b(&ctx);
+  AffineExpr dim0, dim1;
+  bindDims(&ctx, dim0, dim1);
+  auto sourceMemref = MemRefType::get({10, 5}, b.getIndexType(),
+                                      AffineMap::get(2, 0, 1000 * dim0 + dim1));
+  auto reducedType = SubViewOp::inferRankReducedResultType(
+      /*resultRank=*/0, sourceMemref, {2, 3}, {1, 1}, {1, 1});
+  auto expectedType =
+      MemRefType::get({}, b.getIndexType(),
+                      AffineMap::get(0, 0, b.getAffineConstantExpr(2003)));
+  EXPECT_EQ(reducedType, expectedType);
+}

utils/bazel/llvm-project-overlay/mlir/unittests/BUILD.bazel

cc_test(
     deps = [
         "//llvm:TestingSupport",
         "//llvm:gtest_main",
         "//mlir:Dialect",
     ],
 )
 
 cc_test(
+    name = "memref_tests",
+    size = "small",
+    srcs = glob([
+        "Dialect/MemRef/*.cpp",
+        "Dialect/MemRef/*.h",
+    ]),
+    deps = [
+        "//llvm:TestingSupport",
+        "//llvm:gtest_main",
+        "//mlir:MemRefDialect",
+    ],
+)
+
+cc_test(
     name = "quantops_tests",
     size = "small",
     srcs = glob([
         "Dialect/Quant/*.cpp",
         "Dialect/Quant/*.h",
     ]),
     deps = [
         "//llvm:TestingSupport",