Systems and methods for data placement for in-memory-compute

1. A memory module comprising:

a memory die comprising a dynamic random access memory (DRAM) bank comprising:

an array of DRAM cells arranged in a plurality of pages; and

an in-memory compute (IMC) module comprising:

an arithmetic logic unit (ALU); and

a memory controller configured to:

receive, from a host processor, an operand and an instruction;

determine, based on the instruction, a data layout from a plurality of data layouts, the data layout specifying placement of the operand among the pages of the array of DRAM cells and the IMC;

supply the operand to the DRAM bank in accordance with the data layout; and

control the IMC module of the DRAM bank to perform an ALU operation on the operand in accordance with the instruction.

2. The memory module of claim 1 , wherein the data layouts comprise:

a one operand (1OP) data layout, wherein a first operand is read from the DRAM cells and the operand comprises a second operand, the second operand being supplied directly from the host processor to the IMC of the DRAM bank.

3. The memory module of claim 2 , wherein the first operand is divided into a plurality of first tiles and the second operand is divided into a plurality of second tiles, each tile comprising a plurality of values,

wherein the IMC module further comprises an operand register, and

wherein the memory controller is further configured to:

store a first tile of the first operand in the operand register; and

perform the ALU operation on the first operand stored in the operand register and each of a plurality of second tiles of the second operand.

4. The memory module of claim 1 , wherein the operand is divided into a plurality of first tiles and a second operand is divided into a plurality of second tiles, each tile comprising a plurality of values, and

wherein the data layouts comprise a same page (SR) data layout, wherein the memory controller stores one or more of the first tiles and one or more of the second tiles in a same page of the DRAM cells.

5. The memory module of claim 4 , wherein the IMC module further comprises an operand register, and

wherein the memory controller is further configured to:

store a first tile of the one or more of the first tiles in the operand register; and

perform the ALU operation on the first tile stored in the operand register and each of the one or more second tiles stored in the same page of the array of DRAM cells as the first tile.

6. The memory module of claim 5 , wherein the IMC module of the DRAM bank further comprises an accumulator, the accumulator comprising an accumulator register configured to store an accumulated value, the accumulator being configured to:

receive the output computed by the ALU; and

update the accumulator register with the sum of the accumulated value and the output, and

wherein the instruction comprises computing an inner product of the operand and the second operand, wherein the first tile of the first tiles stores row data and a second tile of the second tiles comprises column data.

7. The memory module of claim 5 , wherein the first tile has a first number of values and the second tile has a second number of values,

wherein the IMC module of the DRAM bank comprises an output buffer, the output buffer having a size for storing greater than or equal to the product of the first number of values and the second number of values, and

wherein the instruction comprises computing an outer product of the operand and the second operand, wherein the first tile of the first tiles stores row data and a second tile of the second tiles comprises column data.

8. The memory module of claim 5 , wherein the first tile has a first number of values and the second tile has a second number of values,

wherein the IMC module of the DRAM bank comprises an output buffer, the output buffer having a size for storing greater than or equal to the larger of the first number of values and the second number of values, and

wherein the instruction comprises computing a tensor product of the operand and the second operand, wherein the first tile of the first tiles stores row data and a second tile of the second tiles comprises column data.

9. The memory module of claim 1 , wherein the operand is divided into a plurality of first tiles and a second operand is divided into a plurality of second tiles, each tile comprising a plurality of values, and

wherein the data layouts comprise a different page (DR) data layout wherein the memory controller stores a subset of the first tiles in a first page of the array of DRAM cells and a subset of the second tiles in a second page of the array of DRAM cells.

10. The memory module of claim 9 , wherein the IMC module further comprises an operand register, and

wherein the memory controller is further configured to:

store a first tile of the operand from the first page in the operand register; and

perform the ALU operation on the first tile of the operand stored in the operand register and each of a plurality of second tiles of the second operand from the second page.

11. The memory module of claim 1 , the IMC module of the DRAM bank further comprising a hardware buffer configured to buffer the output computed by the ALU.

12. The memory module of claim 11 , wherein the hardware buffer is at least four times the size of a result register of the IMC module.

13. The memory module of claim 1 , the IMC module of the DRAM bank further comprising an accumulator, the accumulator comprising an accumulator register configured to store an accumulated value, the accumulator being configured to:

receive the output computed by the ALU; and

update the accumulator register with the sum of the accumulated value and the output.

14. The memory module of claim 1 , wherein the memory module is a high bandwidth memory (HBM) module comprising a stack of memory dies connected by through silicon vias, the stack of memory dies comprising the memory die.

15. A method for performing computations in-memory comprising:

receiving, by a memory controller of a memory module, an operand and an instruction;

determining, by the memory controller, based on the instruction, a data layout from a plurality of data layouts;

supplying the operand to a dynamic random access memory (DRAM) bank of the memory module in accordance with the data layout, the DRAM bank comprising:

an array of DRAM cells arranged in a plurality of pages; and

an in-memory compute (IMC) module comprising an arithmetic logic unit (ALU); and

controlling the IMC module of the DRAM bank to perform an ALU operation on the operand in accordance with the instruction,

wherein the data layout specifies placement of the operand among the pages of the array of DRAM cells and the IMC.

16. The method of claim 15 , wherein the data layouts comprise:

a one operand (1OP) data layout, wherein a first operand is written to the DRAM cells and the operand comprises a second operand, the second operand being supplied directly from a host processor to the IMC module of the DRAM bank.

17. The method of claim 16 , wherein the first operand is divided into a plurality of first tiles and the second operand is divided into a plurality of second tiles, each tile comprising a plurality of values,

wherein the IMC module further comprises an operand register, and

wherein the memory controller is further configured to:

store a first tile of the first operand in the operand register; and

perform the ALU operation on the first operand stored in the operand register and each of a plurality of second tiles of the second operand.

18. The method of claim 15 , wherein the operand is divided into a plurality of first tiles and a second operand is divided into a plurality of second tiles, each tile comprising a plurality of values, and

wherein the data layouts comprise a same page (SR) data layout, wherein the memory controller stores one or more of the first tiles and one or more of the second tiles in a same page of the DRAM cells.

19. The method of claim 18 , wherein the IMC module further comprises an operand register, and

wherein the memory controller is further configured to:

store a first tile of the one or more of the first tiles in the operand register; and

perform the ALU operation on the first tile stored in the operand register and each of the one or more second tiles stored in the same page of the array of DRAM cells as the first tile.

20. The method of claim 19 , wherein the IMC module of the DRAM bank further comprises an accumulator, the accumulator comprising an accumulator register configured to store an accumulated value, the accumulator being configured to:

receive the output computed by the ALU; and

update the accumulator register with the sum of the accumulated value and the output, and

wherein the instruction comprises computing an inner product of the operand and the second operand, wherein the first tile of the first tiles stores row data and a second tile of the second tiles comprises column data.

21. The method of claim 19 , wherein the first tile has a first number of values and the second tile has a second number of values,

wherein the IMC module of the DRAM bank comprises an output buffer, the output buffer having a size for storing greater than or equal to the product of the first number of values and the second number of values, and

wherein the instruction comprises computing an outer product of the operand and the second operand, wherein the first tile of the first tiles stores row data and a second tile of the second tiles comprises column data.

22. The method of claim 19 , wherein the first tile has a first number of values and the second tile has a second number of values,

wherein the IMC module of the DRAM bank comprises an output buffer, the output buffer having a size for storing greater than or equal to the larger of the first number of values and the second number of values, and

wherein the instruction comprises computing a tensor product of the operand and the second operand, wherein the first tile of the first tiles stores row data and a second tile of the second tiles comprises column data.

23. The method of claim 15 , wherein the operand is divided into a plurality of first tiles and a second operand is divided into a plurality of second tiles, each tile comprising a plurality of values, and

wherein the data layouts comprise a different page (DR) data layout wherein the memory controller stores a subset of the first tiles in a first page of the array of DRAM cells and a subset of the second tiles in a second page of the array of DRAM cells.

24. The method of claim 23 , wherein the IMC module further comprises an operand register, and

wherein the memory controller is further configured to:

store a first tile of the operand from the first page in the operand register; and

perform the ALU operation on the first tile stored in the operand register and each of a plurality of second tiles of the second operand from the second page.