Method encode

tests/python/relax/test_codegen_cutlass.py:1300–1365

(
            A: T.Buffer((T.int64(128), T.int64(64)), "float16"),
            w_gathered: T.Buffer((T.int64(64), T.int64(64)), "int8"),
            compute: T.Buffer((T.int64(128),), "float16"),
        )

        @T.prim_func(s_tir=True)
        def encode(
            A: T.Buffer((T.int64(128), T.int64(64)), "float16"),
            w_gathered: T.Buffer((T.int64(64), T.int64(64)), "int8"),
            compute: T.Buffer((T.int64(128),), "float16"),
        ):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            max_abs_value = T.sblock_alloc_buffer((T.int64(128),), "float16")
            scale = T.sblock_alloc_buffer((T.int64(128),))
            for i, k in T.grid(T.int64(128), T.int64(64)):
                with T.sblock("max_abs_value"):
                    v_i, v_k = T.axis.remap("SR", [i, k])
                    T.reads(A[v_i, v_k])
                    T.writes(max_abs_value[v_i])
                    with T.init():
                        max_abs_value[v_i] = T.float16(-65504)
                    max_abs_value[v_i] = T.max(max_abs_value[v_i], T.fabs(A[v_i, v_k]))
            for i in range(T.int64(128)):
                with T.sblock("scale"):
                    v_i = T.axis.spatial(T.int64(128), i)
                    T.reads(max_abs_value[v_i])
                    T.writes(scale[v_i])
                    scale[v_i] = T.max(
                        T.Cast("float32", max_abs_value[v_i]), T.float32(0.0001)
                    ) * T.float32(0.125)
            for j, i, k in T.grid(T.int64(64), T.int64(64), T.int64(2)):
                with T.sblock("w_gathered"):
                    v_j, v_i, v_k = T.axis.remap("SSR", [j, i, k])
                    T.reads(A[v_i * T.int64(2) + v_k, v_j], scale[v_i * T.int64(2) + v_k])
                    T.writes(w_gathered[v_j, v_i])
                    with T.init():
                        w_gathered[v_j, v_i] = T.int8(0)
                    w_gathered[v_j, v_i] = T.bitwise_or(
                        w_gathered[v_j, v_i],
                        T.if_then_else(
                            v_i * T.int64(2) + v_k < T.int64(128),
                            T.shift_left(
                                T.bitwise_and(
                                    T.Cast(
                                        "int8",
                                        T.min(
                                            T.max(
                                                T.round(
                                                    T.Cast(
                                                        "float32", A[v_i * T.int64(2) + v_k, v_j]
                                                    )
                                                    / scale[v_i * T.int64(2) + v_k]
                                                ),
                                                T.float32(-8),
                                            ),
                                            T.float32(7),
                                        ),
                                    ),
                                    T.int8(15),
                                ),
                                T.Cast("int8", v_k) * T.int8(4),
                            ),
                            T.int8(0),