Function get_relax_matmul_dequantize_module

tests/python/relax/test_codegen_cublas.py:108–142

Create a matmul op followd by dequantize operations.

(
    x_shape,
    y_shape,
    in_dtype,
    out_dtype,
    transposed_y=False,
    scale_const=1.0,
    zero_point_const=0.0,
)

def get_relax_matmul_dequantize_module(
    x_shape,
    y_shape,
    in_dtype,
    out_dtype,
    transposed_y=False,
    scale_const=1.0,
    zero_point_const=0.0,
):
    """Create a matmul op followd by dequantize operations."""
    with IRBuilder() as builder:
        with relax_builder.function():
            R.func_name("main")
            x = R.arg("x", R.Tensor(x_shape, in_dtype))
            y = R.arg("y", R.Tensor(y_shape, in_dtype))

            with R.dataflow() as frame:
                if transposed_y:
                    axes = list(range(len(y_shape) - 2)) + [-1, -2]
                    y = R.emit(R.permute_dims(y, axes=axes))
                result = R.emit(R.matmul(x, y, out_dtype="float32"))
                result = R.emit(
                    R.dequantize(
                        result,
                        scale=R.const(scale_const, "float16"),
                        zero_point=R.const(zero_point_const, "float16"),
                        axis=-1,
                        out_dtype=out_dtype,
                    )
                )
                R.output(result)
            R.func_ret_value(frame.output_vars[0])

    func = builder.get()
    return tvm.IRModule({"main": func})


def get_relax_matmul_multiply_module(