[Mlir-commits] [mlir] [mlir][arith] Expand the scaling ops with the value of the scale (PR #217892)

[Hung-Kuan Tseng]

Tim096 wrote:

Following up on the `ArithToAMDGPU` paragraph in the description: the tree
answers that question, so I should not have left it open.

`mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td` describes every `cvt.scalef32` op
this path lowers to as multiplying or dividing by *the exponent part of* the
scale -- 16 occurrences, all consistent. For the two the scaling ops reach:

- `:2952`, `cvt.scalef32.pk.{fp8,bf8}.f32`, which `amdgpu.packed_scaled_trunc`
  becomes: "dividing by the exponent in `scale`".
- `:2996`, `cvt.scalef32.pk.{f32,f16,bf16}.{fp8,bf8}`, which
  `amdgpu.scaled_ext_packed` becomes: "multiplying by the exponent in `scale`".

And the conversion places no restriction on the scale type. There is no
`f8E8M0FNU` check anywhere in `ArithToAMDGPU.cpp` -- contrast `ExpandOps.cpp:682`,
which bails -- the scale is cast to `f32` and forwarded verbatim (`:592`), and
there are tests that pass a plain `f32` scale through
(`mlir/test/Conversion/ArithToAMDGPU/scaling-extf.mlir:265`, `:277`). The
patterns are added only for `*maybeChipset == kGfx950` (`:729`).

That has a consequence worth stating plainly, because it cuts against this
patch: under `in / scale` the two lowerings disagree on the **value** of a
non-power-of-two scale rather than on its rounding, and this patch makes that
gap wider rather than narrower. The generic expansion starts using the whole
scale; the hardware path keeps using its exponent. It is the divergence #215295
was opened about, now fixed on one side only.

Two ways to close it, and I would rather be told which than pick:

1. @krzysz00's suggestion on the issue -- `ArithToAMDGPU` takes the instruction's
   fast path only when the scale already is `f8E8M0FNU`, or when the IR contains
   an explicit `arith.truncf ... toward_zero` for it, and otherwise emits the
   multiply or divide. This keeps the op's meaning single.
2. Treat the exponent-only hardware path as an approximation the target is
   allowed to make, and say so in the op documentation.

Either is a separate patch. I raise it on this one because a reviewer should get
to weigh "narrows the gap in `arith`, widens it against the target" while
reviewing, rather than find it afterwards.


https://github.com/llvm/llvm-project/pull/217892