[Mlir-commits] [mlir] [mlir][arith] Expand the scaling ops with the value of the scale (PR #217892)
Hung-Kuan Tseng
llvmlistbot at llvm.org
Fri Aug 21 05:09:46 PDT 2026
https://github.com/Tim096 created https://github.com/llvm/llvm-project/pull/217892
`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.
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. For a scale that already is `f8E8M0FNU` that cast is the same
widening as before, so those expansions are emitted unchanged.
`arith.extf` and `arith.truncf` each require a strict change in width, so a
scale as wide as the operand but of a different type has no conversion to
spell. The previous code did not check for that and built an op that does not
verify; this reports a match failure instead.
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. Four are added: an `f8E5M3FNU`
scale for each op, and the equal-width 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 in this patch: `ArithToAMDGPU` casts any scale type to `f32` and hands it
to `amdgpu::PackedScaledTruncOp`. If that instruction reads only the exponent
bits, that path drops the mantissa under this reading too, but the answer is
in hardware documentation rather than in the tree, so it is a separate change.
cc @tgymnich @krzysz00 @umangyadav @kuhar
>From bf5128a96ff8b84da81f8211aecf2a2e82604567 Mon Sep 17 00:00:00 2001
From: Hung-Kuan Tseng <tseng.tim096 at gmail.com>
Date: Fri, 21 Aug 2026 15:39:56 +0800
Subject: [PATCH] [mlir][arith] Expand the scaling ops with the value of the
scale
`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.
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. For a scale that already is `f8E8M0FNU` that cast is the same
widening as before, so those expansions are emitted unchanged.
`arith.extf` and `arith.truncf` each require a strict change in width, so a
scale as wide as the operand but of a different type has no conversion to
spell. The previous code did not check for that and built an op that does not
verify; this reports a match failure instead.
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. Four are added: an `f8E5M3FNU`
scale for each op, and the equal-width case each op now rejects.
The op descriptions carry the same truncation in their example lowerings and
are updated with it.
---
.../include/mlir/Dialect/Arith/IR/ArithOps.td | 21 ++--
.../Dialect/Arith/Transforms/ExpandOps.cpp | 87 ++++++++--------
mlir/test/Dialect/Arith/expand-ops.mlir | 99 +++++++++++++++----
3 files changed, 139 insertions(+), 68 deletions(-)
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..af014dea5aad5 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ExpandOps.cpp
@@ -662,37 +662,48 @@ struct F8E8M0TruncFOpConverter : public OpRewritePattern<arith::TruncFOp> {
}
};
+/// Casts `value` to `targetTy`, which must be a float type of the same shape.
+/// `arith.extf` and `arith.truncf` both require a strict change in width, so a
+/// conversion between two distinct types of equal width has no spelling and is
+/// reported as a failure rather than built.
+static FailureOr<Value> castFloatValue(ImplicitLocOpBuilder &b, Value value,
+ Type targetTy,
+ arith::FastMathFlagsAttr fastmath) {
+ Type sourceTy = value.getType();
+ if (sourceTy == targetTy)
+ return value;
+
+ unsigned sourceWidth = getElementTypeOrSelf(sourceTy).getIntOrFloatBitWidth();
+ unsigned targetWidth = getElementTypeOrSelf(targetTy).getIntOrFloatBitWidth();
+ if (sourceWidth == targetWidth)
+ return failure();
+ if (sourceWidth < targetWidth)
+ return Value(arith::ExtFOp::create(b, targetTy, value, fastmath));
+ return Value(arith::TruncFOp::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 cannot be cast to the result type: neither arith.extf nor "
+ "arith.truncf spells a conversion between types of equal width");
}
- 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 +711,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 +721,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 cannot be cast to the input type: neither arith.extf nor "
+ "arith.truncf spells a conversion between types of equal width");
+ }
+ 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..1659d3b846d55 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,51 @@ 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]]
+
+// -----
+
+// arith.extf and arith.truncf both require a strict change in width, so a scale
+// as wide as the input but of a different type has no conversion to spell.
+func.func @invalid_scaling_truncf_equal_width_scale(%arg0: f8E4M3FN, %arg1 : f8E5M2FNUZ) -> f4E2M1FN {
+ // expected-error at +1 {{failed to legalize operation 'arith.scaling_truncf' that was explicitly marked illegal}}
+ %0 = arith.scaling_truncf %arg0, %arg1 : f8E4M3FN, f8E5M2FNUZ to f4E2M1FN
+ return %0 : f4E2M1FN
+}
+
// -----
func.func @extf_f8E8M0FNU_to_f32(%arg0 : f8E8M0FNU) -> f32 {
@@ -480,20 +517,46 @@ 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 @invalid_scaling_extf_equal_width_scale(%arg0: f4E2M1FN, %arg1 : f8E5M2FNUZ) -> f8E4M3FN {
+ // expected-error at +1 {{failed to legalize operation 'arith.scaling_extf' that was explicitly marked illegal}}
+ %0 = arith.scaling_extf %arg0, %arg1 : f4E2M1FN, f8E5M2FNUZ to f8E4M3FN
+ return %0 : f8E4M3FN
+}
+
// -----
func.func @scaling_extf_vector_to_f32(%arg0: vector<4xf4E2M1FN>, %arg1 : vector<4xf8E8M0FNU>) -> vector<4xf32> {
@@ -541,8 +604,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 +618,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]]
More information about the Mlir-commits
mailing list