112 lines · cpp
1//===- LayoutTest.cpp - unit tests related to memref layout ---------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/MemRef/IR/MemRef.h"10#include "mlir/IR/AffineMap.h"11#include "mlir/IR/Builders.h"12#include "mlir/IR/BuiltinTypes.h"13#include "gtest/gtest.h"14 15using namespace mlir;16using namespace mlir::memref;17 18//19// Test the correctness of `memref::getNumContiguousTrailingDims`20//21TEST(MemRefLayout, numContigDim) {22 MLIRContext ctx;23 OpBuilder b(&ctx);24 25 const int64_t _ = ShapedType::kDynamic;26 const FloatType f32 = b.getF32Type();27 auto strided = [&ctx](ArrayRef<int64_t> s) {28 return StridedLayoutAttr::get(&ctx, 0, s);29 };30 31 // Special case for identity maps and no explicit `strided` attribute - the32 // memref is entirely contiguous even if the strides cannot be determined33 // statically.34 35 // memref<?x?x?xf32>36 auto m0 = MemRefType::get({_, _, _}, f32);37 EXPECT_EQ(m0.getNumContiguousTrailingDims(), 3);38 39 // Conservatively assume memref is sparse everywhere if cannot get the40 // strides.41 42 // memref<2x2x2xf32, (i,j,k)->(i,k,j)>43 auto m1 = MemRefType::get(44 {2, 2, 2}, f32,45 AffineMap::getPermutationMap(ArrayRef<int64_t>{0, 2, 1}, &ctx));46 EXPECT_EQ(m1.getNumContiguousTrailingDims(), 0);47 48 // A base cases of a fixed memref with the usual strides.49 50 // memref<2x2x2xf32, strided<[4, 2, 1]>>51 auto m3 = MemRefType::get({2, 2, 2}, f32, strided({4, 2, 1}));52 EXPECT_EQ(m3.getNumContiguousTrailingDims(), 3);53 54 // A fixed memref with a discontinuity in the rightmost dimension.55 56 // memref<2x2x2xf32, strided<[8, 4, 2]>>57 auto m4 = MemRefType::get({2, 2, 2}, f32, strided({8, 4, 2}));58 EXPECT_EQ(m4.getNumContiguousTrailingDims(), 0);59 60 // A fixed memref with a discontinuity in the "middle".61 62 // memref<2x2x2xf32, strided<[8, 2, 1]>>63 auto m5 = MemRefType::get({2, 2, 2}, f32, strided({8, 2, 1}));64 EXPECT_EQ(m5.getNumContiguousTrailingDims(), 2);65 66 // A dynamic memref where the dynamic dimension breaks continuity.67 68 // memref<2x?x2xf32, strided<[4, 2, 1]>>69 auto m6 = MemRefType::get({2, _, 2}, f32, strided({4, 2, 1}));70 EXPECT_EQ(m6.getNumContiguousTrailingDims(), 2);71 72 // A edge case of a dynamic memref where the dynamic dimension is the first73 // one.74 75 // memref<?x2x2xf32, strided<[4, 2, 1]>>76 auto m7 = MemRefType::get({2, _, 2}, f32, strided({4, 2, 1}));77 EXPECT_EQ(m7.getNumContiguousTrailingDims(), 2);78 79 // A memref with a unit dimension. Unit dimensions do not affect continuity,80 // even if the corresponding stride is dynamic.81 82 // memref<2x1x2xf32, strided<[2,?,1]>>83 auto m8 = MemRefType::get({2, 1, 2}, f32, strided({2, _, 1}));84 EXPECT_EQ(m8.getNumContiguousTrailingDims(), 3);85}86 87//88// Test the member function `memref::areTrailingDimsContiguous`89//90TEST(MemRefLayout, contigTrailingDim) {91 MLIRContext ctx;92 OpBuilder b(&ctx);93 94 const int64_t _ = ShapedType::kDynamic;95 const FloatType f32 = b.getF32Type();96 auto strided = [&ctx](ArrayRef<int64_t> s) {97 return StridedLayoutAttr::get(&ctx, 0, s);98 };99 100 // A not-entirely-continuous, not-entirely-discontinuous memref.101 // ensure `areTrailingDimsContiguous` returns `true` for the value102 // returned by `getNumContiguousTrailingDims` and `false` for the next bigger103 // number.104 105 // memref<2x?x2xf32, strided<[?,2,1]>>106 auto m = MemRefType::get({2, _, 2}, f32, strided({_, 2, 1}));107 int64_t n = m.getNumContiguousTrailingDims();108 EXPECT_TRUE(m.areTrailingDimsContiguous(n));109 ASSERT_TRUE(n + 1 <= m.getRank());110 EXPECT_FALSE(m.areTrailingDimsContiguous(n + 1));111}112