| 1903 | class After: |
| 1904 | @T.prim_func(private=True, s_tir=True) |
| 1905 | def fused_func( |
| 1906 | input_ids: T.Buffer((T.int64(1),), "int32"), |
| 1907 | input_embeds: T.Buffer((T.int64(4096), T.int64(4096)), "float16"), |
| 1908 | T_take: T.Buffer((T.int64(1), T.int64(4096)), "float16"), |
| 1909 | ): |
| 1910 | T.func_attr({"tirx.noalias": True}) |
| 1911 | Out_handle_intermediate = T.sblock_alloc_buffer( |
| 1912 | (T.int64(4096), T.int64(4096)), "float16" |
| 1913 | ) |
| 1914 | for i, j in T.grid(T.int64(4096), T.int64(4096)): |
| 1915 | with T.sblock("add"): |
| 1916 | vi, vj = T.axis.remap("SS", [i, j]) |
| 1917 | Out_handle_intermediate[vi, vj] = input_embeds[vi, vj] + T.float16(1) |
| 1918 | for ax0, ax1 in T.grid(T.int64(1), T.int64(4096)): |
| 1919 | with T.sblock("T_take"): |
| 1920 | v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) |
| 1921 | T_take[v_ax0, v_ax1] = Out_handle_intermediate[input_ids[v_ax0], v_ax1] |
| 1922 | |
| 1923 | @R.function |
| 1924 | def main( |