chop.passes.module.transform.pim

chop.passes.module.transform.pim#

pim_matmul_transform_pass#

chop.passes.module.transforms.pim.pim_matmul_transform_pass(model, q_config={}, lora_config=None)[source]#

Apply PIM (Process-in-Memory) transformation to the given nn.Module.

This pass replaces supported layers (Linear, Conv2d) with their PIM-aware counterparts (PIMLinear, PIMConv2d) or LoRA-enabled PIM layers (LoraPIMLinear).

Parameters:

  • model (torch.nn.Module) – The input network to be transformed.

  • q_config (dict, optional) – Configuration for the PIM transformation, specifying how to match modules and their parameters.

  • lora_config (dict, optional) – Configuration for LoRA if applying LoRA-enabled PIM transformation.

Example q_config:

q_config = {
    "by": "type",
    "linear": {
        "config": {
            "tile_type": "pcm",
            "core_size": 256,
            "num_bits": 8,
            "programming_noise": True,
            "read_noise": True,
            "ir_drop": True,
            "out_noise": True,
        }
    },
}
Returns:

A tuple containing the transformed model and an empty dictionary (for consistency with other passes).

Return type:

tuple

Contents