Processing-in-memory (pim) architecture supporting multiple contexts with efficient context switching

An apparatus for supporting multiple contexts using a processing-in-memory (PIM) device, the apparatus comprising:

• a memory comprising one or more registers and configured to store PIM contexts;
• and
• a memory controller coupled to the memory, and configured to perform operations, the operations comprising:
• initiating execution of a first task on a first virtual PIM hosted by a first SoC client of a plurality of SoC clients;
• receiving, based on an interruption of the execution of the first task, a first PIM context associated with the interrupted execution of the first task;
• saving the first PIM context associated with the interrupted execution of the first task to the memory;
• and
• restoring, using the first PIM context saved to the memory, execution of the first task on the first virtual PIM.

A system for supporting multiple contexts in a processing-in memory (PIM) device, the system comprising:

• a plurality of system on a chip (SoC) clients, wherein each SoC client is configured to host a respective virtual PIM, wherein the plurality of SoC clients includes at least a first SoC client hosting a first virtual PIM and a second SoC client hosting a second virtual PIM;
• a memory comprising one or more registers and configured to store PIM contexts;
• and
• a memory controller coupled to the memory, and configured to perform operations comprising:
• receiving a plurality of tasks including a first task and a second task;
• initiating execution of the first task on the first virtual PIM;
• interrupting the execution of the first task based on a context switching protocol for switching between tasks of the plurality of tasks;
• saving a first PIM context associated with the interrupted execution of the first task to the memory by reading the first PIM context from a first register of the one or more registers;
• initiating execution of the second task on the second virtual PIM;
• saving a second PIM context associated with a completed execution of the second task to the memory by reading the second PIM context from a second register of the one or more registers;
• and
• restoring, using the first PIM context saved to the memory, the execution of the first task on the first virtual PIM.