tenzr::sharding¶

Classes¶

Functions¶

Functions Documentation¶

function AnalyzeSharding¶

TENZR_APIShardingAnalysis AnalyzeSharding(
    const NamedSharding & sharding,
    const Shape & tensor_shape,
    std::string_view var_name ="%0",
    std::string_view tensor_type ="tensor<f32>"
)

Analyze a NamedSharding configuration for both Shardy MLIR dialect and GSPMD HLO targets.

Parameters:

• sharding Sharding to analyze.
• tensor_shape Shape of the value being sharded.
• var_name Variable name used in the constraint op.
• tensor_type Tensor type string used in the constraint op.

Return: Populated analysis record.

Lower a NamedSharding to its Shardy and GSPMD textual forms for a value.

function FormatShardyModuleHeader¶

TENZR_API std::string FormatShardyModuleHeader(
    const Mesh & mesh,
    std::string_view func_name,
    const std::vector< NamedSharding > & arg_shardings,
    const std::vector< Shape > & arg_shapes,
    std::string_view mesh_sym ="mesh"
)

Format complete Shardy MLIR module header including mesh declarations and function signatures.

Parameters:

• mesh Device mesh to declare.
• func_name Module/function name.
• arg_shardings Per-argument shardings.
• arg_shapes Per-argument shapes.
• mesh_sym Symbol name for the mesh.

Return: Partial module header text.

Emit the opening of an MLIR module with a Shardy mesh declaration and main signature.