(a: T.handle, c: T.handle)
| 857 | |
| 858 | @T.prim_func(s_tir=True) |
| 859 | def wmma_load_impl(a: T.handle, c: T.handle) -> None: |
| 860 | s1 = T.int32() |
| 861 | s0 = T.int32() |
| 862 | d1 = T.int32() |
| 863 | d0 = T.int32() |
| 864 | A = T.match_buffer( |
| 865 | a, |
| 866 | (frag_m, frag_n), |
| 867 | dtype, |
| 868 | align=64, |
| 869 | offset_factor=offset_factor, |
| 870 | scope=shared_scope, |
| 871 | strides=[s1, s0], |
| 872 | ) |
| 873 | C = T.match_buffer( |
| 874 | c, |
| 875 | (frag_m, frag_n), |
| 876 | dtype, |
| 877 | align=64, |
| 878 | offset_factor=offset_factor, |
| 879 | scope=wmma_fragment_scope, |
| 880 | strides=[d1, d0], |
| 881 | ) |
| 882 | with T.sblock("root"): |
| 883 | T.reads(A[0:frag_m, 0:frag_n]) |
| 884 | T.writes(C[0:frag_m, 0:frag_n]) |
| 885 | T.evaluate( |
| 886 | T.tvm_load_matrix_sync( |
| 887 | C.data, |
| 888 | m_dim, |
| 889 | n_dim, |
| 890 | k_dim, |
| 891 | get_wmma_fragment_index(C, d1, frag_m, frag_n), |
| 892 | A.access_ptr("r"), |
| 893 | s1, |
| 894 | layout, |
| 895 | dtype="handle", |
| 896 | ) |
| 897 | ) |
| 898 | |
| 899 | return wmma_load_desc, wmma_load_impl |
| 900 |
nothing calls this directly
no test coverage detected
searching dependent graphs…