8272b6bd61
Fixes an issue where the `SimpleAffineExprFlattener` would simplify `lhs % rhs` to just `-(lhs floordiv rhs)` instead of `lhs - (lhs floordiv rhs)` if `lhs` happened to be equal to `lhs floordiv rhs`. The reported failure case was `(d0, d1) -> (((d1 - (d1 + 2)) floordiv 8) % 8)` from https://github.com/llvm/llvm-project/issues/114654. Note that many paths that simplify AffineMaps (e.g. the AffineApplyOp folder and canonicalization) would not observe this bug because of of slightly different paths taken by the code. Slightly different grouping of the terms could also result in avoiding the bug. Resolves https://github.com/llvm/llvm-project/issues/114654.
150 lines
5.1 KiB
C++
150 lines
5.1 KiB
C++
//===- AffineExprTest.cpp - unit tests for affine expression API ----------===//
|
|
//
|
|
// 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 <cstdint>
|
|
#include <limits>
|
|
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "gtest/gtest.h"
|
|
|
|
using namespace mlir;
|
|
|
|
static std::string toString(AffineExpr expr) {
|
|
std::string s;
|
|
llvm::raw_string_ostream ss(s);
|
|
ss << expr;
|
|
return s;
|
|
}
|
|
|
|
// Test creating AffineExprs using the overloaded binary operators.
|
|
TEST(AffineExprTest, constructFromBinaryOperators) {
|
|
MLIRContext ctx;
|
|
OpBuilder b(&ctx);
|
|
|
|
auto d0 = b.getAffineDimExpr(0);
|
|
auto d1 = b.getAffineDimExpr(1);
|
|
|
|
auto sum = d0 + d1;
|
|
auto difference = d0 - d1;
|
|
auto product = d0 * d1;
|
|
auto remainder = d0 % d1;
|
|
|
|
ASSERT_EQ(sum.getKind(), AffineExprKind::Add);
|
|
ASSERT_EQ(difference.getKind(), AffineExprKind::Add);
|
|
ASSERT_EQ(product.getKind(), AffineExprKind::Mul);
|
|
ASSERT_EQ(remainder.getKind(), AffineExprKind::Mod);
|
|
}
|
|
|
|
TEST(AffineExprTest, constantFolding) {
|
|
MLIRContext ctx;
|
|
OpBuilder b(&ctx);
|
|
auto cn1 = b.getAffineConstantExpr(-1);
|
|
auto c0 = b.getAffineConstantExpr(0);
|
|
auto c1 = b.getAffineConstantExpr(1);
|
|
auto c2 = b.getAffineConstantExpr(2);
|
|
auto c3 = b.getAffineConstantExpr(3);
|
|
auto c6 = b.getAffineConstantExpr(6);
|
|
auto cmax = b.getAffineConstantExpr(std::numeric_limits<int64_t>::max());
|
|
auto cmin = b.getAffineConstantExpr(std::numeric_limits<int64_t>::min());
|
|
|
|
ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::Add, c1, c2), c3);
|
|
ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::Mul, c2, c3), c6);
|
|
ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::FloorDiv, c3, c2), c1);
|
|
ASSERT_EQ(getAffineBinaryOpExpr(AffineExprKind::CeilDiv, c3, c2), c2);
|
|
|
|
// Test division by zero:
|
|
auto c3ceildivc0 = getAffineBinaryOpExpr(AffineExprKind::CeilDiv, c3, c0);
|
|
ASSERT_EQ(c3ceildivc0.getKind(), AffineExprKind::CeilDiv);
|
|
|
|
auto c3floordivc0 = getAffineBinaryOpExpr(AffineExprKind::FloorDiv, c3, c0);
|
|
ASSERT_EQ(c3floordivc0.getKind(), AffineExprKind::FloorDiv);
|
|
|
|
auto c3modc0 = getAffineBinaryOpExpr(AffineExprKind::Mod, c3, c0);
|
|
ASSERT_EQ(c3modc0.getKind(), AffineExprKind::Mod);
|
|
|
|
// Test overflow:
|
|
auto cmaxplusc1 = getAffineBinaryOpExpr(AffineExprKind::Add, cmax, c1);
|
|
ASSERT_EQ(cmaxplusc1.getKind(), AffineExprKind::Add);
|
|
|
|
auto cmaxtimesc2 = getAffineBinaryOpExpr(AffineExprKind::Mul, cmax, c2);
|
|
ASSERT_EQ(cmaxtimesc2.getKind(), AffineExprKind::Mul);
|
|
|
|
auto cminceildivcn1 =
|
|
getAffineBinaryOpExpr(AffineExprKind::CeilDiv, cmin, cn1);
|
|
ASSERT_EQ(cminceildivcn1.getKind(), AffineExprKind::CeilDiv);
|
|
|
|
auto cminfloordivcn1 =
|
|
getAffineBinaryOpExpr(AffineExprKind::FloorDiv, cmin, cn1);
|
|
ASSERT_EQ(cminfloordivcn1.getKind(), AffineExprKind::FloorDiv);
|
|
}
|
|
|
|
TEST(AffineExprTest, divisionSimplification) {
|
|
MLIRContext ctx;
|
|
OpBuilder b(&ctx);
|
|
auto cn6 = b.getAffineConstantExpr(-6);
|
|
auto c6 = b.getAffineConstantExpr(6);
|
|
auto d0 = b.getAffineDimExpr(0);
|
|
auto d1 = b.getAffineDimExpr(1);
|
|
|
|
ASSERT_EQ(c6.floorDiv(-1), cn6);
|
|
ASSERT_EQ((d0 * 6).floorDiv(2), d0 * 3);
|
|
ASSERT_EQ((d0 * 6).floorDiv(4).getKind(), AffineExprKind::FloorDiv);
|
|
ASSERT_EQ((d0 * 6).floorDiv(-2), d0 * -3);
|
|
ASSERT_EQ((d0 * 6 + d1).floorDiv(2), d0 * 3 + d1.floorDiv(2));
|
|
ASSERT_EQ((d0 * 6 + d1).floorDiv(-2), d0 * -3 + d1.floorDiv(-2));
|
|
ASSERT_EQ((d0 * 6 + d1).floorDiv(4).getKind(), AffineExprKind::FloorDiv);
|
|
|
|
ASSERT_EQ(c6.ceilDiv(-1), cn6);
|
|
ASSERT_EQ((d0 * 6).ceilDiv(2), d0 * 3);
|
|
ASSERT_EQ((d0 * 6).ceilDiv(4).getKind(), AffineExprKind::CeilDiv);
|
|
ASSERT_EQ((d0 * 6).ceilDiv(-2), d0 * -3);
|
|
}
|
|
|
|
TEST(AffineExprTest, modSimplificationRegression) {
|
|
MLIRContext ctx;
|
|
OpBuilder b(&ctx);
|
|
auto d0 = b.getAffineDimExpr(0);
|
|
auto sum = d0 + d0.floorDiv(3).floorDiv(-3);
|
|
ASSERT_EQ(sum.getKind(), AffineExprKind::Add);
|
|
}
|
|
|
|
TEST(AffineExprTest, divisorOfNegativeFloorDiv) {
|
|
MLIRContext ctx;
|
|
OpBuilder b(&ctx);
|
|
ASSERT_EQ(b.getAffineDimExpr(0).floorDiv(-1).getLargestKnownDivisor(), 1);
|
|
}
|
|
|
|
TEST(AffineExprTest, d0PlusD0FloorDivNeg2) {
|
|
// Regression test for a bug where this was rewritten to d0 mod -2. We do not
|
|
// support a negative RHS for mod in LowerAffinePass.
|
|
MLIRContext ctx;
|
|
OpBuilder b(&ctx);
|
|
auto d0 = b.getAffineDimExpr(0);
|
|
auto sum = d0 + d0.floorDiv(-2) * 2;
|
|
ASSERT_EQ(toString(sum), "d0 + (d0 floordiv -2) * 2");
|
|
}
|
|
|
|
TEST(AffineExprTest, simpleAffineExprFlattenerRegression) {
|
|
|
|
// Regression test for a bug where mod simplification was not handled
|
|
// properly when `lhs % rhs` was happened to have the property that `lhs
|
|
// floordiv rhs = lhs`.
|
|
MLIRContext ctx;
|
|
OpBuilder b(&ctx);
|
|
|
|
auto d0 = b.getAffineDimExpr(0);
|
|
|
|
// Manually replace variables by constants to avoid constant folding.
|
|
AffineExpr expr = (d0 - (d0 + 2)).floorDiv(8) % 8;
|
|
AffineExpr result = mlir::simplifyAffineExpr(expr, 1, 0);
|
|
|
|
ASSERT_TRUE(isa<AffineConstantExpr>(result));
|
|
ASSERT_EQ(cast<AffineConstantExpr>(result).getValue(), 7);
|
|
}
|