Skip to content

Commit b7bc15a

Browse files
lhutton1tatwaichong
lhutton1
and
tatwaichong
committed
[mlir][tosa] Add support for cast_from/to_block_scaled
This commit adds support for the cast_from/to_block_scaled operations from the ext-mxfp extension. This includes: - Operation definition in TosaOps.td - Micro-scaling supported types definition - Shape inference and verifiers - Validation pass checks to ensure usage is only valid when the target environment includes ext-mxfp and at least v1.1.draft of the specification. Note: currently it excludes support for mxint8. This will be added in a later commit. Note: this commit adds support as defined in the spec in https://review.mlplatform.org/c/tosa/specification/+/15362. EXT_MXFP extension is considered experimental and subject to breaking change. Co-authored-by: Tat Wai Chong <[email protected]> Change-Id: I490645ce99b7ccd7021ed06acaf1530b4fbf6dfd
1 parent 47f8fd1 commit b7bc15a

15 files changed

+526
-34
lines changed

‎mlir/include/mlir/Dialect/Tosa/IR/TosaComplianceData.h.inc‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -748,6 +748,32 @@ extensionComplianceMap = {
748748
{{fp8e5m2T, fp32T}, SpecificationVersion::V_1_0},
749749
{{fp16T, fp8e5m2T}, SpecificationVersion::V_1_0},
750750
{{fp32T, fp8e5m2T}, SpecificationVersion::V_1_0}}}}},
751+
{"tosa.cast_from_block_scaled",
752+
{{{Extension::bf16, Extension::mxfp},
753+
{{{fp4e2m1T, fp8ue8m0T, bf16T}, SpecificationVersion::V_1_1_DRAFT},
754+
{{fp6e2m3T, fp8ue8m0T, bf16T}, SpecificationVersion::V_1_1_DRAFT},
755+
{{fp6e3m2T, fp8ue8m0T, bf16T}, SpecificationVersion::V_1_1_DRAFT},
756+
{{fp8e4m3T, fp8ue8m0T, bf16T}, SpecificationVersion::V_1_1_DRAFT},
757+
{{fp8e5m2T, fp8ue8m0T, bf16T}, SpecificationVersion::V_1_1_DRAFT}}, allOf},
758+
{{Extension::mxfp},
759+
{{{fp4e2m1T, fp8ue8m0T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
760+
{{fp6e2m3T, fp8ue8m0T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
761+
{{fp6e3m2T, fp8ue8m0T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
762+
{{fp8e4m3T, fp8ue8m0T, fp32T}, SpecificationVersion::V_1_1_DRAFT},
763+
{{fp8e5m2T, fp8ue8m0T, fp32T}, SpecificationVersion::V_1_1_DRAFT}}}}},
764+
{"tosa.cast_to_block_scaled",
765+
{{{Extension::mxfp},
766+
{{{bf16T, fp4e2m1T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
767+
{{fp32T, fp4e2m1T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
768+
{{fp32T, fp6e2m3T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
769+
{{fp32T, fp6e3m2T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
770+
{{fp32T, fp8e4m3T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
771+
{{fp32T, fp8e5m2T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT}}},
772+
{{Extension::bf16, Extension::mxfp},
773+
{{{bf16T, fp6e2m3T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
774+
{{bf16T, fp6e3m2T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
775+
{{bf16T, fp8e4m3T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT},
776+
{{bf16T, fp8e5m2T, fp8ue8m0T}, SpecificationVersion::V_1_1_DRAFT}}, allOf}}},
751777
{"tosa.rescale",
752778
{{{Extension::int16},
753779
{{{i48T, i48T, i8T, i8T}, SpecificationVersion::V_1_0},

‎mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2472,6 +2472,69 @@ def Tosa_CastOp: Tosa_Op<"cast", [Pure, SameOperandsAndResultShape,
24722472
let hasFolder = 1;
24732473
}
24742474

2475+
//===----------------------------------------------------------------------===//
2476+
// Operator: cast_from_block_scaled
2477+
//===----------------------------------------------------------------------===//
2478+
def Tosa_CastFromBlockScaledOp: Tosa_InferShapedTypeOp<"cast_from_block_scaled"> {
2479+
let summary = "Apply scales from a scale tensor to the values in a value tensor";
2480+
2481+
let description = [{
2482+
Apply the scales from a scale tensor to the values in a value tensor, casting
2483+
the result to the output type. The block dimension must be the last dimension
2484+
of the tensor.
2485+
}];
2486+
2487+
let arguments = (ins
2488+
Tosa_MXFPDataTensorAtLeast1D:$input_data,
2489+
Tosa_MXFPScaleTensorAtLeast1D:$input_scale,
2490+
Tosa_BlockSizeAttr:$block_size
2491+
);
2492+
2493+
let results = (outs
2494+
Tosa_TensorAtLeast1D: $output_data
2495+
);
2496+
2497+
list<Availability> availability = [
2498+
Profile<[Tosa_PRO_FP]>,
2499+
Extension<[Tosa_EXT_BF16, Tosa_EXT_MXFP]>,
2500+
];
2501+
2502+
let hasVerifier = 1;
2503+
let hasCustomAssemblyFormat = 1;
2504+
}
2505+
2506+
//===----------------------------------------------------------------------===//
2507+
// Operator: cast_to_block_scaled
2508+
//===----------------------------------------------------------------------===//
2509+
def Tosa_CastToBlockScaledOp : Tosa_InferShapedTypeOp<"cast_to_block_scaled"> {
2510+
let summary = "Calculate scale tensor values per block, output to separate scale and data tensors.";
2511+
2512+
let description = [{
2513+
Calculate a scale value per block of input values and use that to calculate
2514+
scaled data values from an input tensor. The output tensors are cast to the
2515+
specified scale and value types. The block dimension will be the last dimension
2516+
of the tensor.
2517+
}];
2518+
2519+
let arguments = (ins
2520+
Tosa_TensorAtLeast1D:$input_data,
2521+
Tosa_BlockSizeAttr:$block_size
2522+
);
2523+
2524+
let results = (outs
2525+
Tosa_MXFPDataTensorAtLeast1D:$output_data,
2526+
Tosa_MXFPScaleTensorAtLeast1D:$output_scale
2527+
);
2528+
2529+
list<Availability> availability = [
2530+
Profile<[Tosa_PRO_FP]>,
2531+
Extension<[Tosa_EXT_BF16, Tosa_EXT_MXFP]>
2532+
];
2533+
2534+
let hasVerifier = 1;
2535+
let hasCustomAssemblyFormat = 1;
2536+
}
2537+
24752538
//===----------------------------------------------------------------------===//
24762539
// Operator: rescale
24772540
//===----------------------------------------------------------------------===//

‎mlir/include/mlir/Dialect/Tosa/IR/TosaProfileCompliance.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ class ProfileInfoDepot {
7979

8080
LogicalResult populatationDispatch(Operation *op);
8181

82-
LogicalResult populateProfileInfo(ValueRange operands, Value output);
82+
LogicalResult populateProfileInfo(ValueRange operands, ValueRange output);
8383

8484
// Base
8585
template <typename T>

‎mlir/include/mlir/Dialect/Tosa/IR/TosaTypesBase.td‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,16 @@ def Tosa_MXFPScaleTensor3D : AnyTypeOf<[
199199
TosaUnrankedTensorOf<[Tosa_MXFPScaleNumber]>,
200200
TosaTensorRankOf<[Tosa_MXFPScaleNumber], [3]>
201201
]>;
202+
def Tosa_MXFPDataTensorAtLeast1D : AnyTypeOf<[
203+
TosaUnrankedTensorOf<[Tosa_MXFPNumber]>,
204+
TosaRankedTensorOf<[Tosa_MXFPNumber], [AtLeastRankOne]>],
205+
"tosa-conformant tensor of at least rank 1", "::mlir::TensorType"
206+
>;
207+
def Tosa_MXFPScaleTensorAtLeast1D : AnyTypeOf<[
208+
TosaUnrankedTensorOf<[Tosa_MXFPScaleNumber]>,
209+
TosaRankedTensorOf<[Tosa_MXFPScaleNumber], [AtLeastRankOne]>],
210+
"tosa-conformant tensor of at least rank 1", "::mlir::TensorType"
211+
>;
202212

203213
//===----------------------------------------------------------------------===//
204214
// Generic scalar, vector, or tensor of a particular type.

‎mlir/lib/Dialect/Tosa/IR/TosaOps.cpp‎

Lines changed: 158 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -370,7 +370,7 @@ ParseResult parseWithEnumHandling(OpAsmParser &parser, OperationState &result) {
370370
result.operands)))
371371
return failure();
372372

373-
result.addTypes(fnTy.getResult(0));
373+
result.addTypes(fnTy.getResults());
374374
result.addAttributes(attrs);
375375

376376
return success();
@@ -532,6 +532,24 @@ void MatmulTBlockScaledOp::print(OpAsmPrinter &parser) {
532532
printWithEnumHandling(parser, *this);
533533
}
534534

535+
ParseResult CastFromBlockScaledOp::parse(OpAsmParser &parser,
536+
OperationState &result) {
537+
return parseWithEnumHandling<tosa::BlockSize>(parser, result);
538+
}
539+
540+
void CastFromBlockScaledOp::print(OpAsmPrinter &parser) {
541+
printWithEnumHandling(parser, *this);
542+
}
543+
544+
ParseResult CastToBlockScaledOp::parse(OpAsmParser &parser,
545+
OperationState &result) {
546+
return parseWithEnumHandling<tosa::BlockSize>(parser, result);
547+
}
548+
549+
void CastToBlockScaledOp::print(OpAsmPrinter &parser) {
550+
printWithEnumHandling(parser, *this);
551+
}
552+
535553
//===----------------------------------------------------------------------===//
536554
// Tosa utilities.
537555
//===----------------------------------------------------------------------===//
@@ -3944,6 +3962,145 @@ LogicalResult RescaleOp::inferReturnTypeComponents(
39443962
return success();
39453963
}
39463964

3965+
LogicalResult CastFromBlockScaledOp::inferReturnTypeComponents(
3966+
MLIRContext *context, ::std::optional<Location> location,
3967+
CastFromBlockScaledOp::Adaptor adaptor,
3968+
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
3969+
const ShapeAdaptor inputShape(adaptor.getInputData().getType());
3970+
inferredReturnShapes.push_back(ShapedTypeComponents(inputShape));
3971+
return success();
3972+
}
3973+
3974+
LogicalResult CastFromBlockScaledOp::verify() {
3975+
const Type inputDataType = getInputData().getType();
3976+
const Type outputDataType = getResult().getType();
3977+
if (failed(verifyCompatibleShape(inputDataType, outputDataType)))
3978+
return emitOpError() << "require compatible shapes for input_data ("
3979+
<< inputDataType << ") and "
3980+
<< "output_data (" << outputDataType << ")";
3981+
3982+
const ShapeAdaptor inputDataShape = ShapeAdaptor(inputDataType);
3983+
3984+
if (inputDataShape.hasRank()) {
3985+
const unsigned int blockSize =
3986+
BlockSizeAttr::getBlockSizeValue(getBlockSize());
3987+
const int64_t inputDataLastDim =
3988+
inputDataShape.getDimSize(inputDataShape.getRank() - 1);
3989+
if (inputDataLastDim % blockSize != 0)
3990+
return emitOpError() << "expect last dimension of input_data ("
3991+
<< inputDataLastDim
3992+
<< ") to be divisible by block_size (" << blockSize
3993+
<< ")";
3994+
3995+
const Type inputScaleType = getInputScale().getType();
3996+
const ShapeAdaptor inputScaleShape = ShapeAdaptor(inputScaleType);
3997+
3998+
if (inputScaleShape.hasRank()) {
3999+
SmallVector<int64_t> inputDataDims, inputScaleDims;
4000+
inputDataShape.getDims(inputDataDims);
4001+
inputScaleShape.getDims(inputScaleDims);
4002+
4003+
if (inputDataDims.size() != inputScaleDims.size() ||
4004+
failed(verifyCompatibleShape(
4005+
ArrayRef<int64_t>(inputDataDims).drop_back(1),
4006+
ArrayRef<int64_t>(inputScaleDims).drop_back(1))))
4007+
return emitOpError() << "require compatible shapes for input_data ("
4008+
<< inputDataType << ") and "
4009+
<< "input_scale (" << inputScaleType
4010+
<< ") except for the last dimension";
4011+
4012+
const SmallVector<int64_t, 2> dimsToCheck{inputDataLastDim / blockSize,
4013+
inputScaleDims.back()};
4014+
if (ShapedType::isStatic(inputDataLastDim) &&
4015+
failed(verifyCompatibleDims(dimsToCheck)))
4016+
return emitOpError()
4017+
<< "expect last dimension of input_scale ("
4018+
<< inputScaleDims.back()
4019+
<< ") to be equal to last dimension of input_data / block_size ("
4020+
<< inputDataDims.back() / blockSize << ")";
4021+
}
4022+
}
4023+
4024+
return success();
4025+
}
4026+
4027+
LogicalResult CastToBlockScaledOp::inferReturnTypeComponents(
4028+
MLIRContext *context, ::std::optional<Location> location,
4029+
CastToBlockScaledOp::Adaptor adaptor,
4030+
SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
4031+
const ShapeAdaptor inputShape(adaptor.getInputData().getType());
4032+
inferredReturnShapes.push_back(ShapedTypeComponents(inputShape));
4033+
if (!inputShape.hasRank())
4034+
return success();
4035+
4036+
// Calculate output_scale shape if ranked input provided
4037+
SmallVector<int64_t> outputScaleShape;
4038+
inputShape.getDims(outputScaleShape);
4039+
const int64_t lastDimLoc = inputShape.getRank() - 1;
4040+
const int64_t lastDimSize = inputShape.getDimSize(lastDimLoc);
4041+
if (ShapedType::isStatic(lastDimSize)) {
4042+
const unsigned int blockSize =
4043+
BlockSizeAttr::getBlockSizeValue(adaptor.getBlockSize());
4044+
outputScaleShape[lastDimLoc] = lastDimSize / blockSize;
4045+
}
4046+
inferredReturnShapes.push_back(ShapedTypeComponents(outputScaleShape));
4047+
return success();
4048+
}
4049+
4050+
LogicalResult CastToBlockScaledOp::verify() {
4051+
const Type inputDataType = getInputData().getType();
4052+
const Type outputDataType = getResult(0).getType();
4053+
if (failed(verifyCompatibleShape(inputDataType, outputDataType)))
4054+
return emitOpError() << "require compatible shapes for input_data ("
4055+
<< inputDataType << ") and "
4056+
<< "output_data (" << outputDataType << ")";
4057+
4058+
const unsigned int blockSize =
4059+
BlockSizeAttr::getBlockSizeValue(getBlockSize());
4060+
const ShapeAdaptor inputDataShape = ShapeAdaptor(inputDataType);
4061+
if (inputDataShape.hasRank()) {
4062+
const int64_t inputDataLastDim =
4063+
inputDataShape.getDimSize(inputDataShape.getRank() - 1);
4064+
if (ShapedType::isStatic(inputDataLastDim) &&
4065+
inputDataLastDim % blockSize != 0)
4066+
return emitOpError() << "expect last dimension of input_data ("
4067+
<< inputDataLastDim
4068+
<< ") to be divisible by block_size (" << blockSize
4069+
<< ")";
4070+
}
4071+
4072+
const ShapeAdaptor outputDataShape = ShapeAdaptor(outputDataType);
4073+
const Type outputScaleType = getResult(1).getType();
4074+
const ShapeAdaptor outputScaleShape = ShapeAdaptor(outputScaleType);
4075+
if (outputDataShape.hasRank() && outputScaleShape.hasRank()) {
4076+
SmallVector<int64_t> outputDataDims, outputScaleDims;
4077+
outputDataShape.getDims(outputDataDims);
4078+
outputScaleShape.getDims(outputScaleDims);
4079+
4080+
if (outputDataDims.size() != outputScaleDims.size() ||
4081+
failed(verifyCompatibleShape(
4082+
ArrayRef<int64_t>(outputDataDims).drop_back(1),
4083+
ArrayRef<int64_t>(outputScaleDims).drop_back(1))))
4084+
return emitOpError() << "require compatible shapes for output_data ("
4085+
<< outputDataType << ") and "
4086+
<< "output_scale (" << outputScaleType
4087+
<< ") except for the last dimension";
4088+
4089+
const int64_t outputDataLastDim = outputDataDims.back();
4090+
const SmallVector<int64_t, 2> dimsToCheck{outputDataLastDim / blockSize,
4091+
outputScaleDims.back()};
4092+
if (ShapedType::isStatic(outputDataLastDim) &&
4093+
failed(verifyCompatibleDims(dimsToCheck)))
4094+
return emitOpError()
4095+
<< "expect last dimension of output_scale ("
4096+
<< outputScaleDims.back()
4097+
<< ") to be equal to last dimension of output_data / block_size ("
4098+
<< outputDataDims.back() / blockSize << ")";
4099+
}
4100+
4101+
return success();
4102+
}
4103+
39474104
LogicalResult IfOp::inferReturnTypeComponents(
39484105
MLIRContext *context, ::std::optional<Location> location,
39494106
IfOp::Adaptor adaptor,

‎mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp‎

Lines changed: 9 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,11 @@ TosaProfileCompliance::getProfileComplianceMap() {
5050

5151
// Base populating function
5252
LogicalResult ProfileInfoDepot::populateProfileInfo(ValueRange operands,
53-
Value output) {
54-
for (auto operand : operands)
53+
ValueRange outputs) {
54+
for (const auto &operand : operands)
5555
addValue(operand);
56-
addValue(output);
56+
for (const auto &output : outputs)
57+
addValue(output);
5758
return success();
5859
}
5960

@@ -175,23 +176,6 @@ LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::ResizeOp op) {
175176
return success();
176177
}
177178

178-
template <>
179-
LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::FFT2dOp op) {
180-
addValue(op.getInputReal());
181-
addValue(op.getInputImag());
182-
addValue(op.getOutputReal());
183-
addValue(op.getOutputImag());
184-
return success();
185-
}
186-
187-
template <>
188-
LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::RFFT2dOp op) {
189-
addValue(op.getInputReal());
190-
addValue(op.getOutputReal());
191-
addValue(op.getOutputImag());
192-
return success();
193-
}
194-
195179
template <>
196180
LogicalResult ProfileInfoDepot::populateProfileInfo(tosa::SelectOp op) {
197181
addValue(op.getOnTrue());
@@ -245,7 +229,7 @@ LogicalResult ProfileInfoDepot::populatationDispatch(Operation *op) {
245229
// This helper function populates the info for all operands.
246230
#define POPULATE_PROFILE_INFO_COMMON(tosaOp) \
247231
if (isa<tosa::tosaOp##Op>(op)) { \
248-
return populateProfileInfo(op->getOperands(), op->getResult(0)); \
232+
return populateProfileInfo(op->getOperands(), op->getResults()); \
249233
}
250234

251235
// Skip irrelevant operands when they are independent and not tied to any
@@ -256,8 +240,6 @@ LogicalResult ProfileInfoDepot::populatationDispatch(Operation *op) {
256240
POPULATE_PROFILE_INFO_CUSTOM(Conv3D)
257241
POPULATE_PROFILE_INFO_CUSTOM(DepthwiseConv2D)
258242
POPULATE_PROFILE_INFO_CUSTOM(Mul)
259-
POPULATE_PROFILE_INFO_CUSTOM(FFT2d)
260-
POPULATE_PROFILE_INFO_CUSTOM(RFFT2d)
261243
POPULATE_PROFILE_INFO_CUSTOM(Concat)
262244
POPULATE_PROFILE_INFO_CUSTOM(Pad)
263245
POPULATE_PROFILE_INFO_CUSTOM(Reshape)
@@ -276,7 +258,11 @@ LogicalResult ProfileInfoDepot::populatationDispatch(Operation *op) {
276258
// For the most of tosa operators, all operands are profile/extension related
277259
// and hence are all considered in this profile-based compilance check.
278260
POPULATE_PROFILE_INFO_COMMON(MatmulTBlockScaled)
261+
POPULATE_PROFILE_INFO_COMMON(FFT2d)
262+
POPULATE_PROFILE_INFO_COMMON(RFFT2d)
279263
POPULATE_PROFILE_INFO_COMMON(Cast)
264+
POPULATE_PROFILE_INFO_COMMON(CastFromBlockScaled)
265+
POPULATE_PROFILE_INFO_COMMON(CastToBlockScaled)
280266
POPULATE_PROFILE_INFO_COMMON(Const)
281267
POPULATE_PROFILE_INFO_COMMON(ArgMax)
282268
POPULATE_PROFILE_INFO_COMMON(Sub)

‎mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -635,6 +635,8 @@ LogicalResult TosaValidation::levelCheckRanksAndSizes(Operation *op) {
635635
CHECK_RANKS_AND_SIZES(Transpose);
636636
// Type Conversion
637637
CHECK_RANKS_AND_SIZES(Cast);
638+
CHECK_RANKS_AND_SIZES(CastFromBlockScaled);
639+
CHECK_RANKS_AND_SIZES(CastToBlockScaled);
638640
CHECK_RANKS_AND_SIZES(Rescale);
639641
// Control Flow Operators
640642
CHECK_RANKS_AND_SIZES(If);

0 commit comments

Comments
 (0)