[Solved(?)] Calling external functions and passing arguments

[aca88]

Hello all,

Background Info

I am trying to do the following:

1. Use TVM to generate the AST of a certain tensor function (which one it is makes no difference, but as an example it is a copy operation)
2. Replace part of the generated AST nodes with calls to function which I have defined externally in Python (using the @tvm.register_func decorator)
3. Run the module, which should trigger the execution of the functions in 2.

Use TVM to generate the AST of a certain tensor function

Assume I have some tvm.compute(...) which gives the following (printed) AST

for (h.outer, 0, 160) {
  for (i1, 0, 6) {
     for (i2, 0, 642) {
        for (i3, 0, 64) {
          //the line below should be replaced
          dest_array[some_pointer_arith] = src_array[some_poiner_arith]
        }
     }
  }
}

This is the kind of AST which would be generated for a copy operation of a tensor which is is split only in the H dimension in 160 tiles.

Replace part of the generated AST nodes with calls to function which I have defined in Python

Following the Tutorial:External Tensor Functions, I also have some Python file with

@tvm.register_func("tvm.my_function")
def my_function(iteration_number):
    print(iteration_number)

So in this case, lets say I just want to print all the iteration indexes for all lines. There is no specific reason I want to do this (i.e. print). But I want to be able to define external functions which use internal iteration variables as arguments.

So I use the tvm.ir_pass.IRTransform(stmt, None, _replace_stmt,["Store"]) to find the place in the AST which I want to replace doing the following

def _replace_stmt(stmt):
    irb = tvm.ir_builder.create()
    _func_call_node = tvm.make.Evaluate(tvm.call_packed("tvm.my_function",tvm.expr.Var('h.outer','int')))
    #NOTE1: I also tried tvm.make.Evaluate(tvm.call_packed("tvm.my_function")), but this cant work since this would theoretically be a function without parameters
    irb.emit(_func_call_node)
    return irb.get()

Which should produce the following (printed) AST (it actually does if simple_mode=True)

for (h.outer, 0, 160) {
  for (i1, 0, 6) {
     for (i2, 0, 642) {
          for (i3, 0, 64) {
            tvm_call_packed("tvm.my_function", h.outer)
          }
      }
  }
}

Run the module

So at some point in my code, I have:

f = tvm.lower(some_schedule, some_place_holders, name="some_name")
m = tvm.build(f, target="llvm")
#some lines to fill the place_holders with values
m(some_tvm_nd_arrays)

While executing m(some_tvm_nd_arrays) I expect the module execution to jump into tvm.my_function(...).

Questions

1. There are a couple of “external functions” in TVM. When should we use which and am I using the right kind?
   Example: the first time I tried something similar to what I described before, I actually didnt pass arguments and used tvm.call_extern(...) instead of tvm.call_packed(...) in replace_stmt() and it didn’t work, but the VTA examples use tvm.call_extern(), but they have a .so which I dont.
   If I dont pass arguments and use tvm.call_packed(...) (i.e. _func_call_node = tvm.make.Evaluate(tvm.call_packed("tvm.my_function")) everything works (even running the module)

As described in the problem statement, the f=tvm.lower(...) line actually fails because

File "~/tvm/src/pass/make_api.cc", line 188
TVMError: Not all Vars are passed in api_args:  'h.outer' does not appeared in api_args

2. How do I pass variables which are present in the AST representation to external functions?
   I mean those which are not really tensor of the tvm.compute(...), but for example part of the iteration variables (and more specifically those after all loop splits and fusions)

What if my_function() is actually a class method:

class MyClass():
    @tvm.register_func("tvm.my_function")
    def my_function(self,iteration_number):
        print(iteration_numer) 

3. How to call it?
   Example: Note here that the self parameter is part of the argument list but it is not visible from the module’s perspective, which actually leads to some compiling error because the argument list doesnt match

4. Am I overcomplicating myself?
   Example: should I be using the hybrid script? if so how?

Thanks for the help

Small Update #1

I tried (not at the same time):

• some_place_holders.append(tvm.var(name='h.outer', dtype='int32'))
• some_place_holders.append(tvm.placeholder((1,), name='h.outer'))

before calling tvm.lower(...) but this also didnt work

Last EDIT:
I didnt want to explicitly tag anyone, but seeing as the question remains unanswered I will make one last edit with tags.
Those people tagged are the Top-5 in the user list for the quarter 28th of May - 28th of August, in an attempt to tag people who have been most active recently.

I apologize in advance:
@vinx13 @tqchen @thierry @yzhliu @FrozenGene

[aca88]

Here a minimal example of what I am trying to achieve

import numpy as np
import tvm

#Required Functions 
def __override__build_config(**kwargs):
    pass_list = []
    #Add our custom IR pass routine in 3rd stage
    pass_list.append((3,__find_store))
        
    return tvm.build_config(add_lower_pass=pass_list,
                            **kwargs)

def __find_store(stmt):
    def __inject_extern(stmt):
        #Replace the a [...] = b[...] by an external call
            irb = tvm.ir_builder.create()
            func_create_node = tvm.make.Evaluate(tvm.call_packed("tvm.my_func",tvm.expr.Var('i1','int32')))
            irb.emit(func_create_node)
            return irb.get()
    #Find the store nodes and define the postorder function
    stmt_out = tvm.ir_pass.IRTransform(stmt, None, __inject_extern,["Store"])
    return stmt_out

@tvm.register_func("tvm.my_func")
def __my_external_func(iteration_idx):
    #External Pyton function with hook
    print(iteration_idx)


#Start of Code
input_shape = (3,3,3)

#Define Placeholders
input_ph = tvm.placeholder(input_shape,name='Input')
output_ph = tvm.placeholder(input_shape,name='Output')
#Define the Compute Operation
output_t = tvm.compute(input_shape,lambda i1, i2, i3: input_ph[i1,i2,i3], name = 'CopyOperation')
#Create the schedule
sched = tvm.create_schedule(output_t.op)
#Use the overwriten build_config()
with __override__build_config():
    #The output of this following print is what I would expect the code to look like
    print('Expected Output Code')
    print(tvm.lower(sched,[input_ph,output_ph], simple_mode=True))
    
    #Lower & Build
    f = tvm.lower(sched,[input_ph,output_ph], name='Copy') 
    #The above line throws TVMError: Not all Vars are passed in api_args:  'i1'  does not appeared in api_args
    # doing  f = tvm.lower(sched,[input_ph,output_ph, tvm.var(name='i1')], name='Copy')  does not help
    target = "llvm"
    dtype = "float32"
    m = tvm.build(f, target=target)
    assert m
    ctx = tvm.context(target,0)
    #Create input and output ND arrays
    input_np = np.random.rand(3,3,3).astype(dtype)
    input_nd = tvm.nd.array(input_np, ctx)
    output_nd = tvm.nd.array(np.zeros((3,3,3),dtype=dtype),ctx)
    #Call the Module
    m(input_nd,output_nd)

Why do I think something like this is important?

Well imagine I have a longer schedule, with some internal tensor internal_tensor=tvm.compute(...). By internal tensor, I mean one which is not part of the placeholder list (i.e. I allow TVM to generate it’s size and handle the allocation).
Now imagine that I would like to insert a certain fault at specific indexes of internal_tensor for a specific sched = tvm.create_schedule(...).
How would I do this?

[aca88]

Although I don’t think I solved the general problem (i.e. passing any kind of variables to an external function), I do think I have an “OK” solution for my problem.
TLDR of my problem: Wanted to call external Python declared functions with iteration variables as arguments.

My Solution:

import numpy as np
import tvm

#Required Functions 
@tvm.register_func("tvm.my_func")
def __my_external_func(it_idx_0):
    #External Pyton function with hook
    print(it_idx_0)

#Start of Code
#This is just one toy example case so that we can tile in the 0th dimension
input_shape = (4,3,4)

#Define Placeholders
input_ph = tvm.placeholder(input_shape,name='Input')
output_ph = tvm.placeholder(input_shape,name='Output')

#Define the Compute Operation
output_t = tvm.compute(input_shape,lambda  i1, i2,i3: tvm.call_packed("tvm.my_func",i1),name="extern")
#The above line is the most important for the solution

#Create the schedule
sched = tvm.create_schedule(output_t.op)
#Tile the function's parameter axis
sched[output_t].split(sched[output_t].op.axis[0],factor=2)

#with __override__build_config(): Commented out since now we dont need a customized IR_pass
print(tvm.lower(sched,[input_ph,output_ph], simple_mode=True))
#Lower & Build
f = tvm.lower(sched,[input_ph,output_ph], name='Copy')
target = "llvm"
dtype = "float32"
m = tvm.build(f, target=target)
assert m
ctx = tvm.context(target,0)
#Create input and output ND arrays
input_np = np.random.rand(4,3,4).astype(dtype)
input_nd = tvm.nd.array(input_np, ctx)
output_nd = tvm.nd.array(np.ones((4,3,4),dtype=dtype),ctx)
#Call the Module
m(input_nd,output_nd)
print(output_nd)

Interestingly enough the (printed) AST looks as follows:

// attr [extern] storage_scope = "global"
allocate extern[int32 * 48]
produce extern {
  for (i1.outer, 0, 2) {
    for (i1.inner, 0, 2) {
      for (i2, 0, 3) {
        for (i3, 0, 4) {
          extern[((((((i1.outer*2) + i1.inner)*3) + i2)*4) + i3)] = tvm_call_packed("tvm.my_func", ((i1.outer*2) + i1.inner))
        }
      }
    }
  }
}

Which nicely shows splitting effects on the wanted axis.
In addition, a customized ir_pass can be declared to clean the arguments if one does not care for either .outer or .inner variables.

So in essence:

2. How do I pass variables which are present in the AST representation to external functions?
   I mean those which are not really tensor of the tvm.compute(...) , but for example part of the iteration variables (and more specifically those after all loop splits and fusions)

and

4. Am I overcomplicating myself?

Have been partially answered.