Introduce Flow InferNumericNarrow, OptimizeNumerics and CleanupNumericNarrowing passes (#7975) Together, these passes: * Infer where it is safe to use a narrower, low precision type in place of floating point arithmetic. * Use this information to do targeted rewrites of linalg math ops. * Propagate casts. When combined with constant hoisting and eval, this will perform most of the work needed for low-precision optimization of inference workloads. Also adds a new Util::NumericCastOpInterface and applies it to the casting ops in the arith dialect. I'm going to put some more mileage on this and then will likely propose it for upstream.
diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD index e52a14d..b9f93b6 100644 --- a/iree/compiler/Dialect/Flow/Transforms/BUILD +++ b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -31,6 +31,7 @@ cc_library( name = "Transforms", srcs = [ + "CleanupNumericNarrowing.cpp", "ConvertConv2D1x1ToMatmulPass.cpp", "ConvertConv2DToImg2ColPass.cpp", "ConvertLinalgMatmulToMmt4D.cpp", @@ -40,8 +41,10 @@ "DispatchLinalgOnTensors.cpp", "ExportBenchmarkFuncs.cpp", "FusionOfTensorOps.cpp", + "InferNumericNarrowing.cpp", "InjectDispatchTracing.cpp", "InterchangeGenericOps.cpp", + "OptimizeNumerics.cpp", "OutlineDispatchRegions.cpp", "PadLinalgOps.cpp", "PadTensorToSubTensorInsert.cpp", @@ -64,6 +67,9 @@ "//iree/compiler/Dialect/Flow/Conversion/TensorToFlow", "//iree/compiler/Dialect/Flow/IR", "//iree/compiler/Dialect/HAL/IR", + "//iree/compiler/Dialect/Util/Analysis", + "//iree/compiler/Dialect/Util/Analysis/Attributes", + "//iree/compiler/Dialect/Util/Analysis/DFX", "//iree/compiler/Dialect/Util/IR", "//iree/compiler/Dialect/Util/Transforms", "//iree/compiler/Utils", @@ -71,8 +77,10 @@ "//llvm-external-projects/iree-dialects:IREELinalgExtTransforms", "@llvm-project//llvm:Support", "@llvm-project//mlir:Affine", + "@llvm-project//mlir:ArithmeticDialect", "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:IR", + "@llvm-project//mlir:LinalgInterfaces", "@llvm-project//mlir:LinalgOps", "@llvm-project//mlir:LinalgTransforms", "@llvm-project//mlir:MemRefDialect",
diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt index e381f09..60e371b 100644 --- a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt +++ b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -28,6 +28,7 @@ "Passes.h.inc" "TypeConverter.h" SRCS + "CleanupNumericNarrowing.cpp" "ConvertConv2D1x1ToMatmulPass.cpp" "ConvertConv2DToImg2ColPass.cpp" "ConvertLinalgMatmulToMmt4D.cpp" @@ -37,8 +38,10 @@ "DispatchLinalgOnTensors.cpp" "ExportBenchmarkFuncs.cpp" "FusionOfTensorOps.cpp" + "InferNumericNarrowing.cpp" "InjectDispatchTracing.cpp" "InterchangeGenericOps.cpp" + "OptimizeNumerics.cpp" "OutlineDispatchRegions.cpp" "PadLinalgOps.cpp" "PadTensorToSubTensorInsert.cpp" @@ -55,6 +58,7 @@ IREELinalgExtPasses LLVMSupport MLIRAffine + MLIRArithmetic MLIRIR MLIRLinalg MLIRLinalgTransforms @@ -71,6 +75,9 @@ iree::compiler::Dialect::Flow::Conversion::TensorToFlow iree::compiler::Dialect::Flow::IR iree::compiler::Dialect::HAL::IR + iree::compiler::Dialect::Util::Analysis + iree::compiler::Dialect::Util::Analysis::Attributes + iree::compiler::Dialect::Util::Analysis::DFX iree::compiler::Dialect::Util::IR iree::compiler::Dialect::Util::Transforms iree::compiler::Utils
diff --git a/iree/compiler/Dialect/Flow/Transforms/CleanupNumericNarrowing.cpp b/iree/compiler/Dialect/Flow/Transforms/CleanupNumericNarrowing.cpp new file mode 100644 index 0000000..4b200e8 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/CleanupNumericNarrowing.cpp
@@ -0,0 +1,37 @@ +// Copyright 2021 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "iree/compiler/Dialect/Flow/Transforms/PassDetail.h" +#include "iree/compiler/Dialect/Flow/Transforms/Passes.h" +#include "iree/compiler/Dialect/Util/IR/UtilOps.h" + +namespace mlir { +namespace iree_compiler { +namespace IREE { +namespace Flow { + +namespace { + +class CleanupNumericNarrowingPass + : public CleanupNumericNarrowingBase<CleanupNumericNarrowingPass> { + void runOnOperation() override { + getOperation()->walk([](IREE::Util::NumericOptionalNarrowOp op) { + op.getResult().replaceAllUsesWith(op.getOperand()); + op->erase(); + }); + } +}; + +} // namespace + +std::unique_ptr<Pass> createCleanupNumericNarrowingPass() { + return std::make_unique<CleanupNumericNarrowingPass>(); +} + +} // namespace Flow +} // namespace IREE +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/InferNumericNarrowing.cpp b/iree/compiler/Dialect/Flow/Transforms/InferNumericNarrowing.cpp new file mode 100644 index 0000000..26c8539 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/InferNumericNarrowing.cpp
@@ -0,0 +1,142 @@ +// Copyright 2021 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "iree/compiler/Dialect/Flow/Transforms/PassDetail.h" +#include "iree/compiler/Dialect/Flow/Transforms/Passes.h" +#include "iree/compiler/Dialect/Util/Analysis/Attributes/Range.h" +#include "iree/compiler/Dialect/Util/Analysis/DFX/Solver.h" +#include "iree/compiler/Dialect/Util/Analysis/DFX/State.h" +#include "iree/compiler/Dialect/Util/Analysis/Explorer.h" +#include "iree/compiler/Dialect/Util/IR/UtilDialect.h" +#include "iree/compiler/Dialect/Util/IR/UtilOps.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/Support/Debug.h" +#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h" + +using llvm::SmallPtrSet; + +namespace mlir { +namespace iree_compiler { +namespace IREE { +namespace Flow { + +namespace { + +IntegerType deriveIntegerTypeFromRange(MLIRContext *context, int64_t minValue, + int64_t maxValue) { + // Clamp min/max to span 0. + const int64_t zero = 0; + minValue = std::min(zero, minValue); + maxValue = std::max(zero, maxValue); + bool isSigned; + if (minValue < 0) { + // For signed, make symmetric from -N:N-1 + isSigned = true; + maxValue = std::max(std::abs(minValue) - 1, maxValue); + minValue = std::min(-maxValue - 1, minValue); + } else { + isSigned = false; + } + int64_t n = maxValue - minValue + 1; + int64_t numBits = std::ceil(std::log2(n)); + + return IntegerType::get(context, numBits, + isSigned + ? IntegerType::SignednessSemantics::Signed + : IntegerType::SignednessSemantics::Unsigned); +} + +class InferNumericNarrowingPass + : public InferNumericNarrowingBase<InferNumericNarrowingPass> { + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert<IREE::Util::UtilDialect>(); + } + + void runOnOperation() override { + auto probePoints = collectProbePoints(); + + Explorer explorer(getOperation(), TraversalAction::SHALLOW); + llvm::BumpPtrAllocator allocator; + DFX::Solver solver(explorer, allocator); + + // Prime with probe points. + for (Value probePoint : probePoints) { + solver.getOrCreateElementFor<IREE::Util::FloatRangeValueElement>( + Position::forValue(probePoint)); + } + + // Solve. + if (failed(solver.run())) { + return signalPassFailure(); + } + + // Annotate. + for (Value probePoint : probePoints) { + auto *elt = solver.lookupElementFor<IREE::Util::FloatRangeValueElement>( + Position::forValue(probePoint)); + if (!elt) { + // Not valid analysis. + continue; + } + + applyAnnotation(probePoint, elt->getKnown()); + } + } + + SmallPtrSet<Value, 8> collectProbePoints() { + SmallPtrSet<Value, 8> probePoints; + getOperation()->walk([&](Operation *op) { + if (auto linalgOp = llvm::dyn_cast<linalg::LinalgOp>(op)) { + for (Value input : linalgOp.inputs()) { + probePoints.insert(input); + } + for (Value output : linalgOp.outputs()) { + probePoints.insert(output); + } + } + }); + return probePoints; + } + + void applyAnnotation(Value probePoint, IREE::Util::FloatRangeStats stats) { + if (stats.isTruncated() && stats.isFinite()) { + // Integer annotation. + applyIntegerAnnotation(probePoint, stats); + } + } + + void applyIntegerAnnotation(Value probePoint, + IREE::Util::FloatRangeStats stats) { + auto context = probePoint.getContext(); + auto minValue = static_cast<int64_t>(stats.minValue); + auto maxValue = static_cast<int64_t>(stats.maxValue); + IntegerType type = + deriveIntegerTypeFromRange(probePoint.getContext(), minValue, maxValue); + + // Insert the annotation. + OpBuilder builder(context); + builder.setInsertionPointAfterValue(probePoint); + Optional<std::pair<int64_t, int64_t>> range; + // i0 values cannot parse any values so omit. + if (type.getWidth() != 0) { + range = std::make_pair(minValue, maxValue); + } + auto annotationOp = builder.create<IREE::Util::NumericOptionalNarrowOp>( + probePoint.getLoc(), probePoint, type, range); + probePoint.replaceAllUsesExcept(annotationOp, {annotationOp}); + } +}; + +} // namespace + +std::unique_ptr<Pass> createInferNumericNarrowingPass() { + return std::make_unique<InferNumericNarrowingPass>(); +} + +} // namespace Flow +} // namespace IREE +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/OptimizeNumerics.cpp b/iree/compiler/Dialect/Flow/Transforms/OptimizeNumerics.cpp new file mode 100644 index 0000000..4c9ac77 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/OptimizeNumerics.cpp
@@ -0,0 +1,283 @@ +// Copyright 2021 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "iree/compiler/Dialect/Flow/Transforms/PassDetail.h" +#include "iree/compiler/Dialect/Flow/Transforms/Passes.h" +#include "iree/compiler/Dialect/Util/IR/UtilOps.h" +#include "llvm/Support/Debug.h" +#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +namespace mlir { +namespace iree_compiler { +namespace IREE { +namespace Flow { + +namespace { + +int getNextPotBitWidth(int bitWidth, int minBitWidth = 8) { + for (int i = minBitWidth;; i *= 2) { + if (i >= bitWidth) return i; + } +} + +Type withNewElementType(Type origType, Type elementType) { + if (auto st = origType.dyn_cast<ShapedType>()) { + return st.clone(elementType); + } else { + return elementType; + } +} + +Type makeLowPType(Type origType, int bitWidth) { + auto *context = origType.getContext(); + auto elementType = IntegerType::get(context, bitWidth); + return withNewElementType(origType, elementType); +} + +Value castNumeric(Value origValue, Type toType, bool isSigned, + OpBuilder &builder) { + Location loc = origValue.getLoc(); + Type origElementType = getElementTypeOrSelf(origValue.getType()); + Type toElementType = getElementTypeOrSelf(toType); + + if (origElementType.isa<FloatType>() && toElementType.isa<IntegerType>()) { + if (isSigned) { + return builder.create<arith::FPToSIOp>(loc, toType, origValue); + } else { + return builder.create<arith::FPToUIOp>(loc, toType, origValue); + } + } else if (origElementType.isa<IntegerType>() && + toElementType.isa<FloatType>()) { + if (isSigned) { + return builder.create<arith::SIToFPOp>(loc, toType, origValue); + } else { + return builder.create<arith::UIToFPOp>(loc, toType, origValue); + } + } else { + // If we need int<->int and float<->float, implement those cases. Since + // this is just needed for things in this file, it is ok to leave it + // under implemented. + llvm_unreachable("unsupported numeric cast"); + } +} + +struct NarrowParams { + static Optional<NarrowParams> forValue(Value value) { + if (auto narrowOp = + llvm::dyn_cast_or_null<IREE::Util::NumericOptionalNarrowOp>( + value.getDefiningOp())) { + NarrowParams params; + params.producer = narrowOp.operand(); + params.fromType = value.getType(); + params.toElementType = narrowOp.semantic_type(); + params.range = narrowOp.getIntegerRange(); + + return params; + } + return {}; + } + + bool isFromFloat() { return getElementTypeOrSelf(fromType).isa<FloatType>(); } + + bool isToInteger() { return toElementType.isa<IntegerType>(); } + + bool isToSigned() { return toElementType.cast<IntegerType>().isSigned(); } + + int getToBitWidth() { return toElementType.cast<IntegerType>().getWidth(); } + + Value producer; + Type fromType; + Type toElementType; + Optional<std::pair<int64_t, int64_t>> range; +}; + +// Eliminates a cast produced by an init_tensor by just initializing to that +// type directly. +struct LinalgInitTensorCast + : OpInterfaceRewritePattern<IREE::Util::NumericCastOpInterface> { + using OpInterfaceRewritePattern::OpInterfaceRewritePattern; + + LogicalResult matchAndRewrite(IREE::Util::NumericCastOpInterface castOp, + PatternRewriter &rewriter) const override { + auto initTensorOp = castOp.getInput().getDefiningOp<linalg::InitTensorOp>(); + if (!initTensorOp) return failure(); + Type resultType = castOp.getCasted().getType(); + + rewriter.replaceOpWithNewOp<linalg::InitTensorOp>( + castOp, resultType, initTensorOp.sizes(), initTensorOp.static_sizes()); + return success(); + } +}; + +// For a cast produced by a fill, rewrites the cast to be on the fill operands. +struct LinalgFillCast + : public OpInterfaceRewritePattern<IREE::Util::NumericCastOpInterface> { + using OpInterfaceRewritePattern::OpInterfaceRewritePattern; + + LogicalResult matchAndRewrite(IREE::Util::NumericCastOpInterface castOp, + PatternRewriter &rewriter) const override { + auto loc = castOp.getLoc(); + auto fillOp = castOp.getInput().getDefiningOp<linalg::FillOp>(); + if (!fillOp) return failure(); + Type toElementType = getElementTypeOrSelf(castOp.getCastedType()); + + Value fillInput = fillOp.value(); + Value fillInit = fillOp.output(); + fillInput = castOp + .cloneWithInput( + rewriter, + withNewElementType(fillInput.getType(), toElementType), + fillInput) + .getCasted(); + fillInit = + castOp + .cloneWithInput( + rewriter, withNewElementType(fillInit.getType(), toElementType), + fillInit) + .getCasted(); + Value fillResult = + rewriter.create<linalg::FillOp>(loc, fillInput, fillInit).result(); + rewriter.replaceOp(castOp, fillResult); + return success(); + } +}; + +// For narrowable inputs, selects +struct LinalgFpMatmulToLowP : public OpRewritePattern<linalg::MatmulOp> { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(linalg::MatmulOp matmulOp, + PatternRewriter &rewriter) const override { + Location loc = matmulOp.getLoc(); + Type origResultType = matmulOp.getResult(0).getType(); + auto lhsParams = NarrowParams::forValue(matmulOp.inputs()[0]); + auto rhsParams = NarrowParams::forValue(matmulOp.inputs()[1]); + auto accumParams = NarrowParams::forValue(matmulOp.outputs()[0]); + if (!lhsParams || !rhsParams || !accumParams) { + return rewriter.notifyMatchFailure(matmulOp, "no narrowing annotations"); + } + + // TODO(#7987): This could be more flexible, allowing mix and match + // integer/float types. + if (!lhsParams->isFromFloat() || !rhsParams->isFromFloat()) { + return rewriter.notifyMatchFailure(matmulOp, "not from floating point"); + } + + // TODO(#7987): Could support partial conversion to integer. + if (!lhsParams->isToInteger() || !rhsParams->isToInteger() || + !accumParams->isToInteger()) { + return rewriter.notifyMatchFailure(matmulOp, "not to an integer type"); + } + + int lhsBitWidth = lhsParams->getToBitWidth(); + int rhsBitWidth = rhsParams->getToBitWidth(); + + // Handle signed/unsigned mismatch. + // TODO(#7987): Implement a proper unsigned->signed widening. + bool isSigned; + if (lhsParams->isToSigned() != rhsParams->isToSigned()) { + // Mixed signed/unsigned. Promote to signed. + isSigned = true; + if (!lhsParams->isToSigned()) { + lhsBitWidth += 1; + } + if (!rhsParams->isToSigned()) { + rhsBitWidth += 1; + } + } else { + // Uniform signed/unsigned. + isSigned = lhsParams->isToSigned(); + } + + // Round up to a suitable POT width. + lhsBitWidth = getNextPotBitWidth(lhsBitWidth); + rhsBitWidth = getNextPotBitWidth(rhsBitWidth); + + // Promote accumulator to match signedness. + int accumBitWidth = accumParams->getToBitWidth(); + if (isSigned && !accumParams->isToSigned()) { + // TODO(#7987): A proper unsigned widening based on range. + accumBitWidth += 1; + } + + // Determine an appropriate accumulator size. + // TODO(#7987): Apply the clamp of: + // lhsBitWidth + rhsBitWidth + log2_ceil(contraction_dim + 1) to determine + // the accumulator size. Note: Can drop the +1 if one of lhs/rhs is signed + // and symmetric (i.e. does not use the asymmetric lower bound). + if (lhsBitWidth > 8 || rhsBitWidth > 8) { + return rewriter.notifyMatchFailure(matmulOp, "outside of low-p range"); + } + accumBitWidth = getNextPotBitWidth(accumBitWidth, 32); + if (accumBitWidth > 32) { + return rewriter.notifyMatchFailure(matmulOp, "accumulator > 32 bits"); + } + + Type lhsLowPType = makeLowPType(lhsParams->fromType, lhsBitWidth); + Type rhsLowPType = makeLowPType(rhsParams->fromType, rhsBitWidth); + Type accumLowPType = makeLowPType(accumParams->fromType, accumBitWidth); + + // Replace the matmul op. + Value newLhs = + castNumeric(lhsParams->producer, lhsLowPType, isSigned, rewriter); + Value newRhs = + castNumeric(rhsParams->producer, rhsLowPType, isSigned, rewriter); + Value newAccum = + castNumeric(accumParams->producer, accumLowPType, isSigned, rewriter); + Value newResult; + + if (isSigned) { + newResult = rewriter + .create<linalg::MatmulOp>(loc, ValueRange{newLhs, newRhs}, + ValueRange{newAccum}) + .getResult(0); + } else { + newResult = rewriter + .create<linalg::MatmulUnsignedOp>( + loc, ValueRange{newLhs, newRhs}, ValueRange{newAccum}) + .getResult(0); + } + + // Cast back. + newResult = castNumeric(newResult, origResultType, isSigned, rewriter); + rewriter.replaceOp(matmulOp, ValueRange{newResult}); + + return success(); + } +}; + +class OptimizeNumericsPass : public OptimizeNumericsBase<OptimizeNumericsPass> { + void runOnOperation() override { + MLIRContext *context = &getContext(); + RewritePatternSet patterns(context); + + // Precision reduction. + patterns.insert<LinalgFpMatmulToLowP>(context); + + // Cast propagation. + patterns.insert<LinalgInitTensorCast>(context); + patterns.insert<LinalgFillCast>(context); + + if (failed(applyPatternsAndFoldGreedily(getOperation(), + std::move(patterns)))) { + return signalPassFailure(); + } + } +}; + +} // namespace + +std::unique_ptr<Pass> createOptimizeNumericsPass() { + return std::make_unique<OptimizeNumericsPass>(); +} + +} // namespace Flow +} // namespace IREE +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.h b/iree/compiler/Dialect/Flow/Transforms/Passes.h index 340ae1d..fb9854a 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.h +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -57,6 +57,10 @@ // Input canonicalization and legalization //===----------------------------------------------------------------------===// +// Cleans up any numeric narrowing ops inserted by +// iree-flow-infer-numeric-narrowing. +std::unique_ptr<Pass> createCleanupNumericNarrowingPass(); + /// Creates a pass to convert linalg convolution ops with 1x1 kernels into /// linalg.matmul std::unique_ptr<Pass> createConvertConv2D1x1ToMatmulPass(); @@ -76,6 +80,10 @@ /// Creates a pass to fuse Linalg operations on tensors. std::unique_ptr<Pass> createFusionOfTensorOpsPass(); +/// Infers and inserts util.numeric.optional_narrow ops at points that may be +/// beneficial. +std::unique_ptr<Pass> createInferNumericNarrowingPass(); + /// Create a pass to interchange generic ops to force the reduction loop to be /// the most inner loops. std::unique_ptr<Pass> createInterchangeGenericOpsPass(); @@ -87,6 +95,10 @@ // equivalent flow ops. std::unique_ptr<Pass> createConvertToFlowAfterDispatchFormation(); +// Optimizes numerics given annotations added via +// iree-flow-infer-numeric-narrowing. +std::unique_ptr<Pass> createOptimizeNumericsPass(); + // Promote I1 tensor constants to I8 tensors to match later operations. std::unique_ptr<OperationPass<mlir::FuncOp>> createPromoteI1ToI8Pass();
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.td b/iree/compiler/Dialect/Flow/Transforms/Passes.td index 2d9644d..a1e0775 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.td +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -9,6 +9,12 @@ include "mlir/Pass/PassBase.td" +def CleanupNumericNarrowing : + Pass<"iree-flow-cleanup-numeric-narrowing", ""> { + let summary = "Cleans up any numeric narrowing ops inserted by iree-flow-infer-numeric-narrowing"; + let constructor = "mlir::iree_compiler::IREE::Flow::createCleanupNumericNarrowingPass()"; +} + def ConvertConv2D1x1ConvToMatmul : Pass<"iree-flow-convert-conv2d-1x1-to-matmul", ""> { let summary = "Convert linalg convolution ops with 1x1 kernels into linalg matrix multiplication ops."; @@ -57,6 +63,12 @@ let constructor = "mlir::iree_compiler::IREE::Flow::createFusionOfTensorOpsPass()"; } +def InferNumericNarrowing : + Pass<"iree-flow-infer-numeric-narrowing", ""> { + let summary = "Infers and inserts util.numeric.optional_narrow ops at points that may be beneficial"; + let constructor = "mlir::iree_compiler::IREE::Flow::createInferNumericNarrowingPass()"; +} + def InjectDispatchTracing : Pass<"iree-flow-inject-dispatch-tracing", ""> { let summary = "Injects dispatch region tracing."; @@ -69,6 +81,12 @@ let constructor = "mlir::iree_compiler::IREE::Flow::createInterchangeGenericOpsPass()"; } +def OptimizeNumerics : + Pass<"iree-flow-optimize-numerics", ""> { + let summary = "Optimizes numerics given annotations added via iree-flow-infer-numeric-narrowing"; + let constructor = "mlir::iree_compiler::IREE::Flow::createOptimizeNumericsPass()"; +} + def OutlineDispatchRegions : Pass<"iree-flow-outline-dispatch-regions", "mlir::ModuleOp"> { let summary = "Outlines dispatch regions into executables";
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/BUILD b/iree/compiler/Dialect/Flow/Transforms/test/BUILD index b66b655..6423883 100644 --- a/iree/compiler/Dialect/Flow/Transforms/test/BUILD +++ b/iree/compiler/Dialect/Flow/Transforms/test/BUILD
@@ -17,6 +17,7 @@ name = "lit", srcs = enforce_glob( [ + "cleanup_numeric_narrowing.mlir", "conv1x1_to_matmul.mlir", "conv2d_to_img2col.mlir", "convert_linalg_tensor_ops_after.mlir", @@ -26,9 +27,11 @@ "dispatch_linalg_on_tensors_elementwise.mlir", "dispatch_linalg_on_tensors_fusion.mlir", "export_benchmark_funcs.mlir", + "infer_numeric_narrowing.mlir", "inject_dispatch_tracing.mlir", "interchange_generic_ops.mlir", "matmul_to_mmt4d.mlir", + "optimize_numerics.mlir", "outline_dispatch_regions.mlir", "pad_linalg_ops.mlir", "pad_tensor_to_tensor.mlir",
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt index 205634c..9183bd0 100644 --- a/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt +++ b/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -14,6 +14,7 @@ NAME lit SRCS + "cleanup_numeric_narrowing.mlir" "conv1x1_to_matmul.mlir" "conv2d_to_img2col.mlir" "convert_linalg_tensor_ops_after.mlir" @@ -23,9 +24,11 @@ "dispatch_linalg_on_tensors_elementwise.mlir" "dispatch_linalg_on_tensors_fusion.mlir" "export_benchmark_funcs.mlir" + "infer_numeric_narrowing.mlir" "inject_dispatch_tracing.mlir" "interchange_generic_ops.mlir" "matmul_to_mmt4d.mlir" + "optimize_numerics.mlir" "outline_dispatch_regions.mlir" "pad_linalg_ops.mlir" "pad_tensor_to_tensor.mlir"
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/cleanup_numeric_narrowing.mlir b/iree/compiler/Dialect/Flow/Transforms/test/cleanup_numeric_narrowing.mlir new file mode 100644 index 0000000..2d2ea4b --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/test/cleanup_numeric_narrowing.mlir
@@ -0,0 +1,8 @@ +// RUN: iree-opt -iree-flow-cleanup-numeric-narrowing %s | IreeFileCheck %s + +// CHECK-LABEL: @remove_inferences +func @remove_inferences(%arg0 : tensor<5x3xf32>) -> tensor<5x3xf32> { + %0 = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as ui3 {max_value = 5 : ui3, min_value = 5 : ui3} + // CHECK: return %arg0 + return %0 : tensor<5x3xf32> +}
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/infer_numeric_narrowing.mlir b/iree/compiler/Dialect/Flow/Transforms/test/infer_numeric_narrowing.mlir new file mode 100644 index 0000000..47bf850 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/test/infer_numeric_narrowing.mlir
@@ -0,0 +1,75 @@ +// RUN: iree-opt -iree-flow-infer-numeric-narrowing %s | IreeFileCheck %s +// This does not test all of the analysis logic, just that the annotations +// are inserted at proper points in the right way. Probe points checked: +// - Every operand of a LinalgOp + +// CHECK-LABEL: @probe_linalg_op +// Checks as a by-product: +// - Infering ui0 for [0, 0] range +// - Infering unsigned for >= 0 range +func @probe_linalg_op(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + // CHECK-DAG: %[[RHS:.*]] = arith.constant dense + // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 + // CHECK-DAG: util.numeric.optional_narrow %[[ZERO]] : f32 as ui0 + // CHECK-DAG: util.numeric.optional_narrow %[[RHS]] : tensor<3x1xf32> as ui7 {max_value = 127 : ui7, min_value = 0 : ui7} + // CHECK-DAG: %[[FILL:.*]] = linalg.fill + // CHECK-DAG: util.numeric.optional_narrow %[[FILL]] : tensor<5x1xf32> as ui0 + %rhs = arith.constant dense< + [[3.900000e+01], [0.000000e+00], [1.270000e+02]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @infer_symmetric_signed +// CHECK: util.numeric.optional_narrow %{{.*}} : tensor<3x1xf32> as si8 {max_value = 127 : si8, min_value = -39 : si8} +func @infer_symmetric_signed(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + %rhs = arith.constant dense< + [[-3.900000e+01], [0.000000e+00], [1.270000e+02]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @infer_i1_signed +// Signed i1 is a silly boundary condition worth checking. +// CHECK: util.numeric.optional_narrow %{{.*}} : tensor<3x1xf32> as si1 {max_value = 0 : si1, min_value = -1 : si1} +func @infer_i1_signed(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + %rhs = arith.constant dense< + [[0.000000e+00], [0.000000e+00], [-1.000000e+00]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @infer_positive_non_straddling_zero +// A range that does not straddle zero is a special case in the code. +// CHECK: util.numeric.optional_narrow %{{.*}} : tensor<3x1xf32> as ui2 {max_value = 2 : ui2, min_value = 1 : ui2} +func @infer_positive_non_straddling_zero(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + %rhs = arith.constant dense< + [[1.000000e+00], [1.000000e+00], [2.000000e+00]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @infer_negative_non_straddling_zero +// A range that does not straddle zero is a special case in the code. +// CHECK: util.numeric.optional_narrow %{{.*}} : tensor<3x1xf32> as si2 {max_value = -1 : si2, min_value = -2 : si2} +func @infer_negative_non_straddling_zero(%arg0 : tensor<5x3xf32>) -> tensor<5x1xf32> { + %rhs = arith.constant dense< + [[-1.000000e+00], [-1.000000e+00], [-2.000000e+00]]> : tensor<3x1xf32> + %init_value = arith.constant 0.000000e+00 : f32 + %0 = linalg.init_tensor [5, 1] : tensor<5x1xf32> + %1 = linalg.fill(%init_value, %0) : f32, tensor<5x1xf32> -> tensor<5x1xf32> + %2 = linalg.matmul ins(%arg0, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%1 : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +}
diff --git a/iree/compiler/Dialect/Flow/Transforms/test/optimize_numerics.mlir b/iree/compiler/Dialect/Flow/Transforms/test/optimize_numerics.mlir new file mode 100644 index 0000000..cc8e368 --- /dev/null +++ b/iree/compiler/Dialect/Flow/Transforms/test/optimize_numerics.mlir
@@ -0,0 +1,77 @@ +// RUN: iree-opt -iree-flow-optimize-numerics %s | IreeFileCheck %s + +// CHECK-LABEL: @matmul_i8_i8_i32_unsigned +func @matmul_i8_i8_i32_unsigned(%arg0 : tensor<5x3xf32>, %arg1 : tensor<3x1xf32>, %arg2 : tensor<5x1xf32>) -> tensor<5x1xf32> { + // CHECK: %[[LHS:.*]] = arith.fptoui %arg0 : tensor<5x3xf32> to tensor<5x3xi8> + // CHECK: %[[RHS:.*]] = arith.fptoui %arg1 : tensor<3x1xf32> to tensor<3x1xi8> + // CHECK: %[[INIT:.*]] = arith.fptoui %arg2 : tensor<5x1xf32> to tensor<5x1xi32> + %lhs = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as ui7 {max_value = 127 : ui7, min_value = 0 : ui7} + %rhs = util.numeric.optional_narrow %arg1 : tensor<3x1xf32> as ui7 {max_value = 127 : ui7, min_value = 0 : ui7} + %init = util.numeric.optional_narrow %arg2 : tensor<5x1xf32> as ui0 + // CHECK: %[[RESULT:.*]] = linalg.matmul_unsigned ins(%[[LHS]], %[[RHS]] : tensor<5x3xi8>, tensor<3x1xi8>) outs(%[[INIT]] : tensor<5x1xi32>) + %2 = linalg.matmul ins(%lhs, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%init : tensor<5x1xf32>) -> tensor<5x1xf32> + // CHECK: arith.uitofp %[[RESULT]] : tensor<5x1xi32> to tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @matmul_i8_i8_i32_signed +func @matmul_i8_i8_i32_signed(%arg0 : tensor<5x3xf32>, %arg1 : tensor<3x1xf32>, %arg2 : tensor<5x1xf32>) -> tensor<5x1xf32> { + // CHECK: %[[LHS:.*]] = arith.fptosi %arg0 : tensor<5x3xf32> to tensor<5x3xi8> + // CHECK: %[[RHS:.*]] = arith.fptosi %arg1 : tensor<3x1xf32> to tensor<3x1xi8> + // CHECK: %[[INIT:.*]] = arith.fptosi %arg2 : tensor<5x1xf32> to tensor<5x1xi32> + %lhs = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as ui7 {max_value = 127 : ui7, min_value = 0 : ui7} + %rhs = util.numeric.optional_narrow %arg1 : tensor<3x1xf32> as si8 {max_value = 127 : si8, min_value = -127 : si8} + %init = util.numeric.optional_narrow %arg2 : tensor<5x1xf32> as ui0 + // CHECK: %[[RESULT:.*]] = linalg.matmul ins(%[[LHS]], %[[RHS]] : tensor<5x3xi8>, tensor<3x1xi8>) outs(%[[INIT]] : tensor<5x1xi32>) + %2 = linalg.matmul ins(%lhs, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%init : tensor<5x1xf32>) -> tensor<5x1xf32> + // CHECK: arith.sitofp %[[RESULT]] : tensor<5x1xi32> to tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @matmul_i4_i4_i32_signed +// For now we clamp this to i8 +func @matmul_i4_i4_i32_signed(%arg0 : tensor<5x3xf32>, %arg1 : tensor<3x1xf32>, %arg2 : tensor<5x1xf32>) -> tensor<5x1xf32> { + // CHECK: %[[LHS:.*]] = arith.fptosi %arg0 : tensor<5x3xf32> to tensor<5x3xi8> + // CHECK: %[[RHS:.*]] = arith.fptosi %arg1 : tensor<3x1xf32> to tensor<3x1xi8> + // CHECK: %[[INIT:.*]] = arith.fptosi %arg2 : tensor<5x1xf32> to tensor<5x1xi32> + %lhs = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as si4 {max_value = 7 : si4, min_value = -7 : si4} + %rhs = util.numeric.optional_narrow %arg1 : tensor<3x1xf32> as si4 {max_value = 3 : si4, min_value = -7 : si4} + %init = util.numeric.optional_narrow %arg2 : tensor<5x1xf32> as ui0 + // CHECK: %[[RESULT:.*]] = linalg.matmul ins(%[[LHS]], %[[RHS]] : tensor<5x3xi8>, tensor<3x1xi8>) outs(%[[INIT]] : tensor<5x1xi32>) + %2 = linalg.matmul ins(%lhs, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%init : tensor<5x1xf32>) -> tensor<5x1xf32> + // CHECK: arith.sitofp %[[RESULT]] : tensor<5x1xi32> to tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @matmul_reject_gt_8bit +// We may relax this restriction at some point but for right now we have it +// because less analysis is needed to prove safety. +// CHECK-NOT: fptosi +func @matmul_reject_gt_8bit(%arg0 : tensor<5x3xf32>, %arg1 : tensor<3x1xf32>, %arg2 : tensor<5x1xf32>) -> tensor<5x1xf32> { + %lhs = util.numeric.optional_narrow %arg0 : tensor<5x3xf32> as ui9 {max_value = 312 : ui9, min_value = 0 : ui9} + %rhs = util.numeric.optional_narrow %arg1 : tensor<3x1xf32> as si8 {max_value = 127 : si8, min_value = -127 : si8} + %init = util.numeric.optional_narrow %arg2 : tensor<5x1xf32> as ui0 + // CHECK: linalg.matmul {{.*}} -> tensor<5x1xf32> + %2 = linalg.matmul ins(%lhs, %rhs : tensor<5x3xf32>, tensor<3x1xf32>) outs(%init : tensor<5x1xf32>) -> tensor<5x1xf32> + return %2 : tensor<5x1xf32> +} + +// CHECK-LABEL: @cast_fill +func @cast_fill(%arg0 : f32, %arg1 : tensor<3xf32>) -> tensor<3xi8> { + // CHECK: %[[SCALAR:.*]] = arith.fptosi %arg0 : f32 to i8 + // CHECK: %[[INIT:.*]] = arith.fptosi %arg1 : tensor<3xf32> to tensor<3xi8> + // CHECK: %[[RESULT:.*]] = linalg.fill(%[[SCALAR]], %[[INIT]]) : i8, tensor<3xi8> -> tensor<3xi8> + // CHECK: return %[[RESULT]] + %0 = linalg.fill(%arg0, %arg1) : f32, tensor<3xf32> -> tensor<3xf32> + %1 = arith.fptosi %0 : tensor<3xf32> to tensor<3xi8> + return %1 : tensor<3xi8> +} + +// CHECK-LABEL: @cast_init +func @cast_init() -> tensor<5x9xi8> { + // CHECK: %[[RESULT:.*]] = linalg.init_tensor [5, 9] : tensor<5x9xi8> + // CHECK: return %[[RESULT]] + %0 = linalg.init_tensor [5, 9] : tensor<5x9xf32> + %1 = arith.fptosi %0 : tensor<5x9xf32> to tensor<5x9xi8> + return %1 : tensor<5x9xi8> +}
diff --git a/iree/compiler/Dialect/Util/Analysis/Attributes/Range.h b/iree/compiler/Dialect/Util/Analysis/Attributes/Range.h index aca1cb8..d8c640e 100644 --- a/iree/compiler/Dialect/Util/Analysis/Attributes/Range.h +++ b/iree/compiler/Dialect/Util/Analysis/Attributes/Range.h
@@ -56,6 +56,12 @@ return static_cast<TruncationFlag>(std::max(lhs, rhs)); } + // Whether the range is known to only contain values that have been + // truncated to exclude fractional bits. + bool isTruncated() { return truncationFlag == TRUNC; } + + bool isFinite() { return std::isfinite(minValue) && std::isfinite(maxValue); } + // Reset to initial state. void reset() { *this = FloatRangeStats(); }
diff --git a/iree/compiler/Dialect/Util/IR/UtilDialect.cpp b/iree/compiler/Dialect/Util/IR/UtilDialect.cpp index 1b3e2a4..58b4e7a 100644 --- a/iree/compiler/Dialect/Util/IR/UtilDialect.cpp +++ b/iree/compiler/Dialect/Util/IR/UtilDialect.cpp
@@ -143,6 +143,45 @@ results.insert<FoldDimOp<tensor::DimOp>>(getContext()); } +//===----------------------------------------------------------------------===// +// Interface external models +//===----------------------------------------------------------------------===// + +namespace { + +// Since all details of the interface are provided via default implementations, +// we can just have one templated external model to apply per op, vs one +// explicit model per op. +struct GenericNumericCastExternalModel { + template <typename OpTy> + struct ExternalModel + : public NumericCastOpInterface::ExternalModel<ExternalModel<OpTy>, + OpTy> {}; + + template <typename OpTy> + static void add(DialectRegistry ®istry) { + registry.addOpInterface<OpTy, ExternalModel<OpTy>>(); + } + + template <typename OpTy1, typename OpTy2, typename... More> + static void add(DialectRegistry ®istry) { + add<OpTy1>(registry); + add<OpTy2, More...>(registry); + } +}; + +} // namespace + +void registerUtilExternalModels(DialectRegistry ®istry) { + // Must ensure that any dependent dialects are registered. + registry.insert<arith::ArithmeticDialect>(); + + GenericNumericCastExternalModel::add< + arith::BitcastOp, arith::ExtFOp, arith::ExtUIOp, arith::ExtSIOp, + arith::FPToSIOp, arith::FPToUIOp, arith::IndexCastOp, arith::TruncFOp, + arith::TruncIOp, arith::SIToFPOp, arith::UIToFPOp>(registry); +} + } // namespace Util } // namespace IREE } // namespace iree_compiler
diff --git a/iree/compiler/Dialect/Util/IR/UtilDialect.h b/iree/compiler/Dialect/Util/IR/UtilDialect.h index 8a07a6e..8e08e28 100644 --- a/iree/compiler/Dialect/Util/IR/UtilDialect.h +++ b/iree/compiler/Dialect/Util/IR/UtilDialect.h
@@ -36,6 +36,8 @@ void registerTypes(); }; +void registerUtilExternalModels(DialectRegistry& registry); + } // namespace Util } // namespace IREE } // namespace iree_compiler
diff --git a/iree/compiler/Dialect/Util/IR/UtilInterfaces.td b/iree/compiler/Dialect/Util/IR/UtilInterfaces.td index c85926f..b415059 100644 --- a/iree/compiler/Dialect/Util/IR/UtilInterfaces.td +++ b/iree/compiler/Dialect/Util/IR/UtilInterfaces.td
@@ -118,6 +118,101 @@ } //===----------------------------------------------------------------------===// +// IREE::Util::NumericCastOpInterface +//===----------------------------------------------------------------------===// + +def Util_NumericCastOpInterface : OpInterface<"NumericCastOpInterface"> { + let cppNamespace = "::mlir::iree_compiler::IREE::Util"; + + let description = [{ + Applied to numeric casting ops which can convert between different numeric + types or shaped-types thereof. Example ops include `fptosi`, `trunci`, etc. + Treating these generically allows us to perform various cast movement + optimizations. + + Conforming operations must: + * Have no attributes. + * Have one operand and one result. + * Be able to operate on supported scalar types: IntegerType, FloatType, + IndexType. + * Be able to operate on tensors/vectors of supported scalar types. + * Have a builder that takes (Type, Value). + }]; + + let methods = [ + InterfaceMethod< + /*desc=*/[{ + Gets the input value. + }], + /*retTy=*/"Value", + /*methodName=*/"getInput", + /*args=*/(ins), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op->getOperand(0); + }] + >, + + InterfaceMethod< + /*desc=*/[{ + Gets the input type. + }], + /*retTy=*/"Type", + /*methodName=*/"getInputType", + /*args=*/(ins), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op->getOperand(0).getType(); + }] + >, + + InterfaceMethod< + /*desc=*/[{ + Gets the result casted value. + }], + /*retTy=*/"Value", + /*methodName=*/"getCasted", + /*args=*/(ins), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op->getResult(0); + }] + >, + InterfaceMethod< + /*desc=*/[{ + Gets the result casted type. + }], + /*retTy=*/"Type", + /*methodName=*/"getCastedType", + /*args=*/(ins), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return $_op->getResult(0).getType(); + }] + >, + InterfaceMethod< + /*desc=*/[{ + Clones the operation with a new result type and input. + Note that it is generally not legal to generically change the scalar + type of the input or the result, but it is legal to transform between + a scalar type and a tensor/vector with an element type of the original + scalar type (and vica-versa). + }], + /*retTy=*/"NumericCastOpInterface", + /*methodName=*/"cloneWithInput", + /*args=*/(ins "OpBuilder &":$builder, "Type":$resultType, "Value":$input), + /*methodBody=*/[{}], + /*defaultImplementation=*/[{ + return llvm::cast<NumericCastOpInterface>( + builder.create<ConcreteOp>($_op->getLoc(), resultType, input) + .getOperation()); + }] + >, + ]; +} + + +//===----------------------------------------------------------------------===// // IREE::Util::TiedOpInterface //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Util/IR/UtilOps.cpp b/iree/compiler/Dialect/Util/IR/UtilOps.cpp index f555f63..9e0b3ff 100644 --- a/iree/compiler/Dialect/Util/IR/UtilOps.cpp +++ b/iree/compiler/Dialect/Util/IR/UtilOps.cpp
@@ -663,6 +663,24 @@ } //===----------------------------------------------------------------------===// +// Numeric ops +//===----------------------------------------------------------------------===// + +Optional<std::pair<int64_t, int64_t>> +NumericOptionalNarrowOp::getIntegerRange() { + if (!min_value() || !max_value()) return {}; + bool signExtend = isSigned(); + // Note: Cannot sign extend 0 bit values. + int64_t minValue = signExtend && min_value()->getBitWidth() > 0 + ? min_value()->getSExtValue() + : min_value()->getZExtValue(); + int64_t maxValue = signExtend && max_value()->getBitWidth() > 0 + ? max_value()->getSExtValue() + : max_value()->getZExtValue(); + return std::make_pair(minValue, maxValue); +} + +//===----------------------------------------------------------------------===// // Structural ops //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Util/IR/UtilOps.td b/iree/compiler/Dialect/Util/IR/UtilOps.td index 7b5131a..f3561de 100644 --- a/iree/compiler/Dialect/Util/IR/UtilOps.td +++ b/iree/compiler/Dialect/Util/IR/UtilOps.td
@@ -96,6 +96,81 @@ } //===----------------------------------------------------------------------===// +// Data type conversions +//===----------------------------------------------------------------------===// + +def Util_NumericOptionalNarrowOp : Util_PureOp<"numeric.optional_narrow", [ + SameOperandsAndResultType +]> { + let summary = "memorializes an optional numeric narrowing that is valid"; + let description = [{ + Serves as a placeholder for points in the computation where an optional + numeric narrowing can be performed without loss of information. Such ops + can guide optimization passes wishing to perform precision reduction. + + In addition to the operand and result type, this op takes an additional + `semantic_type` attribute representing the semantic target type which can + be: + * FloatType + * Signed IntegerType + * Unsigned IntegerType + + Note that this `semantic_type` must be a sign-carrying integer if using an + integer type and cannot be IndexType (i.e. it can be used to indicate a + possible narrowing of an IndexType to a specific integer). + + If the operand is a TensorType, then the result must be a TensorType. The + `semantic_type` constrains the element type. + + Optionally, the minimum and maximum integer values (for integer semantic + types) are tracked if known. + }]; + + let arguments = (ins + AnyTypeOf<[Util_Element, Util_Tensor]>:$operand, + TypeAttr:$semantic_type, + OptionalAttr<APIntAttr>:$min_value, + OptionalAttr<APIntAttr>:$max_value + ); + let results = (outs + AnyTypeOf<[Util_Element, Util_Tensor]>:$result + ); + + let assemblyFormat = [{ + $operand `:` type($operand) `as` $semantic_type attr-dict + }]; + + let builders = [ + OpBuilder<(ins + "Value":$operand, + "Type":$type, + "Optional<std::pair<int64_t, int64_t>>":$integerRange + ), + [{ + IntegerAttr minValueAttr; + IntegerAttr maxValueAttr; + if (integerRange) { + minValueAttr = $_builder.getIntegerAttr(type, integerRange->first); + maxValueAttr = $_builder.getIntegerAttr(type, integerRange->second); + } + build($_builder, $_state, operand.getType(), operand, TypeAttr::get(type), + minValueAttr, maxValueAttr); + }]>, + ]; + + let extraClassDeclaration = [{ + bool isSigned() { + if (auto integerType = getType().dyn_cast<IntegerType>()) { + return !integerType.isUnsigned(); + } + return true; + } + + Optional<std::pair<int64_t, int64_t>> getIntegerRange(); + }]; +} + +//===----------------------------------------------------------------------===// // Range arithmetic //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/Util/IR/test/BUILD b/iree/compiler/Dialect/Util/IR/test/BUILD index 457b1e6..84f17fb 100644 --- a/iree/compiler/Dialect/Util/IR/test/BUILD +++ b/iree/compiler/Dialect/Util/IR/test/BUILD
@@ -26,6 +26,7 @@ "hint_folding.mlir", "hint_ops.mlir", "list_ops.mlir", + "numeric_ops.mlir", "range_folding.mlir", "range_ops.mlir", "structural_folding.mlir",
diff --git a/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt b/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt index 3fa6522..c2f61f8 100644 --- a/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt +++ b/iree/compiler/Dialect/Util/IR/test/CMakeLists.txt
@@ -23,6 +23,7 @@ "hint_folding.mlir" "hint_ops.mlir" "list_ops.mlir" + "numeric_ops.mlir" "range_folding.mlir" "range_ops.mlir" "structural_folding.mlir"
diff --git a/iree/compiler/Dialect/Util/IR/test/numeric_ops.mlir b/iree/compiler/Dialect/Util/IR/test/numeric_ops.mlir new file mode 100644 index 0000000..0481412 --- /dev/null +++ b/iree/compiler/Dialect/Util/IR/test/numeric_ops.mlir
@@ -0,0 +1,19 @@ +// RUN: iree-opt -split-input-file %s | IreeFileCheck %s + +func @optional_convert_scalar(%arg0 : i32) -> i32 { + // CHECK: util.numeric.optional_narrow %arg0 : i32 as si8 + %0 = util.numeric.optional_narrow %arg0 : i32 as si8 + return %0 : i32 +} + +func @optional_convert_tensor(%arg0 : tensor<f32>) -> tensor<f32> { + // CHECK: util.numeric.optional_narrow %arg0 : tensor<f32> as si8 + %0 = util.numeric.optional_narrow %arg0 : tensor<f32> as si8 + return %0 : tensor<f32> +} + +func @optional_convert_zero(%arg0 : i32) -> i32 { + // CHECK: util.numeric.optional_narrow %arg0 : i32 as ui0 + %0 = util.numeric.optional_narrow %arg0 : i32 as ui0 + return %0 : i32 +}
diff --git a/iree/tools/init_iree_dialects.h b/iree/tools/init_iree_dialects.h index e4c8277..0aa7cb1 100644 --- a/iree/tools/init_iree_dialects.h +++ b/iree/tools/init_iree_dialects.h
@@ -47,6 +47,7 @@ // clang-format on IREE::LinalgExt::registerTiledOpInterfaceExternalModels(registry); + IREE::Util::registerUtilExternalModels(registry); registerCodegenInterfaces(registry); }