mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
Lei Wang
0f9bbd759c
* [Enhancement] Improve layout inference in ParallelOpNode * Enhanced the layout inference mechanism in ParallelOpNode to utilize two strategies: compute_loop_layout_from_buffer and PlanLoopPartition, selecting the one that minimizes replication while ensuring compatibility. * Updated the logic to choose the best candidate layout based on replication size and containment checks. * Refactored the HasKnownLayoutAnchor function to clarify its purpose in prioritizing buffer layouts. * Added a new test case to validate the layout inference behavior, ensuring the correct fragments are generated in the output. This update aims to optimize layout inference for parallel operations, improving performance and resource utilization. * lint fix * bug fix and refactor * lint fix
37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def _tilelang_issue_layout_free_inference_choose_smallest_replication():
|
|
@T.prim_func
|
|
def main(A: T.Tensor((128, 4), T.float), B: T.Tensor((128, 4), T.float)):
|
|
with T.Kernel(1, threads=128) as _:
|
|
A_frag = T.alloc_fragment((128, 4), T.float)
|
|
B_frag = T.alloc_fragment((128, 4), T.float)
|
|
S_frag = T.alloc_fragment((4,), T.float)
|
|
T.annotate_layout(
|
|
{
|
|
A_frag: T.Fragment(A_frag.shape, lambda i, j: (i, j)),
|
|
}
|
|
)
|
|
for i, j in T.Parallel(128, 4):
|
|
A_frag[i, j] = S_frag[j]
|
|
for i, j in T.Parallel(128, 4):
|
|
B_frag[i, j] = S_frag[j]
|
|
|
|
return main
|
|
|
|
|
|
def test_tilelang_issue_layout_free_inference_choose_smallest_replication():
|
|
kernel = _tilelang_issue_layout_free_inference_choose_smallest_replication()
|
|
source = kernel.get_kernel_source()
|
|
assert "float S_frag[4];" in source, "S_frag is not in the source"
|
|
assert "float B_frag[4];" in source, "B_frag is not in the source"
|
|
assert "float A_frag[4];" in source, "A_frag is not in the source"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|