[TIR] add loop partition hint pragma

[wrongtest-intellif]

Currently LoopPartition pass will try to partition loops assiociated with condition in likely tag, it would be great if developers can take control of which loop to partition, no-matter whether the condition to eliminate is "likely" tagged or not.

The PR add a pragma attr key loop_partition_hint, which can be tagged explicitly in schedule phase. The loop partition pass will consider all arith conditions for hinted loop var.

Below are two examples of how explicit controlled loop partition benefits, the target is on Ubuntu20.08 i7-7700, with llvm version 11.0 :

• For max pooling with padding inlined, which create conditional buffer accesses

  data = te.placeholder([1, 128, 56, 56], name="x")
  out = topi.nn.pool2d(data, kernel=[5, 5], stride=[1, 1], padding=[2, 2, 2, 2], pool_type="max", dilation=[1, 1], layout="NCHW")
  pad = out.op.input_tensors[0]
  x = tvm.nd.array(np.random.randint(0, 64, [1, 128, 56, 56]).astype("float32"))
  
  def test(do_partition):
      s = te.create_schedule([out.op])
      s[pad].compute_inline()
      n, c, h, w = s[out].op.axis
      if do_partition:
          s[out].pragma(h, "loop_partition_hint")
          s[out].pragma(w, "loop_partition_hint")
  
      with tvm.ir.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
          f = tvm.build(s, [data, out], "llvm")
      y = tvm.nd.array(np.zeros([1, 128, 56, 56]).astype("float32"))
      f(x, y)
      result = y.asnumpy()
      print(f.get_source("asm"))
      evaluator = f.time_evaluator(f.entry_name, tvm.cpu(), number=1000)
      print("partition=%s: %.3f millisecs" % (do_partition, evaluator(x, y).mean * 1000))
      return result
  
  r1 = test(do_partition=False)
  r2 = test(do_partition=True)
  testing.assert_allclose(r1, r2, rtol=1e-5)

  The performance I get:

  • no loop partition: 3.708 millisecs
  • with loop partition: 0.975 millisecs

• For tiled matmul following TVM tensor expression tutorial, but with shape not divided by tiling factor. The tir split do not create a likely condition for it now.

  M, N, K = 1025, 1025, 1025
  dtype = "float32"
  dev = tvm.cpu()
  a = tvm.nd.array(np.random.rand(M, K).astype(dtype), dev)
  b = tvm.nd.array(np.random.rand(K, N).astype(dtype), dev)
  k = te.reduce_axis((0, K), "k")
  A = te.placeholder((M, K), name="A")
  B = te.placeholder((K, N), name="B")
  C = te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")
  f = te.create_prim_func([A, B, C])
  s = tvm.tir.Schedule(f)
  
  def evaluate_operation(s, target, optimization):
     with tvm.ir.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
         print(tvm.lower(s.mod["main"], [], simple_mode=True))
         func = tvm.build(s.mod["main"], [], target=target, name="mmult")
         assert func
  
     c = tvm.nd.array(np.zeros((M, N), dtype=dtype), dev)
     func(a, b, c)
     evaluator = func.time_evaluator(func.entry_name, dev, number=10)
     mean_time = evaluator(a, b, c).mean
     print("%s: %f" % (optimization, mean_time))
  
  # no opt
  evaluate_operation(s, target="llvm", optimization="none")
  
  # tiling and vectorize
  x, y, k = s.get_loops(s.get_block("C"))
  xo, xi = s.split(x, factors=[None, 32])
  yo, yi = s.split(y, factors=[None, 32])
  ko, ki = s.split(k, factors=[None, 4])
  s.reorder(xo, yo, ko, ki, xi, yi)
  s.vectorize(yi)
  evaluate_operation(s, target="llvm", optimization="blocking")
  
  # loop partition
  def pragma(s, rv, key):
     sref = s.get_sref(rv)
     loop = sref.stmt
     new_loop = tvm.tir.For(loop.loop_var, loop.min, loop.extent, loop.kind, loop.body, annotations={key: 1})
     s.state.replace(sref, new_loop)
  pragma(s, xo, "pragma_loop_partition_hint")
  pragma(s, yo, "pragma_loop_partition_hint")
  evaluate_operation(s, target="llvm",  optimization="loop_partition")

  The performance I get:

  • no opt: 1.374402
  • with tiling + vectorize: 0.843930
  • with tiling + vectorize + loop partition: 0.272183

[junrushao]

CC @areusch @Hzfengsy @vinx13 @zxybazh @ZihengJiang would you guys review this PR? looks like it's relevant to some of our previous discussion

[wrongtest-intellif]

@Hzfengsy hi~ I resolve the comment issues，can you kindly take another round of review：）？

[Hzfengsy]

Thanks, @wrongtest