From: Uday Bondhugula Date: Tue, 10 Dec 2019 23:49:07 +0000 (-0800) Subject: More affine expr simplifications for floordiv and mod X-Git-Tag: llvmorg-11-init~1466^2~106 X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=36a415bcc543553891af6809c5256e6e2469357d;p=platform%2Fupstream%2Fllvm.git 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 Closes tensorflow/mlir#202 COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/mlir/pull/202 from bondhugula:affine b13fcb2f1c00a39ca5434613a02408e085a80e77 PiperOrigin-RevId: 284866710 --- diff --git a/mlir/include/mlir/IR/AffineExpr.h b/mlir/include/mlir/IR/AffineExpr.h index cca7eac..928ced5 100644 --- a/mlir/include/mlir/IR/AffineExpr.h +++ b/mlir/include/mlir/IR/AffineExpr.h @@ -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; diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp index 95ebc0a..19599a8 100644 --- a/mlir/lib/IR/AffineExpr.cpp +++ b/mlir/lib/IR/AffineExpr.cpp @@ -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(); auto rhsConst = rhs.dyn_cast(); + // 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(); if (lBin && lBin.getKind() == AffineExprKind::Mul) { if (auto lrhs = lBin.getRHS().dyn_cast()) { - // 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(); if (lBin && lBin.getKind() == AffineExprKind::Mul) { if (auto lrhs = lBin.getRHS().dyn_cast()) { - // 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(); auto rhsConst = rhs.dyn_cast(); + // 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(); + 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 { diff --git a/mlir/test/IR/affine-map.mlir b/mlir/test/IR/affine-map.mlir index d867d51..ebbd473 100644 --- a/mlir/test/IR/affine-map.mlir +++ b/mlir/test/IR/affine-map.mlir @@ -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>) diff --git a/mlir/test/Transforms/Vectorize/compose_maps.mlir b/mlir/test/Transforms/Vectorize/compose_maps.mlir index a8afbec..f1826f4 100644 --- a/mlir/test/Transforms/Vectorize/compose_maps.mlir +++ b/mlir/test/Transforms/Vectorize/compose_maps.mlir @@ -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 -} \ No newline at end of file +} diff --git a/mlir/test/Transforms/unroll.mlir b/mlir/test/Transforms/unroll.mlir index 208df58..da2a5e5 100644 --- a/mlir/test/Transforms/unroll.mlir +++ b/mlir/test/Transforms/unroll.mlir @@ -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() {