Rethink on approach to low precision FP types

MLIR
21

That’s correct. If you want to use #my_research.custom_float inside an MLIR vector type or use certain ops such as those from the arith dialect, we need (3) because those types/ops only accept types that upstream MLIR considers to be float types (i.e., FloatType).

Yes, that’s what I implemented in the above-mentioned PR. It comes at a compilation time cost though.

22

Thanks for the analysis. I think this provides a pretty reasonably derived worst case bound to what the cost is.

I don’t know how to make this decision, but the first thing I would want to ask is what the typical cost is for a moderately complicated real world program. My guess is that the cost is probably negligible in such cases. And that is a far easier position to argue a benefit from if it is true.

23

The “heaviest” integration test that we have in MLIR is the sparse compiler test suite. But that’s still a pretty small test.

Benchmark 1: mlir-opt sparse_binary.mlir --sparsifier="enable-runtime-library=false enable-buffer-initialization=true vl=2 reassociate-fp-reductions=true enable-index-optimizations=true"
  BEFORE
  Time (mean ± σ):     229.5 ms ±   4.4 ms    [User: 331.0 ms, System: 103.6 ms]
  Range (min … max):   222.7 ms … 241.3 ms    50 runs

  AFTER
  Time (mean ± σ):     230.8 ms ±   3.6 ms    [User: 332.2 ms, System: 105.8 ms]
  Range (min … max):   225.1 ms … 247.4 ms    50 runs

Do you have some larger examples/models in IREE that I can try? (And instructions how to compile them?) If I remember correctly, the IREE CI can also measure compilation time. Would that be a good indicator or is the data too noisy?

24

@marbre anything come to mind to test?

May also want to ask the flang guys.

25

Just in case it might be tangentially relevant as a proof of existence for the parametrizability of the family of FP types, here are some reference (slow) conversion helpers between f32 and {f16,bf16,f8e4m3,f8e5m2,f8e4m3fnuz,f8e5m2fnuz}. Just to give a taste of what a parameter space might look like.

26

May also want to ask the flang guys.

I did a few runs on big Fortran application using floating points (mainly WRF from SPEC FP 2017 (couple hundred Fortran files that take ~1100s to compile on an AMD EPYC 9334) and aermod_11 from Polyhedron (about 50kloc that takes about 80s to compile), I saw no measurable compile time impact with @matthias-springer patch.

Instrumenting the code, getFloatSemantics current number of calls in flang (at -O3) is roughly linear with the number of lines in these apps, so the extra cost there is not visible when compared to all the other compilation costs.

So I do not see problems with this change on our side.

27

That evidence is sufficient for me. +1 on opening up the type as @matthias-springer recommends.

28

The FloatType interface change has been merged.

I am now looking at a few further cleanups on the MLIR side. Some fundamental types get special treatment in MLIR, and I’m wondering if low-precision FP types should get that treatment. Especially, given that the list is expected to grow further.

Do we need to have all floating-point types in the Builder API?

class Builder {
  // Types.
  FloatType getFloat4E2M1FNType();
  FloatType getFloat6E2M3FNType();
  FloatType getFloat6E3M2FNType();
  FloatType getFloat8E5M2Type();
  FloatType getFloat8E4M3Type();
  FloatType getFloat8E4M3FNType();
  FloatType getFloat8E5M2FNUZType();
  FloatType getFloat8E4M3FNUZType();
  FloatType getFloat8E4M3B11FNUZType();
  FloatType getFloat8E3M4Type();
  FloatType getFloat8E8M0FNUType();
  FloatType getBF16Type();
  FloatType getF16Type();
  FloatType getTF32Type();
  FloatType getF32Type();
  FloatType getF64Type();
  FloatType getF80Type();
  FloatType getF128Type();

Users can always write b.getType<Float8E4M3FNType>() or Float8E4M3FNType::get(b.getContext()) instead of b.getFloat8E4M3FNType(). My thinking is to remove all but the most commonly used floating-point types (and the ones that are valid LLVM types) from the Builder API. I.e., what would be remaining: BF16, F16, TF32, F32, F64, F80, F128.

We could even go as far as removing all get...Type() from the builder API. But that would be a quite drastic change, affecting lots of code. So probably not…

Do we need to cache all floating-point types in MLIRContext?

Or just the most frequently used ones. (Same list as above: BF16, F16, TF32, F32, F64, F80, F128.)

We currently cache these types for faster lookup:

class MLIRContextImpl {
  /// Cached Type Instances.
  Float4E2M1FNType f4E2M1FNTy;
  Float6E2M3FNType f6E2M3FNTy;
  Float6E3M2FNType f6E3M2FNTy;
  Float8E5M2Type f8E5M2Ty;
  Float8E4M3Type f8E4M3Ty;
  Float8E4M3FNType f8E4M3FNTy;
  Float8E5M2FNUZType f8E5M2FNUZTy;
  Float8E4M3FNUZType f8E4M3FNUZTy;
  Float8E4M3B11FNUZType f8E4M3B11FNUZTy;
  Float8E3M4Type f8E3M4Ty;
  Float8E8M0FNUType f8E8M0FNUTy;
  BFloat16Type bf16Ty;
  Float16Type f16Ty;
  FloatTF32Type tf32Ty;
  Float32Type f32Ty;
  Float64Type f64Ty;
  Float80Type f80Ty;
  Float128Type f128Ty;
  IndexType indexTy;
  IntegerType int1Ty, int8Ty, int16Ty, int32Ty, int64Ty, int128Ty;
  NoneType noneType;

Any thoughts?

29

I vote we remove them, they don’t add any real measurable value over using the templated alternative (which is why I added the template in the first place). I cleansed the builder API of a lot of these types of methods a long time ago, it’d be good to keep cleaning up that API.

We should only cache the more fundamental types, the rest can go through the default path (which is already cached, just with a map lookup as opposed to directly through the context).

– River

30

One more idea for allowing user provided semantics would be to have callbacks in fltSemantics which will be guarded in if constexpr so they will not slow down existing semantics but still allow custom paths to be taken. I think the template-ization of initFromIEEEAPInt/convert could make this possible now:

struct fltSemantics {
  // [...]

  // Note: IEEEFloat needed right now, could be APFloat with some internal refactorings
  void (* customInitFromAPInt)(const IEEEFloat&, APInt&) = nullptr;
  void (* customBitcastToAPInt)(const APInt&, IEEEFloat&) = nullptr;
};

Used guarded with an if constexpr like this:

template <const fltSemantics &S>
void IEEEFloat::initFromIEEEAPInt(const APInt &api) {
  if constexpr (S.customInitFromAPInt != nullptr) {
    S.customInitFromAPInt(api, *this);
  }
  else {
    /* existing implementation */
  }
}

This would make new (LLVM API) user provided custom types pay for the dynamic call but guarantee existing types do not get pessimized for the extensibility.

Replacing the function pointer with static member functions could also get rid of the dynamic call but would probably need to expose many functions as templates in the APFloat API which probably is not an option.

We’d also need to allow semantics which are not covered by the built-ins in APFloat. This could be done by calling the custom functions from the end of IEEEFloat::initFromAPInt instead of llvm_unreachable (maybe even instead of the branch in initFromIEEEAPInt above).

This could also be generalized for arithmetic operations. Having one callback for each op might be a bit much but I think having one callback which can fix the result of any operation might be enough to at least cover special conversion rules for NaN/Inf/etc. which seems to not be implemented for all types right now (see APFloat::convert() followed by bitcastToAPInt() crashes when converting infinity to FiniteOnly semantics (e.g. Float4E2M1FN) · Issue #182433 · llvm/llvm-project · GitHub ).

enum fltOperation { Add, Sub, ... };

struct fltSemantics {
  void (* fixAfterArithOp)(const APFloat& result, fltOperation op);
};

It would add an if (S.fixAfterArithOp != null) S.fixAfterArithOp(*this); to the end of all arith operations. Given that they operate on dynamically sized types this might be reasonable overhead but I haven’t measured. Any thoughts?