Method dequantize

tests/python/relax/backend/adreno/mod_utils.py:752–781

(weight: T.handle, scale: T.handle, var_dequantize: T.handle)

        @T.prim_func(s_tir=True)
        def dequantize(weight: T.handle, scale: T.handle, var_dequantize: T.handle):
            T.func_attr({"tirx.noalias": T.bool(True)})
            lm_head_q_weight1 = T.match_buffer(weight, (T.int64(K // 8), T.int64(N)), "uint32")
            lm_head_q_scale1 = T.match_buffer(scale, (T.int64(K // 32), T.int64(N)), "float16")
            dequantize = T.match_buffer(var_dequantize, (T.int64(K), T.int64(N)), "float16")
            # with T.sblock("root"):
            compute = T.alloc_buffer((T.int64(K), T.int64(N)), "float16")
            for i0, i1 in T.grid(T.int64(K), T.int64(N)):
                with T.sblock("compute"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(lm_head_q_weight1[v_i0 // T.int64(8), v_i1])
                    T.writes(compute[v_i0, v_i1])
                    compute[v_i0, v_i1] = T.Cast(
                        "float16",
                        T.bitwise_and(
                            T.shift_right(
                                lm_head_q_weight1[v_i0 // T.int64(8), v_i1],
                                T.Cast("uint32", v_i0 % T.int64(8) * T.int64(4)),
                            ),
                            T.uint32(15),
                        ),
                    )
            for i0, i1 in T.grid(T.int64(K), T.int64(N)):
                with T.sblock("dequantize"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(compute[v_i0, v_i1], lm_head_q_scale1[v_i0 // T.int64(32), v_i1])
                    T.writes(dequantize[v_i0, v_i1])
                    dequantize[v_i0, v_i1] = (
                        compute[v_i0, v_i1] - T.float16(7.0)
                    ) * lm_head_q_scale1[v_i0 // T.int64(32), v_i1]

    return DequantMatmul