[mlir][memref] Add initial Wide Int Emulation pass and patterns

[kuhar]

Summary

Add a new pass and conversions to emulate wide integer operations over memrefs.
The emulation is implemented on top of the existing pass to emulate wide integer arith ops.

Improve naming in the arith pass to avoid potential name clashes.

[kuhar]

Let me know if you know a better way to organize this

[kuhar]

Clean up

[kuhar]

Clean up

[antiagainst]

Nice, thanks!

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

2

Nit: This is missing the header line like L6?

[kuhar]

Fix copyright header. Rebase.

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

 std::unique_ptr<Pass> createArithBufferizePass();
 
 /// Create a pass to bufferize arith.constant ops.
 std::unique_ptr<Pass> createConstantBufferizePass(uint64_t alignment = 0);
 
 /// Adds patterns to emulate wide Arith 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(WideIntEmulationConverter &typeConverter,
-                                      RewritePatternSet &patterns);
+void populateArithWideIntEmulationPatterns(
+    WideIntEmulationConverter &typeConverter, RewritePatternSet &patterns);
 
 /// Add patterns to expand Arith ceil/floor division ops.
 void populateCeilFloorDivExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Add patterns to expand Arith ops.
 void populateArithExpandOpsPatterns(RewritePatternSet &patterns);
 
 /// Create a pass to legalize Arith ops.

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

let description = [{
     comparisons.
   }];
   let constructor = "mlir::arith::createArithUnsignedWhenEquivalentPass()";
 }
 
 def ArithEmulateWideInt : 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.
+    Emulate arith 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_ARITH_TRANSFORMS_PASSES

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

 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 class AffineDialect;
 class ModuleOp;
 
+namespace arith {
+class WideIntEmulationConverter;
+} // namespace arith
+
 namespace func {
 class FuncDialect;
 } // namespace func
 namespace tensor {
 class TensorDialect;
 } // namespace tensor
 namespace vector {
 class VectorDialect;

 /// terms of shapes of its input operands.
 void populateResolveShapedTypeResultDimsPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns for simplifying extract_strided_metadata(other_op) into
 /// easier to analyze constructs.
 void populateSimplifyExtractStridedMetadataOpPatterns(
     RewritePatternSet &patterns);
 
+/// Appends patterns for emulating wide integer memref operations with ops over
+/// narrower integer types.
+void populateMemRefWideIntEmulationPatterns(
+    arith::WideIntEmulationConverter &typeConverter,
+    RewritePatternSet &patterns);
+
+/// Appends type converions for emulating wide integer memref operations with
+/// ops over narrowe integer types.
+void populateMemRefWideIntEmulationConversions(
+    arith::WideIntEmulationConverter &typeConverter);
+
 /// Transformation to do multi-buffering/array expansion to remove dependencies
 /// on the temporary allocation between consecutive loop iterations.
 /// It returns the new allocation if the original allocation was multi-buffered
 /// and returns failure() otherwise.
 /// Example:
 /// ```
 /// %0 = memref.alloc() : memref<4x128xf32>
 /// scf.for %iv = %c1 to %c1024 step %c3 {

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

let description = [{
     from/to the original memref.
   }];
   let constructor = "mlir::memref::createFoldMemRefAliasOpsPass()";
   let dependentDialects = [
       "AffineDialect", "memref::MemRefDialect", "vector::VectorDialect"
   ];
 }
 
+def MemRefEmulateWideInt : Pass<"memref-emulate-wide-int"> {
+  let summary = "Emulate 2*N-bit integer operations using N-bit operations";
+  let description = [{
+    Emulate memref 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.
+
+    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"];
+}
+
 def NormalizeMemRefs : Pass<"normalize-memrefs", "ModuleOp"> {
   let summary = "Normalize memrefs";
    let description = [{
     This pass transforms memref types with a non-trivial
     [layout map](https://mlir.llvm.org/docs/LangRef/#layout-map) into
     memref types with an identity layout map, e.g. (i, j) -> (i, j). This
     pass is inter-procedural, in the sense that it can modify function
     interfaces and call sites that pass memref types. In order to modify

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

auto opLegalCallback = [&typeConverter](Operation *op) {
       return typeConverter.isLegal(op);
     };
     target.addDynamicallyLegalOp<func::CallOp, func::ReturnOp>(opLegalCallback);
     target
         .addDynamicallyLegalDialect<arith::ArithDialect, vector::VectorDialect>(
             opLegalCallback);
 
     RewritePatternSet patterns(ctx);
-    arith::populateWideIntEmulationPatterns(typeConverter, patterns);
+    arith::populateArithWideIntEmulationPatterns(typeConverter, patterns);
 
     if (failed(applyPartialConversion(op, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
addConversion([this](FunctionType ty) -> Optional<Type> {
     SmallVector<Type> results;
     if (failed(convertTypes(ty.getResults(), results)))
       return None;
 
     return FunctionType::get(ty.getContext(), inputs, results);
   });
 }
 
-void arith::populateWideIntEmulationPatterns(
+void arith::populateArithWideIntEmulationPatterns(
     WideIntEmulationConverter &typeConverter, RewritePatternSet &patterns) {
   // Populate `func.*` conversion patterns.
   populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
                                                                  typeConverter);
   populateCallOpTypeConversionPattern(patterns, typeConverter);
   populateReturnOpTypeConversionPattern(patterns, typeConverter);
 
   // Populate `arith.*` conversion patterns.

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

 add_mlir_dialect_library(MLIRMemRefTransforms
   ComposeSubView.cpp
   ExpandOps.cpp
+  EmulateWideInt.cpp
   FoldMemRefAliasOps.cpp
   MultiBuffer.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
   SimplifyExtractStridedMetadata.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRAffineUtils
   MLIRArithDialect
+  MLIRArithTransforms
   MLIRFuncDialect
   MLIRInferTypeOpInterface
   MLIRLoopLikeInterface
   MLIRMemRefDialect
   MLIRPass
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
 )
 

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

+//
+// 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/IR/Arith.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
+
+#include "mlir/Dialect/Arith/Transforms/WideIntEmulationConverter.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/MathExtras.h"
+#include <cassert>
+
+namespace mlir::memref {
+#define GEN_PASS_DEF_MEMREFEMULATEWIDEINT
+#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
+} // namespace mlir::memref
+
+using namespace mlir;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// ConvertMemRefAlloc
+//===----------------------------------------------------------------------===//
+
+struct ConvertMemRefAlloc final : OpConversionPattern<memref::AllocOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::AllocOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type newTy = getTypeConverter()->convertType(op.getType());
+    if (!newTy)
+      return rewriter.notifyMatchFailure(
+          op->getLoc(),
+          llvm::formatv("failed to convert memref type: {0}", op.getType()));
+
+    rewriter.replaceOpWithNewOp<memref::AllocOp>(
+        op, newTy, adaptor.getDynamicSizes(), adaptor.getSymbolOperands(),
+        adaptor.getAlignmentAttr());
+    return success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// ConvertMemRefLoad
+//===----------------------------------------------------------------------===//
+
+struct ConvertMemRefLoad final : OpConversionPattern<memref::LoadOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::LoadOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type newResTy = getTypeConverter()->convertType(op.getType());
+    if (!newResTy)
+      return rewriter.notifyMatchFailure(
+          op->getLoc(), llvm::formatv("failed to convert memref type: {0}",
+                                      op.getMemRefType()));
+
+    rewriter.replaceOpWithNewOp<memref::LoadOp>(
+        op, newResTy, adaptor.getMemref(), adaptor.getIndices());
+    return success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// ConvertMemRefStore
+//===----------------------------------------------------------------------===//
+
+struct ConvertMemRefStore final : OpConversionPattern<memref::StoreOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::StoreOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type newTy = getTypeConverter()->convertType(op.getMemRefType());
+    if (!newTy)
+      return rewriter.notifyMatchFailure(
+          op->getLoc(), llvm::formatv("failed to convert memref type: {0}",
+                                      op.getMemRefType()));
+
+    rewriter.replaceOpWithNewOp<memref::StoreOp>(
+        op, adaptor.getValue(), adaptor.getMemref(), adaptor.getIndices());
+    return success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// Pass Definition
+//===----------------------------------------------------------------------===//
+
+struct EmulateWideIntPass final
+    : memref::impl::MemRefEmulateWideIntBase<EmulateWideIntPass> {
+  using MemRefEmulateWideIntBase::MemRefEmulateWideIntBase;
+
+  void runOnOperation() override {
+    if (!llvm::isPowerOf2_32(widestIntSupported) || widestIntSupported < 2) {
+      signalPassFailure();
+      return;
+    }
+
+    Operation *op = getOperation();
+    MLIRContext *ctx = op->getContext();
+
+    arith::WideIntEmulationConverter typeConverter(widestIntSupported);
+    memref::populateMemRefWideIntEmulationConversions(typeConverter);
+    ConversionTarget target(*ctx);
+    target.addDynamicallyLegalDialect<
+        arith::ArithDialect, memref::MemRefDialect, vector::VectorDialect>(
+        [&typeConverter](Operation *op) { return typeConverter.isLegal(op); });
+
+    RewritePatternSet patterns(ctx);
+    // Add common pattenrs to support contants, functions, etc.
+    arith::populateArithWideIntEmulationPatterns(typeConverter, patterns);
+
+    memref::populateMemRefWideIntEmulationPatterns(typeConverter, patterns);
+
+    if (failed(applyPartialConversion(op, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
+} // end anonymous namespace
+
+//===----------------------------------------------------------------------===//
+// Public Interface Definition
+//===----------------------------------------------------------------------===//
+
+void memref::populateMemRefWideIntEmulationPatterns(
+    arith::WideIntEmulationConverter &typeConverter,
+    RewritePatternSet &patterns) {
+  // Populate `memref.*` conversion patterns.
+  patterns.add<ConvertMemRefAlloc, ConvertMemRefLoad, ConvertMemRefStore>(
+      typeConverter, patterns.getContext());
+}
+
+void memref::populateMemRefWideIntEmulationConversions(
+    arith::WideIntEmulationConverter &typeConverter) {
+  typeConverter.addConversion(
+      [&typeConverter](MemRefType ty) -> Optional<Type> {
+        auto intTy = ty.getElementType().dyn_cast<IntegerType>();
+        if (!intTy)
+          return ty;
+
+        if (intTy.getIntOrFloatBitWidth() <=
+            typeConverter.getMaxTargetIntBitWidth())
+          return ty;
+
+        Type newElemTy = typeConverter.convertType(intTy);
+        if (!newElemTy)
+          return None;
+
+        return ty.cloneWith(None, newElemTy);
+      });
+}

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

+// RUN: mlir-opt --memref-emulate-wide-int="widest-int-supported=32" %s | FileCheck %s
+
+// Expect no conversions, i32 is supported.
+// CHECK-LABEL: func @memref_i32
+// CHECK:         [[M:%.+]] = memref.alloc() : memref<4xi32, 1>
+// CHECK-NEXT:    [[V:%.+]] = memref.load [[M]][{{%.+}}] : memref<4xi32, 1>
+// CHECK-NEXT:    memref.store {{%.+}}, [[M]][{{%.+}}] : memref<4xi32, 1>
+// CHECK-NEXT:    return
+func.func @memref_i32() {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : i32
+    %m = memref.alloc() : memref<4xi32, 1>
+    %v = memref.load %m[%c0] : memref<4xi32, 1>
+    memref.store %c1, %m[%c0] : memref<4xi32, 1>
+    return
+}
+
+// Expect no conversions, f64 is not an integer type.
+// CHECK-LABEL: func @memref_f32
+// CHECK:         [[M:%.+]] = memref.alloc() : memref<4xf32, 1>
+// CHECK-NEXT:    [[V:%.+]] = memref.load [[M]][{{%.+}}] : memref<4xf32, 1>
+// CHECK-NEXT:    memref.store {{%.+}}, [[M]][{{%.+}}] : memref<4xf32, 1>
+// CHECK-NEXT:    return
+func.func @memref_f32() {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1.0 : f32
+    %m = memref.alloc() : memref<4xf32, 1>
+    %v = memref.load %m[%c0] : memref<4xf32, 1>
+    memref.store %c1, %m[%c0] : memref<4xf32, 1>
+    return
+}
+
+// CHECK-LABEL: func @alloc_load_store_i64
+// CHECK:         [[C1:%.+]] = arith.constant dense<[1, 0]> : vector<2xi32>
+// CHECK-NEXT:    [[M:%.+]]  = memref.alloc() : memref<4xvector<2xi32>, 1>
+// CHECK-NEXT:    [[V:%.+]]  = memref.load [[M]][{{%.+}}] : memref<4xvector<2xi32>, 1>
+// CHECK-NEXT:    memref.store [[C1]], [[M]][{{%.+}}] : memref<4xvector<2xi32>, 1>
+// CHECK-NEXT:    return
+func.func @alloc_load_store_i64() {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : i64
+    %m = memref.alloc() : memref<4xi64, 1>
+    %v = memref.load %m[%c0] : memref<4xi64, 1>
+    memref.store %c1, %m[%c0] : memref<4xi64, 1>
+    return
+}

mlir/test/lib/Dialect/Arith/TestEmulateWideInt.cpp

void runOnOperation() override {
     ConversionTarget target(*ctx);
     target
         .addDynamicallyLegalDialect<arith::ArithDialect, vector::VectorDialect>(
             [&typeConverter](Operation *op) {
               return typeConverter.isLegal(op);
             });
 
     RewritePatternSet patterns(ctx);
-    arith::populateWideIntEmulationPatterns(typeConverter, patterns);
+    arith::populateArithWideIntEmulationPatterns(typeConverter, patterns);
     if (failed(applyPartialConversion(op, target, std::move(patterns))))
       signalPassFailure();
   }
 
   Option<std::string> testFunctionPrefix{
       *this, "function-prefix",
       llvm::cl::desc("Prefix of functions to run the emulation pass on"),
       llvm::cl::init("emulate_")};