[Mlir-commits] [mlir] [mlir][arith] Expand the scaling ops with the value of the scale (PR #217892)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 21 10:56:52 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-arith
Author: Hung-Kuan Tseng (Tim096)
<details>
<summary>Changes</summary>
`ScalingExtFOpConverter` and `ScalingTruncFOpConverter` truncate a scale of
16 bits or wider to `f8E8M0FNU` before using it, and reject an 8-bit scale
that is not already `f8E8M0FNU`:
```mlir
// The scale is rounded to a power of two before it is ever used.
%0 = arith.scaling_truncf %in, %scale : f16, f16 to f4E2M1FN
// No generic expansion exists; this fails to legalize.
%1 = arith.scaling_truncf %in, %e5m3 : f16, f8E5M3FNU to f4E2M1FN
```
Both follow from reading the scale as an exponent. The OCP MXFP spec fixes
that reading only for the `E8M0` scale of an MX format, and hardware that
consumes an `f8E5M3` scale uses its mantissa, so the general reading is that
the input is multiplied or divided by the value of the scale. `ArithToAMDGPU`
already reads it that way in the IR: it casts the scale to `f32` by width and
passes the value on, with no `f8E8M0FNU` check anywhere in the file.
The truncation is a wrong value rather than a wrong rounding mode. Constant
folding the IR the expansion emits for the scale gives the divisor it
actually uses:
| scale | divisor used | exact |
|---|---|---|
| `3.0` | `2.0` | `3.0` |
| `1.6` | `1.0` | `1.6` |
| `7.0` | `4.0` | `7.0` |
The error approaches a factor of two just below a power of two.
## Change
Both converters cast the scale to the type the arithmetic happens in and use
it as it is, picking the cast that fits the change in width: `arith.extf` and
`arith.truncf` where the width changes, and `arith.convertf` where it does
not. For a scale that already is `f8E8M0FNU` and arithmetic wider than eight
bits, that cast is the same widening as before, so those expansions are
emitted unchanged.
No cast bridges a scalar scale and a shaped operand, which the op verifier
permits because `ElementwiseMappable` exempts scalars, and `arith` cannot
broadcast one. That combination used to build an `arith.extf` that does not
verify; it is now reported as a match failure.
Eight tests in `expand-ops.mlir` use a scale that is not `f8E8M0FNU` and all
eight change. Six assert the `arith.truncf ... to f8E8M0FNU` step that is
gone. Two assert that a `f8E5M2FNUZ` scale fails to legalize, which it no
longer does, so they become expansion tests. Six are added: an `f8E5M3FNU`
scale and an equal-width scale for each op, a scale narrower than the
arithmetic type, and the scalar-scale case each op now rejects.
The op descriptions carry the same truncation in their example lowerings and
are updated with it.
---
Context: #<!-- -->215295. That thread asks whether `arith.scaling_truncf(in, scale)`
means `in / scale` or `in / 2^exponent(scale)`; this implements the first,
which is the reading @<!-- -->krzysz00 gave with the CDNA5 ISA sections for `f8E5M3`
scales. @<!-- -->tgymnich read it the other way earlier on #<!-- -->215123 and has not said
where he stands since, so this is the patch for one of the two answers rather
than a settled question -- please say so if you still read it as the exponent.
Not closed by this patch: `ArithToAMDGPU` casts any scale type to `f32` and
hands it to `amdgpu::PackedScaledTruncOp`, and `ROCDLOps.td` describes every
`cvt.scalef32` op it becomes as multiplying or dividing by *the exponent part
of* the scale (16 occurrences, e.g. `:2952`, `:2996`). So under this reading
the two lowerings disagree on the value of a non-power-of-two scale, and this
patch widens that gap rather than narrowing it: the generic expansion starts
using the whole scale while the hardware path keeps using its exponent. Closing
it is a separate change, and there is more than one reasonable way -- see the
comments below.
cc @<!-- -->tgymnich @<!-- -->krzysz00 @<!-- -->umangyadav @<!-- -->kuhar
---
Full diff: https://github.com/llvm/llvm-project/pull/217892.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/Arith/IR/ArithOps.td (+13-8)
- (modified) mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp (+54-42)
- (modified) mlir/test/Dialect/Arith/expand-ops.mlir (+124-18)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
index 54481d3232483..e730d7f2ec1e6 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
+++ b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
@@ -1477,15 +1477,18 @@ def Arith_ScalingExtFOp
```mlir
// Cast scale to result type.
- %0 = arith.truncf %1 : f32 to f8E8M0FNU
- %1 = arith.extf %0 : f8E8M0FNU to f16
+ %scale = arith.extf %arg1 : f8E8M0FNU to f16
// Cast input to result type.
- %2 = arith.extf %3 : f4E2M1FN to f16
+ %in = arith.extf %arg0 : f4E2M1FN to f16
// Perform scaling
- %3 = arith.mulf %2, %1 : f16
+ %res = arith.mulf %in, %scale : f16
```
+ The input is multiplied by the value of the scale. A scale type that has a
+ mantissa, such as `f8E5M3FNU`, therefore contributes it: the scale is not
+ reduced to its exponent.
+
It propagates NaN values. Therefore, if either scale or the input element
contains NaN, then the output element value will also be a NaN.
@@ -1668,15 +1671,17 @@ def Arith_ScalingTruncFOp
```mlir
// Cast scale to input type.
- %0 = arith.truncf %1 : f32 to f8E8M0FNU
- %1 = arith.extf %0 : f8E8M0FNU to f16
+ %scale = arith.extf %arg1 : f8E8M0FNU to f16
// Perform scaling.
- %3 = arith.divf %2, %1 : f16
+ %scaled = arith.divf %arg0, %scale : f16
// Cast to result type.
- %4 = arith.truncf %3 : f16 to f4E2M1FN
+ %res = arith.truncf %scaled : f16 to f4E2M1FN
```
+ The input is divided by the value of the scale. A scale type that has a
+ mantissa, such as `f8E5M3FNU`, therefore contributes it: the scale is not
+ reduced to its exponent.
Example:
diff --git a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
index 02ed6ccd87a42..9f61919532bd8 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -662,37 +662,57 @@ struct F8E8M0TruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
}
};
+/// Casts `value` to the float type `targetTy`, picking the cast that fits the
+/// change in width: `arith.extf` and `arith.truncf` each require a strict one,
+/// and `arith.convertf` covers two types of equal width. Fails when the two
+/// have different shapes, which none of the three casts can bridge.
+static FailureOr<Value> castFloatValue(ImplicitLocOpBuilder &b, Value value,
+ Type targetTy,
+ arith::FastMathFlagsAttr fastmath) {
+ Type sourceTy = value.getType();
+ if (sourceTy == targetTy)
+ return value;
+
+ auto sourceShapedTy = dyn_cast<ShapedType>(sourceTy);
+ auto targetShapedTy = dyn_cast<ShapedType>(targetTy);
+ if (static_cast<bool>(sourceShapedTy) != static_cast<bool>(targetShapedTy))
+ return failure();
+ if (sourceShapedTy && sourceShapedTy.getShape() != targetShapedTy.getShape())
+ return failure();
+
+ unsigned sourceWidth = getElementTypeOrSelf(sourceTy).getIntOrFloatBitWidth();
+ unsigned targetWidth = getElementTypeOrSelf(targetTy).getIntOrFloatBitWidth();
+ if (sourceWidth < targetWidth)
+ return Value(arith::ExtFOp::create(b, targetTy, value, fastmath));
+ if (sourceWidth > targetWidth)
+ return Value(
+ arith::TruncFOp::create(b, targetTy, value, nullptr, fastmath));
+ return Value(
+ arith::ConvertFOp::create(b, targetTy, value, nullptr, fastmath));
+}
+
struct ScalingExtFOpConverter : public OpRewritePattern<arith::ScalingExtFOp> {
using Base::Base;
LogicalResult matchAndRewrite(arith::ScalingExtFOp op,
PatternRewriter &rewriter) const final {
ImplicitLocOpBuilder b(op.getLoc(), rewriter);
Value inputOperand = op.getIn();
- Value scaleOperand = op.getScale();
- Type scaleTy = scaleOperand.getType();
- Type scaleETy = getElementTypeOrSelf(scaleOperand);
- // allow implicit exponent extraction from 16/32 bits floats
- if (scaleETy.getIntOrFloatBitWidth() >= 16) {
- scaleETy = b.getF8E8M0Type();
- scaleTy = cloneToShapedType(scaleTy, scaleETy);
- scaleOperand = arith::TruncFOp::create(b, scaleTy, scaleOperand, nullptr,
- op.getFastmathAttr());
- }
- // Catch scale types like f8E5M2.
- if (!llvm::isa<Float8E8M0FNUType>(scaleETy)) {
+ Type resultTy = op.getType();
+ // The scale multiplies the input by its value, so it is cast to the type
+ // the multiplication happens in rather than reduced to its exponent. For a
+ // f8E8M0FNU scale the cast is the same widening this used to perform, and
+ // it keeps propagating NaNs.
+ FailureOr<Value> scaleExt =
+ castFloatValue(b, op.getScale(), resultTy, op.getFastmathAttr());
+ if (failed(scaleExt)) {
return rewriter.notifyMatchFailure(
- op, "scaling_extf is using scales of type which can not be converted "
- "to f8E8M0FNU");
+ op, "scale and result have different shapes, so the scale cannot be "
+ "cast to the type the multiplication happens in");
}
- Type resultTy = op.getType();
- // extf on scale will essentially create floating point number
- // of type resulTy that is 2^scale and will also propagate NaNs
- Value scaleExt =
- arith::ExtFOp::create(b, resultTy, scaleOperand, op.getFastmathAttr());
Value inputExt =
arith::ExtFOp::create(b, resultTy, inputOperand, op.getFastmathAttr());
Value result =
- arith::MulFOp::create(b, inputExt, scaleExt, op.getFastmathAttr());
+ arith::MulFOp::create(b, inputExt, *scaleExt, op.getFastmathAttr());
rewriter.replaceOp(op, result);
return success();
}
@@ -700,8 +720,8 @@ struct ScalingExtFOpConverter : public OpRewritePattern<arith::ScalingExtFOp> {
/*
Expands arith.ScalingTruncFOp(in, scale) into
- scale = arith.truncf(scale) : scaleTy -> f8E8M0FNU
- result = arith.truncf(in / (2^scale))
+ scale = <cast scale to the type of in>
+ result = arith.truncf(in / scale)
*/
struct ScalingTruncFOpConverter
: public OpRewritePattern<arith::ScalingTruncFOp> {
@@ -710,28 +730,20 @@ struct ScalingTruncFOpConverter
PatternRewriter &rewriter) const final {
ImplicitLocOpBuilder b(op.getLoc(), rewriter);
Value inputOperand = op.getIn();
- Value scaleOperand = op.getScale();
- Type scaleTy = scaleOperand.getType();
- Type scaleETy = getElementTypeOrSelf(scaleOperand);
- // allow implicit exponent extraction from 16/32 bits floats
- if (scaleETy.getIntOrFloatBitWidth() >= 16) {
- scaleETy = b.getF8E8M0Type();
- scaleTy = cloneToShapedType(scaleTy, scaleETy);
- scaleOperand = arith::TruncFOp::create(b, scaleTy, scaleOperand, nullptr,
- op.getFastmathAttr());
- }
- if (!llvm::isa<Float8E8M0FNUType>(scaleETy)) {
- return rewriter.notifyMatchFailure(
- op, "scaling_truncf is using scales type which can not be converted "
- "to f8E8M0FNU");
- }
Type resultTy = op.getType();
Type inputTy = inputOperand.getType();
- // this will create a floating point number of type
- // inputTy that is 2^scale and will also propagate NaNs
- scaleOperand =
- arith::ExtFOp::create(b, inputTy, scaleOperand, op.getFastmathAttr());
- Value result = arith::DivFOp::create(b, inputOperand, scaleOperand,
+ // The scale divides the input by its value, so it is cast to the type the
+ // division happens in rather than reduced to its exponent. For a f8E8M0FNU
+ // scale the cast is the same widening this used to perform, and it keeps
+ // propagating NaNs.
+ FailureOr<Value> scaleCast =
+ castFloatValue(b, op.getScale(), inputTy, op.getFastmathAttr());
+ if (failed(scaleCast)) {
+ return rewriter.notifyMatchFailure(
+ op, "scale and input have different shapes, so the scale cannot be "
+ "cast to the type the division happens in");
+ }
+ Value result = arith::DivFOp::create(b, inputOperand, *scaleCast,
op.getFastmathAttr());
Value resultCast = arith::TruncFOp::create(
b, resultTy, result, op.getRoundingmodeAttr(), op.getFastmathAttr());
diff --git a/mlir/test/Dialect/Arith/expand-ops.mlir b/mlir/test/Dialect/Arith/expand-ops.mlir
index 20f00b82505e7..75252a3125813 100644
--- a/mlir/test/Dialect/Arith/expand-ops.mlir
+++ b/mlir/test/Dialect/Arith/expand-ops.mlir
@@ -340,11 +340,12 @@ func.func @scaling_truncf_propagate_rounding_mode_fast_math(%arg0 : vector<4xf16
%0 = arith.scaling_truncf %arg0, %arg1 to_nearest_even fastmath<fast> : vector<4xf16>, vector<4xf16> to vector<4xf6E3M2FN>
return %0 : vector<4xf6E3M2FN>
}
+// The scale already has the type the division happens in, so it is used as it
+// is rather than routed through f8E8M0FNU.
// SCHECK-LABEL: @scaling_truncf_propagate_rounding_mode_fast_math
-// SCHECK: %[[SCALEF8:.+]] = arith.truncf %arg1 fastmath<fast> : vector<4xf16> to vector<4xf8E8M0FNU>
-// SCHECK: %[[SCALEINTY:.+]] = arith.extf %[[SCALEF8]] fastmath<fast> : vector<4xf8E8M0FNU> to vector<4xf16>
-// SCHECK: %[[DIVF:.+]] = arith.divf %arg0, %[[SCALEINTY]] fastmath<fast> : vector<4xf16>
-// SCHECK: %[[TRUNCF:.+]] = arith.truncf [[_:%[a-zA-Z0-9_]+]] to_nearest_even fastmath<fast> : vector<4xf16> to vector<4xf6E3M2FN>
+// SCHECK-NOT: f8E8M0FNU
+// SCHECK: %[[DIVF:.+]] = arith.divf %arg0, %arg1 fastmath<fast> : vector<4xf16>
+// SCHECK: %[[TRUNCF:.+]] = arith.truncf %[[DIVF]] to_nearest_even fastmath<fast> : vector<4xf16> to vector<4xf6E3M2FN>
// SCHECK: return %[[TRUNCF]] : vector<4xf6E3M2FN>
// -----
@@ -354,8 +355,10 @@ func.func @scaling_truncf_f16_to_f4E2M1FN_using_f16_scales(%arg0: f16, %arg1 : f
return %0 : f4E2M1FN
}
// SCHECK-LABEL: @scaling_truncf_f16_to_f4E2M1FN_using_f16_scales
-// SCHECK: %[[SCALETRUNCF:.+]] = arith.truncf %arg1 : f16 to f8E8M0FN
-// SCHECK: return
+// SCHECK-NOT: f8E8M0FNU
+// SCHECK: %[[DIVF:.+]] = arith.divf %arg0, %arg1 : f16
+// SCHECK: %[[TRUNCF:.+]] = arith.truncf %[[DIVF]] : f16 to f4E2M1FN
+// SCHECK: return %[[TRUNCF]]
// -----
func.func @scaling_truncf_vector_f16_to_f4E2M1FN_using_f16_scales(%arg0: vector<4xf16>, %arg1 : vector<4xf16>) -> vector<4xf4E2M1FN> {
@@ -363,17 +366,81 @@ func.func @scaling_truncf_vector_f16_to_f4E2M1FN_using_f16_scales(%arg0: vector<
return %0 : vector<4xf4E2M1FN>
}
// SCHECK-LABEL: @scaling_truncf_vector_f16_to_f4E2M1FN_using_f16_scales
-// SCHECK: %[[SCALETRUNCF:.+]] = arith.truncf %arg1 : vector<4xf16> to vector<4xf8E8M0FNU>
-// SCHECK: return
+// SCHECK-NOT: f8E8M0FNU
+// SCHECK: %[[DIVF:.+]] = arith.divf %arg0, %arg1 : vector<4xf16>
+// SCHECK: %[[TRUNCF:.+]] = arith.truncf %[[DIVF]] : vector<4xf16> to vector<4xf4E2M1FN>
+// SCHECK: return %[[TRUNCF]]
// -----
-func.func @invalid_scaling_truncf_to_f4E2M1FN(%arg0: f16, %arg1 : f8E5M2FNUZ) -> f4E2M1FN {
- // expected-error at +1 {{failed to legalize operation 'arith.scaling_truncf' that was explicitly marked illegal}}
+// A scale that is neither f8E8M0FNU nor as wide as the input is widened to the
+// input type and divided by, rather than rejected.
+func.func @scaling_truncf_f16_to_f4E2M1FN_using_f8E5M2FNUZ_scales(%arg0: f16, %arg1 : f8E5M2FNUZ) -> f4E2M1FN {
%0 = arith.scaling_truncf %arg0, %arg1 : f16, f8E5M2FNUZ to f4E2M1FN
return %0 : f4E2M1FN
}
+// SCHECK-LABEL: @scaling_truncf_f16_to_f4E2M1FN_using_f8E5M2FNUZ_scales
+// SCHECK: %[[SCALE:.+]] = arith.extf %arg1 : f8E5M2FNUZ to f16
+// SCHECK: %[[DIVF:.+]] = arith.divf %arg0, %[[SCALE]] : f16
+// SCHECK: %[[TRUNCF:.+]] = arith.truncf %[[DIVF]] : f16 to f4E2M1FN
+// SCHECK: return %[[TRUNCF]]
+
+// -----
+
+// An f8E5M3FNU scale carries a mantissa, and that mantissa is part of the value
+// the input is divided by.
+func.func @scaling_truncf_f16_to_f4E2M1FN_using_f8E5M3FNU_scales(%arg0: f16, %arg1 : f8E5M3FNU) -> f4E2M1FN {
+ %0 = arith.scaling_truncf %arg0, %arg1 : f16, f8E5M3FNU to f4E2M1FN
+ return %0 : f4E2M1FN
+}
+
+// SCHECK-LABEL: @scaling_truncf_f16_to_f4E2M1FN_using_f8E5M3FNU_scales
+// SCHECK: %[[SCALE:.+]] = arith.extf %arg1 : f8E5M3FNU to f16
+// SCHECK: %[[DIVF:.+]] = arith.divf %arg0, %[[SCALE]] : f16
+// SCHECK: %[[TRUNCF:.+]] = arith.truncf %[[DIVF]] : f16 to f4E2M1FN
+// SCHECK: return %[[TRUNCF]]
+
+// -----
+
+// A scale as wide as the input but of a different type is bridged by
+// arith.convertf, which neither arith.extf nor arith.truncf can spell.
+func.func @scaling_truncf_equal_width_scale(%arg0: f8E4M3FN, %arg1 : f8E5M2FNUZ) -> f4E2M1FN {
+ %0 = arith.scaling_truncf %arg0, %arg1 : f8E4M3FN, f8E5M2FNUZ to f4E2M1FN
+ return %0 : f4E2M1FN
+}
+
+// SCHECK-LABEL: @scaling_truncf_equal_width_scale
+// SCHECK: %[[SCALE:.+]] = arith.convertf %arg1 : f8E5M2FNUZ to f8E4M3FN
+// SCHECK: %[[DIVF:.+]] = arith.divf %arg0, %[[SCALE]] : f8E4M3FN
+// SCHECK: %[[TRUNCF:.+]] = arith.truncf %[[DIVF]] : f8E4M3FN to f4E2M1FN
+// SCHECK: return %[[TRUNCF]]
+
+// -----
+
+// The scale is narrowed to the type the division happens in. The op's rounding
+// mode belongs to the result cast, not to this one; fastmath reaches both.
+func.func @scaling_truncf_f16_to_f4E2M1FN_using_f32_scales(%arg0: f16, %arg1 : f32) -> f4E2M1FN {
+ %0 = arith.scaling_truncf %arg0, %arg1 to_nearest_even fastmath<fast> : f16, f32 to f4E2M1FN
+ return %0 : f4E2M1FN
+}
+
+// SCHECK-LABEL: @scaling_truncf_f16_to_f4E2M1FN_using_f32_scales
+// SCHECK: %[[SCALE:.+]] = arith.truncf %arg1 fastmath<fast> : f32 to f16
+// SCHECK: %[[DIVF:.+]] = arith.divf %arg0, %[[SCALE]] fastmath<fast> : f16
+// SCHECK: %[[TRUNCF:.+]] = arith.truncf %[[DIVF]] to_nearest_even fastmath<fast> : f16 to f4E2M1FN
+// SCHECK: return %[[TRUNCF]]
+
+// -----
+
+// No cast bridges a scalar scale and a shaped operand, and arith cannot
+// broadcast one.
+func.func @invalid_scaling_truncf_scalar_scale(%arg0: vector<4xf16>, %arg1 : f16) -> vector<4xf4E2M1FN> {
+ // expected-error at +1 {{failed to legalize operation 'arith.scaling_truncf' that was explicitly marked illegal}}
+ %0 = arith.scaling_truncf %arg0, %arg1 : vector<4xf16>, f16 to vector<4xf4E2M1FN>
+ return %0 : vector<4xf4E2M1FN>
+}
+
// -----
func.func @extf_f8E8M0FNU_to_f32(%arg0 : f8E8M0FNU) -> f32 {
@@ -480,20 +547,59 @@ func.func @scaling_extf_to_f32_using_f16_scales(%arg0: f4E2M1FN, %arg1 : f16) ->
}
// SCHECK-LABEL: @scaling_extf_to_f32_using_f16_scales
-// SCHECK: %[[TRUNCF_SCALE:.+]] = arith.truncf %arg1 : f16 to f8E8M0FNU
-// SCHECK: %[[EXT_SCALE:.+]] = arith.extf %[[TRUNCF_SCALE]] : f8E8M0FNU to f32
+// SCHECK-NOT: f8E8M0FNU
+// SCHECK: %[[EXT_SCALE:.+]] = arith.extf %arg1 : f16 to f32
// SCHECK: %[[EXT_INPUT:.+]] = arith.extf %arg0 : f4E2M1FN to f32
// SCHECK: %[[RESULT:.+]] = arith.mulf %[[EXT_INPUT]], %[[EXT_SCALE]] : f32
// SCHECK: return %[[RESULT]]
// -----
-func.func @invalid_scaling_extf_to_f32(%arg0: f4E2M1FN, %arg1 : f8E5M2FNUZ) -> f32 {
- // expected-error at +1 {{failed to legalize operation 'arith.scaling_extf' that was explicitly marked illegal}}
+func.func @scaling_extf_to_f32_using_f8E5M2FNUZ_scales(%arg0: f4E2M1FN, %arg1 : f8E5M2FNUZ) -> f32 {
%0 = arith.scaling_extf %arg0, %arg1 : f4E2M1FN, f8E5M2FNUZ to f32
return %0 : f32
}
+// SCHECK-LABEL: @scaling_extf_to_f32_using_f8E5M2FNUZ_scales
+// SCHECK: %[[EXT_SCALE:.+]] = arith.extf %arg1 : f8E5M2FNUZ to f32
+// SCHECK: %[[EXT_INPUT:.+]] = arith.extf %arg0 : f4E2M1FN to f32
+// SCHECK: %[[RESULT:.+]] = arith.mulf %[[EXT_INPUT]], %[[EXT_SCALE]] : f32
+// SCHECK: return %[[RESULT]]
+
+// -----
+
+func.func @scaling_extf_to_f32_using_f8E5M3FNU_scales(%arg0: f4E2M1FN, %arg1 : f8E5M3FNU) -> f32 {
+ %0 = arith.scaling_extf %arg0, %arg1 : f4E2M1FN, f8E5M3FNU to f32
+ return %0 : f32
+}
+
+// SCHECK-LABEL: @scaling_extf_to_f32_using_f8E5M3FNU_scales
+// SCHECK: %[[EXT_SCALE:.+]] = arith.extf %arg1 : f8E5M3FNU to f32
+// SCHECK: %[[EXT_INPUT:.+]] = arith.extf %arg0 : f4E2M1FN to f32
+// SCHECK: %[[RESULT:.+]] = arith.mulf %[[EXT_INPUT]], %[[EXT_SCALE]] : f32
+// SCHECK: return %[[RESULT]]
+
+// -----
+
+func.func @scaling_extf_equal_width_scale(%arg0: f4E2M1FN, %arg1 : f8E5M2FNUZ) -> f8E4M3FN {
+ %0 = arith.scaling_extf %arg0, %arg1 : f4E2M1FN, f8E5M2FNUZ to f8E4M3FN
+ return %0 : f8E4M3FN
+}
+
+// SCHECK-LABEL: @scaling_extf_equal_width_scale
+// SCHECK: %[[EXT_SCALE:.+]] = arith.convertf %arg1 : f8E5M2FNUZ to f8E4M3FN
+// SCHECK: %[[EXT_INPUT:.+]] = arith.extf %arg0 : f4E2M1FN to f8E4M3FN
+// SCHECK: %[[RESULT:.+]] = arith.mulf %[[EXT_INPUT]], %[[EXT_SCALE]] : f8E4M3FN
+// SCHECK: return %[[RESULT]]
+
+// -----
+
+func.func @invalid_scaling_extf_scalar_scale(%arg0: vector<4xf4E2M1FN>, %arg1 : f8E8M0FNU) -> vector<4xf32> {
+ // expected-error at +1 {{failed to legalize operation 'arith.scaling_extf' that was explicitly marked illegal}}
+ %0 = arith.scaling_extf %arg0, %arg1 : vector<4xf4E2M1FN>, f8E8M0FNU to vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
// -----
func.func @scaling_extf_vector_to_f32(%arg0: vector<4xf4E2M1FN>, %arg1 : vector<4xf8E8M0FNU>) -> vector<4xf32> {
@@ -541,8 +647,8 @@ func.func @scaling_extf_vector_to_f32_using_f16_scales(%arg0: vector<4xf4E2M1FN>
}
// SCHECK-LABEL: @scaling_extf_vector_to_f32_using_f16_scales
-// SCHECK: %[[TRUNCF_SCALE:.+]] = arith.truncf %arg1 : vector<4xf16> to vector<4xf8E8M0FNU>
-// SCHECK: %[[EXT_SCALE:.+]] = arith.extf %[[TRUNCF_SCALE]] : vector<4xf8E8M0FNU> to vector<4xf32>
+// SCHECK-NOT: f8E8M0FNU
+// SCHECK: %[[EXT_SCALE:.+]] = arith.extf %arg1 : vector<4xf16> to vector<4xf32>
// SCHECK: %[[EXT_INPUT:.+]] = arith.extf %arg0 : vector<4xf4E2M1FN> to vector<4xf32>
// SCHECK: %[[RESULT:.+]] = arith.mulf %[[EXT_INPUT]], %[[EXT_SCALE]] : vector<4xf32>
// SCHECK: return %[[RESULT]]
@@ -555,8 +661,8 @@ func.func @scaling_extf_vector_to_f32_using_f16_scales_fastmath(%arg0: vector<4x
}
// SCHECK-LABEL: @scaling_extf_vector_to_f32_using_f16_scales_fastmath
-// SCHECK: %[[TRUNCF_SCALE:.+]] = arith.truncf %arg1 fastmath<fast> : vector<4xf16> to vector<4xf8E8M0FNU>
-// SCHECK: %[[EXT_SCALE:.+]] = arith.extf %[[TRUNCF_SCALE]] fastmath<fast> : vector<4xf8E8M0FNU> to vector<4xf32>
+// SCHECK-NOT: f8E8M0FNU
+// SCHECK: %[[EXT_SCALE:.+]] = arith.extf %arg1 fastmath<fast> : vector<4xf16> to vector<4xf32>
// SCHECK: %[[EXT_INPUT:.+]] = arith.extf %arg0 fastmath<fast> : vector<4xf4E2M1FN> to vector<4xf32>
// SCHECK: %[[RESULT:.+]] = arith.mulf %[[EXT_INPUT]], %[[EXT_SCALE]] fastmath<fast> : vector<4xf32>
// SCHECK: return %[[RESULT]]
``````````
</details>
https://github.com/llvm/llvm-project/pull/217892
More information about the Mlir-commits
mailing list