[LLVM] Codegen subroutine call when CallNode::op is GlobalVar

[Lunderberg]

Previously, the CallNode::op must be a known built-in operation. This commit allows LLVM codegen to produce a subroutine call to another function within the same IRModule, with CallNode::op specifying the GlobalVar that represents that function.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• No users to tag found in teams: llvm _{See #10317 for details}

_{Generated by tvm-bot}

[Lunderberg]

This PR is a subset of the functionality in #14862. As that PR has grown to a difficult-to-review size, I'm carving off independent chunks of it into separate PRs, after which #14862 can be rebased to a more manageable size.

[LeiWang1999]

looks like this pr introduced some rocm backend issues.

code to reproduce:

import tvm
from tvm import te
import numpy as np
import tvm.testing
from tvm.script import tir as T
from tvm.tir import TensorIntrin


M = 64
N = 64

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        T.func_attr({"global_symbol": "main"})
        A = T.match_buffer(a, (M, N), dtype="float32")
        B = T.match_buffer(b, (M, N), dtype="float32")
        for i, j in T.grid(M, N):
            with T.block("B"):
                vi, vj = T.axis.remap("SS", [i, j])
                B[vi, vj] = A[vi, vj] * 2.0


ir_module = MyModule
sch = tvm.tir.Schedule(ir_module, debug_mask="all")

block_b = sch.get_block("B")

i, j = sch.get_loops(block_b)
sch.bind(i, "blockIdx.x")
sch.bind(j, "threadIdx.x")

print(sch.mod["main"].script())

ctx = tvm.rocm(0)
with tvm.transform.PassContext():
    rocm_mod = tvm.build(sch.mod, target="rocm")

'''
Traceback (most recent call last):
  File "../../tvm_rocm/memory_copy.py", line 38, in <module>
    rocm_mod = tvm.build(sch.mod, target="rocm -mcpu=gfx90a")
  File "/home/aiscuser/v-leiwang3/tvm/python/tvm/driver/build_module.py", line 281, in build
    rt_mod_host = _driver_ffi.tir_to_runtime(annotated_mods, target_host)
  File "/home/aiscuser/v-leiwang3/tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 238, in __call__
    raise get_last_ffi_error()
tvm._ffi.base.TVMError: Traceback (most recent call last):
  7: TVMFuncCall
  6: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::runtime::Module (tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)>::AssignTypedLambda<tvm::{lambda(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)#6}>(tvm::{lambda(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target)#6}, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, tvm::runtime::TVMRetValue)
  5: tvm::TIRToRuntime(tvm::runtime::Map<tvm::Target, tvm::IRModule, void, void> const&, tvm::Target const&)
  4: tvm::codegen::Build(tvm::IRModule, tvm::Target)
  3: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::TypedPackedFunc<tvm::runtime::Module (tvm::IRModule, tvm::Target)>::AssignTypedLambda<tvm::runtime::Module (*)(tvm::IRModule, tvm::Target)>(tvm::runtime::Module (*)(tvm::IRModule, tvm::Target), std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  2: tvm::codegen::BuildAMDGPU(tvm::IRModule, tvm::Target)
  1: tvm::codegen::CodeGenLLVM::Finish()
  0: tvm::codegen::CodeGenLLVM::Verify() const
  File "/home/aiscuser/v-leiwang3/tvm/src/target/llvm/codegen_llvm.cc", line 361
TVMError: LLVM module verification failed with the following errors: 
Calling convention requires void return type
ptr @main_kernel
Function return type does not match operand type of return inst!
  ret void
 i32
'''

this can be resolved through:

// src/target/llvm/codegen_amdgpu.cc
class CodeGenAMDGPU : public CodeGenLLVM {
 public:
  CodeGenAMDGPU() = default;
  virtual ~CodeGenAMDGPU() = default;

  llvm::Function* DeclareFunction(const GlobalVar& gvar, const PrimFunc& f) {
    // amd gpu kernel function should be void return
    return this->DeclareFunctionInternal(gvar, f, true);
  }

because amd gpu kernels should be void return but by default it wll call a llvm codegen implementation which returns int by default, Previously there were no issues, until this code snippet was added.

// src/target/llvm/codegen_llvm.cc 
  if (auto it = functions_.find(gvar.get()); it != functions_.end()) {
    return it->second;
  }

and if we wanna leverage shared memory, there's another issue:

TVMError: LLVM module verification failed with the following errors: 
Global is external, but doesn't have external or weak linkage!
ptr addrspace(3) @shmem
Global is external, but doesn't have external or weak linkage!
ptr addrspace(3) @shmem.1

I still have no idea how to resolve it, please help me take a look @Lunderberg

[tqchen]

We should not support external global in the device codegen

[tqchen]

Following up on this. Device side codegen follows a different convention and likely we cannot simply have a call from the host to device via the external call. Instead, we need to do them through packed func calls (and we need to do them via PackedFunc convention)

So I wonder if we should revert some of the changes in the device codegen(nvptx, rocm) and only support subroutine rewriting in host side. cc @Lunderberg

[Lunderberg]

Agreed, the device-side codegen should not handle host-to-device calls. If any such calls exist (e.g. resulting from SplitHostDevice), they are lowered to PackedFunc API in LowerDeviceKernelLaunch.

I'm investigating the issue at the moment, and can reproduce it locally. (Thank you @LeiWang1999 for the test case!) It looks like it is a discrepancy between the PrimFunc::ret_type and the CallNode::dtype, currently checking where that discrepancy gets introduced.

[Lunderberg]

@LeiWang1999 It looks like the discrepancy is specifically between the return type of the PrimFunc (DataType::Void()) and the return type provided to LLVM-IR (i32). I can confirm that your fix (adding CodeGenAMDGPU::DeclareFunction) does avoid the issue in the test script, but I can't reproduce the shared memory issue. Is that a different test case for that failure mode?

Long-term, rather than providing the ret_void = true argument, it may be more robust to use the PrimFunc::ret_type directly. That is, because the PrimFunc::ret_type is void, then the ret_void argument duplicates information without adding expressibility. There's a TODO in the comments (from PR#5107) stating that this was an original intention of the code.

[Lunderberg]

So I wonder if we should revert some of the changes in the device codegen(nvptx, rocm) and only support subroutine rewriting in host side. cc @Lunderberg

@tqchen Forgot to respond to this part. That is the intent, that the LLVM subroutines are supported only on the host side. The primary change to nvptx/rocm were (a) to make sure they could provide the GlobalVar when delegating to CodeGenLLVM::AddFunction and (b) to call DeclareFunction prior to AddFunction. If we want to explicitly prevent TIR-to-TIR calls in nvptx/rocm, we could revert the changes of type (b) without impacting the changes of type(a). I suppose we could also move the callee resolution into CodeGenLLVM::Finish, since that would avoid the need for DeclareFunction altogether, but that would be a larger change.

[Lunderberg]

@LeiWang1999 Can you try running your test case with PR #15127 ? This is a draft PR that removes the ret_void argument altogether, instead using prim_func->ret_type. This resolved the failure mode in the test case you provided, though I haven't tested it against the shared memory issue you mentioned.

[LeiWang1999]

Thanks @Lunderberg , I will take a look later, and to reproduce the shared memory external global issue, the test case can be accessed through https://gist.github.com/LeiWang1999/d0db6e802f358a2a42fca32b5cd42441

[Lunderberg]

@LeiWang1999 Thank you for the test case. I can reproduce the shmem error, though it looks like it predates the LLVM changes in this PR. Unfortunately, that means there isn't a clear before/after that could be compared and fixed. Do you know if the shared memory on rocm has worked in the past?

FYI, I also condensed the test case down the the following test case, which reproduces the same error:

import tvm
from tvm.script import tir as T

@T.prim_func
def func():
    A = T.alloc_buffer(1, scope="shared")

    for i in T.thread_binding(1, thread="threadIdx.x"):
        A[0] = 0.0

tvm.build(func, target="rocm")

[LeiWang1999]

@Lunderberg at least from this commit 1d98634 the rocm shared memory should work.

[LeiWang1999]

@LeiWang1999 Can you try running your test case with PR #15127 ? This is a draft PR that removes the ret_void argument altogether, instead using prim_func->ret_type. This resolved the failure mode in the test case you provided, though I haven't tested it against the shared memory issue you mentioned.

I think ret void argument issue has been fixed through the related prs, thanks @Lunderberg ! but from my test the external global function still remains.

[zhangxiao-stack]

@LeiWang1999 Thank you for the test case. I can reproduce the shmem error, though it looks like it predates the LLVM changes in this PR. Unfortunately, that means there isn't a clear before/after that could be compared and fixed. Do you know if the shared memory on rocm has worked in the past?

FYI, I also condensed the test case down the the following test case, which reproduces the same error:

import tvm
from tvm.script import tir as T

@T.prim_func
def func():
    A = T.alloc_buffer(1, scope="shared")

    for i in T.thread_binding(1, thread="threadIdx.x"):
        A[0] = 0.0

tvm.build(func, target="rocm")

I am facing the same issue on the ROCm platform. Is there a solution yet?

[LeiWang1999]

@zhangxiao-stack checkout this branch https://github.com/LeiWang1999/tvm/blob/lei/feat-hip/tests/python/unittest/test_tir_schedule_tensorize_hip_mfma.py#L112 this is a alter hip source code generation, it might be one possible impermanent solution on rocm.