mscclpp.language.loop

Classes

LoopIterationContext(unit, num_chunks)

A context manager for defining pipelined loop operations in MSCCL++ programs.

class mscclpp.language.loop.LoopIterationContext(unit, num_chunks)

Bases: object

A context manager for defining pipelined loop operations in MSCCL++ programs.

LoopIterationContext provides a way to define operations that will be executed in a pipelined manner across multiple iterations, where each pipeline iteration processes a specific chunk size (unit) of data. The pipeline allows overlapping execution of operations enabling efficient data processing.

unit

The unit size for pipeline operations.

Type:: int

num_chunks

The number of chunks to process in the pipeline.

Type:: int

operations

List of operations to be pipelined.

Type:: list

add_operation(rank, tb, operation)

Add an operation to be included in the pipeline.

This method is called internally to collect operations that should be pipelined together. Operations added here will be grouped and converted into pipeline operations when the context exits.

Parameters:

rank (int) – The rank where this operation will be executed.
tb (int) – The thread block ID that will execute this operation.
operation – The operation object to be added to the pipeline.