main

GGitHub[BACKEND] Allow backend to specify special rules for membar insertion (#4675)

8a3fb7e3创建于 2024年9月9日历史提交

文件	最后提交记录	最后更新时间
CMakeLists.txt	[BUILD][FRONTEND] working 3P backend (#2896) AMD is enabled by default, but not ripe for usage (not tested). Lots of work will be necessary to make everything robust and maintainable.	2 年前
TestAlias.cpp	[BACKEND] Refactor shared memory representation in TTGIR (#3261) Existing shared memory representation in TTGIR had some semantic problems as it was mixing value semantic with memory semantic. In order to solve that this PR moves shared memory representation to memory semantic only. This means shared memory is now not represented as tensors but as allocations that may or may not be mutable. tensor cannot have shared encoding now. Convert_layout #distributed -> #shared becomes triton_gpu.local_alloc %init tensor<#distributed> -> !tt.memdesc<#shared> Convert_layout #shared -> #distributed becomes triton_gpu.local_load %mem !tt.memdesc<#shared> -> tensor<#distributed> Insert_slice_async becomes async_copy_global_to_local	2 年前
TestAllocation.cpp	[FRONTEND][BACKEND] Add the noinline annotation for triton.jit (#1568) # Introducing the noinline Parameter for Triton JIT Decorator We're excited to introduce a new parameter, noinline, that can be added to the jit decorator in Triton. This parameter allows developers to specify that a particular Triton function should not be inlined into its callers. In this post, we'll dive into the syntax, purpose, and implementation details of this new feature. ## Syntax To use the noinline parameter, simply add noinline=True to the jit decorator for the function that you don't want to be inlined. Here's an example: ```python @triton.jit(noinline=True) def device_fn(x, y, Z): z = x + y tl.store(Z, z) def test_noinline(): @triton.jit def kernel(X, Y, Z): x = tl.load(X) y = tl.load(Y) device_fn(x, y, Z) ``` In this example, the device_fn function is decorated with @triton.jit(noinline=True), indicating that it should not be inlined into its caller, kernel. ## Purpose The noinline parameter serves several key purposes: - Reducing code size: By preventing inlining, we can reduce the size of the compiled code. - Facilitating debugging: Keeping functions separate can make it easier to debug the code. - Avoiding common subexpression elimination (CSE) in certain cases: CSE can sometimes be avoided by using the noinline parameter to reduce register pressure. - Enabling dynamic linking: This parameter makes it possible to dynamically link Triton functions. ## Implementation The implementation of the noinline parameter involves significant changes to three analysis modules in Triton: *Allocation*, *Membar*, and *AxisInfo*. Prior to this update, these modules assumed that all Triton functions had been inlined into the root kernel function. With the introduction of non-inlined functions, we've had to rework these assumptions and make corresponding changes to the analyses. ### Call Graph and Limitations <div style="text-align: center;"> <img src="https://user-images.githubusercontent.com/2306281/234663904-12864247-3412-4405-987b-6991cdf053bb.png" alt="figure 1" width="200" height="auto"> </div> To address the changes, we build a call graph and perform all the analyses on the call graph instead of a single function. The call graph is constructed by traversing the call edges and storing them in an edge map. Roots are extracted by checking nodes with no incoming edges. The call graph has certain limitations: - It does not support recursive function calls, although this could be implemented in the future. - It does not support dynamic function calls, where the function name is unknown at compilation time. ### Allocation <div style="text-align: center;"> <img src="https://user-images.githubusercontent.com/2306281/234665110-bf6a2660-06fb-4648-85dc-16429439e72d.png" alt="figure 2" width="400" height="auto"> </div> In Triton, shared memory allocation is achieved through two operations: triton_gpu.convert_layout and triton_gpu.alloc_tensor. The convert_layout operation allocates an internal tensor, which we refer to as a *scratch* buffer, while the alloc_tensor operation returns an allocated tensor and is thus known as an *explicit* buffer. To accommodate the introduction of function calls, we are introducing a third type of buffer called a *virtual* buffer. Similar to scratch buffers, virtual buffers are allocated internally within the scope of a function call, and the buffers allocated by the called functions remain invisible to subsequent operations in the calling function. However, virtual buffers are distinct from scratch buffers in that the call operation itself does not allocate memory—instead, it specifies the total amount of memory required by all the child functions being called. The actual allocation of buffers is performed by individual operations within these child functions. For example, when invoking edge e1, no memory is allocated, but the total amount of memory needed by function B is reserved. Notably, the amount of shared memory used by function B remains fixed across its call sites due to the consideration of dynamic control flows within each function. An additional challenge to address is the calculation of shared memory offsets for functions within a call graph. While we can assume a shared memory offset starting at 0 for a single root function, this is not the case with a call graph, where we must determine each function's starting offset based on the call path. Although each function has a fixed memory consumption, the starting offset may vary. For instance, in Figure 2, the starting offset of function C through edges e1->e2 differs from that through edges e2->e4. To handle this, we accumulate the starting offset at each call site and pass it as an argument to the called function. Additionally, we amend both the function declaration and call sites by appending an offset variable. ### Membar <div style="text-align: center;"> <img src="https://user-images.githubusercontent.com/2306281/234665157-844dd66f-5028-4ef3-bca2-4ca74b8f969d.png" alt="figure 3" width="300" height="auto"> </div> The membar pass is dependent on the allocation analysis. Once the offset and size of each buffer are known, we conduct a post-order traversal of the call graph and analyze each function on an individual basis. Unlike previous analyses, we now return buffers that remain unsynchronized at the end of functions, allowing the calling function to perform synchronization in cases of overlap. ### AxisInfo <div style="text-align: center;"> <img src="https://user-images.githubusercontent.com/2306281/234665183-790a11ac-0ba1-47e1-98b1-e356220405a3.png" alt="figure 4" width="400" height="auto"> </div> The AxisInfo analysis operates differently from both membar and allocation, as it traverses the call graph in topological order. This is necessary because function arguments may contain axis information that will be utilized by callee functions. As we do not implement optimizations like function cloning, each function has a single code base, and the axis information for an argument is determined as a conservative result of all axis information passed by the calling functions. --------- Co-authored-by: Philippe Tillet <phil@openai.com>	3 年前
TestAxisInfo.cpp	[Backend] Clean up AxisInfo (NFC). (#3240) [Backend] Clean up AxisInfo (NFC). - Remove unused member variables. - Clarify and reformat comments. - Move code out of header. - Add namespaces.	2 年前
TestMembar.cpp	[BACKEND] Allow backend to specify special rules for membar insertion (#4675) With block level kind of operations like TMA it is possible that some ops access the shared memory but don't require barriers. This adds a lambda that backends can pass to explicitly skip barriers in between some ops.	1 年前