List of Enabled XLA Hlo Passes

Hlo Pass definitions

class hloenv.HloPass.AlgebraicSimplifier(self: hloenv.hlo_pass.AlgebraicSimplifier, options: dict) → None

class hloenv.HloPass.AllGatherBroadcastReorder(self: hloenv.hlo_pass.AllGatherBroadcastReorder) → None

class hloenv.HloPass.AllGatherCombiner(self: hloenv.hlo_pass.AllGatherCombiner, combine_threshold_in_bytes: int, combine_threshold_count: int) → None

class hloenv.HloPass.AllReduceBlueConnect(self: hloenv.hlo_pass.AllReduceBlueConnect, num_devices_per_host: int) → None

class hloenv.HloPass.AllReduceCombiner(self: hloenv.hlo_pass.AllReduceCombiner, combine_threshold_in_bytes: int, combine_threshold_count: int) → None

class hloenv.HloPass.AllReduceContiguous(self: hloenv.hlo_pass.AllReduceContiguous) → None

class hloenv.HloPass.AllReduceFolder(self: hloenv.hlo_pass.AllReduceFolder) → None

class hloenv.HloPass.AllReduceReassociate(self: hloenv.hlo_pass.AllReduceReassociate) → None

class hloenv.HloPass.AllToAllDecomposer(self: hloenv.hlo_pass.AllToAllDecomposer, decompose_to_tuple: bool = True, min_array_rank: int = 0) → None

class hloenv.HloPass.AsyncCollectiveCreator(self: hloenv.hlo_pass.AsyncCollectiveCreator) → None

class hloenv.HloPass.BFloat16Normalization(self: hloenv.hlo_pass.BFloat16Normalization, arg0: bool) → None

class hloenv.HloPass.BatchNormExpander(self: hloenv.hlo_pass.BatchNormExpander, rewrite_training_op: bool = False, rewrite_inference_op: bool = False, rewrite_grad_op: bool = False) → None

class hloenv.HloPass.BitcastDtypesExpander(self: hloenv.hlo_pass.BitcastDtypesExpander) → None

class hloenv.HloPass.CallInliner(self: hloenv.hlo_pass.CallInliner, single_call_site: bool = False, update_domain: bool = False) → None

class hloenv.HloPass.CollectivesScheduleLinearizer(self: hloenv.hlo_pass.CollectivesScheduleLinearizer) → None

class hloenv.HloPass.ComparisonExpander(self: hloenv.hlo_pass.ComparisonExpander) → None

class hloenv.HloPass.ConditionalCanonicalizer(self: hloenv.hlo_pass.ConditionalCanonicalizer) → None

class hloenv.HloPass.ConditionalSimplifier(self: hloenv.hlo_pass.ConditionalSimplifier) → None

class hloenv.HloPass.Convolution4DExpander(self: hloenv.hlo_pass.Convolution4DExpander) → None

class hloenv.HloPass.CublasPadForGemms(self: hloenv.hlo_pass.CublasPadForGemms, datatype: xla::PrimitiveType, pad_to_multiple_of: int) → None

BF16 = <PrimitiveType.BF16: 16>

C128 = <PrimitiveType.C128: 18>

C64 = <PrimitiveType.C64: 15>

F16 = <PrimitiveType.F16: 10>

F32 = <PrimitiveType.F32: 11>

F64 = <PrimitiveType.F64: 12>

OPAQUE_TYPE = <PrimitiveType.OPAQUE_TYPE: 14>

PRED = <PrimitiveType.PRED: 1>

PRIMITIVE_TYPE_INVALID = <PrimitiveType.PRIMITIVE_TYPE_INVALID: 0>

class PrimitiveType(self: hloenv.hlo_pass.CublasPadForGemms.PrimitiveType, value: int) → None

Members:

PRIMITIVE_TYPE_INVALID

PRED

S8

S16

S32

S64

U8

U16

U32

U64

F16

F32

BF16

F64

C64

C128

TUPLE

OPAQUE_TYPE

TOKEN

BF16 = <PrimitiveType.BF16: 16>

C128 = <PrimitiveType.C128: 18>

C64 = <PrimitiveType.C64: 15>

F16 = <PrimitiveType.F16: 10>

F32 = <PrimitiveType.F32: 11>

F64 = <PrimitiveType.F64: 12>

OPAQUE_TYPE = <PrimitiveType.OPAQUE_TYPE: 14>

PRED = <PrimitiveType.PRED: 1>

PRIMITIVE_TYPE_INVALID = <PrimitiveType.PRIMITIVE_TYPE_INVALID: 0>

S16 = <PrimitiveType.S16: 3>

S32 = <PrimitiveType.S32: 4>

S64 = <PrimitiveType.S64: 5>

S8 = <PrimitiveType.S8: 2>

TOKEN = <PrimitiveType.TOKEN: 17>

TUPLE = <PrimitiveType.TUPLE: 13>

U16 = <PrimitiveType.U16: 7>

U32 = <PrimitiveType.U32: 8>

U64 = <PrimitiveType.U64: 9>

U8 = <PrimitiveType.U8: 6>

name

value

S16 = <PrimitiveType.S16: 3>

S32 = <PrimitiveType.S32: 4>

S64 = <PrimitiveType.S64: 5>

S8 = <PrimitiveType.S8: 2>

TOKEN = <PrimitiveType.TOKEN: 17>

TUPLE = <PrimitiveType.TUPLE: 13>

U16 = <PrimitiveType.U16: 7>

U32 = <PrimitiveType.U32: 8>

U64 = <PrimitiveType.U64: 9>

U8 = <PrimitiveType.U8: 6>

class hloenv.HloPass.CudnnFusedConvRewriter(self: hloenv.hlo_pass.CudnnFusedConvRewriter) → None

class hloenv.HloPass.CudnnPadForConvolutions(self: hloenv.hlo_pass.CudnnPadForConvolutions) → None

class hloenv.HloPass.CudnnVectorizeConvolutions(self: hloenv.hlo_pass.CudnnVectorizeConvolutions) → None

class hloenv.HloPass.DotDecomposer(self: hloenv.hlo_pass.DotDecomposer) → None

class hloenv.HloPass.DotMerger(self: hloenv.hlo_pass.DotMerger, max_size_to_merge: int) → None

class hloenv.HloPass.DynamicDimensionSimplifier(self: hloenv.hlo_pass.DynamicDimensionSimplifier) → None

class hloenv.HloPass.DynamicIndexSplitter(self: hloenv.hlo_pass.DynamicIndexSplitter) → None

class hloenv.HloPass.DynamicPadder(self: hloenv.hlo_pass.DynamicPadder, arg0: dict) → None

class ShapeCheckMode(self: hloenv.hlo_pass.DynamicPadder.ShapeCheckMode, value: int) → None

Members:

kInvalid

kCompileTime

kRuntimeTime

kIgnore

kCompileTime = <ShapeCheckMode.kCompileTime: 1>

kIgnore = <ShapeCheckMode.kIgnore: 3>

kInvalid = <ShapeCheckMode.kInvalid: 0>

kRuntimeTime = <ShapeCheckMode.kRuntimeTime: 2>

name

value

kCompileTime = <ShapeCheckMode.kCompileTime: 1>

kIgnore = <ShapeCheckMode.kIgnore: 3>

kInvalid = <ShapeCheckMode.kInvalid: 0>

kRuntimeTime = <ShapeCheckMode.kRuntimeTime: 2>

class hloenv.HloPass.EighExpander(self: hloenv.hlo_pass.EighExpander) → None

class hloenv.HloPass.FlattenCallGraph(self: hloenv.hlo_pass.FlattenCallGraph) → None

class hloenv.HloPass.FusionBitcastLift(self: hloenv.hlo_pass.FusionBitcastLift) → None

class hloenv.HloPass.FusionMerger(self: hloenv.hlo_pass.FusionMerger) → None

class hloenv.HloPass.GatherExpander(self: hloenv.hlo_pass.GatherExpander, m: xla::GatherExpander::Mode) → None

class Mode(self: hloenv.hlo_pass.GatherExpander.Mode, value: int) → None

Members:

kEliminateAllGathers

kEliminateSimpleGathers

kEliminateAllGathers = <Mode.kEliminateAllGathers: 0>

kEliminateSimpleGathers = <Mode.kEliminateSimpleGathers: 1>

name

value

kEliminateAllGathers = <Mode.kEliminateAllGathers: 0>

kEliminateSimpleGathers = <Mode.kEliminateSimpleGathers: 1>

class hloenv.HloPass.GemmAlgorithmPicker(self: hloenv.hlo_pass.GemmAlgorithmPicker) → None

class hloenv.HloPass.GemmBroadcastFoldingRewriter(self: hloenv.hlo_pass.GemmBroadcastFoldingRewriter) → None

class hloenv.HloPass.GemmRewriter(self: hloenv.hlo_pass.GemmRewriter) → None

class hloenv.HloPass.GeneralFusion(self: hloenv.hlo_pass.GeneralFusion) → None

class hloenv.HloPass.GpuConvAlgorithmPicker(self: hloenv.hlo_pass.GpuConvAlgorithmPicker) → None

class hloenv.HloPass.GpuConvPaddingLegalization(self: hloenv.hlo_pass.GpuConvPaddingLegalization) → None

class hloenv.HloPass.GpuConvRewriter(self: hloenv.hlo_pass.GpuConvRewriter) → None

class hloenv.HloPass.GpuHorizontalInputFusion(self: hloenv.hlo_pass.GpuHorizontalInputFusion) → None

class hloenv.HloPass.GpuHorizontalLoopFusion(self: hloenv.hlo_pass.GpuHorizontalLoopFusion) → None

class hloenv.HloPass.GpuInstructionFusion(self: hloenv.hlo_pass.GpuInstructionFusion, may_duplicate: bool) → None

class hloenv.HloPass.GpuLayoutAssignment(self: hloenv.hlo_pass.GpuLayoutAssignment, arg0: hloenv.HloModule) → None

class hloenv.HloPass.GpuMultiOutputFusion(self: hloenv.hlo_pass.GpuMultiOutputFusion) → None

class hloenv.HloPass.GpuScatterExpander(self: hloenv.hlo_pass.GpuScatterExpander) → None

class hloenv.HloPass.GpuTreeReductionRewriter(self: hloenv.hlo_pass.GpuTreeReductionRewriter) → None

class hloenv.HloPass.GpusolverRewriter(self: hloenv.hlo_pass.GpusolverRewriter) → None

class hloenv.HloPass.HloCSE(self: hloenv.hlo_pass.HloCSE, is_layout_sensitive: bool, only_fusion_computations: bool = False) → None

class hloenv.HloPass.HloConstantFolding(self: hloenv.hlo_pass.HloConstantFolding) → None

class hloenv.HloPass.HloDCE(*args, **kwargs)

Overloaded function.

1. __init__(self: hloenv.hlo_pass.HloDCE) -> None
2. __init__(self: hloenv.hlo_pass.HloDCE, remove_cross_partition_collective_ops: bool) -> None

class hloenv.HloPass.HloPassInterface

class hloenv.HloPass.HloVerifier(self: hloenv.hlo_pass.HloVerifier, layout_sensitive: bool, allow_mixed_precision: bool) → None

class hloenv.HloPass.LogisticExpander(self: hloenv.hlo_pass.LogisticExpander, expansion_type: xla::LogisticExpansionType) → None

class LogisticExpansionType(self: hloenv.hlo_pass.LogisticExpander.LogisticExpansionType, value: int) → None

Members:

kTanh

kExp

kExp = <LogisticExpansionType.kExp: 1>

kTanh = <LogisticExpansionType.kTanh: 0>

name

value

kExp = <LogisticExpansionType.kExp: 1>

kTanh = <LogisticExpansionType.kTanh: 0>

class hloenv.HloPass.OperandUpcaster(self: hloenv.hlo_pass.OperandUpcaster) → None

class hloenv.HloPass.OptimizationBarrierExpander(self: hloenv.hlo_pass.OptimizationBarrierExpander) → None

class hloenv.HloPass.QrExpander(self: hloenv.hlo_pass.QrExpander) → None

class hloenv.HloPass.RealImagExpander(self: hloenv.hlo_pass.RealImagExpander) → None

class hloenv.HloPass.ReduceScatterCombiner(self: hloenv.hlo_pass.ReduceScatterCombiner, combine_threshold_in_bytes: int, combine_threshold_count: int) → None

class hloenv.HloPass.ReduceScatterCreator(self: hloenv.hlo_pass.ReduceScatterCreator) → None

class hloenv.HloPass.ReductionDegenerateDimRemover(self: hloenv.hlo_pass.ReductionDegenerateDimRemover) → None

class hloenv.HloPass.ReductionDimensionGrouper(self: hloenv.hlo_pass.ReductionDimensionGrouper) → None

class hloenv.HloPass.ReductionLayoutNormalizer(self: hloenv.hlo_pass.ReductionLayoutNormalizer) → None

class hloenv.HloPass.ReductionSplitter(self: hloenv.hlo_pass.ReductionSplitter) → None

class hloenv.HloPass.ReshapeMover(self: hloenv.hlo_pass.ReshapeMover) → None

class hloenv.HloPass.ResultCaster(self: hloenv.hlo_pass.ResultCaster) → None

class hloenv.HloPass.RngBitGeneratorExpander(self: hloenv.hlo_pass.RngBitGeneratorExpander, default_algorithm: xla::RandomAlgorithm) → None

RNG_DEFAULT = <RandomAlgorithm.RNG_DEFAULT: 0>

RNG_PHILOX = <RandomAlgorithm.RNG_PHILOX: 2>

RNG_THREE_FRY = <RandomAlgorithm.RNG_THREE_FRY: 1>

class RandomAlgorithm(self: hloenv.hlo_pass.RngBitGeneratorExpander.RandomAlgorithm, value: int) → None

Members:

RNG_DEFAULT

RNG_THREE_FRY

RNG_PHILOX

RNG_DEFAULT = <RandomAlgorithm.RNG_DEFAULT: 0>

RNG_PHILOX = <RandomAlgorithm.RNG_PHILOX: 2>

RNG_THREE_FRY = <RandomAlgorithm.RNG_THREE_FRY: 1>

name

value

class hloenv.HloPass.RngExpander(self: hloenv.hlo_pass.RngExpander) → None

class hloenv.HloPass.ScatterExpander(self: hloenv.hlo_pass.ScatterExpander, m: xla::ScatterExpander::Mode) → None

class Mode(self: hloenv.hlo_pass.ScatterExpander.Mode, value: int) → None

Members:

kEliminateAllScatters

kEliminateSimpleScatters

kEliminateAllScatters = <Mode.kEliminateAllScatters: 0>

kEliminateSimpleScatters = <Mode.kEliminateSimpleScatters: 1>

name

value

kEliminateAllScatters = <Mode.kEliminateAllScatters: 0>

kEliminateSimpleScatters = <Mode.kEliminateSimpleScatters: 1>

class hloenv.HloPass.ShardingPropagation(self: hloenv.hlo_pass.ShardingPropagation, is_spmd: bool = False, propagate_metadata: bool = False, allow_spmd_sharding_propagation_to_output: bool = False, cse_prevention_only: bool = False) → None

class hloenv.HloPass.ShardingRemover(self: hloenv.hlo_pass.ShardingRemover) → None

class hloenv.HloPass.SliceSinker(self: hloenv.hlo_pass.SliceSinker) → None

class hloenv.HloPass.SortSimplifier(self: hloenv.hlo_pass.SortSimplifier) → None

class hloenv.HloPass.StableSortExpander(self: hloenv.hlo_pass.StableSortExpander) → None

class hloenv.HloPass.StatefulRngSpmdPartitioner(self: hloenv.hlo_pass.StatefulRngSpmdPartitioner, num_partitions: int, num_replicas: int) → None

class hloenv.HloPass.TransposeFolding(self: hloenv.hlo_pass.TransposeFolding) → None

class hloenv.HloPass.TriangularSolveRewriter(self: hloenv.hlo_pass.TriangularSolveRewriter) → None

class hloenv.HloPass.TupleSimplifier(*args, **kwargs)

Overloaded function.

1. __init__(self: hloenv.hlo_pass.TupleSimplifier) -> None
2. __init__(self: hloenv.hlo_pass.TupleSimplifier, exclude_entry_computation: bool) -> None

class hloenv.HloPass.VariadicOpSplitter(self: hloenv.hlo_pass.VariadicOpSplitter) → None

class hloenv.HloPass.WhileLoopConstantSinking(self: hloenv.hlo_pass.WhileLoopConstantSinking) → None

class hloenv.HloPass.WhileLoopSimplifier(self: hloenv.hlo_pass.WhileLoopSimplifier) → None

class hloenv.HloPass.WhileLoopTripCountAnnotator(self: hloenv.hlo_pass.WhileLoopTripCountAnnotator) → None

class hloenv.HloPass.ZeroSizedHloElimination(self: hloenv.hlo_pass.ZeroSizedHloElimination) → None