[mlir][arith] Add initial integer bitwidth narrowing pass

[kuhar]

Summary

This pass reduces the logical complexity of arith ops by choosing
narrowest supported operand bitwidth. On some targets like mobile GPUs,
narrower bitwidths also bring better runtime performance.

The first batch of rewrites handles a simple case of arith.sitofp
and arith.uitofp with zero/sign-extended inputs. In future revisions,
I plan to extend it with the following:

• Propagating sign/zero-extensions through bit-pattern-preserving ops, e.g., vector transpose, broadcast, insertions/extractions.
• Handling linalg.index using the ValueBounds interface.
• Handling more arith ops.

[antiagainst]

Cool, stuff. LGTM for the first version as a basis!

mlir/test/Dialect/Arith/int-narrowing.mlir

2

Also add tests for cases where we have:

• Consecutive extension ops of the same kind
• Consecutive extension ops of different kinds

to check the recursive aspect

[kuhar]

Add test cases with consecutive extensions

[springerm]

mlir/test/Dialect/Arith/int-narrowing.mlir

12

Put // ----- between tests, otherwise --split-input-file has no effect.

[kuhar]

Drop --split-file

mlir/include/mlir/Dialect/Arith/Transforms/Passes.h

 
 /// Add patterns for int range based optimizations.
 void populateIntRangeOptimizationsPatterns(RewritePatternSet &patterns,
                                            DataFlowSolver &solver);
 
 /// Create a pass which do optimizations based on integer range analysis.
 std::unique_ptr<Pass> createIntRangeOptimizationsPass();
 
+/// Add patterns for integer bitwidth narrowing.
+void populateArithIntNarrowingPatterns(RewritePatternSet &patterns,
+                                       const ArithIntNarrowingOptions &options);
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
 
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITH_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/Arith/Transforms/Passes.td

def ArithEmulateWideInt : Pass<"arith-emulate-wide-int"> {
   }];
   let options = [
     Option<"widestIntSupported", "widest-int-supported", "unsigned",
            /*default=*/"32", "Widest integer type supported by the target">,
   ];
   let dependentDialects = ["vector::VectorDialect"];
 }
 
+def ArithIntNarrowing : Pass<"arith-int-narrowing"> {
+  let summary = "Reduce integer operation bitwidth";
+  let description = [{
+    Reduce bitwidths of integer types used in arith operations. This pass
+    prefers the narrowest available integer bitwidths that are guaranteed to
+    produce the same results.
+  }];
+  let options = [
+    ListOption<"bitwidthsSupported", "int-bitwidths-supported", "unsigned",
+               "Integer bitwidths supported">,
+  ];
+ }
+
 #endif // MLIR_DIALECT_ARITH_TRANSFORMS_PASSES

mlir/lib/Dialect/Arith/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRArithTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   EmulateWideInt.cpp
   ExpandOps.cpp
+  IntNarrowing.cpp
   IntRangeOptimizations.cpp
   ReifyValueBounds.cpp
   UnsignedWhenEquivalent.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith/Transforms
 
   DEPENDS

mlir/lib/Dialect/Arith/Transforms/IntNarrowing.cpp

+//===- IntNarrowing.cpp - Integer bitwidth reduction optimizations --------===//
+//
+// Part of the LLVM Project, 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 "mlir/Dialect/Arith/Transforms/Passes.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include <cassert>
+#include <cstdint>
+
+namespace mlir::arith {
+#define GEN_PASS_DEF_ARITHINTNARROWING
+#include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
+} // namespace mlir::arith
+
+namespace mlir::arith {
+namespace {
+//===----------------------------------------------------------------------===//
+// Common Helpers
+//===----------------------------------------------------------------------===//
+
+/// The base for integer bitwidth narrowing patterns.
+template <typename SourceOp>
+struct NarrowingPattern : OpRewritePattern<SourceOp> {
+  NarrowingPattern(MLIRContext *ctx, const ArithIntNarrowingOptions &options,
+                   PatternBenefit benefit = 1)
+      : OpRewritePattern<SourceOp>(ctx, benefit),
+        supportedBitwidths(options.bitwidthsSupported.begin(),
+                           options.bitwidthsSupported.end()) {
+    assert(!supportedBitwidths.empty() && "Invalid options");
+    assert(!llvm::is_contained(supportedBitwidths, 0) && "Invalid bitwidth");
+    llvm::sort(supportedBitwidths);
+  }
+
+  FailureOr<unsigned>
+  getNarrowestCompatibleBitwidth(unsigned bitsRequired) const {
+    for (unsigned candidate : supportedBitwidths)
+      if (candidate >= bitsRequired)
+        return candidate;
+
+    return failure();
+  }
+
+  /// Returns the narrowest supported type that fits `bitsRequired`.
+  FailureOr<Type> getNarrowType(unsigned bitsRequired, Type origTy) const {
+    assert(origTy);
+    FailureOr<unsigned> bestBitwidth =
+        getNarrowestCompatibleBitwidth(bitsRequired);
+    if (failed(bestBitwidth))
+      return failure();
+
+    auto elemTy = getElementTypeOrSelf(origTy);
+    if (!isa<IntegerType>(elemTy))
+      return failure();
+
+    auto newElemTy = IntegerType::get(origTy.getContext(), bitsRequired);
+    if (newElemTy == elemTy)
+      return failure();
+
+    if (origTy == elemTy)
+      return newElemTy;
+
+    if (auto shapedTy = dyn_cast<ShapedType>(origTy))
+      if (auto elemTy = dyn_cast<IntegerType>(shapedTy.getElementType()))
+        return shapedTy.clone(shapedTy.getShape(), newElemTy);
+
+    return failure();
+  }
+
+private:
+  // Supported integer bitwidths in the ascending order.
+  llvm::SmallVector<unsigned, 6> supportedBitwidths;
+};
+
+/// Returns the integer bitwidth required to represent `type`.
+FailureOr<unsigned> calculateBitsRequired(Type type) {
+  assert(type);
+  if (auto intTy = dyn_cast<IntegerType>(getElementTypeOrSelf(type)))
+    return intTy.getWidth();
+
+  return failure();
+}
+
+enum class ExtensionKind { Sign, Zero };
+
+/// Returns the integer bitwidth required to represent `value`.
+/// Looks through either sign- or zero-extension as specified by
+/// `lookThroughExtension`.
+FailureOr<unsigned> calculateBitsRequired(Value value,
+                                          ExtensionKind lookThroughExtension) {
+  if (lookThroughExtension == ExtensionKind::Sign) {
+    if (auto sext = value.getDefiningOp<arith::ExtSIOp>())
+      return calculateBitsRequired(sext.getIn().getType());
+  } else if (lookThroughExtension == ExtensionKind::Zero) {
+    if (auto zext = value.getDefiningOp<arith::ExtUIOp>())
+      return calculateBitsRequired(zext.getIn().getType());
+  }
+
+  // If nothing else worked, return the type requirements for this element type.
+  return calculateBitsRequired(value.getType());
+}
+
+//===----------------------------------------------------------------------===//
+// *IToFPOp Patterns
+//===----------------------------------------------------------------------===//
+
+template <typename IToFPOp, ExtensionKind Extension>
+struct IToFPPattern final : NarrowingPattern<IToFPOp> {
+  using NarrowingPattern<IToFPOp>::NarrowingPattern;
+
+  LogicalResult matchAndRewrite(IToFPOp op,
+                                PatternRewriter &rewriter) const override {
+    FailureOr<unsigned> narrowestWidth =
+        calculateBitsRequired(op.getIn(), Extension);
+    if (failed(narrowestWidth))
+      return failure();
+
+    FailureOr<Type> narrowTy =
+        this->getNarrowType(*narrowestWidth, op.getIn().getType());
+    if (failed(narrowTy))
+      return failure();
+
+    Value newIn = rewriter.createOrFold<arith::TruncIOp>(op.getLoc(), *narrowTy,
+                                                         op.getIn());
+    rewriter.replaceOpWithNewOp<IToFPOp>(op, op.getType(), newIn);
+    return success();
+  }
+};
+using SIToFPPattern = IToFPPattern<arith::SIToFPOp, ExtensionKind::Sign>;
+using UIToFPPattern = IToFPPattern<arith::UIToFPOp, ExtensionKind::Zero>;
+
+//===----------------------------------------------------------------------===//
+// Pass Definitions
+//===----------------------------------------------------------------------===//
+
+struct ArithIntNarrowingPass final
+    : impl::ArithIntNarrowingBase<ArithIntNarrowingPass> {
+  using ArithIntNarrowingBase::ArithIntNarrowingBase;
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    MLIRContext *ctx = op->getContext();
+    RewritePatternSet patterns(ctx);
+    populateArithIntNarrowingPatterns(
+        patterns, ArithIntNarrowingOptions{bitwidthsSupported});
+    if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Public API
+//===----------------------------------------------------------------------===//
+
+void populateArithIntNarrowingPatterns(
+    RewritePatternSet &patterns, const ArithIntNarrowingOptions &options) {
+  patterns.add<SIToFPPattern, UIToFPPattern>(patterns.getContext(), options);
+}
+
+} // namespace mlir::arith

mlir/test/Dialect/Arith/int-narrowing.mlir

+// RUN: mlir-opt --arith-int-narrowing="int-bitwidths-supported=1,8,16,32" \
+// RUN:          --verify-diagnostics --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: func.func @sitofp_extsi_i16
+// CHECK-SAME:    (%[[ARG:.+]]: i16)
+// CHECK-NEXT:    %[[RET:.+]] = arith.sitofp %[[ARG]] : i16 to f16
+// CHECK-NEXT:    return %[[RET]] : f16
+func.func @sitofp_extsi_i16(%a: i16) -> f16 {
+  %b = arith.extsi %a : i16 to i32
+  %f = arith.sitofp %b : i32 to f16
+  return %f : f16
+}
+
+// CHECK-LABEL: func.func @sitofp_extsi_vector_i16
+// CHECK-SAME:    (%[[ARG:.+]]: vector<3xi16>)
+// CHECK-NEXT:    %[[RET:.+]] = arith.sitofp %[[ARG]] : vector<3xi16> to vector<3xf16>
+// CHECK-NEXT:    return %[[RET]] : vector<3xf16>
+func.func @sitofp_extsi_vector_i16(%a: vector<3xi16>) -> vector<3xf16> {
+  %b = arith.extsi %a : vector<3xi16> to vector<3xi32>
+  %f = arith.sitofp %b : vector<3xi32> to vector<3xf16>
+  return %f : vector<3xf16>
+}
+
+// CHECK-LABEL: func.func @sitofp_extsi_tensor_i16
+// CHECK-SAME:    (%[[ARG:.+]]: tensor<3x?xi16>)
+// CHECK-NEXT:    %[[RET:.+]] = arith.sitofp %[[ARG]] : tensor<3x?xi16> to tensor<3x?xf16>
+// CHECK-NEXT:    return %[[RET]] : tensor<3x?xf16>
+func.func @sitofp_extsi_tensor_i16(%a: tensor<3x?xi16>) -> tensor<3x?xf16> {
+  %b = arith.extsi %a : tensor<3x?xi16> to tensor<3x?xi32>
+  %f = arith.sitofp %b : tensor<3x?xi32> to tensor<3x?xf16>
+  return %f : tensor<3x?xf16>
+}
+
+// Narrowing to i64 is not enabled in pass options.
+//
+// CHECK-LABEL: func.func @sitofp_extsi_i64
+// CHECK-SAME:    (%[[ARG:.+]]: i64)
+// CHECK-NEXT:    %[[EXT:.+]] = arith.extsi %[[ARG]] : i64 to i128
+// CHECK-NEXT:    %[[RET:.+]] = arith.sitofp %[[EXT]] : i128 to f32
+// CHECK-NEXT:    return %[[RET]] : f32
+func.func @sitofp_extsi_i64(%a: i64) -> f32 {
+  %b = arith.extsi %a : i64 to i128
+  %f = arith.sitofp %b : i128 to f32
+  return %f : f32
+}
+
+// CHECK-LABEL: func.func @uitofp_extui_i16
+// CHECK-SAME:    (%[[ARG:.+]]: i16)
+// CHECK-NEXT:    %[[RET:.+]] = arith.uitofp %[[ARG]] : i16 to f16
+// CHECK-NEXT:    return %[[RET]] : f16
+func.func @uitofp_extui_i16(%a: i16) -> f16 {
+  %b = arith.extui %a : i16 to i32
+  %f = arith.uitofp %b : i32 to f16
+  return %f : f16
+}
+
+// CHECK-LABEL: func.func @sitofp_extsi_extsi_i8
+// CHECK-SAME:    (%[[ARG:.+]]: i8)
+// CHECK-NEXT:    %[[RET:.+]] = arith.sitofp %[[ARG]] : i8 to f16
+// CHECK-NEXT:    return %[[RET]] : f16
+func.func @sitofp_extsi_extsi_i8(%a: i8) -> f16 {
+  %b = arith.extsi %a : i8 to i16
+  %c = arith.extsi %b : i16 to i32
+  %f = arith.sitofp %c : i32 to f16
+  return %f : f16
+}
+
+// CHECK-LABEL: func.func @uitofp_extui_extui_i8
+// CHECK-SAME:    (%[[ARG:.+]]: i8)
+// CHECK-NEXT:    %[[RET:.+]] = arith.uitofp %[[ARG]] : i8 to f16
+// CHECK-NEXT:    return %[[RET]] : f16
+func.func @uitofp_extui_extui_i8(%a: i8) -> f16 {
+  %b = arith.extui %a : i8 to i16
+  %c = arith.extui %b : i16 to i32
+  %f = arith.uitofp %c : i32 to f16
+  return %f : f16
+}
+
+// CHECK-LABEL: func.func @uitofp_extsi_extui_i8
+// CHECK-SAME:    (%[[ARG:.+]]: i8)
+// CHECK-NEXT:    %[[EXT:.+]] = arith.extsi %[[ARG]] : i8 to i16
+// CHECK-NEXT:    %[[RET:.+]] = arith.uitofp %[[EXT]] : i16 to f16
+// CHECK-NEXT:    return %[[RET]] : f16
+func.func @uitofp_extsi_extui_i8(%a: i8) -> f16 {
+  %b = arith.extsi %a : i8 to i16
+  %c = arith.extui %b : i16 to i32
+  %f = arith.uitofp %c : i32 to f16
+  return %f : f16
+}
+
+// CHECK-LABEL: func.func @uitofp_trunci_extui_i8
+// CHECK-SAME:    (%[[ARG:.+]]: i16)
+// CHECK-NEXT:    %[[TR:.+]]  = arith.trunci %[[ARG]] : i16 to i8
+// CHECK-NEXT:    %[[RET:.+]] = arith.uitofp %[[TR]] : i8 to f16
+// CHECK-NEXT:    return %[[RET]] : f16
+func.func @uitofp_trunci_extui_i8(%a: i16) -> f16 {
+  %b = arith.trunci %a : i16 to i8
+  %c = arith.extui %b : i8 to i32
+  %f = arith.uitofp %c : i32 to f16
+  return %f : f16
+}
+
+// This should not be folded because arith.extui changes the signed
+// range of the number. For example:
+//  extsi -1 : i16 to i32 ==> -1
+//  extui -1 : i16 to i32 ==> U16_MAX
+//
+/// CHECK-LABEL: func.func @sitofp_extui_i16
+// CHECK-SAME:    (%[[ARG:.+]]: i16)
+// CHECK-NEXT:    %[[EXT:.+]] = arith.extui %[[ARG]] : i16 to i32
+// CHECK-NEXT:    %[[RET:.+]] = arith.sitofp %[[EXT]] : i32 to f16
+// CHECK-NEXT:    return %[[RET]] : f16
+func.func @sitofp_extui_i16(%a: i16) -> f16 {
+  %b = arith.extui %a : i16 to i32
+  %f = arith.sitofp %b : i32 to f16
+  return %f : f16
+}
+
+// This should not be folded because arith.extsi changes the unsigned
+// range of the number. For example:
+//  extsi -1 : i16 to i32 ==> U32_MAX
+//  extui -1 : i16 to i32 ==> U16_MAX
+//
+// CHECK-LABEL: func.func @uitofp_extsi_i16
+// CHECK-SAME:    (%[[ARG:.+]]: i16)
+// CHECK-NEXT:    %[[EXT:.+]] = arith.extsi %[[ARG]] : i16 to i32
+// CHECK-NEXT:    %[[RET:.+]] = arith.uitofp %[[EXT]] : i32 to f16
+// CHECK-NEXT:    return %[[RET]] : f16
+func.func @uitofp_extsi_i16(%a: i16) -> f16 {
+  %b = arith.extsi %a : i16 to i32
+  %f = arith.uitofp %b : i32 to f16
+  return %f : f16
+}