Support Inline Assembly Microkernel

This issue arises through discussion with several people @eqy @adityaatluri @ajtulloch @cowanmeg So far we can use IR to embed a sequence of LLVM intrinsics to generate microkernel. But this may not be ideal, in certain cases, we need to make use of hand-crafted micro-kernels, or utilize assembly generator such as Peachey to generate and embed them into the code. This issue is used to discuss possible API proposals and ways to do so.

Notably, llvm already support inline assembly, so the only gap is to add support in the IR intrinsic side. There are two possible ways to make this happen

Link to a module with assembly intrinsic

This is the easiest approach, @eqy already have some hacked up version that does this. The idea is that load the asm module as a separate LLVM module (wrapped in llvm ir), set the property so that this module must be force inlined, and load this module in code generation phase and link.

To implement this, we can specify the desired linked external module (possibly via arguments, or a thread-local scope context setting function that set the linked modules)

import tvm

ll_asm= """
asm code wrapped in llvm format, expose a function my_function
"""
# In the IR construction, add a call_extern that calls my_function.
with tvm.build_config(extern_llvm_modules=[ll_asm]):
    # build code here.    

Ideally, we want to support inline asm string, as well as a path to the asm file.

One key characteristic of this approach is that all the asm are wrapped as functions. The LLVM’s SSA property means it might be hard to do certain things, like the array of registers have to be passed by address and we need to rely on the register promotion pass. It is not sure if there are cases we could not handle

Action Items

• Modify codegen_llvm.cc to support two modes of pragma_import.
  • pragma import /path/to/file.ll
  • pragma import str_llvm
  • Related pragma handling code
  • Put things into a unordered_set<std::string>(deduplication)
  • Before optimize, import those modules as llvm module, mark them to be force inline and link together
  • Write testcases about pragma_import(schedule, and tensor intrinsic)
    • We can write sch[xx].pragma("import", "path/to/file.ll")
    • We can also directly embed things in tensor intrinsics
• Support contrib.clang.create_llvm to make it easy to generate llvm from embedded c asm.

Can be in a separate PR

• Write a tutorial on how to embedd inline asm. (might also need tensorization tutorial).

[ajtulloch]

Ah sorry @tqchen, I was out the last few weeks. This looks great, thanks so much. I’m restarted working on these tensorized implementations in https://github.com/ajtulloch/tvm/blob/43f50eba9eab45ec505e92ce372c45d91227778a/tensorize/gemm__neon.c + https://github.com/ajtulloch/tvm/blob/43f50eba9eab45ec505e92ce372c45d91227778a/tensorize/gemm__avx2.c, and there’s so pretty promising results on ARMv7 in addition to AVX2 already.

[ajtulloch]

@tqchen sure, me (or someone I loop in) will work on that part.