Why is PrimFunc.buffer_map cleared in MakePackedAPI

Development
#1

There is a benefit of having the association between a Var and the corresponding Buffer throughout the optimizations, including codegen. Buffer is not a PrimExpr, so it cannot be passed around directly, but it contains a lot of useful information (such as shape etc.). If we had the buffer_var -> Buffer mapping, we could rely on that to associate any Var with its buffer (if there is any).

However, the buffer_map is cleared in MakePackedAPI, and the LLVM codegen expects it to be empty. What are the reasons for not having it? Can we keep it?

#2

The main reason is that buffer map represented a specific semantics of variable defs(eg the vars in the buffer).

MakePackedAPI desugars that into a low level form that uses load and intrinsics to define those vars. As a result, the buffer map information can now be stale(or duplicated with the desugared impl).

I would try to defer MakePackedAPI as late as possible and do most of the transformations before that, which allows us to take benefit of the high level mapping info.

#3

The problem is that MakePackedAPI doesn’t simply change the form of the information (e.g. from high level to low level), it just uses what it needs at the moment, and drop the rest. After MakePackedAPI there is no way to synthesize the buffer information back. For example, the shape of a buffer may only be used in assertions, but if assertions are turned off, it will be lost. Even if assertions are present, analyzing them to extract the buffer shape information seems like a fragile (and complicated) approach.

I think we should consider having some mechanism to preserve this information. The absence of this data should not affect correctness of the generated code, but it can have a great effect on the performance. In that sense it fits the “optional metadata” definition. Maybe we could have attributes for that?

#4

I agree that such information is useful, the main Q though is shall we consider defer MakePackedAPI even further if we need to perform the additional transformations that would require the buffer info?

One argument could be we hold until codegen, and I wonder if it makes sense to do argument binding generation as a util and let codegen invoke them, so we do not need to call MakePackedAPI.

1 Like
#5

I think that having arg binding as a utility would be great.

The use case for codegen having shape information (that I had in mind) was a case where codegen itself tries to vectorize scalar code. In some situations generating a vector load could access more memory than the original scalar code (which would normally be a thing to avoid), but with the shape information we’d not only know if it’s a legal thing to do (i.e. it won’t introduce an exception), but we’d also know what data was loaded. For example, if the buffer is narrow, and it’s accessed row by row, a single vector load may load two rows at the time, and we’d only know that if we know the buffer’s shapes/strides.

#6

We already have tvm/arg_binder.h at main · apache/tvm · GitHub but i guess we could go one step further to introduce something like PackedABIBinder that pragmatically generate the fragments, feel free to take a stab at this.