[mlir][arith] Add wide integer emulation pass

[kuhar]

Summary

In this first patch in a series to add wide integer emulation:

• Set up the initial pass structure
• Add a custom type converter
• Handle func ops

The initial implementation supports power-of-two integers types only. We
emulate wide integer operations by splitting original i2N integer types
into two iN halves

My immediate use case is to emulate i64 operations using i32 ones
on mobile GPUs that do not support i64.

[krzysz00]

As a general question, why limit yourself to the next power of 2? (and also, that limitation didn't feel documented)

Would it be too much bother to allow, for example, i128 -> 4xi32?

[kuhar]

In D133135#3765224, @krzysz00 wrote:

As a general question, why limit yourself to the next power of 2?

Would it be too much bother to allow, for example, i128 -> 4xi32?

I had a few reasons:

1. In my use case, I don't need that level of generality.
2. Specializing for the next power of 2 makes the implementation and testing simpler. I'd expect a lot more corner cases to appear once we allow arbitrarily large integers. For example, what if someone wants to emulate i64 with i1s?
3. Some ops get quite inefficient when emulated. For example, my conversion pattern for arith.muli` generates ~60 instructions already when splitting the bit width in half.
4. The current implementation provides a workaround: you can run the emulation pass multiple times, gradually decreasing the widest int type supported.

(and also, that limitation didn't feel documented)

The patch mentions this in the pass description and in the comment by the create function. Do you have some recommendations on how to communicate this better?

[kuhar]

(and also, that limitation didn't feel documented)

The patch mentions this in the pass description and in the comment by the create function. Do you have some recommendations on how to communicate this better?

Ah, I think I see what you mean. It does mention i2N -> iN but not that N has to be a power of two.

[krzysz00]

I think it was more that the description gave the impression that this was
the general case when I first saw it.

If this were explicitly a "split wide integer math in half" pass it'd be
less confusing - "wide integer emulation" feels like it'd handle the
general case.

(and if someone wants to autoconvert i64 math to i1, let 'em - I expect the
compiler will translate add with carry in and carry out to the relevant
Boolean operators or tell the user to get lost)

[kuhar]

Update the pass description and comments to be more explicit about supported integer types and how we emulate wide integer operations.

[kuhar]

Fix a comment.

[kuhar]

Thanks for the feedback, @krzysz00.

In principle, I think that generalizing this pass to use a more general emulation scheme could be useful. For the initial implementation, I'd prefer to keep things simple (and sound) and cover the use case we have in mind.

I updated the comments, documentation, and revision description to be more explicit about the supported types and how we emulate wide operations. Does this address your concerns?

[jacquesguan]

I have a little question, what about if some target does support different max integer width between scalar and vector type? For example, in RISCV, RV64 + Zve32x supports i64 integer scalar but only max to i32 integer vector type. What should we set for this case?

[kuhar]

In D133135#3771398, @jacquesguan wrote:

I have a little question, what about if some target does support different max integer width between scalar and vector type? For example, in RISCV, RV64 + Zve32x supports i64 integer scalar but only max to i32 integer vector type. What should we set for this case?

Right now this is not supported, but I think we should be able to handle this. We can add a second option to the pass and type converter for the max vector int width supported so that integers and vectors of integers get converted to different types. As for the types used for emulation, we convert ints to 2-element vectors but don't perform any dynamic operations on them -- we only use vectors to keep two results in one value. If this is not supported by the target, we could instead use a struct with 2 elements or tuple with 2 values.

[Mogball]

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

15–16

33

Do you need this constructor? I think it's auto-generated by ODS

38

Returning a pointer to a TypeConverter is a little odd. The usual pattern is to subclass TypeConverter (e.g. WideIntTypeConverter) and expose it in the header

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

60

I would put efficiency as a TODO

67

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

12–21

31

This sentence would suggest the type converter is capable of converting multiple integer types, but in fact only one type is supported: i2N since all other types are rejected by the converter.

31–37

48

54

64

68

I would expect vector<4xi64> to become vector<4x2xi32> instead of vector<2x4xi32>

111

Why are your asserting false? The program will abort instead of allowing the pass to gracefully fail.

126

Is this necessary for a partial conversion?

135

Can you outline this lambda?

141

This shouldn't be necessary in a partial conversion

mlir/test/Dialect/Arithmetic/emulate-wide-int.mlir

33

Can you add a test for a scalar i64 conversion?

[antiagainst]

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

33

Nit: widestMachineIntBitwidth? widestIntSupported makes me think that the widest integer supported for emulation is 32.. Would also be nice to explain the parameter in the comments.

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

56

... on narrower integer ...

62

.. only power-of two integer bitwidths ...

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

46

Just capture this?

61

Same here.

63

In MLIR the convention is to do early returns.

68

Hmm, I'd expect this to be vector<...x2xiN> to keep the two halves close. The current way it would require transposing.

142

These shouldn't be unconditionally legal right? Actually the dynamic legality specification at L129 should applicable to all ops.

[antiagainst]

+1 to progressively build up the functionality. This way we can make sure we build up sufficient amount of testing to cover various cases too.

[kuhar]

Addressed non-functional issues.

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

33

Thanks, I didn't realize ODS also automates constructors with pass options.

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

126

It seems so to me. I added a test case that demonstrates this. Without these materializations, it fails with:

error: failed to legalize un
resolved materialization from 'i64' to 'vector<2xi32>' that remained live after conversion                                          
    return %i : i64                                               
    ^

[kuhar]

Moved type converter to its own header. Changed vector layout.

[kuhar]

All comments should be addressed now. Could you take a second look, @Mogball and @antiagainst?

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

68

Thanks for the suggestion. Initially, I chose vector<2x..iN> because it seemed easy to construct/destruct vectors with this new outer dim. However, your suggestion makes more sense because it essentially allows for bitcasting between original types and emulated ones.

[Mogball]

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

29

Please move these all out of the anonymous namespace and make them static.

[kuhar]

Made helper functions static. Rebased.

[Mogball]

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

27

I would just inline this function in addDynamicallyLegalOp

41

Wouldn't this allow i0 or i1?

59

can you split this up with target.addDynamicallyLegalOp<func::funcOp>(...) and target.addDynamicallyLegalOp<func::CallOp, func::ReturnOp>(...) so that isOpLegal doesn't have to switch on func::FuncOp?

73–75

126

This seems like a test case that *should* fail. If an i64 operation could not be emulated (because it lacks an implementation, for example), then this pass should fail.

131

[kuhar]

Cleanups. Removed support for unhandled arith ops.

[kuhar]

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

41

0 is not considered a power of 2. Emulating i2 with i1 is supported, although I haven't tested it.

59

thanks for the suggestion, this simplifies things

[Mogball]

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

41

Right, thanks.

[Mogball]

LGTM

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

58

Can you drop this since you have using namespace mlir declared at the top of the file?

[kuhar]

Dropped the mlir namespace.

[antiagainst]

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

34

I think you can just return signalPassFailure() here.

[kuhar]

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

34

The return type is void, so this seems unusual to me. Not sure what the benefit would be apart from saving one line of code.

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

 //===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_ARITHMETIC_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_ARITHMETIC_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+class TypeConverter;
 namespace arith {
 
 #define GEN_PASS_DECL_ARITHMETICBUFFERIZE
+#define GEN_PASS_DECL_ARITHMETICEMULATEWIDEINT
 #define GEN_PASS_DECL_ARITHMETICEXPANDOPS
 #define GEN_PASS_DECL_ARITHMETICUNSIGNEDWHENEQUIVALENT
 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h.inc"
 
 /// Create a pass to bufferize Arithmetic ops.
 std::unique_ptr<Pass> createArithmeticBufferizePass();
 
 /// Create a pass to bufferize arith.constant ops.
 std::unique_ptr<Pass> createConstantBufferizePass(uint64_t alignment = 0);
 
+/// Creates a type converter that splits wide power-of-two integer types i2N
+/// into two iN halves.
+std::unique_ptr<TypeConverter>
+createWideIntEmulationTypeConverter(unsigned widestIntSupported = 32);
+
+/// Adds patterns to emulate wide Arithmetic and Function ops over integer
+/// types into supported ones. This is done by splitting original power-of-two
+/// i2N integer types into two iN halves.
+void populateWideIntEmulationPatterns(TypeConverter &typeConverter,
+                                      RewritePatternSet &patterns);
+
 /// Add patterns to expand Arithmetic ops for LLVM lowering.
 void populateArithmeticExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Create a pass to legalize Arithmetic ops for LLVM lowering.
 std::unique_ptr<Pass> createArithmeticExpandOpsPass();
 
 /// Create a pass to replace signed ops with unsigned ones where they are proven
 /// equivalent.

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

let description = [{
     range [0, signed_max(type)].
 
     The affect ops include division, remainder, shifts, min, max, and integer
     comparisons.
   }];
   let constructor = "mlir::arith::createArithmeticUnsignedWhenEquivalentPass()";
 }
 
+def ArithmeticEmulateWideInt : Pass<"arith-emulate-wide-int"> {
+  let summary = "Emulate 2*N-bit integer operations using N-bit operations";
+  let description = [{
+    Emulate integer operations that use too wide integer types with equivalent
+    operations on supported narrow integer types. This is done by splitting
+    original integer values into two halves.
+
+    This pass is intended preserve semantics but not necessarily provide the
+    most efficient implementation.
+    TODO: Optimize op emulation.
+
+    Currently, only power-of-two integer bitwidths are supported.
+  }];
+  let options = [
+    Option<"widestIntSupported", "widest-int-supported", "unsigned",
+           /*default=*/"32", "Widest integer type supported by the target">,
+  ];
+  let dependentDialects = ["vector::VectorDialect"];
+}
+
 #endif // MLIR_DIALECT_ARITHMETIC_TRANSFORMS_PASSES

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

 add_mlir_dialect_library(MLIRArithmeticTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
+  EmulateWideInt.cpp
   ExpandOps.cpp
   UnsignedWhenEquivalent.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arithmetic/Transforms
 
   DEPENDS
   MLIRArithmeticTransformsIncGen

mlir/lib/Dialect/Arithmetic/Transforms/EmulateWideInt.cpp

+//===- EmulateWideInt.cpp - Pass to emulate wide integer operations ------===//
+//
+// 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/Arithmetic/Transforms/Passes.h"
+
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/MathExtras.h"
+#include <cassert>
+
+namespace mlir::arith {
+#define GEN_PASS_DEF_ARITHMETICEMULATEWIDEINT
+#include "mlir/Dialect/Arithmetic/Transforms/Passes.h.inc"
+} // namespace mlir::arith
+
+using namespace mlir;
+
+namespace {
+/// Converts integer types that are too wide for the target by splitting them in
+/// two halves and thus turning into supported ones, i.e., i2*N --> iN, where N
+/// is the widest integer bitwidth supported by the target.
+/// Currently, we only handle power-of-two integer types and support conversions
+/// of integers twice as wide as the maxium supported by the target. Wide
+/// integers are represented as vectors, e.g., i64 --> vector<2xi32>, where the
+/// first element is the low half of the original integer, and the second
+/// element the high half.
+class WideIntEmulationConverter final : public TypeConverter {
+public:
+  explicit WideIntEmulationConverter(unsigned widestIntSupportedByTarget)
+      : maxIntWidth(widestIntSupportedByTarget) {
+    assert(llvm::isPowerOf2_32(widestIntSupportedByTarget) &&
+           "Only power-of-two integers are supported");
+
+    // Scalar case.
+    addConversion([this](IntegerType ty) -> Optional<Type> {
+      unsigned width = ty.getWidth();
+      if (width <= maxIntWidth)
+        return ty;
+
+      // i2N --> vector<2xiN>
+      if (width == 2 * maxIntWidth)
+        return VectorType::get(2,
+                               IntegerType::get(ty.getContext(), maxIntWidth));
+
+      return None;
+    });
+
+    // Vector case.
+    addConversion([this](VectorType ty) -> Optional<Type> {
+      auto intTy = ty.getElementType().dyn_cast<IntegerType>();
+      if (!intTy)
+        return ty;
+
+      unsigned width = intTy.getWidth();
+      if (width <= maxIntWidth)
+        return ty;
+
+      // vector<...xi2N> --> vector<2x...xiN>
+      if (width == 2 * maxIntWidth) {
+        SmallVector<int64_t> newShape = {2};
+        llvm::append_range(newShape, ty.getShape());
+        return VectorType::get(newShape,
+                               IntegerType::get(ty.getContext(), maxIntWidth));
+      }
+
+      return None;
+    });
+
+    // Function case.
+    addConversion([this](FunctionType ty) -> Optional<Type> {
+      // Convert inputs and results, e.g.:
+      //   (i2N, i2N) -> i2N --> (vector<2xiN>, vector<2xiN>) -> vector<2xiN>
+      SmallVector<Type> inputs;
+      if (failed(convertTypes(ty.getInputs(), inputs)))
+        return None;
+
+      SmallVector<Type> results;
+      if (failed(convertTypes(ty.getResults(), results)))
+        return None;
+
+      return FunctionType::get(ty.getContext(), inputs, results);
+    });
+  }
+
+  unsigned getMaxIntegerWidth() const { return maxIntWidth; }
+
+private:
+  unsigned maxIntWidth;
+};
+
+bool isOpLegal(Operation *op, WideIntEmulationConverter &converter) {
+  if (auto func = dyn_cast<func::FuncOp>(op))
+    return converter.isLegal(func.getFunctionType());
+
+  return converter.isLegal(op);
+}
+
+struct EmulateWideIntPass final
+    : arith::impl::ArithmeticEmulateWideIntBase<EmulateWideIntPass> {
+  using ArithmeticEmulateWideIntBase::ArithmeticEmulateWideIntBase;
+
+  void runOnOperation() override {
+    if (!llvm::isPowerOf2_32(widestIntSupported)) {
+      signalPassFailure();
+      return;
+    }
+
+    Operation *op = getOperation();
+    MLIRContext *ctx = op->getContext();
+
+    WideIntEmulationConverter typeConverter(widestIntSupported);
+    auto addUnrealizedCast = [](OpBuilder &builder, Type type,
+                                ValueRange inputs, Location loc) {
+      auto cast = builder.create<UnrealizedConversionCastOp>(loc, type, inputs);
+      return Optional<Value>(cast.getResult(0));
+    };
+    typeConverter.addSourceMaterialization(addUnrealizedCast);
+    typeConverter.addTargetMaterialization(addUnrealizedCast);
+
+    ConversionTarget target(*ctx);
+    target.addDynamicallyLegalOp<func::FuncOp, func::CallOp, func::ReturnOp>(
+        [&typeConverter](Operation *op) {
+          return isOpLegal(op, typeConverter);
+        });
+
+    RewritePatternSet patterns(ctx);
+    arith::populateWideIntEmulationPatterns(typeConverter, patterns);
+
+    if (failed(applyPartialConversion(op, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // end anonymous namespace
+
+namespace mlir::arith {
+
+void populateWideIntEmulationPatterns(TypeConverter &typeConverter,
+                                      RewritePatternSet &patterns) {
+  // Populate `func.*` conversion patterns.
+  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                 typeConverter);
+  populateCallOpTypeConversionPattern(patterns, typeConverter);
+  populateReturnOpTypeConversionPattern(patterns, typeConverter);
+}
+
+std::unique_ptr<TypeConverter>
+createWideIntEmulationTypeConverter(unsigned widestIntSupported) {
+  return std::make_unique<WideIntEmulationConverter>(widestIntSupported);
+}
+
+} // namespace mlir::arith

mlir/test/Dialect/Arithmetic/emulate-wide-int.mlir

+// RUN: mlir-opt --arith-emulate-wide-int="widest-int-supported=32" %s | FileCheck %s
+
+// Expect no conversions, i32 is supported.
+// CHECK-LABEL: func @addi_same_i32
+// CHECK-SAME:    ([[ARG:%.+]]: i32) -> i32
+// CHECK-NEXT:    [[X:%.+]] = arith.addi [[ARG]], [[ARG]] : i32
+// CHECK-NEXT:    return [[X]] : i32
+func.func @addi_same_i32(%a : i32) -> i32 {
+    %x = arith.addi %a, %a : i32
+    return %x : i32
+}
+
+// Expect no conversions, i32 is supported.
+// CHECK-LABEL: func @addi_same_vector_i32
+// CHECK-SAME:    ([[ARG:%.+]]: vector<2xi32>) -> vector<2xi32>
+// CHECK-NEXT:    [[X:%.+]] = arith.addi [[ARG]], [[ARG]] : vector<2xi32>
+// CHECK-NEXT:    return [[X]] : vector<2xi32>
+func.func @addi_same_vector_i32(%a : vector<2xi32>) -> vector<2xi32> {
+    %x = arith.addi %a, %a : vector<2xi32>
+    return %x : vector<2xi32>
+}
+
+// CHECK-LABEL: func @identity_scalar
+// CHECK-SAME:     ([[ARG:%.+]]: vector<2xi32>) -> vector<2xi32>
+// CHECK-NEXT:     return [[ARG]] : vector<2xi32>
+func.func @identity_scalar(%x : i64) -> i64 {
+    return %x : i64
+}
+
+// CHECK-LABEL: func @identity_vector
+// CHECK-SAME:     ([[ARG:%.+]]: vector<2x4xi32>) -> vector<2x4xi32>
+// CHECK-NEXT:     return [[ARG]] : vector<2x4xi32>
+func.func @identity_vector(%x : vector<4xi64>) -> vector<4xi64> {
+    return %x : vector<4xi64>
+}
+
+// CHECK-LABEL: func @identity_vector2d
+// CHECK-SAME:     ([[ARG:%.+]]: vector<2x3x4xi32>) -> vector<2x3x4xi32>
+// CHECK-NEXT:     return [[ARG]] : vector<2x3x4xi32>
+func.func @identity_vector2d(%x : vector<3x4xi64>) -> vector<3x4xi64> {
+    return %x : vector<3x4xi64>
+}
+
+// CHECK-LABEL: func @call
+// CHECK-SAME:     ([[ARG:%.+]]: vector<2x4xi32>) -> vector<2x4xi32>
+// CHECK-NEXT:     [[RES:%.+]] = call @identity_vector([[ARG]]) : (vector<2x4xi32>) -> vector<2x4xi32>
+// CHECK-NEXT:     return [[RES]] : vector<2x4xi32>
+func.func @call(%a : vector<4xi64>) -> vector<4xi64> {
+    %res = func.call @identity_vector(%a) : (vector<4xi64>) -> vector<4xi64>
+    return %res : vector<4xi64>
+}
+
+// CHECK-LABEL: func @unrealized_conversion_cast
+// CHECK-SAME:     ([[ARG:%.+]]: vector<2xi32>) -> vector<2xi32>
+// CHECK-NEXT:     [[CAST0:%.+]] = builtin.unrealized_conversion_cast [[ARG]] : vector<2xi32> to i64
+// CHECK-NEXT:     [[FLOAT:%.+]] = arith.sitofp [[CAST0]] : i64 to f64
+// CHECK-NEXT:     [[INT:%.+]]   = arith.fptosi [[FLOAT]] : f64 to i64
+// CHECK-NEXT:     [[CAST1:%.+]] = builtin.unrealized_conversion_cast [[INT]] : i64 to vector<2xi32>
+// CHECK-NEXT:     return [[CAST1]] : vector<2xi32>
+func.func @unrealized_conversion_cast(%a : i64) -> i64 {
+    %f = arith.sitofp %a : i64 to f64
+    %i = arith.fptosi %f : f64 to i64
+    return %i : i64
+}