Add FloatType to MLIR Python bindings

MLIR
1

The Python bindings do not currently have a common base class for floating point types. This makes it difficult to quickly check if a given ir.Type is floating point. Instead, you have to do something like

_FP_TYPES = (ir.F16Type, ir.F32Type, ir.F64Type, ...)

def is_float_type(t: ir.Type) -> bool:
    return isinstance(t, _FP_TYPES)

Wdyt about introducing such a type? It can be done in a somewhat ad-hoc way via metaclasses (see below), but of course it would be better to make all floating point types to inherit from a common base class in Python.

For completeness, the metaclass hack

_FLOAT_WIDTH = {
    ir.F16Type: 16,
    ir.F32Type: 32,
    ir.F64Type: 64,
    # ...
}
_FLOAT_TYPES = tuple(_FLOAT_WIDTH)


class FloatTypeMeta(type):

    def __instancecheck__(cls, instance: object) -> bool:
        return isinstance(instance, _FLOAT_TYPES)

    def __subclasscheck__(cls, subclass: type[object]) -> bool:
        return issubclass(subclass, _FLOAT_TYPES)


class FloatType(metaclass=FloatTypeMeta):
    """Fake base class for MLIR floating point types."""

    def __init__(self, t: ir.Type) -> None:
        self.t = t

    @property
    def width(self) -> int:
        return _FLOAT_WIDTH[type(self.t)]
1 Like
2

I’ve needed this a number of times but have hacked around it locally and not gotten to it.

I think step 1 is so expose it in the capi similar to how ShapedType is defined.

I’d need to stare more at the python side but I think this can be retrofitted without meta classes. The Tensor type hierarchy may be a guide.

3 
bool mlirTypeIsAFloat(MlirType t) {
  return llvm::isa<FloatType>(unwrap(t));
}

class PyFloatType : public PyConcreteType<PyFloatType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat;
  ...
}
...
class PyF16Type : public PyConcreteType<PyF16Type, PyFloatType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF16;
}

Rest you get for free because these types have for a long time (some patch I made last year - don’t remember) already supported isinstance just by virtue of being cast correctly in all the right places when being shipped from C++ → Python.

4

We probably want the super-type to also give access to the bit width. I have often wanted that at the super-type.

5

Yea sure - there are lots/many/several things that can go into PyFloatType - I was just confirming that there’s no meta hijinks necessary (as you suggested).

6

Sure, the metaclass is only necessary to emulate the base class.

Maks, how do you feel about converting your sketch into a PR? :slight_smile:

7

I feel I am stretched very thin lately (working pretty hard to try to finish up dissertation stuff so I can graduate…) but I’d be happy to advise/steer/review. You can DM me on discord under the same username/handle if you want a more in-real-time discussion.

8

Thanks, Maks! I’ve sent a PR: [MLIR][Python] Added a base class to all builtin floating point types by superbobry · Pull Request #81720 · llvm/llvm-project · GitHub.

2 Likes