Tensorize in multithreaded environment

taknevski · August 16, 2018, 9:25pm

Hi,
I have a small test case where I do a 2-d matrix addition.
A = tvm.placeholder((4,16), name=‘input1’)
B = tvm.placeholder((4,16), name=‘input2’)
C = tvm.compute((4,16),
lambda m,n: A[m,n] + B[m,n],
name=‘output’)

I tile on the leading dimension by 2 and parallelize as follows:

m,n = s[C].op.axis

m_outer,m_inner = s[C].split(m,factor=2)
s[C].reorder(m_outer,m_inner,n)
s[func].parallel(m_outer)

I then match the inner matrix add with a tensorize routine:
And I add my own c function that does the matrix add as a packed func.
I print the address of the pointer passed into the routine. When there is a single thread(TVM_NUM_THREADS=1) the pointers are updated correctly prior to being passed to the function
But when there are 2 threads are more, the input pointer is not offset for threads other than the first one.

taknevski · August 16, 2018, 9:30pm

Please look here for the sample code
https://github.com/dmlc/tvm/files/2295210/test.tar.gz

tqchen · August 20, 2018, 5:16pm

If we need the input of tensorization to have non-zero offset, we might need to redefine the buffer abstraction(by having buffer_alignment)

taknevski · August 20, 2018, 5:33pm

Hi Tianqi,

Thanks for the response, is it possible to pass in the loop index variables as arguments to the tensorized function?

taknevski · August 20, 2018, 5:54pm

So in the output of tvm.lower after tensorize:

After tensorize
produce output {
parallel (m.outer, 0, 2) {
tvm_call_packed(“test_intrin”, tvm_address_of(input1[(m.outer*32)]), tvm_address_of(input2[(m.outer32)]), tvm_address_of(output[(m.outer32)]))
}
}

We see that the tvm_address_of is pointing to the right offset in the buffer. Can the implementation of tvm_address_of be modifed to yield the right offset, taking into account parallelized loops?

tqchen · August 20, 2018, 7:07pm

OK, this might have something to do with a bug in TVM code generation when supporting packed function calling in the parallel body, Can you try swap call_packed by call_extern for now, and directly provide an extern “C” function with the same signature? This will get around this

taknevski · August 20, 2018, 8:53pm

Hi,

Can you give an example of how to use call_extern please? I’m passing the output in a s a pointer to the external function so I have a void return type, is this supported?

tqchen · August 20, 2018, 10:22pm

You can simply use int as return type and not use the values for void functions, and it should work

tqchen · August 20, 2018, 10:23pm

see https://github.com/dmlc/tvm/blob/master/python/tvm/intrin.py#L134

taknevski · August 21, 2018, 1:50am

Some qns:
Do you use TVM_REGISTER_GLOBAL to register the extern function?
How do you export the extern function ,by loading a .so?
do you need get_global_func?
are the arguments to the function C++ datatypes or TVMArgValues as for packed func?

tqchen · August 21, 2018, 6:15am

For normal extern function, you don’t need to register them, just make sure they are exported as extern C in the runtime is fine. You cannot get them through tvm.get_global as they are typed c functions

taknevski · August 21, 2018, 2:22pm

I get the following error:

LLVM ERROR: Program used external function ‘test_intrin’ which could not be resolved!

I declare the function as extern “C” in the .so file.

f = tvm.module.load("./test.so")

def intrin_func(ins, outs):
body = tvm.call_extern (“int32”,“test_intrin”,ins[0].access_ptr(“rw”),ins[1].access_ptr(“rw”),outs[0].access_ptr(“rw”))

and here is my function

extern “C” int test_intrin(void * input, void * input2, void * output){
printf(“input %.8X\n”, input);
};

tqchen · August 21, 2018, 9:28pm

github.com/apache/tvm

[CODEGEN] Enable inline llvm asm code

master ← tqchen:master

opened 05:10AM - 25 Jul 18 UTC

tqchen

+243 -0

Ref #1276 ## Usage We can have a pragma attribute pragma_import_llvm, which… can be used in either ir builder or schedule ## Example ```python def test_llvm_import(): # extern "C" is necessary to get the correct signature # we can put inline asm here cc_code = """ extern "C" float my_add(float x, float y) { return x + y; } """ n = 10 A = tvm.placeholder((n,), name='A') B = tvm.compute((n,), lambda *i: tvm.call_pure_extern("float32", "my_add", A(*i), 1.0), name='B') def check_llvm(use_file): if not tvm.module.enabled("llvm"): return if not clang.find_clang(required=False): print("skip because clang is not available") return temp = util.tempdir() ll_path = temp.relpath("temp.ll") # Create LLVM ir from c source code ll_code = clang.create_llvm(cc_code, output=ll_path) s = tvm.create_schedule(B.op) if use_file: s[B].pragma(s[B].op.axis[0], "import_llvm", ll_path) else: s[B].pragma(s[B].op.axis[0], "import_llvm", ll_code) # BUILD and invoke the kernel. f = tvm.build(s, [A, B], "llvm") ctx = tvm.cpu(0) # launch the kernel. a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), ctx) b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), ctx) f(a, b) np.testing.assert_allclose( b.asnumpy(), a.asnumpy() + 1.0) check_llvm(use_file=True) ```

taknevski · October 15, 2018, 10:36am

Hi,

Thank you for the response. The issue I’m facing is that I want to use call_extern to invoke a function in a dynamic .so library which I link by using tvm.module.load.

In your approach you implement the extern function as llvm inline assembly, but for my case the function is implemented in the dynamic .so library and I want to call that… Is there a way to do this?

Thanks,
Anand

tqchen · October 15, 2018, 4:01pm

You can, as long as the .so library is loaded prior to the tvm.module.load in the tvm runtime

taknevski · October 15, 2018, 11:38pm

Hi,

The .so library is an external library which I have compiled using a c++ compiler. Now I load it using tvm.module.load to expose it to tvm runtime. But the external function implemented in the library is not resolved as indicated by the error I reported earlier.

Thanks,
Anand

tqchen · November 22, 2018, 12:55am

Please check out the example here https://github.com/dmlc/tvm/pull/2156, as I said in my last post, it is fine as long as .so library is loaded in the tvm runtime.

However, we do need to load the so as RTLD_GLOBAL object, so it is visible to other so library. Or you can re-expose it in the tvm runtime and have tvm runtime link against your dll

taknevski · November 22, 2018, 2:41am

Thanks, this looks like what I need, I’ll confirm if it worked for me.