[mlir][arith] Support wide integer constant emulation

[Mogball]

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

46

Why not cast to VectorType?

47

72

Why do you need to cast this to size_t?

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

68

I would expect this to be [[0, 1], [0, 1], [0, 1]]

[antiagainst]

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

84

As commented in the previous patch, I think it's more natural to put these two halves adjacent. (Eventually we'd need to support memory types and load/store ops; putting them adjacent would also mean easier to handle there and better load/store behavior.) Any reason transposing it like this way?

96

Just make this as a static helper function?

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

68

Can use {LITERAL} to avoid all the escapes. https://llvm.org/docs/CommandGuide/FileCheck.html#directive-modifiers

[kuhar]

Addressed comments

[kuhar]

Removed an unnecessary include

[Mogball]

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

39

Please add a short doc

[kuhar]

Added documentation for the getHalves helper.

[Mogball]

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

39

Can you add code headers?

//===---------------------------------------------===//
// ConvertConstant
//===---------------------------------------------===//

[kuhar]

Added conversion pattern header. Rebased.

[Mogball]

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

38–42

95

Can you add a header here as well, something like "Pass Definition"

[kuhar]

Added headers. Rebased.

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

95

Done. Are these headers standardized / covered by the coding convention in this part of MLIR, or more of a personal preference?
For me, these don't seem very useful.

[Mogball]

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

95

They're nice to have when files start getting really long and op/pattern/pass definitions start to blur together. They also clearly indicate where new code should go that belong to the category (e.g. helper functions for patterns, etc.)

[Mogball]

LGTM just nit on the spacing around the headers

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

 //===- EmulateWideInt.cpp - Wide integer operation emulation ----*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
 
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Transforms/WideIntEmulationConverter.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/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;
 
+// Returns N bottom and N top bits from `value`, where N = `newBitWidth`.
+// Treats `value` as a 2*N bits-wide integer.
+// The bottom bits are returned in the first pair element, while the top bits in
+// the second one.
+static std::pair<APInt, APInt> getHalves(const APInt &value,
+                                         unsigned newBitWidth) {
+  APInt low = value.extractBits(newBitWidth, 0);
+  APInt high = value.extractBits(newBitWidth, newBitWidth);
+  return {std::move(low), std::move(high)};
+}
+
 namespace {
+//===----------------------------------------------------------------------===//
+// ConvertConstant
+//===----------------------------------------------------------------------===//
+struct ConvertConstant final : OpConversionPattern<arith::ConstantOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(arith::ConstantOp op, OpAdaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type oldType = op.getType();
+    auto newType = getTypeConverter()->convertType(oldType).cast<VectorType>();
+    unsigned newBitWidth = newType.getElementTypeBitWidth();
+    Attribute oldValue = op.getValueAttr();
+
+    if (auto intAttr = oldValue.dyn_cast<IntegerAttr>()) {
+      auto [low, high] = getHalves(intAttr.getValue(), newBitWidth);
+      auto newAttr = DenseElementsAttr::get(newType, {low, high});
+      rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);
+      return success();
+    }
+
+    if (auto splatAttr = oldValue.dyn_cast<SplatElementsAttr>()) {
+      auto [low, high] =
+          getHalves(splatAttr.getSplatValue<APInt>(), newBitWidth);
+      int64_t numSplatElems = splatAttr.getNumElements();
+      SmallVector<APInt> values;
+      values.reserve(numSplatElems * 2);
+      for (int64_t i = 0; i < numSplatElems; ++i) {
+        values.push_back(low);
+        values.push_back(high);
+      }
+
+      auto attr = DenseElementsAttr::get(newType, values);
+      rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, attr);
+      return success();
+    }
+
+    if (auto elemsAttr = oldValue.dyn_cast<DenseElementsAttr>()) {
+      int64_t numElems = elemsAttr.getNumElements();
+      SmallVector<APInt> values;
+      values.reserve(numElems * 2);
+      for (const APInt &origVal : elemsAttr.getValues<APInt>()) {
+        auto [low, high] = getHalves(origVal, newBitWidth);
+        values.push_back(std::move(low));
+        values.push_back(std::move(high));
+      }
+
+      auto attr = DenseElementsAttr::get(newType, values);
+      rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, attr);
+      return success();
+    }
+
+    return rewriter.notifyMatchFailure(op.getLoc(),
+                                       "unhandled constant attribute");
+  }
+};
+
 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();
 
     arith::WideIntEmulationConverter typeConverter(widestIntSupported);
     ConversionTarget target(*ctx);
     target.addDynamicallyLegalOp<func::FuncOp>([&typeConverter](Operation *op) {
       return typeConverter.isLegal(cast<func::FuncOp>(op).getFunctionType());
     });
-    target.addDynamicallyLegalOp<func::CallOp, func::ReturnOp>(
+    target.addDynamicallyLegalOp<
+        // `func.*` ops
+        func::CallOp, func::ReturnOp,
+        // `arith.*` ops
+        arith::ConstantOp>(
         [&typeConverter](Operation *op) { return typeConverter.isLegal(op); });
 
     RewritePatternSet patterns(ctx);
     arith::populateWideIntEmulationPatterns(typeConverter, patterns);
 
     if (failed(applyPartialConversion(op, target, std::move(patterns))))
       signalPassFailure();
   }

 
 void arith::populateWideIntEmulationPatterns(
     WideIntEmulationConverter &typeConverter, RewritePatternSet &patterns) {
   // Populate `func.*` conversion patterns.
   populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                  typeConverter);
   populateCallOpTypeConversionPattern(patterns, typeConverter);
   populateReturnOpTypeConversionPattern(patterns, typeConverter);
+
+  // Populate `arith.*` conversion patterns.
+  patterns.add<ConvertConstant>(typeConverter, patterns.getContext());
 }
 
 } // namespace mlir

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

 // CHECK-LABEL: func @call
 // CHECK-SAME:     ([[ARG:%.+]]: vector<4x2xi32>) -> vector<4x2xi32>
 // CHECK-NEXT:     [[RES:%.+]] = call @identity_vector([[ARG]]) : (vector<4x2xi32>) -> vector<4x2xi32>
 // CHECK-NEXT:     return [[RES]] : vector<4x2xi32>
 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 @constant_scalar
+// CHECK-SAME:     () -> vector<2xi32>
+// CHECK-NEXT:     [[C0:%.+]] = arith.constant dense<0> : vector<2xi32>
+// CHECK-NEXT:     [[C1:%.+]] = arith.constant dense<[0, 1]> : vector<2xi32>
+// CHECK-NEXT:     [[C2:%.+]] = arith.constant dense<[-7, -1]> : vector<2xi32>
+// CHECK-NEXT:     return [[C0]] : vector<2xi32>
+func.func @constant_scalar() -> i64 {
+    %c0 = arith.constant 0 : i64
+    %c1 = arith.constant 4294967296 : i64
+    %c2 = arith.constant -7 : i64
+    return %c0 : i64
+}
+
+// CHECK-LABEL: func @constant_vector
+// CHECK-SAME:     () -> vector<3x2xi32>
+// CHECK-NEXT:     [[C0:%.+]] = arith.constant dense
+// CHECK-SAME{LITERAL}:                             <[[0, 1], [0, 1], [0, 1]]> : vector<3x2xi32>
+// CHECK-NEXT:     [[C1:%.+]] = arith.constant dense
+// CHECK-SAME{LITERAL}:                             <[[0, 0], [1, 0], [-2, -1]]> : vector<3x2xi32>
+// CHECK-NEXT:     return [[C0]] : vector<3x2xi32>
+func.func @constant_vector() -> vector<3xi64> {
+    %c0 = arith.constant dense<4294967296> : vector<3xi64>
+    %c1 = arith.constant dense<[0, 1, -2]> : vector<3xi64>
+    return %c0 : vector<3xi64>
+}