test_arith_canonical_simplify.py 11 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
17 18 19 20 21 22 23 24
import tvm

class CanonicalChecker:
    def __init__(self):
        self.analyzer = tvm.arith.Analyzer()

    def verify(self, data, expected):
        res = self.analyzer.canonical_simplify(data)
25
        assert tvm.ir_pass.Equal(res, expected), "\ndata={}\nres={}\nexpected={}".format(data, res, expected)
26 27 28 29 30


def test_mul_sum_simplify():
    ck = CanonicalChecker()
    x, y, z = tvm.var("x"), tvm.var("y"), tvm.var("z")
31

32 33 34 35
    ck.verify(2 + (3 * x + z + y + 1) * 4 + x,
              x * 13 + z * 4 + y * 4 +6)
    ck.verify(x * 3 - 4 * x + 1, 1 - x)
    ck.verify(y + x * 3 - 5 * x + 1 + y, y * 2 + 1 - x * 2)
36 37
    tdiv = tvm.truncdiv
    tmod = tvm.truncmod
38
    # trucdiv
39 40
    ck.verify(tdiv(x + y + x + y * 3, 2), y * 2 + x)
    ck.verify(tmod(x + y + x + y * 3, 2), 0)
41 42 43 44 45 46 47 48

    # floordiv
    fld = tvm.floordiv
    flm = tvm.floormod
    ck.verify(flm(x + x + y * 3, 2), flm(y * 3, 2))
    ck.verify(fld(x + y + x + y * 3, 2), y * 2 + x)
    ck.verify(flm(x + y + x + y * 3, 2), 0)
    ck.verify(fld(x + x + y * 3, 2), fld(y * 3, 2) + x)
49 50 51 52 53 54


def test_split_index_simplify():
    ck = CanonicalChecker()
    x, y, z = tvm.var("x"), tvm.var("y"), tvm.var("z")

55
    # trucdiv
56 57 58
    tdiv = tvm.truncdiv
    tmod = tvm.truncmod

59
    # split div const
60 61 62 63 64 65
    ck.verify(tdiv(x, 3) *3 + tmod(x, 3), x)
    ck.verify(tdiv(x, 6) * 6 + tmod(tdiv(x, 3), 2) * 3 + tmod(x, 3), x)
    ck.verify(tdiv(tdiv(tmod(x, 16), 2) * 2, 4), tdiv(tmod(x, 16), 4))
    ck.verify(tdiv(tmod(x, 2), 8), 0)
    ck.verify(tdiv(tmod(x, 2), 7), 0)
    ck.verify(tdiv(tdiv(tmod(x, 16), 2) * 2, 6), tdiv(tmod(x, 16), 6))
66 67

    # split mod const
68 69
    ck.verify(tmod((x * 8), 16), tmod(x, 2) * 8)
    ck.verify(tmod(x * 8, 2), 0)
70 71 72 73

    # simplify then fold
    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 1000))
    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 1000))
74
    ck.verify(tdiv(x * 4 + y, 2) * 2 + tmod(x * 4 + y, 2), x * 4 + y)
75
    # complex fold
76
    ck.verify(tdiv(z * 9 + y, 2) * 2 + tmod(z * 9 + y, 2), z * 9 + y)
77

78 79
    ck.analyzer.update(x, tvm.arith.ConstIntBound(-100, 1000), True)
    ck.analyzer.update(y, tvm.arith.ConstIntBound(-100, 1000), True)
80
    ck.verify(tdiv(x * 4 + y, 2) * 2 + tmod(x * 4 + y, 2), x * 4 + y)
81 82 83 84 85 86 87 88 89 90 91 92

    # floordiv
    fld = tvm.floordiv
    flm = tvm.floormod
    ck.verify(fld(x, 3) * 3 + flm(x, 3), x)
    ck.verify(fld(x, 6) * 6 + flm(fld(x, 3), 2) * 3 + flm(x, 3), x)
    ck.verify(fld(fld(flm(x, 16), 2) * 2, 4), fld(flm(x, 16), 4))
    ck.verify(fld(flm(x, 2), 8), 0)
    ck.verify(fld(flm(x, 2), 7), 0)
    ck.verify(fld(fld(flm(x, 16), 2) * 2, 6), fld(flm(x, 16), 6))

    # cannot simplify mixed case, unless we canonicalize into one mode.
93
    ck.verify(tdiv(x,6) * 2 + tmod(fld(x,3), 2), tdiv(x,6) * 2 + tmod(fld(x,3), 2))
94 95 96 97 98


def test_div_simplify():
    ck = CanonicalChecker()
    x = tvm.var("x")
99
    tdiv = tvm.truncdiv
100 101

    # truc div
102
    ck.verify(tdiv(16+48*x,16), x*3 + 1)
103 104
    # (17+48*x)/16 is not simplifiable for arbitrary x because when 17+48*x<0
    # (17+48*x)/16 != 1+3*x
105
    ck.verify(tdiv(17 + 48 * x, 16), tdiv(x * 48 + 17, 16))
106 107
    # However, when x >= 0, then 17+48*x >= 0 and (17+48*x)/16 can be simplified
    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 10))
108
    ck.verify(tdiv(17 + 48 * x, 16), x * 3 + 1)
109
    # Trying expressions that are not simplifiable for any values of the variables
110
    ck.verify(tdiv(17 + 47 * x, 16), tdiv(x * 47 + 17, 16))
111

112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127
    # floordiv
    fld = tvm.floordiv
    ck.analyzer.update(x, tvm.arith.ConstIntBound(-1000, 10000), True)
    ck.verify(fld(16+48*x, 16), x*3 + 1)
    ck.verify(fld(17+48*x, 16), x * 3 + 1)
    ck.verify(fld(17+47*x, 16), fld(x * 47 + 17, 16))


def test_floormod_simplify():
    ck = CanonicalChecker()
    flm = tvm.floormod
    x, y = tvm.var("x"), tvm.var("y")
    ck.verify(flm(flm((x*4) + y  - 466036, 24528) - 24512,  16),
              flm((x*4) + y  + 12, 16))


128 129 130 131 132

def test_canonical_mixed():
    ck = CanonicalChecker()
    x = tvm.var("x")
    z = tvm.const(3, "int32")
133 134 135 136
    tdiv = tvm.truncdiv
    tmod = tvm.truncmod
    ck.verify(tdiv(x, (z*z)) - tdiv(x, (z*z)), 0)
    ck.verify(tdiv(x, (z+z)) - tdiv(x, (z+z)), 0)
137 138 139 140
    ck.verify(x - 2 < 3, x < 5)
    ck.verify(tvm.max(x, 1) - tvm.max(x, 1), 0)
    ck.verify(tvm.min(x, 1) - tvm.min(x, 1), 0)
    ck.verify(x * x - x * x, 0)
141

142 143 144 145
    fld = tvm.floordiv
    ck.verify(fld(x, (z*z)) - fld(x, (z*z)), 0)
    ck.verify(fld(x, (z+z)) - fld(x, (z+z)), 0)

146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217

def test_reduce_combiner_simplify():
    ck = CanonicalChecker()
    dummy = tvm.var('dummy')
    comm_reducer = tvm.comm_reducer
    prod = comm_reducer(lambda x, y: x*y, lambda t0: tvm.const(1, t0))

    sum_or_prod = comm_reducer(
        lambda x, y: tvm.expr.Select(dummy < 0,
                                     x + y, x*y),
        lambda t0: tvm.expr.Select(dummy < 0,
                                   tvm.const(0, t0), tvm.const(1, t0)))
    sum_and_prod = comm_reducer(
        lambda x, y: (x[0] + y[0],
                      x[1]*y[1]),
        lambda t0, t1: (tvm.const(0, t0),
                        tvm.const(5, t0) - tvm.const(4, t0)))
    some_reducer1 = comm_reducer(
        lambda x, y: (x[0] + y[0],
                      x[0] + y[0] + x[1] + y[1],
                      x[0]*y[2] + y[0]*x[2],
                      x[1] + y[2],
                    4.0),
        lambda t0, t1, t2, t3, t4: (tvm.const(0, t0),
                                    tvm.const(1, t1),
                                    tvm.const(2, t2),
                                    tvm.const(3, t3),
                                    tvm.const(4, t4)))

    k = tvm.reduce_axis((0, 10), name="k")
    A = tvm.placeholder((10,), name='A')
    # Test that SimplifyCombiner makes use of vranges
    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(-10, -4))
    ck.verify(sum_or_prod(A[k], k), tvm.sum(A[k], k))
    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(5, 9), True)
    ck.verify(sum_or_prod(A[k], k), prod(A[k], k))
    ck.analyzer.update(dummy, tvm.arith.ConstIntBound(-10, 100), True)
    ck.verify(sum_and_prod((A[k], A[10-k]), k)[0], tvm.sum(A[k], k))
    ck.verify(sum_and_prod((A[k], A[10-k]), k)[1], prod(A[10-k], k))

    reference_simplified_sources = [[A[0]],
                                    [A[0], A[1]],
                                    [A[0], A[2]],
                                    [A[0], A[1], A[2], A[3]],
                                    [A[4]]]
    for j in range(5):
        # Here we use the j-th component of the result, so only it and the components it
        # depends on are left.
        simplified = ck.analyzer.canonical_simplify(
            some_reducer1((A[0], A[1], A[2], A[3], A[4]), k)[j])

        # Check that the remaining components are the expected ones.
        for lhs, rhs in zip(simplified.source, reference_simplified_sources[j]):
            assert tvm.ir_pass.Equal(lhs, rhs)

    # Test that components with side effects are not removed
    side_effect = lambda *xs: tvm.make.Call("int32", "dummy", xs, tvm.expr.Call.Intrinsic, None, 0)
    ck.verify(sum_and_prod((A[k], side_effect(A[10-k])), k)[0],
             sum_and_prod((A[k], side_effect(A[10-k])), k)[0])
    ck.verify(sum_and_prod((side_effect(A[k]), A[10-k]), k)[0],
              tvm.sum(side_effect(A[k]), k))


def test_reduce_simplify():
    ck = CanonicalChecker()
    k = tvm.reduce_axis((0, 10), name="k")
    j = tvm.reduce_axis((-5, 3), name="j")
    A = tvm.placeholder((10,), name='A')
    ck.verify(tvm.sum(tvm.expr.Select(k + j < 12, k + j, 0), [k, j]),
              tvm.sum(k + j, [k, j]))
    ck.verify(tvm.sum(A[3], []), A[3])
    # The rule below is not typical, removed for now
218
    ck.verify(tvm.sum(tvm.div(k, 10), k), tvm.sum(tvm.const(0, "int32"), k))
219 220


221 222 223 224
def test_simplify_if_then_else():
    ck = CanonicalChecker()
    x = tvm.var("x")
    y = tvm.var("y")
225 226
    tdiv = tvm.truncdiv
    tmod = tvm.truncmod
227 228
    # simplification that takes condition into account.
    res = tvm.if_then_else((x * 4 + y) >= 466036,
229 230
                           tvm.if_then_else(24512 <= tmod(((x*4) + y) - 466036, 24528),
                                            tmod(tmod(((x*4) + y)  - 466036, 24528) -24512, 16),
231
                                            x), y)
232 233

    res2 = tvm.if_then_else((x * 4) >= 466036 - y,
234 235
                           tvm.if_then_else(24512 <= tmod(((x*4) + y) - 466036, 24528),
                                            tmod(tmod(((x*4) + y)  - 466036, 24528) -24512, 16),
236
                                            x), y)
237 238
    expected = tvm.if_then_else(
        tvm.expr.LE(466036, (x * 4 + y)),
239 240
        tvm.if_then_else(tvm.expr.LE(24512, tmod(((x*4) + y) - 4, 24528)),
                                     tmod(((x*4) + y)  - 4, 16),
241 242
                         x), y)
    ck.verify(res, expected)
243
    ck.verify(res2, expected)
244
    # can only simplify if condition
245 246
    res = tvm.expr.Select(tvm.all(x >= -1, y >= 0), tmod(x + y + 100, 3), tmod(x + 100, 3))
    expected = tvm.expr.Select(tvm.all(x >= -1, y >= 0), tmod(x + y + 1, 3), tmod(x + 100, 3))
247 248 249
    ck.verify(res, ck.analyzer.canonical_simplify(expected))

    res = tvm.expr.Select(x >= 10,
250
                          tvm.if_then_else(tdiv(x, 3) > 2, x, 0), 0)
251 252 253 254
    expected = tvm.expr.Select(x >= 10, x, 0)
    ck.verify(res, ck.analyzer.canonical_simplify(expected))

    res = tvm.expr.Select(x >= 10,
255
                          tvm.if_then_else(tdiv(x, 3) < 2, x, 0), 0)
256 257 258
    ck.verify(res, 0)


259 260 261 262
def test_complex_cases():
    ck = CanonicalChecker()
    x = tvm.var("x")
    y = tvm.var("y")
263 264 265 266 267
    tdiv = tvm.truncdiv
    tmod = tvm.truncmod
    res2 = (tdiv(tdiv(tmod(x*128 + y, 1296),36)*2 + 1,2)*36 +
            tdiv(tmod((x*128) + y, 36)*2 + 1,2)
            - tmod((x*128) + y, 1296) + 1)
268 269 270 271
    ck.analyzer.update(x, tvm.arith.ConstIntBound(0, 5))
    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 127))
    ck.verify(res2, 1)

272
    ck.analyzer.update(y, tvm.arith.ConstIntBound(0, 1024), True)
273 274 275 276
    res3 = (tdiv(x*1024 + y,65536) + tdiv(tmod(x*1024 + y, 65536),256)
            + tdiv(tmod(x*1024 + y, 256),16) + tmod(x*1024 + y, 16) - tdiv(y,256) -
            tdiv(tmod(y, 256),16) - tmod(y, 16) - (x*4))
    ck.verify(res3, tdiv((x*1024) + y, 256) - tdiv(y,256) - (x*4))
277 278


279
if __name__ == "__main__":
280 281
    test_floormod_simplify()
    test_mul_sum_simplify()
282
    test_simplify_if_then_else()
283 284 285
    test_div_simplify()
    test_reduce_simplify()
    test_reduce_combiner_simplify()
286

287 288
    test_split_index_simplify()
    test_canonical_mixed()
289
    test_complex_cases()