[Bug] BuiltinLower has hard-coded heuristic for alloca which is not appropriate for all kDLCPU target devices

[mbs-octoml]

tvm/src/tir/transforms/lower_tvm_builtin.cc

Line 116 in 2aebd33

if (device_type_.defined()) {

This particular code has a complex history but it's purpose is to avoid pool allocations for small tensors on CPU targets (a big perf win). However, not all devices of type 'CPU' support alloca since the stack may not be hardware addressable. In particular, the EthosU target requires all mem ops to be rewritten to the pool allocation primitives for later handling.

So the logic which fires this early exit heuristic (so that a later pass will rewrite allocs to alloca) is too hard coded.

Any easy fix would be to make the kMaxStackAlloca threshold zero on target which do not support alloca. That can be done by moving the value to a TargetKind attribute.

[tqchen]

@mbs-octoml To give a bit of context

In the context of CPU, we want to preserve small alloca until the code generation point. And then the code will generate the stack alloca in an explicit way. Only when memory is big enough(bigger than a constant), we will use an opaque allocation instead.

Stack allocation is important for the performance of the CPU code. In the case of TVM, we do not have explicit concept of registers in most cases. Instead we need to rely on LLVM's mem2reg pass to transform a set of constant indexing into stack allocation and turn them into registers, so the code can run effectively. So removing this code path can complicate the code generator side optimization by quite a bit and slow down the CPU code.

Of course this can be a target specific thing. LowerTVMBuiltin right now has the assumption to only run on host(CPU) code.

• Allocate always prefers (native) stack allocation when possible, but also allows other means of opaque allocation(as long as the allocation is fulfilled)
• There are however, cases when stack allocation is not possible
  • When the size of memory requested is too big, stack alloca will explode the stack space(That is why there is a size check in the CPU case and the use of global opaque was meant as a fallback to avoid stackoverflow in models with big intermediate temp space)
  • LowerTVMBuiltin was originally designed to run on the host side, which means as soon as the allocation is about device side memory, it will need to call onto a (host side) device API to allocate the memory instead

So rationales for the specific CPU side logic:

• We want to have stack alloca on host when possible(to gain mem2reg optimization)
• When the requested size is too large, we fallback to opaque workspace allocation on heap to allow the code to safely handle code with big temp memory requests as well as dynamic size allocation requests.

My guess is we need to look into why VM cannot work with code that allocates on stack in the multiple target case

[mbs-octoml]

Thanks so much for for the context. I'll try to capture that in a comment.
This is a 'bug' only in the sense that this heuristic is not working for the EthosU AOT codegen, I think because they are expecting to intercept the workpool ops downstream? But it does suggest a very simple way forward: make kMaxStackAlloca a TargetKind attribute so they can force it to zero.
@manupa-arm can you chime in here?

[tqchen]

Right, this the gets to the target dependent generation regime where TargetKind attribute is indeed the right solution. We should also send a PR to add comments to that code block so we could have more context in the future.

[manupak]

@tqchen @mbs-octoml ,

This is not specific to Arm(R) Ethos(TM)-U codegen and its generally applicable for any micro controller where we would want to avoid creating allocation of memories in the stack but rather service them via platform abstraction that is handled via TVMBackendWorkspaceAlloc --> TVMPlatformAllocate.

This only showed up in Arm(R) Ethos(TM)-U codegen because we use TVMPlatformAllocate allocate memory from a buffer placed in a memory that is both accessible by CPU and NPU. Thus, it makes this a functional bug.
However, with this change current main produces code that have much higher stack allocation for micros -- that is not desired.

cc : @u99127 @areusch

Stack allocation is important for the performance of the CPU code. In the case of TVM, we do not have explicit concept of registers in most cases. Instead we need to rely on LLVM's mem2reg pass to transform a set of constant indexing into stack allocation and turn them into registers, so the code can run effectively. So removing this code path can complicate the code generator side optimization by quite a bit and slow down the CPU code.

The correct way represent this seems to be using tir.allocates with storage_scope="local" for device=CPU to go into the stack. For targets that needs this behavior, there should be an explicit pass to convert them to local to make them placed to the stack rather than assuming this default behaviour.

Of course this can be a target specific thing. LowerTVMBuiltin right now has the assumption to only run on host(CPU) code.

Allocate always prefers (native) stack allocation when possible, but also allows other means of opaque allocation(as long as the allocation is fulfilled)
There are however, cases when stack allocation is not possible
When the size of memory requested is too big, stack alloca will explode the stack space(That is why there is a size check in the CPU case and the use of global opaque was meant as a fallback to avoid stackoverflow in models with big intermediate temp space)
LowerTVMBuiltin was originally designed to run on the host side, which means as soon as the allocation is about device side memory, it will need to call onto a (host side) device API to allocate the memory instead

Clearly, this definition of CPU leaves out micros.
It feels wrong to print out allocates with "global" storage_scope directly into CPU PrimFunc that gets printed as a stack allocation rather it should be serviced via TVMBAW call which moves the responsibility for the runtime/application layer.

So rationales for the specific CPU side logic:

We want to have stack alloca on host when possible(to gain mem2reg optimization)
When the requested size is too large, we fallback to opaque workspace allocation on heap to allow the code to safely handle code with big temp memory requests as well as dynamic size allocation requests.

This certainly sounds like we could use an optimization pass to convert the tir.allocate's storage_scope for targets that requires this rather than making that the default behaviour for tir.allocates with "global" storage scope.

cc : @tom-gall @mbaret @Mousius

[manupak]

Our proposal is to add a check to that loop whether it has 'local' storage_scope before we place them into the stack as it is the solution that works for the wider definition of the CPU, rather than performing a hidden optimization in the codegen that is applicable for a subset of CPUs.

[tqchen]

@manupa-arm in cpu we do not necessarily differentiate local from global for now as they are from the same namespace.

I can understand the need from the micro side, and I believe this can be resolved by making TVMLowerBuiltin target dependent, and query the max stack alloca property from the target. We can set the value to be zero for micro related targets.

[manupak]

I feel that is a workaround for an optimization that certain set of CPUs require it.

Moreover, TVMBAWs allows an abstraction to use an Arena that could be used in any memory for non-micro cases. Therefore it feels more like the general solution we want unless a specific target/option requires it placed on the stack.

Landing the above change that provides an optimization for certain CPUs in certain cases where creating functional problems for other backend does not seem like the right thing to do. I'd rather revert the change for now and land the correct one that works for the general case.

[tqchen]

Thanks for the discussions. Before we suggest a resolution, it would be helpful
to summarize the discussions so far.

Semantics of Allocate and storage_scope

Allocate at the moment does not specify the way to allocate the memory.
The lowering phase can feel free to pick a runtime dependent one or a
more platform native one(e.g. stack alloca or direct declaration of shared memory).

In practie, we see two different kinds of needs:

• N0: The need for optimizing "device" level code.
• N1: The need to be able to allocate memory from the host side, possibly
  in a runtime dependent way.

N0 is needed to get best performing kernel, since a native way of allocation
that can be understood by the compiler would give the best chance for followup
optimizations. This is the case for CPU related optimizations. Note that such
optimization is also needed for micro setting, when we use TIR to generate kernel
code that requested a temp scracth memory for output tiling.

N1 is needed for "driving" the set of computations. In the case of GPU, we have
to rely on host side driver function. In the case of micro AOT, it is desirable
to allocate using platform specific memory.

One possible approach is to try to ask user to differentiate these two kinds of
allocations. However, to give user maximum convenience, TVM allows a mixed host/device
programming model like CUDA and hide that complexity from the user. The boundary
of N0 and N1 can also be blurred, e.g. a kernel can request to allocate from platform
dependent memory.

In specific specialized target devices, we also see an additional need:

• N2: Allocating memory with special semantics.

For example, an unified device pinned memory that is accessible from both NPU and CPU.
A specialized texture memory or shared memory. The request in N2 is quite different
and brings additional requirement to how we allocate the memory and how the memory can be
used and represented in codegen.

Right now N2 can be covered by introducing a special memory tag -- "global.npu_unified".
This will usually allow the TVM to know that the memory needed to be lowered by possibly
special target dependent passes, and leave them out. It is important to do this because
the compiler makes no assumption that "global" can be accessed by other types of devices.

Possible Resolutions

Based on discussions so far, there are two suggested resolutions.

• R0: Separate out a "local" scope that carries the stack allocation heavior. (proposed by @manupa-arm )
• R1: Keep "global" scope as it is, introduce a special tagged scope "global.workspace"
  that represents a global memory request specifically for workspace memory.
  And introduce lowerings for them. For specialized memory(e.g. NPU unified), introduce
  separate memory scopes.
• R3: Introduce target specific attribute that marks the possible stack alloca size for
  lowering the R0("local") or R1("global"). Note R3 can be done in addition to R0 or R1.

[tqchen]

So in the above post I tried to summarize the state. Now let me try to share some of my thoughts based on the summary.

First of all, R0 and R1 are not that different in nature. Both tries to introduce two separate scopes that brings different behavior. The main questions boils down to how can we name the "global" scope.

Per allocate semantics, we treats "global" as normal CPU memory which can come from stack or platform specific allocation. The system can choose the best way of doing such lowering. Always lowering to TBAW is indeed more general for the need of N1. However, the need N0 would favor stack allocation when possible. Note that we will likely need a related behavior for micro devices as well when generating operator kernels.

While it is OK to differentiate stack allocated memory from a platform specific one, doing so would bring additional burden to the user and would require significant refactor of the operator implementations.

The main requests so far comes from need of N1. In that case, it would be easy for AOT generator to allocate memory with special tags("global.workspace"), that enforces workspace allocation since in this setting there is a single expected behavior.

So my suggestion would be R1+R2, as it helps to resolve the need in a way that is compatible with the current semantics and usecases. It will also open doors for more future scope dependent optimizations

[manupak]

Thanks @tqchen for summarizing the ideas and presenting possible resolutions.

The two needs seems very valid.

For N0, The operators should really be tagged with 'local' storage scope for the needs of N0 as they are quite local to the operator primfunc and they benefit of further optimizations within and beyond TVM -- i.e. follow up C compiler / LLVM.

For N1, we could use the 'global' tag to give the responsibility for the application/runtime layer to service the allocation.

Therefore, the actual fix should have been tagging the allocates that are expected to be optimized to be 'local' to the PrimFunc, rather than making the 'global' allocates to CPU being treated as local.

N0 is needed to get best performing kernel, since a native way of allocation
that can be understood by the compiler would give the best chance for followup
optimizations. This is the case for CPU related optimizations. Note that such
optimization is also needed for micro setting, when we use TIR to generate kernel
code that requested a temp scracth memory for output tiling.

I feel we are incorrectly tagging the storage_scopes here. They should really be 'local' for this specific usecase.

First of all, R0 and R1 are not that different in nature. Both tries to introduce two separate scopes that brings different behavior. The main questions boils down to how can we name the "global" scope.

In the solution R1, I still that as a workaround for incorrect treatment of 'global' scoped memories where we create an override of an actual 'global' what we declare as 'global.workspace'. In shared memory SoCs, it would be un-scalable explosion of tags if we want to keep tagging memories for devices which have access to them. I would think we'd want to treat memories to devices having a many-to-many relationship.

The lowering we had until few days back was general enough so they were serviced by a runtime/application layer routine and that were more aligned with what we call as 'global' (with respect to codegen) scoped storage.

Per allocate semantics, we treats "global" as normal CPU memory which can come from stack or platform specific allocation

Can you explain what you define as 'normal CPU memory' ? A CPU can technically have access to many memories.

However, the need N0 would favor stack allocation when possible. Note that we will likely need a related behavior for micro devices as well when generating operator kernels.

It would have been nice to have a RFC (Apologize in advance if I missed this if there was a one already) to discuss before we move from TVMBAW style allocation which I find more generic than just stack allocations. It almost feel the schedules should have tagged them 'local' if this was the expectation rather than assuming a combined logic : 'global' and CPU.

One possible approach is to try to ask user to differentiate these two kinds of allocations.

Wouldn't it be simpler if tag allocations for N0 to be 'local' and N1 to be 'global' ?

N2: Allocating memory with special semantics. For example, an unified device pinned memory that is accessible from both NPU and CPU. A specialized texture memory or shared memory. The request in N2 is quite different and brings additional requirement to how we allocate the memory and how the memory can be used and represented in codegen.

Its a memory where multiple Target have access to which the runtime/application could provide via TVMBAW with a specialized global.<pool_name>.

It is important to do this because the compiler makes no assumption that "global" can be accessed by other types of devices.

Hmmmm, this argument seems counter-intuitive to me. I think we should assume the 'global' to be accessible unless they are explicitly specified to be restricted. i.e. global.<pool_name>. Otherwise, the terminology is confusing.

R0: Separate out a "local" scope that carries the stack allocation heavior. (proposed by @manupa-arm )
R1: Keep "global" scope as it is, introduce a special tagged scope "global.workspace"
that represents a global memory request specifically for workspace memory.
And introduce lowerings for them. For specialized memory(e.g. NPU unified), introduce
separate memory scopes.
R3: Introduce target specific attribute that marks the possible stack alloca size for
lowering the R0("local") or R1("global"). Note R3 can be done in addition to R0 or R1.

Ideally, the allocates destined to be end up in stack should have been 'local'. Moreover, at which point we can decide not to tag allocates that exceed the target specific attribute rather than dealing this in the codegen or lower_builtin_tvm pass.

I would propose :

R4 : R0 + I think its cleaner if we just introduce a optional pass to tag memories with 'local'. At which point, we should only tag them according to the target specific attribute max stack alloca size -- that could work until we fix the schedules that wants them in stack to have 'local' storage scope for the need N0 -- that is a conscious decision the schedule writer / autoscheduler takes.

[tqchen]

Thanks @manupa-arm . Trying to capture some of the discussions.

• Right now the "global" scope translate to something that can be accessed by CPU, and there was no requirement of allocating from TVMBAW. This is consistent with allocate node for other kinds of scopes as well, such as shared memory in a GPU.
• I can understand the general proposal R4: "global" to "local" would introduce a semantic change of the "global" scope, and can cause backward compatibility issues.

The main reason I would suggest R1 is because this is a backward compatible change(no changes to topi and other libraries needed). Additionally, it reduces the users' mental cost overall to think and choose another kind of memory.
While I can totally get the other type of interpretation as well(by defining "local" to be stack allocation and "global" to be absolutely not), it is not the current intended semantics.

Note that introducing an explicit tag for TVMBAW via "global.workspace" won't necessarily mean we need to introduce tags for all kinds of memory. It does places constraints on how this kind of memory can be allocated. So things can work out of box

[manupak]

Thanks @tqchen .

In the R4, I was not suggesting to change TOPI but saying we could just do a Pass to change the storage scope. What are your thoughts about making the pass to make the storage_scope 'global.stack' after the pass? This should be backward compatible as it just performs an explicit specialization transparently in the IR. We could even do the max alloca size check in this particular pass.

This makes it much clear and specific. Moreover, global.stack will not require further specialization unlike global.workspace where we would want to choose between multiple workspace buffers. The reason being I'd like to keep the tag to denote between global memories the compute devices have access to.