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

Questions
3

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:

and

Have been partially answered.