hand

venv/lib/python3.12/site-packages/jaxlib/mosaic/dialect/gpu/_mosaic_gpu_gen_enums.py

@@ -0,0 +1,200 @@
+
+# Autogenerated by mlir-tblgen; don't manually edit.
+
+from enum import IntEnum, auto, IntFlag
+from jaxlib.mlir.dialects._ods_common import _cext as _ods_cext
+from jaxlib.mlir.ir import register_attribute_builder
+_ods_ir = _ods_cext.ir
+
+class AtomicOpType(IntEnum):
+    """Atomic store operation type"""
+
+    Add = 0
+    Min = 1
+    Max = 2
+    And = 3
+    Or = 4
+    Xor = 5
+
+    def __str__(self):
+        if self is AtomicOpType.Add:
+            return "add"
+        if self is AtomicOpType.Min:
+            return "min"
+        if self is AtomicOpType.Max:
+            return "max"
+        if self is AtomicOpType.And:
+            return "and"
+        if self is AtomicOpType.Or:
+            return "or"
+        if self is AtomicOpType.Xor:
+            return "xor"
+        raise ValueError("Unknown AtomicOpType enum entry.")
+
+
+
+@register_attribute_builder("MosaicGPU_AtomicOpType", allow_existing=True)
+def _mosaicgpu_atomicoptype(x, context):
+    return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
+
+class Dimension(IntEnum):
+    """a dimension, either 'x', 'y', or 'z'"""
+
+    x = 0
+    y = 1
+    z = 2
+
+    def __str__(self):
+        if self is Dimension.x:
+            return "x"
+        if self is Dimension.y:
+            return "y"
+        if self is Dimension.z:
+            return "z"
+        raise ValueError("Unknown Dimension enum entry.")
+
+
+
+@register_attribute_builder("MosaicGPU_Dimension", allow_existing=True)
+def _mosaicgpu_dimension(x, context):
+    return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
+
+class MultimemLoadReductionType(IntEnum):
+    """Reduction operation for multimem reduction."""
+
+    Add = 0
+    Min = 1
+    Max = 2
+    And = 3
+    Or = 4
+    Xor = 5
+    Umin = 6
+    Umax = 7
+    Smin = 8
+    Smax = 9
+
+    def __str__(self):
+        if self is MultimemLoadReductionType.Add:
+            return "add"
+        if self is MultimemLoadReductionType.Min:
+            return "min"
+        if self is MultimemLoadReductionType.Max:
+            return "max"
+        if self is MultimemLoadReductionType.And:
+            return "and"
+        if self is MultimemLoadReductionType.Or:
+            return "or"
+        if self is MultimemLoadReductionType.Xor:
+            return "xor"
+        if self is MultimemLoadReductionType.Umin:
+            return "umin"
+        if self is MultimemLoadReductionType.Umax:
+            return "umax"
+        if self is MultimemLoadReductionType.Smin:
+            return "smin"
+        if self is MultimemLoadReductionType.Smax:
+            return "smax"
+        raise ValueError("Unknown MultimemLoadReductionType enum entry.")
+
+
+
+@register_attribute_builder("MosaicGPU_MultimemLoadReductionType", allow_existing=True)
+def _mosaicgpu_multimemloadreductiontype(x, context):
+    return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
+
+class OOBFillMode(IntEnum):
+    """Out of bounds fill mode for TMA."""
+
+    kUndefined = 0
+    kPromiseInBounds = 1
+    kZeros = 2
+
+    def __str__(self):
+        if self is OOBFillMode.kUndefined:
+            return "undefined"
+        if self is OOBFillMode.kPromiseInBounds:
+            return "promise_in_bounds"
+        if self is OOBFillMode.kZeros:
+            return "zeros"
+        raise ValueError("Unknown OOBFillMode enum entry.")
+
+
+
+@register_attribute_builder("MosaicGPU_OOBFillMode", allow_existing=True)
+def _mosaicgpu_oobfillmode(x, context):
+    return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
+
+class SwizzlingMode(IntEnum):
+    """What swizzling to use for a memory access."""
+
+    kNoSwizzle = 16
+    k32ByteSwizzle = 32
+    k64ByteSwizzle = 64
+    k128ByteSwizzle = 128
+
+    def __str__(self):
+        if self is SwizzlingMode.kNoSwizzle:
+            return "kNoSwizzle"
+        if self is SwizzlingMode.k32ByteSwizzle:
+            return "k32ByteSwizzle"
+        if self is SwizzlingMode.k64ByteSwizzle:
+            return "k64ByteSwizzle"
+        if self is SwizzlingMode.k128ByteSwizzle:
+            return "k128ByteSwizzle"
+        raise ValueError("Unknown SwizzlingMode enum entry.")
+
+
+
+@register_attribute_builder("MosaicGPU_SwizzlingMode", allow_existing=True)
+def _mosaicgpu_swizzlingmode(x, context):
+    return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
+
+class TMAReduction(IntEnum):
+    """Reduction operation for TMA."""
+
+    Add = 0
+    Min = 1
+    Max = 2
+    Inc = 3
+    Dec = 4
+    And = 5
+    Or = 6
+    Xor = 7
+    Umin = 8
+    Umax = 9
+    Smin = 10
+    Smax = 11
+
+    def __str__(self):
+        if self is TMAReduction.Add:
+            return "add"
+        if self is TMAReduction.Min:
+            return "min"
+        if self is TMAReduction.Max:
+            return "max"
+        if self is TMAReduction.Inc:
+            return "inc"
+        if self is TMAReduction.Dec:
+            return "dec"
+        if self is TMAReduction.And:
+            return "and"
+        if self is TMAReduction.Or:
+            return "or"
+        if self is TMAReduction.Xor:
+            return "xor"
+        if self is TMAReduction.Umin:
+            return "umin"
+        if self is TMAReduction.Umax:
+            return "umax"
+        if self is TMAReduction.Smin:
+            return "smin"
+        if self is TMAReduction.Smax:
+            return "smax"
+        raise ValueError("Unknown TMAReduction enum entry.")
+
+
+
+@register_attribute_builder("MosaicGPU_TMAReduction", allow_existing=True)
+def _mosaicgpu_tmareduction(x, context):
+    return _ods_ir.IntegerAttr.get(_ods_ir.IntegerType.get_signless(32, context=context), int(x))
+