[Codegen] codegen with minimum if condition

If condition often slows things down, and it often occurs when split / tile factors do not divide the shape. For example, split an axis of length 1023 by 32 yields a if condition, while split an axis of length 1024 by 32 does not. Note that split an axis of length tvm.var('m') also yields a if condition, as 32 may not divide m .

What if we specify the shape as (32 * tvm.var('M'), 32 * tvm.var('N')) to give more hint to the compiler? I tried it out in a matrix multiplication, but unfortunately the if condition is not eliminated.

import tvm
import numpy
import timeit

M = tvm.const(32) * tvm.var('M')
K = tvm.const(32) * tvm.var('K')
N = tvm.const(32) * tvm.var('N')
M_num = 1024
K_num = 1024
N_num = 1024

# The default tensor type in tvm
dtype = "float32"

target = 'llvm'
ctx = tvm.context(target, 0)

# Random generated tensor for testing
a = tvm.nd.array(numpy.random.rand(M_num, K_num).astype(dtype), ctx)
b = tvm.nd.array(numpy.random.rand(K_num, N_num).astype(dtype), ctx)

# Algorithm
k = tvm.reduce_axis((0, K), 'k')
A = tvm.placeholder((M, K), name='A')
B = tvm.placeholder((K, N), name='B')
C = tvm.compute(
           (M, N),
           lambda x, y: tvm.sum(A[x, k] * B[k, y], axis=k),
           name='C')

bn = 32
s = tvm.create_schedule(C.op)

# Blocking by loop tiling
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
k, = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)

# Hoist reduction domain outside the blocking loop
s[C].reorder(xo, yo, ko, ki, xi, yi)

print(tvm.lower(s, [A, B, C], simple_mode=True))

The code generated:

produce C {
  for (x.outer, 0, (((M*32) + 31)/32)) {
    for (y.outer, 0, (((N*32) + 31)/32)) {
      for (x.inner.init, 0, 32) {
        for (y.inner.init, 0, 32) {
          if (likely((((x.outer*32) + x.inner.init) < (M*32)))) {
            if (likely((((y.outer*32) + y.inner.init) < (N*32)))) {
              C[(((y.outer*32) + (((x.outer*32) + x.inner.init)*(N*32))) + y.inner.init)] = 0f
            }
          }
        }
      }
      for (k.outer, 0, (((K*32) + 3)/4)) {
        for (k.inner, 0, 4) {
          for (x.inner, 0, 32) {
            for (y.inner, 0, 32) {
              if (likely((((x.outer*32) + x.inner) < (M*32)))) {
                if (likely((((y.outer*32) + y.inner) < (N*32)))) {
                  if (likely((((k.outer*4) + k.inner) < (K*32)))) {
                    if (likely((((x.outer*32) + x.inner) < (M*32)))) {
                      if (likely((((y.outer*32) + y.inner) < (N*32)))) {
                        if (likely((((k.outer*4) + k.inner) < (K*32)))) {
                          C[(((y.outer*32) + (((x.outer*32) + x.inner)*(N*32))) + y.inner)] = (C[(((y.outer*32) + (((x.outer*32) + x.inner)*(N*32))) + y.inner)] + (A[(((k.outer*4) + (((x.outer*32) + x.inner)*(K*32))) + k.inner)]*B[(((y.outer*32) + (((k.outer*4) + k.inner)*(N*32))) + y.inner)]))
                        }
                      }
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
  }
}

As can be seen, 32 | 32 * m is not recognized and unnecessary if conditions are generated.

Is there any better ways to avoid the if condition and speed ops up? Note that I want to deal with inputs with general shapes (or at least a family of shapes, like the multiples of 32), instead of specific shapes like (1024, 1024).

1 Like

Thank @yzhliu for guidance.

@hzfan would you mind to share how you resolved this issue?
I am thinking this could be related to Expr Simplifier for tvm.var.

Thanks.

@comaniac actually the issue has not been resolved yet. I agree that the root cause of this issue is similar to the Expr Simplifier for tvm.var