Make arith converter configurable by user

[yzhang93]

Summary

Depends on D151826

[kuhar]

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

89–94

if you expect the converted type to be something specific, pass it to the convert functions: getTypeConverter()->convertType<MemRefType>();. This way the code won't explode if the type converter returns something undesirable.

mlir/test/lib/Dialect/MemRef/TestEmulateNarrowInt.cpp

107

Is this 4 in the flag name intentional?

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

 
 namespace mlir::arith {
 /// Converts integer types that are too narrow and not supported by the target
 /// hardware. Currently, we only handle power-of-two integer types and convert
 /// them to wider integers that are equal or larger than 8 bits.
 class NarrowIntEmulationConverter : public TypeConverter {
 public:
   explicit NarrowIntEmulationConverter(unsigned targetWideInt);
-
-private:
   unsigned targetBitwidth;
 };
 } // namespace mlir::arith
 
 #endif // MLIR_DIALECT_ARITH_NARROW_INT_EMULATION_CONVERTER_H_

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

 // Public Interface Definition
 //===----------------------------------------------------------------------===//
 
 arith::NarrowIntEmulationConverter::NarrowIntEmulationConverter(
     unsigned targetWideInt)
     : targetBitwidth(targetWideInt) {
   assert(llvm::isPowerOf2_32(targetWideInt) &&
          "Only power-of-two integers are supported");
-  assert(targetWideInt >= 8 && "Target integer type too narrow");
 
   // Allow unknown types.
   addConversion([](Type ty) -> std::optional<Type> { return ty; });
 
-  // Scalar case.
-  addConversion([this](IntegerType ty) -> std::optional<Type> {
-    unsigned width = ty.getWidth();
-    if (width >= targetBitwidth)
-      return ty;
-    else
-      return IntegerType::get(ty.getContext(), targetBitwidth);
-
-    return std::nullopt;
-  });
-
-  // Vector case.
-  addConversion([this](VectorType ty) -> std::optional<Type> {
-    auto intTy = dyn_cast<IntegerType>(ty.getElementType());
-    if (!intTy)
-      return ty;
-
-    unsigned width = intTy.getWidth();
-    if (width >= targetBitwidth)
-      return ty;
-    else
-      return VectorType::get(to_vector(ty.getShape()),
-                             IntegerType::get(ty.getContext(), targetBitwidth));
-
-    return std::nullopt;
-  });
-
   // Function case.
   addConversion([this](FunctionType ty) -> std::optional<Type> {
     SmallVector<Type> inputs;
     if (failed(convertTypes(ty.getInputs(), inputs)))
       return std::nullopt;
 
     SmallVector<Type> results;
     if (failed(convertTypes(ty.getResults(), results)))

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

 //===----------------------------------------------------------------------===//
 
 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)
+    Type newTy = getTypeConverter()->convertType(op.getType());
+    if (!newTy)
       return rewriter.notifyMatchFailure(
           op->getLoc(), llvm::formatv("failed to convert memref type: {0}",
                                       op.getMemRefType()));
+    if (op.getMemRefType() == cast<MemRefType>(newTy))
+      return success();
 
     auto loc = op.getLoc();
     Value source = adaptor.getMemref();
     auto sourceType = cast<MemRefType>(source.getType());
+    auto srcElementType = sourceType.getElementType();
     unsigned sourceRank = sourceType.getRank();
+
+    auto oldElementType =
+        cast<MemRefType>(op.getMemref().getType()).getElementType();
+    int srcBits = oldElementType.getIntOrFloatBitWidth();
+    int dstBits = srcElementType.getIntOrFloatBitWidth();
+    assert(dstBits % srcBits == 0);
+
+    // The emulation only works on 1D memref types. To make this work on N-D
+    // memref, we need to linearize the offset.
+    // Specifically, %0 = memref.load %0[%v0][%v1] :
+    // memref<?x?xi4, strided<[?, ?], offset = ?>> can be replaced with
+    // %b, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %0
+    // %linearized_offset = %v0 * %stride#0 + %scaled_v1 * %stride#1
+    // where %scaled_v1 = v1 / targetBits * sourceBits
+    // %linearized_size = %size0 * %size1
+    // %linearized = memref.reinterpret_cast %b, offset = [%offset], sizes =
+    // [%linearized_size], strides = [%stride#1] %load = memref.load
+    // %linearized[%linearized_offset] : memref<?xi4, strided<?, offset = ?>>
     auto stridedMetadata =
         rewriter.create<memref::ExtractStridedMetadataOp>(loc, source);
     auto baseBuffer = stridedMetadata.getBaseBuffer();
     auto baseSizes = stridedMetadata.getSizes();
     auto baseStrides = stridedMetadata.getStrides();
     auto baseOffset = stridedMetadata.getOffset();
 
-    int srcBits = cast<MemRefType>(op.getMemref().getType())
-                      .getElementType()
-                      .getIntOrFloatBitWidth();
-    int dstBits = sourceType.getElementType().getIntOrFloatBitWidth();
-    assert(dstBits % srcBits == 0);
-
     SmallVector<Value> indices = adaptor.getIndices();
     assert(indices.size() == baseStrides.size());
     assert(indices.size() == sourceRank);
 
     // Only the last index is modified to load the bits needed.
     IndexType targetType = rewriter.getIndexType();
     IntegerAttr attr = rewriter.getIndexAttr(dstBits / srcBits);
     auto scaler = rewriter.create<arith::ConstantOp>(loc, targetType, attr);
     indices.back() =
         rewriter.create<arith::DivUIOp>(loc, indices.back(), scaler);
 
     SmallVector<Value> adjustOffsets;
     for (unsigned i = 0; i < sourceRank; ++i) {
       adjustOffsets.push_back(
           rewriter.create<arith::MulIOp>(loc, indices[i], baseStrides[i]));
     }
 
     // Linearize offset and sizes.
-    Value linearized_offset = adjustOffsets[0];
-    Value linearized_size = baseSizes[0];
+    Value linearizedOffset = adjustOffsets[0];
+    Value linearizedSize = baseSizes[0];
     for (unsigned i = 1; i < sourceRank; ++i) {
-      linearized_offset = rewriter.create<arith::AddIOp>(loc, linearized_offset,
-                                                         adjustOffsets[i]);
-      linearized_size =
-          rewriter.create<arith::MulIOp>(loc, linearized_size, baseSizes[i]);
+      linearizedOffset = rewriter.create<arith::AddIOp>(loc, linearizedOffset,
+      linearizedSize =
+          rewriter.create<arith::MulIOp>(loc, linearizedSize, baseSizes[i]);
     }
 
     // Flatten n-D MemRef to 1-D MemRef.
-    StridedLayoutAttr layoutAttr = StridedLayoutAttr::get(
+    auto layoutAttr = StridedLayoutAttr::get(
         sourceType.getContext(), ShapedType::kDynamic, {ShapedType::kDynamic});
     int64_t staticShape = sourceType.hasStaticShape()
                               ? sourceType.getNumElements()
                               : ShapedType::kDynamic;
-    auto flattenMemrefType =
-        MemRefType::get(staticShape, sourceType.getElementType(), layoutAttr,
-                        sourceType.getMemorySpace());
+    auto flattenMemrefType = MemRefType::get(
+        staticShape, srcElementType, layoutAttr, sourceType.getMemorySpace());
 
     auto reinterpret = rewriter.create<memref::ReinterpretCastOp>(
-        loc, flattenMemrefType, baseBuffer, baseOffset, linearized_size,
+        loc, flattenMemrefType, baseBuffer, baseOffset, linearizedSize,
         baseStrides.back());
 
-    auto newLoad =
-        rewriter.create<memref::LoadOp>(loc, newResTy, reinterpret.getResult(),
-                                        linearized_offset, op.getNontemporal());
+    auto newLoad = rewriter.create<memref::LoadOp>(
+        loc, srcElementType, reinterpret.getResult(), linearizedOffset,
+        op.getNontemporal());
 
     // Get the offset and shift the bits to the rightmost.
     auto lastIdx = rewriter.create<arith::IndexCastUIOp>(
-        loc, sourceType.getElementType(), adaptor.getIndices().back());
+        loc, srcElementType, adaptor.getIndices().back());
     Value BitwidthOffset =
         getOffsetForBitwidth(loc, lastIdx, srcBits, dstBits, rewriter);
     auto bitsLoad =
         rewriter.create<arith::ShRSIOp>(loc, newLoad, BitwidthOffset);
 
-    // Apply the mask to extract corresponding bits.
-    auto mask = rewriter.create<arith::ConstantOp>(
-        loc, sourceType.getElementType(),
-        rewriter.getIntegerAttr(sourceType.getElementType(),
-                                (1 << srcBits) - 1));
-    auto result = rewriter.create<arith::AndIOp>(loc, bitsLoad, mask);
-
+    // Get the low bits by truncating the result.
+    auto result =
+        rewriter.create<arith::TruncIOp>(loc, oldElementType, bitsLoad);
     rewriter.replaceOp(op, result.getResult());
+
     return success();
   }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // Public Interface Definition
 //===----------------------------------------------------------------------===//

 void memref::populateMemRefNarrowIntEmulationConversions(
     arith::NarrowIntEmulationConverter &typeConverter) {
   typeConverter.addConversion(
       [&typeConverter](MemRefType ty) -> std::optional<Type> {
         auto intTy = dyn_cast<IntegerType>(ty.getElementType());
         if (!intTy)
           return ty;
 
-        Type newElemTy = typeConverter.convertType(intTy);
+        unsigned width = intTy.getWidth();
+        if (width >= typeConverter.targetBitwidth)
+          return ty;
+        else {
+          Type newElemTy =
+              IntegerType::get(ty.getContext(), typeConverter.targetBitwidth,
+                               intTy.getSignedness());
           if (!newElemTy)
             return std::nullopt;
-
           return ty.cloneWith(std::nullopt, newElemTy);
+        }
       });
 }

mlir/test/Dialect/Arith/emulate-narrow-int.mlir

-// RUN: mlir-opt --test-emulate-narrow-int="target-wide-int=8" %s | FileCheck %s
+// RUN: mlir-opt --test-emulate-narrow-int="int4-arith-bitwidth=8" %s | FileCheck %s
 
 // Expect no conversions, f32 is not an integer type.
 // CHECK-LABEL: func @identity_f32
 // CHECK-SAME:    ([[ARG:%.+]]: f32) -> f32
 // CHECK-NEXT:    return [[ARG]] : f32
 func.func @identity_f32(%a : f32) -> f32 {
     return %a : f32
 }

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

-// RUN: mlir-opt --test-emulate-narrow-int="target-wide-int=8" %s | FileCheck %s
+// RUN: mlir-opt --test-emulate-narrow-int="int4-arith-bitwidth=4 memref-target-bits=8" %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() {

 // CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[OFFSET]]], sizes: [%[[SIZES]]], strides: [%[[STRIDES]]] : memref<i8> to memref<4xi8, strided<[?], offset: ?>>
 // CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<4xi8, strided<[?], offset: ?>>
 // CHECK-NEXT:    %[[I:.*]] = arith.index_castui %[[ARG]] : index to i8
 // CHECK-NEXT:    %[[C2:.*]] = arith.constant 2 : i8
 // CHECK-NEXT:    %[[C4:.*]] = arith.constant 4 : i8
 // CHECK-NEXT:    %[[REM:.*]] = arith.remui %[[I]], %[[C2]] : i8
 // CHECK-NEXT:    %[[STEP:.*]] = arith.muli %[[REM]], %[[C4]] : i8
 // CHECK-NEXT:    %[[SHIFT:.*]] = arith.shrsi %[[LOAD]], %[[STEP]] : i8
-// CHECK-NEXT:    %[[MASK:.*]] = arith.constant 15 : i8
-// CHECK-NEXT:    %[[RES:.*]] = arith.andi %[[SHIFT]], %[[MASK]] : i8
+// CHECK-NEXT:    %[[RES:.*]] = arith.trunci %[[SHIFT]] : i8 to i4
 // CHECK-NEXT:    return
 func.func @memref_load_i4(%arg0: index) {
     %0 = memref.alloc() : memref<4xi4>
     %1 = memref.load %0[%arg0] : memref<4xi4>
     return
 }
 
 

 // CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[OFFSET]]], sizes: [%[[LSIZE]]], strides: [%[[STRIDES]]#1] : memref<i8> to memref<16xi8, strided<[?], offset: ?>>
 // CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<16xi8, strided<[?], offset: ?>>
 // CHECK-NEXT:    %[[I:.*]] = arith.index_castui %[[ARG1]] : index to i8
 // CHECK-NEXT:    %[[C2:.*]] = arith.constant 2 : i8
 // CHECK-NEXT:    %[[C4:.*]] = arith.constant 4 : i8
 // CHECK-NEXT:    %[[REM:.*]] = arith.remui %[[I]], %[[C2]] : i8
 // CHECK-NEXT:    %[[STEP:.*]] = arith.muli %[[REM]], %[[C4]] : i8
 // CHECK-NEXT:    %[[SHIFT:.*]] = arith.shrsi %[[LOAD]], %[[STEP]] : i8
-// CHECK-NEXT:    %[[MASK:.*]] = arith.constant 15 : i8
-// CHECK-NEXT:    %[[RES:.*]] = arith.andi %[[SHIFT]], %[[MASK]] : i8
+// CHECK-NEXT:    %[[RES:.*]] = arith.trunci %[[SHIFT]] : i8 to i4
 // CHECK-NEXT:    return
 func.func @memref_load_i4_rank2(%arg0: index, %arg1: index) {
     %0 = memref.alloc() : memref<4x4xi4>
     %1 = memref.load %0[%arg0,%arg1] : memref<4x4xi4>
     return
-}
+}
 

mlir/test/lib/Dialect/MemRef/TestEmulateNarrowInt.cpp

 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
+
 struct TestEmulateNarrowIntPass
     : public PassWrapper<TestEmulateNarrowIntPass,
                          OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestEmulateNarrowIntPass)
 
   TestEmulateNarrowIntPass() = default;
   TestEmulateNarrowIntPass(const TestEmulateNarrowIntPass &pass)
       : PassWrapper(pass) {}
if (!llvm::isPowerOf2_32(targetWideInt) || targetWideInt < 8) {
       signalPassFailure();
       return;
     }
 
     Operation *op = getOperation();
     MLIRContext *ctx = op->getContext();
 
     arith::NarrowIntEmulationConverter typeConverter(targetWideInt);
+
+    // Convert scalar type.
+    typeConverter.addConversion([this](IntegerType ty) -> std::optional<Type> {
+      unsigned width = ty.getWidth();
+      if (width >= arithBitwidth)
+        return ty;
+      else
+        return IntegerType::get(ty.getContext(), arithBitwidth);
+
+      return std::nullopt;
+    });
+
+    // Convert vector type.
+    typeConverter.addConversion([this](VectorType ty) -> std::optional<Type> {
+      auto intTy = dyn_cast<IntegerType>(ty.getElementType());
+      if (!intTy)
+        return ty;
+
+      unsigned width = intTy.getWidth();
+      if (width >= arithBitwidth)
+        return ty;
+      else
+        return VectorType::get(
+            to_vector(ty.getShape()),
+            IntegerType::get(ty.getContext(), arithBitwidth));
+
+      return std::nullopt;
+    });
+
     memref::populateMemRefNarrowIntEmulationConversions(typeConverter);
     ConversionTarget target(*ctx);
     target.addDynamicallyLegalOp<func::FuncOp>([&typeConverter](Operation *op) {
       return typeConverter.isLegal(cast<func::FuncOp>(op).getFunctionType());
     });
     auto opLegalCallback = [&typeConverter](Operation *op) {
       return typeConverter.isLegal(op);
     };
     target.addDynamicallyLegalOp<func::CallOp, func::ReturnOp>(opLegalCallback);
     target.addDynamicallyLegalDialect<
         arith::ArithDialect, vector::VectorDialect, memref::MemRefDialect>(
         [&typeConverter](Operation *op) { return typeConverter.isLegal(op); });
 
     RewritePatternSet patterns(ctx);
 
-    // Add common pattenrs to support contants, functions, etc.
     arith::populateArithNarrowIntEmulationPatterns(typeConverter, patterns);
     memref::populateMemRefNarrowIntEmulationPatterns(typeConverter, patterns);
 
     if (failed(applyPartialConversion(op, target, std::move(patterns))))
       signalPassFailure();
   }
 
-  Option<unsigned> targetWideInt{*this, "target-wide-int",
+  Option<unsigned> targetWideInt{*this, "memref-target-bits",
                                  llvm::cl::desc("Target integer bit width"),
                                  llvm::cl::init(8)};
+
+  Option<unsigned> arithBitwidth{*this, "int4-arith-bitwidth",
+                                 llvm::cl::desc("Target integer bit width"),
+                                 llvm::cl::init(4)};
 };
 } // namespace
 
 namespace mlir::test {
 void registerTestEmulateNarrowIntPass() {
   PassRegistration<TestEmulateNarrowIntPass>();
 }
 } // namespace mlir::test