maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

More affine expr simplifications for floordiv and mod 

Add one more simplification for floordiv and mod affine expressions.
Examples:
 (2*d0 + 1) floordiv 2 is simplified to d0
 (8*d0 + 4*d1 + d2) floordiv 4 simplified to 4*d0 + d1 + d2 floordiv 4.
 etc.

 Similarly, (4*d1 + 1) mod 2 is simplified to 1,
            (2*d0 + 8*d1) mod 8 simplified to 2*d0 mod 8.

Change getLargestKnownDivisor to return int64_t to be consistent and
to avoid casting at call sites (since the return value is used in expressions
of int64_t/index type).

Signed-off-by: Uday Bondhugula <uday@polymagelabs.com>

Closes tensorflow/mlir#202

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/mlir/pull/202 from bondhugula:affine b13fcb2f1c00a39ca5434613a02408e085a80e77
PiperOrigin-RevId: 284866710
This commit is contained in:
Uday Bondhugula 2019-12-10 15:49:07 -08:00 committed by A. Unique TensorFlower
parent d1213ae51d
commit 36a415bcc5
5 changed files with 48 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mlir/include/mlir/IR/AffineExpr.h
View File

@ -114,8 +114,9 @@ public:
/// floordiv, ceildiv, and mod is only allowed w.r.t constants.
bool isPureAffine() const;
/// Returns the greatest known integral divisor of this affine expression.
uint64_t getLargestKnownDivisor() const;
/// Returns the greatest known integral divisor of this affine expression. The
/// result is always positive.
int64_t getLargestKnownDivisor() const;
/// Return true if the affine expression is a multiple of 'factor'.
bool isMultipleOf(int64_t factor) const;

43
mlir/lib/IR/AffineExpr.cpp
View File

@ -160,7 +160,7 @@ bool AffineExpr::isPureAffine() const {
}
// Returns the greatest known integral divisor of this affine expression.
uint64_t AffineExpr::getLargestKnownDivisor() const {
int64_t AffineExpr::getLargestKnownDivisor() const {
AffineBinaryOpExpr binExpr(nullptr);
switch (getKind()) {
case AffineExprKind::SymbolId:
@ -444,6 +444,7 @@ static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs) {
auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
// mlir floordiv by zero or negative numbers is undefined and preserved as is.
if (!rhsConst || rhsConst.getValue() < 1)
return nullptr;
@ -453,18 +454,32 @@ static AffineExpr simplifyFloorDiv(AffineExpr lhs, AffineExpr rhs) {
// Fold floordiv of a multiply with a constant that is a multiple of the
// divisor. Eg: (i * 128) floordiv 64 = i * 2.
if (rhsConst.getValue() == 1)
if (rhsConst == 1)
return lhs;
// Simplify (expr * const) floordiv divConst when expr is known to be a
// multiple of divConst.
auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
if (lBin && lBin.getKind() == AffineExprKind::Mul) {
if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) {
// rhsConst is known to be positive if a constant.
// rhsConst is known to be a positive constant.
if (lrhs.getValue() % rhsConst.getValue() == 0)
return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue());
}
}
// Simplify (expr1 + expr2) floordiv divConst when either expr1 or expr2 is
// known to be a multiple of divConst.
if (lBin && lBin.getKind() == AffineExprKind::Add) {
int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor();
int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor();
// rhsConst is known to be a positive constant.
if (llhsDiv % rhsConst.getValue() == 0 ||
lrhsDiv % rhsConst.getValue() == 0)
return lBin.getLHS().floorDiv(rhsConst.getValue()) +
lBin.getRHS().floorDiv(rhsConst.getValue());
}
return nullptr;
}
@ -497,10 +512,12 @@ static AffineExpr simplifyCeilDiv(AffineExpr lhs, AffineExpr rhs) {
if (rhsConst.getValue() == 1)
return lhs;
// Simplify (expr * const) ceildiv divConst when const is known to be a
// multiple of divConst.
auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
if (lBin && lBin.getKind() == AffineExprKind::Mul) {
if (auto lrhs = lBin.getRHS().dyn_cast<AffineConstantExpr>()) {
// rhsConst is known to be positive if a constant.
// rhsConst is known to be a positive constant.
if (lrhs.getValue() % rhsConst.getValue() == 0)
return lBin.getLHS() * (lrhs.getValue() / rhsConst.getValue());
}
@ -526,6 +543,7 @@ static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs) {
auto lhsConst = lhs.dyn_cast<AffineConstantExpr>();
auto rhsConst = rhs.dyn_cast<AffineConstantExpr>();
// mod w.r.t zero or negative numbers is undefined and preserved as is.
if (!rhsConst || rhsConst.getValue() < 1)
return nullptr;
@ -539,11 +557,20 @@ static AffineExpr simplifyMod(AffineExpr lhs, AffineExpr rhs) {
if (lhs.getLargestKnownDivisor() % rhsConst.getValue() == 0)
return getAffineConstantExpr(0, lhs.getContext());
// Simplify (expr1 + expr2) mod divConst when either expr1 or expr2 is
// known to be a multiple of divConst.
auto lBin = lhs.dyn_cast<AffineBinaryOpExpr>();
if (lBin && lBin.getKind() == AffineExprKind::Add) {
int64_t llhsDiv = lBin.getLHS().getLargestKnownDivisor();
int64_t lrhsDiv = lBin.getRHS().getLargestKnownDivisor();
// rhsConst is known to be a positive constant.
if (llhsDiv % rhsConst.getValue() == 0)
return lBin.getRHS() % rhsConst.getValue();
if (lrhsDiv % rhsConst.getValue() == 0)
return lBin.getLHS() % rhsConst.getValue();
}
return nullptr;
// TODO(bondhugula): In general, this can be simplified more by using the GCD
// test, or in general using quantifier elimination (add two new variables q
// and r, and eliminate all variables from the linear system other than r. All
// of this can be done through mlir/Analysis/'s FlatAffineConstraints.
}
AffineExpr AffineExpr::operator%(uint64_t v) const {

8
mlir/test/IR/affine-map.mlir
View File

@ -156,7 +156,7 @@
#map48 = (i, j, k) -> (i * 64 floordiv 64, i * 512 floordiv 128, 4 * j mod 4, 4*j*4 mod 8)
// Simplifications for mod using known GCD's of the LHS expr.
// CHECK: #map{{[0-9]+}} = (d0, d1)[s0] -> (0, 0, 0, (d0 * 4 + 3) mod 2)
// CHECK: #map{{[0-9]+}} = (d0, d1)[s0] -> (0, 0, 0, 1)
#map49 = (i, j)[s0] -> ( (i * 4 + 8) mod 4, 32 * j * s0 * 8 mod 256, (4*i + (j * (s0 * 2))) mod 2, (4*i + 3) mod 2)
// Floordiv, ceildiv divide by one.
@ -180,6 +180,9 @@
// CHECK: #map{{[0-9]+}} = () -> ()
#map55 = () -> ()
// CHECK: #map{{[0-9]+}} = (d0, d1) -> (d0, d0 * 2 + d1 * 4 + 2, 1, 2, (d0 * 4) mod 8)
#map56 = (d0, d1) -> ((4*d0 + 2) floordiv 4, (4*d0 + 8*d1 + 5) floordiv 2, (2*d0 + 4*d1 + 3) mod 2, (3*d0 - 4) mod 3, (4*d0 + 8*d1) mod 8)
// Single identity maps are removed.
// CHECK: func @f0(memref<2x4xi8, 1>)
func @f0(memref<2x4xi8, #map0, 1>)
@ -355,3 +358,6 @@ func @f54(memref<10xi32, #map54>)
// CHECK: "foo.op"() {map = #map{{[0-9]+}}} : () -> ()
"foo.op"() {map = #map55} : () -> ()
// CHECK: func @f56(memref<1x1xi8, #map{{[0-9]+}}>)
func @f56(memref<1x1xi8, #map56>)

4
mlir/test/Transforms/Vectorize/compose_maps.mlir
View File

@ -78,7 +78,7 @@ func @simple5c() {
}
func @simple5d() {
// CHECK: Composed map: (d0) -> ((d0 * 4 + 24) floordiv 3)
// CHECK: Composed map: (d0) -> ((d0 * 4) floordiv 3 + 8)
"test_affine_map"() { affine_map = (d0) -> (d0 - 1) } : () -> ()
"test_affine_map"() { affine_map = (d0) -> (d0 + 7) } : () -> ()
"test_affine_map"() { affine_map = (d0) -> (d0 * 4) } : () -> ()
@ -128,4 +128,4 @@ func @multi_symbols() {
"test_affine_map"() { affine_map = (d0)[s0] -> (d0 + s0, d0 - s0) } : () -> ()
"test_affine_map"() { affine_map = (d0, d1)[s0, s1] -> (d0 + 1 + s1, d1 - 1 - s0) } : () -> ()
return
}
}

2
mlir/test/Transforms/unroll.mlir
View File

@ -21,7 +21,7 @@
// UNROLL-BY-4-DAG: [[MAP5:#map[0-9]+]] = (d0)[s0] -> (d0 + s0 + 1)
// UNROLL-BY-4-DAG: [[MAP6:#map[0-9]+]] = (d0, d1) -> (d0 * 16 + d1)
// UNROLL-BY-4-DAG: [[MAP11:#map[0-9]+]] = (d0) -> (d0)
// UNROLL-BY-4-DAG: [[MAP_TRIP_COUNT_MULTIPLE_FOUR:#map[0-9]+]] = ()[s0, s1, s2] -> (s0 + ((-s0 + s1) floordiv 4) * 4, s0 + ((-s0 + s2) floordiv 4) * 4, s0 + ((-s0 + 1024) floordiv 4) * 4)
// UNROLL-BY-4-DAG: [[MAP_TRIP_COUNT_MULTIPLE_FOUR:#map[0-9]+]] = ()[s0, s1, s2] -> (s0 + ((-s0 + s1) floordiv 4) * 4, s0 + ((-s0 + s2) floordiv 4) * 4, s0 + ((-s0) floordiv 4) * 4 + 1024)
// UNROLL-FULL-LABEL: func @loop_nest_simplest() {
func @loop_nest_simplest() {