[mlir][arith] Support wide integer constant emulation

[Mogball]

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

44

Why not cast to VectorType?

45

70

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

82

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?

94

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

37

Please add a short doc

[kuhar]

Added documentation for the getHalves helper.

[Mogball]

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

37

Can you add code headers?

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

[kuhar]

Added conversion pattern header. Rebased.

[Mogball]

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

28–97

85

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

[kuhar]

Added headers. Rebased.

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

85

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

85

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/ADT/APInt.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 {
 
 bool isOpLegal(Operation *op, arith::WideIntEmulationConverter &converter) {
   if (auto func = dyn_cast<func::FuncOp>(op))
     return converter.isLegal(func.getFunctionType());
 
   return converter.isLegal(op);
 }
 
+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)};
+}
+
+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);
     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>(
+    target.addDynamicallyLegalOp<func::FuncOp, func::CallOp, func::ReturnOp,
+                                 arith::ConstantOp>(
         [&typeConverter](Operation *op) {
           return isOpLegal(op, typeConverter);
         });
 
     RewritePatternSet patterns(ctx);
     arith::populateWideIntEmulationPatterns(typeConverter, patterns);
 
     if (failed(applyPartialConversion(op, target, std::move(patterns))))
if (failed(convertTypes(ty.getResults(), results)))
       return None;
 
     return FunctionType::get(ty.getContext(), inputs, results);
   });
 }
 
 void populateWideIntEmulationPatterns(WideIntEmulationConverter &typeConverter,
                                       RewritePatternSet &patterns) {
+  patterns.add<ConvertConstant>(typeConverter, patterns.getContext());
+
   // Populate `func.*` conversion patterns.
   populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                  typeConverter);
   populateCallOpTypeConversionPattern(patterns, typeConverter);
   populateReturnOpTypeConversionPattern(patterns, typeConverter);
 }
 
 } // namespace mlir::arith

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

 // 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
 }
+
+// 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>
+}