Text format for Halide IR from a subset of Python

Development pre-RFC
#1

I am planning to build another front end for TVM, because:

  1. some patterns, like triangular memory access, and fine-grained dependences, cannot be easily represented by TVM high-level language
  2. the current ir-builder, although works, makes code redundant

I propose a Python subset which can be lowered to Halide IR, and so far it only uses:

  1. for loop
for i in range(5):
    # something
  1. TVM specialized array (Some TVM interface should be added)
a = tvm.allocate_array_with_shape((5, 5))
  1. if-then-else (no elif)
if cond1:
    #then clause
else:
    #else clause
  1. arithmetic operations (including sqrt, log, and etc)
  2. to simplify the type specification, maybe type annotation is required:
def func_module(a: tvm.Tensor, b: tvm.Tensor, c: tvm.Tensor) -> int:

Something I am not sure yet:
Halide IR does not support variable write, all the variables to be written are arrays with length one. I have 2 solutions to this issue so far:

  1. Let the compiler automatically detect the variable r/w and replace it with array declaration
  2. Add new function call intrinsics for variable declaration and write. The write function should be a non-side-effect-free intrinsic

All the comments and suggestions will be appreciated.

1 Like
Complex matrix inversion?
#2

Bump on this, finally have sometime after the deadlines. https://github.com/dmlc/tvm/pull/1168 is a good starting point. I now feel it is important enough that we should define a clear API around it. Here are some of my initial example thoughts

import tvm

# a good name?
@tvm.hybrid_script
def mycross(a, b):
  # t is a special namespace 
  # to indicate TVM ir constructs
  c = t.array((a.shape[0], b.shape[0]))
  for i in t.unroll_range(a.shape[0]):
    for j in t.vectorize_range(b.shape[0]):
       c[i, j] = a[i] + b[j]
  return c

x = tvm.placeholder((10,))
y = tvm.placeholder((20,))

# returns an extern that creates IR
# Just in time dispatch according to input types, allow:
#   - Symbolic types(tensor, var): return extern(compute op) expressions
#   - Python constant: intepret as constant in expression
#   - concrete data type(numpy array), execute the python code simply
#
z = mycross(x, y)
s = tvm.create_schedule(z.op)
print(tvm.lower(s, [x,y,z]))

x = np.zeros((10,))
y = np.zeros((20,))
print(mycross(x, y))