Conditional Loop Partitioning - Extending to remove if conditions by anijain2305 · Pull Request #1797 · apache/tvm

anijain2305 · 2018-10-02T23:23:34Z

The following PR accomplishes 2 items in the context of loop partitioning pass

Removing if conditions from the AST if they can be removed
Prev implementation was not partitioning loops on multiple loop variables because the other loop variable candidates were becoming stale as we partition for just one loop variable

Illustrative example
Input

realize compute([0, 60]) {
  produce compute {
    for (i.outer, 0, 4) {
      for (i.inner, 0, 16) {
        if (likely(((i.inner + (i.outer*16)) < 60))) {
          if (likely(((i.inner + (i.outer*16)) < 60))) {
            compute((i.inner + (i.outer*16))) =(A((i.inner + (i.outer*16))) + B((i.inner + (i.outer*16))))
          }
        }
      }
    }
  }
}

Output

// attr [compute(compute, 0x7f83a88706c0)] realize_scope = ""
realize compute([0, 60]) {
  produce compute {
    for (i.outer, 0, 3) {
      for (i.inner, 0, 16) {
        compute((i.inner + (i.outer*16))) =(A((i.inner + (i.outer*16))) + B((i.inner + (i.outer*16))))
      }
    }
    for (i.inner, 0, 12) {
      compute((i.inner + 48)) =(A((i.inner + 48)) + B((i.inner + 48)))
    }
  }
}

Example with multiple if conditions
Input

// attr [compute(compute, 0x7ffd9e4a0970)] realize_scope = ""
realize compute([0, 94], [0, 62]) {
  produce compute {
    for (i.outer, 0, 6) {
      for (j.outer, 0, 4) {
        for (i.inner, 0, 16) {
          for (j.inner, 0, 16) {
            if (likely(((i.inner + (i.outer*16)) < 94))) {
              if (likely(((i.inner + (i.outer*16)) < 94))) {
                if (likely(((j.inner + (j.outer*16)) < 62))) {
                  if (likely(((j.inner + (j.outer*16)) < 62))) {
                    compute((i.inner + (i.outer*16)), (j.inner + (j.outer*16))) =(A((i.inner + (i.outer*16)), (j.inner + (j.outer*16))) + B((i.inner + (i.outer*16)), (j.inner + (j.outer*16))))
                  }
                }
              }
            }
          }
        }
      }
    }
  }
}

Output

// attr [compute(compute, 0x7ffd9e4a0970)] realize_scope = ""
realize compute([0, 94], [0, 62]) {
  produce compute {
    for (i.outer, 0, 5) {
      for (j.outer, 0, 3) {
        for (i.inner, 0, 16) {
          for (j.inner, 0, 16) {
            compute((i.inner + (i.outer*16)), (j.inner + (j.outer*16))) =(A((i.inner + (i.outer*16)), (j.inner + (j.outer*16))) + B((i.inner + (i.outer*16)), (j.inner + (j.outer*16))))
          }
        }
      }
      for (i.inner, 0, 16) {
        for (j.inner, 0, 14) {
          compute((i.inner + (i.outer*16)), (j.inner + 48)) =(A((i.inner + (i.outer*16)), (j.inner + 48)) + B((i.inner + (i.outer*16)), (j.inner + 48)))
        }
      }
    }
    for (j.outer, 0, 3) {
      for (i.inner, 0, 14) {
        for (j.inner, 0, 16) {
          compute((i.inner + 80), (j.inner + (j.outer*16))) =(A((i.inner + 80), (j.inner + (j.outer*16))) + B((i.inner + 80), (j.inner + (j.outer*16))))
        }
      }
    }
    for (i.inner, 0, 14) {
      for (j.inner, 0, 14) {
        compute((i.inner + 80), (j.inner + 48)) =(A((i.inner + 80), (j.inner + 48)) + B((i.inner + 80), (j.inner + 48)))
      }
    }
  }
}

anijain2305 · 2018-10-02T23:24:42Z

@xqdan @tqchen @ZihengJiang

anijain2305 · 2018-10-05T20:09:32Z

Ping

xqdan · 2018-10-09T03:01:23Z

We are testing for this patch, will give feedback soon.

anijain2305 · 2018-10-09T03:22:41Z

Thanks @xqdan, please feel free to provide more testcases if something is of your interest and doesn't work with this patch.

xqdan · 2018-10-09T06:56:06Z

@anijain2305 could you try this case.
https://discuss.tvm.ai/t/simplify-if-stmt-within-for/830

xqdan · 2018-10-09T07:14:32Z

more cases:
test_cce_loop_1: fail to partition
test_cce_loop_2: ok
test_cce_loop_3: fail to partition

import tvm

def test_cce_loop_1():
  ib = tvm.ir_builder.create()
  len = 112
  tile = 32
  loop = (len + tile - 1) // tile
  with ib.for_range(0, loop, 'i') as i:
    head = i * tile
    tail = tvm.select(head + tile > len, len, head + tile)
    ib.emit(tvm.call_extern('float32', "cce_intrisic", head, tail))

  stmt = ib.get()
  stmt = tvm.ir_pass.LoopPartition(stmt, True)
  stmt = tvm.ir_pass.Simplify(stmt)
  print stmt

def test_cce_loop_2():
  ib = tvm.ir_builder.create()
  len = 112
  tile = 32
  loop = (len + tile - 1) // tile
  with ib.for_range(0, loop, 'i') as i:
    head = i * tile
    with ib.if_scope(ib.likely(head + tile > len)):
      tail = len
      ib.emit(tvm.call_extern('float32', "cce_intrisic", head, tail))
    with ib.else_scope():
      tail = head + tile
      ib.emit(tvm.call_extern('float32', "cce_intrisic", head, tail))

  stmt = ib.get()
  stmt = tvm.ir_pass.LoopPartition(stmt, True)
  stmt = tvm.ir_pass.Simplify(stmt)
  print stmt

def test_cce_loop_3():
  ib = tvm.ir_builder.create()
  len = 112
  tile = 32
  loop = (len + tile - 1) // tile
  with ib.for_range(0, loop, 'i') as i:
    head = i * tile
    with ib.if_scope(ib.likely(head + tile > len)):
      tail = len
      with ib.if_scope(ib.likely(i > 0)):
        init = 0
        ib.emit(tvm.call_extern('float32', "cce_intrisic", head, tail, init))
      with ib.else_scope():
        init = 1
        ib.emit(tvm.call_extern('float32', "cce_intrisic", head, tail, init))
    with ib.else_scope():
      tail = head + tile
      with ib.if_scope(ib.likely(i > 0)):
        init = 0
        ib.emit(tvm.call_extern('float32', "cce_intrisic", head, tail, init))
      with ib.else_scope():
        init = 1
        ib.emit(tvm.call_extern('float32', "cce_intrisic", head, tail, init))

  stmt = ib.get()
  stmt = tvm.ir_pass.LoopPartition(stmt, True)
  stmt = tvm.ir_pass.Simplify(stmt)
  print stmt

if __name__ == "__main__":
  test_cce_loop_1()
  test_cce_loop_2()
  test_cce_loop_3()

anijain2305 · 2018-10-09T20:00:06Z

@xqdan Thanks for providing the test cases. Current situation is

https://discuss.tvm.ai/t/simplify-if-stmt-within-for/830 --> If you add ib.likely in the conditions, it gets partitioned. Currently, the loop partitioning only works for likely primitives.
cce_loop_1 --> select not supported for loop partitioning
cce_loop_2 --> Seems good
cce_loop_3 --> Will need to dig deeper. I think 2 conditions dependent on i back to back is throwing code off.

  with ib.for_range(0, loop, 'i') as i:
    head = i * tile
    with ib.if_scope(ib.likely(head + tile > len)):

RichardFido · 2018-10-10T23:18:45Z

hi,
here's another case as bellow fail to partition, please consider:

import tvm

def test():
ib = tvm.ir_builder.create()
loop1=4
loop2=9998
tile=39991
with ib.for_range(0,loop2,'i') as i:
with ib.for_range(0,loop1,'j') as j:
head1 = i
head2 = j
with ib.if_scope(ib.likely(head1loop1 < tile - head2)):
with ib.if_scope(ib.likely(head1loop1 < tile - head2)):
ib.emit(tvm.call_extern('float16',"cce_intrinsic",head1))

stmt = ib.get()
print stmt
print "==="
stmt = tvm.ir_pass.LoopPartition(stmt,True)
print stmt

test()

MarisaKirisame · 2018-10-10T23:21:28Z

-        Stmt post_body = Substitute(body, {{Var{var}, var + post_doubt_begin}});
-        post_stmt = MakeFor(node, max - post_doubt_begin + 1, post_body);
+        Stmt post_body;
+        // If the loop is going from 0 to 1, then replace the loop var with min value


remove then

anijain2305 · 2018-10-12T21:31:50Z

@RichardFido Can you please edit the code to print it nicely? It doesn't compile as of now

RichardFido · 2018-10-15T02:43:09Z

sorry about that. @anijain2305
please check it again:


def test():
    ib = tvm.ir_builder.create()
    loop1 = 4 
    loop2 = 9998
    tile = 39991
    with ib.for_range(0,loop2,'i') as i:
        with ib.for_range(0,loop1,'j') as j:
            head1 = i 
            head2 = j 
            with ib.if_scope(ib.likely(head1*loop1 < tile - head2)):
                with ib.if_scope(ib.likely(head1*loop1 < tile - head2)):
                    ib.emit(tvm.call_extern('float16',"cce_intrisic",head1))

    stmt = ib.get()
    stmt = tvm.ir_pass.LoopPartition(stmt,True)
    print stmt

test()

yzhliu · 2018-10-15T18:56:42Z

+            if (*as_const_int(max) == *as_const_int(post_doubt_begin)) {
+                post_body = Substitute(body, {{Var{var}, post_doubt_begin}});
+                post_stmt = post_body;
+            }


just trying to understand, why this need to be handled separately?

An example usecase is as follows. This is what statement looks in between the CLP is running

for (i = 0, 1) for(j = 0, 8) if(i + j < 4) Do_something

If the loop bounds are going from 0 to 1, then its better to replace i with 0 because then it simplifies the conditions, resulting in better partitioning.

anijain2305 · 2018-10-16T00:56:49Z

@RichardFido The example doesn't work because of the rigid structure of how bounds are deduced for the conditionals. I do not understand that portion of the code. There is an easy workaround for now.

Line causing issue - with ib.if_scope(ib.likely(head1*loop1 < tile - head2)):
Replace  - with ib.if_scope(ib.likely(head1*loop1 + head2 < tile)):

If the constant is not present on the right side alone, then the deduce_bound function does not work correctly. Solving for any combination of conditions will require decent changes in deduce_bounds functions or Simplify pass. I do not plan to work on that in this PR.

xqdan · 2018-10-29T09:00:44Z

I suggest we can merge this PR first, we can send more PR to support other cases someone may need.

tqchen · 2018-10-29T22:41:23Z

OK, reviewer, please take another round of look and https://docs.tvm.ai/contribute/code_review.html#approve-and-request-changes-explicitly

tqchen · 2018-10-30T02:45:03Z

THanks @xqdan @MarisaKirisame @anijain2305 @yzhliu this is now merged

…ache#1797)

Conditional Loop Partitioning - Extending to remove if conditions

ee5a2f0

tqchen added the status: need review label Oct 3, 2018

tqchen assigned ZihengJiang Oct 6, 2018

MarisaKirisame reviewed Oct 10, 2018

View reviewed changes

Cleanup. Added more testcases

c52f85c

yzhliu reviewed Oct 15, 2018

View reviewed changes

Additional test case

1b429b6

MarisaKirisame approved these changes Oct 29, 2018

View reviewed changes

xqdan approved these changes Oct 29, 2018

View reviewed changes

tqchen merged commit dd1558a into apache:master Oct 30, 2018

tqchen added status: accepted and removed status: need review labels Oct 30, 2018

FrozenGene pushed a commit to FrozenGene/tvm that referenced this pull request Dec 27, 2018

Conditional Loop Partitioning - Extending to remove if conditions (ap…

ea74668

…ache#1797)

wweic pushed a commit to neo-ai/tvm that referenced this pull request Feb 20, 2019

Conditional Loop Partitioning - Extending to remove if conditions (ap…

dabf5c6

…ache#1797)

wweic pushed a commit to neo-ai/tvm that referenced this pull request Feb 20, 2019

Conditional Loop Partitioning - Extending to remove if conditions (ap…

d5b3f25

…ache#1797)

anijain2305 deleted the conditional_loop_partitioning branch June 14, 2019 18:18

tqchen unassigned ZihengJiang Nov 4, 2019

Conversation

anijain2305 commented Oct 2, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

anijain2305 commented Oct 2, 2018

Uh oh!

anijain2305 commented Oct 5, 2018

Uh oh!

xqdan commented Oct 9, 2018

Uh oh!

anijain2305 commented Oct 9, 2018

Uh oh!

xqdan commented Oct 9, 2018

Uh oh!

xqdan commented Oct 9, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

anijain2305 commented Oct 9, 2018

Uh oh!

RichardFido commented Oct 10, 2018

Uh oh!

MarisaKirisame Oct 10, 2018

Choose a reason for hiding this comment

Uh oh!

anijain2305 commented Oct 12, 2018

Uh oh!

RichardFido commented Oct 15, 2018

Uh oh!

yzhliu Oct 15, 2018

Choose a reason for hiding this comment

Uh oh!

anijain2305 Oct 16, 2018

Choose a reason for hiding this comment

Uh oh!

anijain2305 commented Oct 16, 2018

Uh oh!

xqdan commented Oct 29, 2018

Uh oh!

tqchen commented Oct 29, 2018

Uh oh!

tqchen commented Oct 30, 2018

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

anijain2305 commented Oct 2, 2018 •

edited

Loading

xqdan commented Oct 9, 2018 •

edited

Loading