[Mlir-commits] [mlir] [MLIR][MemRef] Emulate `ExtractStriedMetadata` narrow types (PR #196967)
Alan Li
llvmlistbot at llvm.org
Thu May 14 04:58:46 PDT 2026
https://github.com/lialan updated https://github.com/llvm/llvm-project/pull/196967
>From 4ebcf1d93d519a7490a817938510b38301c2b468 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Mon, 11 May 2026 07:40:25 -0700
Subject: [PATCH 1/4] [MLIR][MemRef] Emulate more narrow types
* Add `ConvertExtractStridedMetadata` when the source memref has
a sub-byte element type. It calls `extract_strided_metadata` on the
converted i8 source, scales the runtime offset back to emulated-element
units, and returns the emulated-element strides/sizes as constants.
* Added tests to cover the end-to-end memref narrow-type story:
`cf.br` carrying a sub-byte memref block-arg, and the `vector.load` path through
`extract_strided_metadata` on a sub-byte source with a dynamic offset.
---
.../MemRef/Transforms/EmulateNarrowType.cpp | 173 +++++++++++++-----
.../MemRef/emulate-narrow-type-cf.mlir | 69 +++++++
...-narrow-type-extract-strided-metadata.mlir | 151 +++++++++++++++
.../Dialect/MemRef/TestEmulateNarrowType.cpp | 40 +++-
4 files changed, 381 insertions(+), 52 deletions(-)
create mode 100644 mlir/test/Dialect/MemRef/emulate-narrow-type-cf.mlir
create mode 100644 mlir/test/Dialect/MemRef/emulate-narrow-type-extract-strided-metadata.mlir
diff --git a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
index a11e14faa5475..9b7c7f96abda7 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
@@ -718,6 +718,80 @@ struct ConvertMemRefExpandShape final
return success();
}
};
+
+//===----------------------------------------------------------------------===//
+// ConvertExtractStridedMetadata
+//===----------------------------------------------------------------------===//
+
+/// Lowers `memref.extract_strided_metadata` on a sub-byte source by
+/// delegating to the i8 container produced by the narrow-type converter and
+/// scaling the runtime offset from i8 units back to emulated-element units.
+struct ConvertExtractStridedMetadata final
+ : OpConversionPattern<memref::ExtractStridedMetadataOp> {
+ ConvertExtractStridedMetadata(
+ const arith::NarrowTypeEmulationConverter &converter, MLIRContext *ctx,
+ PatternBenefit benefit = 1)
+ : OpConversionPattern(converter, ctx, benefit),
+ loadStoreBitwidth(converter.getLoadStoreBitwidth()) {}
+
+ LogicalResult
+ matchAndRewrite(memref::ExtractStridedMetadataOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto srcType = dyn_cast<MemRefType>(op.getSource().getType());
+ if (!srcType)
+ return rewriter.notifyMatchFailure(op, "source is not a MemRefType");
+
+ Type elemTy = srcType.getElementType();
+ if (!elemTy.isIntOrFloat() ||
+ elemTy.getIntOrFloatBitWidth() >= loadStoreBitwidth)
+ return rewriter.notifyMatchFailure(op, "source element not sub-byte");
+
+ unsigned scale = loadStoreBitwidth / elemTy.getIntOrFloatBitWidth();
+
+ Location loc = op.getLoc();
+ auto i8Meta = memref::ExtractStridedMetadataOp::create(rewriter, loc,
+ adaptor.getSource());
+
+ int64_t srcStaticOffset;
+ SmallVector<int64_t> srcStaticStrides;
+ if (failed(srcType.getStridesAndOffset(srcStaticStrides, srcStaticOffset)))
+ return rewriter.notifyMatchFailure(op, "failed to get strides from type");
+
+ Value emulatedOffset;
+ if (srcStaticOffset == ShapedType::kDynamic) {
+ Value scaleCst = arith::ConstantIndexOp::create(rewriter, loc, scale);
+ emulatedOffset =
+ arith::MulIOp::create(rewriter, loc, i8Meta.getOffset(), scaleCst);
+ } else {
+ // The static offset on the original sub-byte type is already in
+ // emulated-element units; no scaling needed.
+ emulatedOffset =
+ arith::ConstantIndexOp::create(rewriter, loc, srcStaticOffset);
+ }
+
+ SmallVector<Value> emulatedSizes;
+ for (int64_t dim : srcType.getShape())
+ emulatedSizes.push_back(
+ arith::ConstantIndexOp::create(rewriter, loc, dim));
+
+ SmallVector<Value> emulatedStrides;
+ for (int64_t stride : srcStaticStrides) {
+ if (stride == ShapedType::kDynamic)
+ return rewriter.notifyMatchFailure(op, "dynamic stride not supported");
+ emulatedStrides.push_back(
+ arith::ConstantIndexOp::create(rewriter, loc, stride));
+ }
+
+ SmallVector<Value> results = {i8Meta.getBaseBuffer(), emulatedOffset};
+ results.append(emulatedSizes);
+ results.append(emulatedStrides);
+ rewriter.replaceOp(op, results);
+ return success();
+ }
+
+private:
+ unsigned loadStoreBitwidth;
+};
} // end anonymous namespace
//===----------------------------------------------------------------------===//
@@ -740,6 +814,8 @@ void memref::populateMemRefNarrowTypeEmulationPatterns(
typeConverter, patterns.getContext(), assumeAligned);
patterns.insert<ConvertMemrefStore>(typeConverter, patterns.getContext(),
disableAtomicRMW);
+ patterns.insert<ConvertExtractStridedMetadata>(typeConverter,
+ patterns.getContext());
memref::populateResolveExtractStridedMetadataPatterns(patterns);
}
@@ -763,53 +839,52 @@ static SmallVector<int64_t> getLinearizedShape(MemRefType ty, int srcBits,
void memref::populateMemRefNarrowTypeEmulationConversions(
arith::NarrowTypeEmulationConverter &typeConverter) {
- typeConverter.addConversion(
- [&typeConverter](MemRefType ty) -> std::optional<Type> {
- Type elementType = ty.getElementType();
- if (!elementType.isIntOrFloat())
- return ty;
-
- unsigned width = elementType.getIntOrFloatBitWidth();
- unsigned loadStoreWidth = typeConverter.getLoadStoreBitwidth();
- if (width >= loadStoreWidth)
- return ty;
-
- // Currently only handle innermost stride being 1, checking
- SmallVector<int64_t> strides;
- int64_t offset;
- if (failed(ty.getStridesAndOffset(strides, offset)))
- return nullptr;
- if (!strides.empty() && strides.back() != 1)
- return nullptr;
-
- auto newElemTy = IntegerType::get(
- ty.getContext(), loadStoreWidth,
- elementType.isInteger()
- ? cast<IntegerType>(elementType).getSignedness()
- : IntegerType::SignednessSemantics::Signless);
- if (!newElemTy)
- return nullptr;
-
- StridedLayoutAttr layoutAttr;
- // If the offset is 0, we do not need a strided layout as the stride is
- // 1, so we only use the strided layout if the offset is not 0.
- if (offset != 0) {
- if (offset == ShapedType::kDynamic) {
- layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
- ArrayRef<int64_t>{1});
- } else {
- // Check if the number of bytes are a multiple of the loadStoreWidth
- // and if so, divide it by the loadStoreWidth to get the offset.
- if ((offset * width) % loadStoreWidth != 0)
- return std::nullopt;
- offset = (offset * width) / loadStoreWidth;
-
- layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
- ArrayRef<int64_t>{1});
- }
- }
-
- return MemRefType::get(getLinearizedShape(ty, width, loadStoreWidth),
- newElemTy, layoutAttr, ty.getMemorySpace());
- });
+ typeConverter.addConversion([&typeConverter](
+ MemRefType ty) -> std::optional<Type> {
+ Type elementType = ty.getElementType();
+ if (!elementType.isIntOrFloat())
+ return ty;
+
+ unsigned width = elementType.getIntOrFloatBitWidth();
+ unsigned loadStoreWidth = typeConverter.getLoadStoreBitwidth();
+ if (width >= loadStoreWidth)
+ return ty;
+
+ // Currently only handle innermost stride being 1, checking
+ SmallVector<int64_t> strides;
+ int64_t offset;
+ if (failed(ty.getStridesAndOffset(strides, offset)))
+ return nullptr;
+ if (!strides.empty() && strides.back() != 1)
+ return nullptr;
+
+ auto newElemTy = IntegerType::get(
+ ty.getContext(), loadStoreWidth,
+ elementType.isInteger() ? cast<IntegerType>(elementType).getSignedness()
+ : IntegerType::SignednessSemantics::Signless);
+ if (!newElemTy)
+ return nullptr;
+
+ StridedLayoutAttr layoutAttr;
+ // If the offset is 0, we do not need a strided layout as the stride is
+ // 1, so we only use the strided layout if the offset is not 0.
+ if (offset != 0) {
+ if (offset == ShapedType::kDynamic) {
+ layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
+ ArrayRef<int64_t>{1});
+ } else {
+ // Check if the number of bytes are a multiple of the loadStoreWidth
+ // and if so, divide it by the loadStoreWidth to get the offset.
+ if ((offset * width) % loadStoreWidth != 0)
+ return std::nullopt;
+ offset = (offset * width) / loadStoreWidth;
+
+ layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
+ ArrayRef<int64_t>{1});
+ }
+ }
+
+ return MemRefType::get(getLinearizedShape(ty, width, loadStoreWidth),
+ newElemTy, layoutAttr, ty.getMemorySpace());
+ });
}
diff --git a/mlir/test/Dialect/MemRef/emulate-narrow-type-cf.mlir b/mlir/test/Dialect/MemRef/emulate-narrow-type-cf.mlir
new file mode 100644
index 0000000000000..a5cebb0110048
--- /dev/null
+++ b/mlir/test/Dialect/MemRef/emulate-narrow-type-cf.mlir
@@ -0,0 +1,69 @@
+// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8 enable-cf-conversion=true" --cse --verify-diagnostics --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8 enable-cf-conversion=true" --split-input-file %s | FileCheck %s --check-prefix=CHECK-NOCSE
+
+// Sub-byte memref type carried through cf.br block args. The
+// BranchOpInterface type-conversion pattern must rewrite both the cf.br
+// operand type and the successor block-arg type to the i8 container, so the
+// downstream uses in the successor block see an i8 source.
+
+// CHECK-LABEL: func.func @cf_br_block_arg_narrow_type
+// CHECK-SAME: %[[ARG:[A-Za-z0-9_]+]]: memref<{{[0-9]+}}xi8>
+// CHECK: cf.br ^[[BB1:.+]](%[[ARG]] : memref<{{[0-9]+}}xi8>)
+// CHECK: ^[[BB1]](%[[BARG:[A-Za-z0-9_]+]]: memref<{{[0-9]+}}xi8>):
+// CHECK: return %[[BARG]]
+// CHECK-NOT: memref<{{[0-9]+}}xi4>
+func.func @cf_br_block_arg_narrow_type(%arg: memref<8xi4>) -> memref<8xi4> {
+ cf.br ^bb1(%arg : memref<8xi4>)
+^bb1(%a: memref<8xi4>):
+ return %a : memref<8xi4>
+}
+
+// -----
+
+// Sub-byte memref with dynamic offset carried through cf.br block-arg,
+// then loaded via vector.load. After FunctionOpInterfaceAllBlocksSignatureConversion
+// converts the block-arg to i8, an unrealized_conversion_cast is inserted.
+// ConvertVectorLoad calls extract_strided_metadata on the original sub-byte
+// source (op.getBase()), which is illegal without ConvertExtractStridedMetadata.
+
+// CHECK-LABEL: func.func @cf_br_block_arg_vector_load_i4
+// CHECK: vector.load {{.*}} : memref<{{[0-9]+}}xi8, strided<[1], offset: ?>>, vector<{{[0-9]+}}xi8>
+// CHECK-NOT: memref<{{[0-9]+}}xi4>
+// CHECK-NOCSE-LABEL: func.func @cf_br_block_arg_vector_load_i4
+// CHECK-NOCSE-NOT: memref<{{[0-9]+}}xi4>
+func.func @cf_br_block_arg_vector_load_i4(%arg: memref<8xi4, strided<[1], offset: ?>>) -> vector<8xi4> {
+ cf.br ^bb1(%arg : memref<8xi4, strided<[1], offset: ?>>)
+^bb1(%a: memref<8xi4, strided<[1], offset: ?>>):
+ %c0 = arith.constant 0 : index
+ %v = vector.load %a[%c0] : memref<8xi4, strided<[1], offset: ?>>, vector<8xi4>
+ return %v : vector<8xi4>
+}
+
+// -----
+
+// Sub-byte memref with static non-zero offset carried through cf.br block-arg,
+// then loaded via vector.load. The static offset (4 i4-elements) is converted
+// to 2 i8-elements in the container type. ConvertExtractStridedMetadata must
+// NOT over-scale the static offset: srcStaticOffset is already in
+// emulated-element (i4) units, so it should be passed through unchanged rather
+// than multiplied by containerBits/emulatedBits (which would give 8, wrong).
+//
+// The CHECK-NOCSE check verifies that the emulated offset constant is 4 (in i4
+// units, as returned by ConvertExtractStridedMetadata) and not 8 (the
+// over-scaled value that would appear with the bug). After CSE the constant
+// is dead and the distinction disappears, so this check requires the no-CSE
+// run line.
+
+// CHECK-LABEL: func.func @cf_br_block_arg_vector_load_i4_static_offset
+// CHECK: vector.load {{.*}} : memref<{{[0-9]+}}xi8, strided<[1], offset: 2>>, vector<{{[0-9]+}}xi8>
+// CHECK-NOT: memref<{{[0-9]+}}xi4>
+// CHECK-NOCSE-LABEL: func.func @cf_br_block_arg_vector_load_i4_static_offset
+// CHECK-NOCSE-NOT: memref<{{[0-9]+}}xi4>
+// CHECK-NOCSE: %c4 = arith.constant 4 : index
+func.func @cf_br_block_arg_vector_load_i4_static_offset(%arg: memref<8xi4, strided<[1], offset: 4>>) -> vector<4xi4> {
+ cf.br ^bb1(%arg : memref<8xi4, strided<[1], offset: 4>>)
+^bb1(%a: memref<8xi4, strided<[1], offset: 4>>):
+ %c0 = arith.constant 0 : index
+ %v = vector.load %a[%c0] : memref<8xi4, strided<[1], offset: 4>>, vector<4xi4>
+ return %v : vector<4xi4>
+}
diff --git a/mlir/test/Dialect/MemRef/emulate-narrow-type-extract-strided-metadata.mlir b/mlir/test/Dialect/MemRef/emulate-narrow-type-extract-strided-metadata.mlir
new file mode 100644
index 0000000000000..fa79f48f82114
--- /dev/null
+++ b/mlir/test/Dialect/MemRef/emulate-narrow-type-extract-strided-metadata.mlir
@@ -0,0 +1,151 @@
+// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8" --cse --verify-diagnostics --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8" --verify-diagnostics --split-input-file %s | FileCheck %s --check-prefix=CHECK-NOCSE
+
+// Tests for `ConvertExtractStridedMetadata` on sub-byte memrefs. The pattern
+// rewrites `memref.extract_strided_metadata` on the original sub-byte source
+// to an `extract_strided_metadata` on the i8 container produced by the
+// type converter, scaling the runtime offset from i8 units back to
+// emulated-element units and returning the original sub-byte sizes/strides
+// as constants.
+
+// -----
+
+// Static zero offset, 1-D i4: emulated offset is constant 0 in i4 units; the
+// size/stride constants match the original sub-byte shape/strides (8 / 1),
+// not the i8 container shape/strides (4 / 1).
+
+// CHECK-LABEL: func.func @extract_strided_metadata_i4
+// CHECK-SAME: %{{.+}}: memref<4xi8>
+// CHECK-DAG: %[[OFF:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[SZ:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[ST:.+]] = arith.constant 1 : index
+// CHECK: return %[[OFF]], %[[SZ]], %[[ST]]
+func.func @extract_strided_metadata_i4(%arg: memref<8xi4>)
+ -> (index, index, index) {
+ %base, %offset, %size, %stride = memref.extract_strided_metadata %arg :
+ memref<8xi4> -> memref<i4>, index, index, index
+ return %offset, %size, %stride : index, index, index
+}
+
+// -----
+
+// Static non-zero offset, 1-D i4: srcStaticOffset on the original sub-byte
+// type is already in emulated-element units, so it must be returned
+// unchanged (8 here). The i8 container offset is 4 (8 i4-elements = 4
+// i8-elements), but the pattern must NOT report 4 nor the over-scaled value
+// 16 (= 8 * 2). Shape 6 keeps every constant (offset, size, stride)
+// distinct so each can be matched independently.
+
+// CHECK-LABEL: func.func @extract_strided_metadata_i4_static_offset
+// CHECK-SAME: %{{.+}}: memref<3xi8, strided<[1], offset: 4>>
+// CHECK-DAG: %[[OFF:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[SZ:.+]] = arith.constant 6 : index
+// CHECK-DAG: %[[ST:.+]] = arith.constant 1 : index
+// CHECK: return %[[OFF]], %[[SZ]], %[[ST]]
+func.func @extract_strided_metadata_i4_static_offset(
+ %arg: memref<6xi4, strided<[1], offset: 8>>) -> (index, index, index) {
+ %base, %offset, %size, %stride = memref.extract_strided_metadata %arg :
+ memref<6xi4, strided<[1], offset: 8>> -> memref<i4>, index, index, index
+ return %offset, %size, %stride : index, index, index
+}
+
+// -----
+
+// Dynamic offset, 1-D i4: the runtime i8 offset returned by the inner
+// `extract_strided_metadata` is multiplied by the scale factor
+// `loadStoreBitwidth / elementBitwidth` (= 2 for i4 @ i8 container) to
+// convert it back to i4-element units.
+
+// CHECK-LABEL: func.func @extract_strided_metadata_i4_dynamic_offset
+// CHECK-SAME: %[[ARG:.+]]: memref<4xi8, strided<[1], offset: ?>>
+// CHECK-DAG: %[[SZ:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[ST:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[SCALE:.+]] = arith.constant 2 : index
+// CHECK-DAG: %{{.+}}, %[[I8OFF:.+]], %{{.+}}, %{{.+}} = memref.extract_strided_metadata %[[ARG]]
+// CHECK-DAG: %[[OFF:.+]] = arith.muli %[[I8OFF]], %[[SCALE]]
+// CHECK: return %[[OFF]], %[[SZ]], %[[ST]]
+func.func @extract_strided_metadata_i4_dynamic_offset(
+ %arg: memref<8xi4, strided<[1], offset: ?>>) -> (index, index, index) {
+ %base, %offset, %size, %stride = memref.extract_strided_metadata %arg :
+ memref<8xi4, strided<[1], offset: ?>> -> memref<i4>, index, index, index
+ return %offset, %size, %stride : index, index, index
+}
+
+// -----
+
+// Dynamic offset, 1-D i2: scale factor is `loadStoreBitwidth /
+// elementBitwidth` = 4. Validates that the scale is derived from the actual
+// element width rather than hardcoded for i4.
+
+// CHECK-LABEL: func.func @extract_strided_metadata_i2_dynamic_offset
+// CHECK-SAME: %[[ARG:.+]]: memref<4xi8, strided<[1], offset: ?>>
+// CHECK-DAG: %[[SZ:.+]] = arith.constant 16 : index
+// CHECK-DAG: %[[ST:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[SCALE:.+]] = arith.constant 4 : index
+// CHECK-DAG: %{{.+}}, %[[I8OFF:.+]], %{{.+}}, %{{.+}} = memref.extract_strided_metadata %[[ARG]]
+// CHECK-DAG: %[[OFF:.+]] = arith.muli %[[I8OFF]], %[[SCALE]]
+// CHECK: return %[[OFF]], %[[SZ]], %[[ST]]
+func.func @extract_strided_metadata_i2_dynamic_offset(
+ %arg: memref<16xi2, strided<[1], offset: ?>>) -> (index, index, index) {
+ %base, %offset, %size, %stride = memref.extract_strided_metadata %arg :
+ memref<16xi2, strided<[1], offset: ?>> -> memref<i2>, index, index, index
+ return %offset, %size, %stride : index, index, index
+}
+
+// -----
+
+// Multi-dim sub-byte source: the converted source is linearized to 1-D, but
+// the pattern still returns the original 2-D sizes/strides as constants
+// derived from the original `MemRefType`. Shape <2x3> keeps strides distinct
+// from sizes (strides [3, 1], sizes [2, 3]) so the per-result mapping is
+// observable.
+
+// CHECK-LABEL: func.func @extract_strided_metadata_i4_2d
+// CHECK-SAME: %{{.+}}: memref<3xi8>
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK: return %[[C0]], %[[C2]], %[[C3]], %[[C3]], %[[C1]]
+func.func @extract_strided_metadata_i4_2d(%arg: memref<2x3xi4>)
+ -> (index, index, index, index, index) {
+ %base, %offset, %sz:2, %st:2 = memref.extract_strided_metadata %arg :
+ memref<2x3xi4> -> memref<i4>, index, index, index, index, index
+ return %offset, %sz#0, %sz#1, %st#0, %st#1 :
+ index, index, index, index, index
+}
+
+// -----
+
+// Base buffer of the rewrite is the i8 container's base, not the original
+// sub-byte base. The no-CSE run line keeps the inner
+// `extract_strided_metadata` so its `base_buffer` result is observable.
+
+// CHECK-NOCSE-LABEL: func.func @extract_strided_metadata_i4_base_buffer
+// CHECK-NOCSE-SAME: %[[ARG:.+]]: memref<4xi8>
+// CHECK-NOCSE: %[[BASE:[A-Za-z0-9_]+]], %{{.+}}, %{{.+}}, %{{.+}} = memref.extract_strided_metadata %[[ARG]] : memref<4xi8> -> memref<i8>
+// CHECK-NOCSE: return %[[BASE]]
+func.func @extract_strided_metadata_i4_base_buffer(%arg: memref<8xi4>)
+ -> memref<i4> {
+ %base, %offset, %size, %stride = memref.extract_strided_metadata %arg :
+ memref<8xi4> -> memref<i4>, index, index, index
+ return %base : memref<i4>
+}
+
+// -----
+
+// Non-sub-byte source: pattern must not fire and the op is left for the
+// existing `populateResolveExtractStridedMetadataPatterns` pipeline (which
+// here resolves it on the unchanged i8 source).
+
+// CHECK-LABEL: func.func @extract_strided_metadata_i8_passthrough
+// CHECK-SAME: %[[ARG:.+]]: memref<8xi8>
+// CHECK: %[[BASE:[A-Za-z0-9_]+]], %[[OFF:.+]], %[[SZ:.+]], %[[ST:.+]] = memref.extract_strided_metadata %[[ARG]] : memref<8xi8>
+// CHECK-NOT: arith.muli
+// CHECK: return %[[OFF]], %[[SZ]], %[[ST]]
+func.func @extract_strided_metadata_i8_passthrough(%arg: memref<8xi8>)
+ -> (index, index, index) {
+ %base, %offset, %size, %stride = memref.extract_strided_metadata %arg :
+ memref<8xi8> -> memref<i8>, index, index, index
+ return %offset, %size, %stride : index, index, index
+}
diff --git a/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp b/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
index bec83a8dcbef9..3e87cd955a627 100644
--- a/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
+++ b/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
@@ -11,11 +11,14 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -34,9 +37,9 @@ struct TestEmulateNarrowTypePass
: PassWrapper(pass) {}
void getDependentDialects(DialectRegistry ®istry) const override {
- registry
- .insert<arith::ArithDialect, func::FuncDialect, memref::MemRefDialect,
- vector::VectorDialect, affine::AffineDialect>();
+ registry.insert<arith::ArithDialect, cf::ControlFlowDialect,
+ func::FuncDialect, memref::MemRefDialect,
+ vector::VectorDialect, affine::AffineDialect>();
}
StringRef getArgument() const final { return "test-emulate-narrow-int"; }
StringRef getDescription() const final {
@@ -96,6 +99,19 @@ struct TestEmulateNarrowTypePass
arith::ArithDialect, vector::VectorDialect, memref::MemRefDialect,
affine::AffineDialect>(opLegalCallback);
+ if (enableCFConversion) {
+ target.addDynamicallyLegalDialect<cf::ControlFlowDialect>(
+ [&typeConverter](Operation *op) -> bool {
+ // Only apply legality check to BranchOpInterface ops; other cf ops
+ // (e.g. cf.assert) have no successor-block-arg conversion concern
+ // and should remain legal.
+ if (!isa<BranchOpInterface>(op))
+ return true;
+ return isLegalForBranchOpInterfaceTypeConversionPattern(
+ op, typeConverter);
+ });
+ }
+
RewritePatternSet patterns(ctx);
arith::populateArithNarrowTypeEmulationPatterns(typeConverter, patterns);
@@ -104,6 +120,17 @@ struct TestEmulateNarrowTypePass
vector::populateVectorNarrowTypeEmulationPatterns(
typeConverter, patterns, disableAtomicRMW, assumeAligned);
+ if (enableCFConversion) {
+ populateBranchOpInterfaceTypeConversionPattern(patterns, typeConverter);
+ // Opt the FunctionOpInterface signature-conversion pattern into the
+ // all-blocks mode so non-entry block-arg types are converted in
+ // lockstep with the function signature and branch operands. Use a
+ // higher benefit so it supersedes the default-benefit pattern
+ // populated by `populateArithNarrowTypeEmulationPatterns`.
+ populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(
+ patterns, typeConverter, /*benefit=*/2, /*convertAllBlocks=*/true);
+ }
+
if (failed(applyPartialConversion(op, target, std::move(patterns))))
signalPassFailure();
}
@@ -133,6 +160,13 @@ struct TestEmulateNarrowTypePass
llvm::cl::desc("assume store offsets are aligned to container element "
"boundaries"),
llvm::cl::init(false)};
+
+ Option<bool> enableCFConversion{
+ *this, "enable-cf-conversion",
+ llvm::cl::desc("register populateBranchOpInterfaceTypeConversionPattern "
+ "and mark cf dialect ops dynamically legal based on "
+ "isLegalForBranchOpInterfaceTypeConversionPattern"),
+ llvm::cl::init(false)};
};
struct TestMemRefFlattenAndVectorNarrowTypeEmulationPass
>From a5759d4b55dce3af7661d09fbcb352ad6e5d0cc4 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Mon, 11 May 2026 13:21:01 -0700
Subject: [PATCH 2/4] [MLIR][MemRef] Simplify rank-0 handling in narrow-type
emulation
Make rank-0 support explicit in `populateMemRefNarrowTypeEmulationConversions`.
* Replace the implicit `!strides.empty()` guard with `ty.getRank() > 0`, and
note that rank-0 memrefs trivially satisfy the innermost-unit-stride
requirement.
* When the original offset is non-zero, build the `StridedLayoutAttr` with
the rank-appropriate stride list (empty for rank-0, `{1}` otherwise),
fixing the previous always-`{1}` construction which was malformed for
rank-0 sources.
* Add rank-0 tests for `ConvertExtractStridedMetadata` (static zero offset
and dynamic offset) covering the `{baseBuffer, offset}` result shape.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../MemRef/Transforms/EmulateNarrowType.cpp | 27 ++++++++------
...-narrow-type-extract-strided-metadata.mlir | 36 +++++++++++++++++++
2 files changed, 53 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
index 9b7c7f96abda7..5634803645f84 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
@@ -850,12 +850,13 @@ void memref::populateMemRefNarrowTypeEmulationConversions(
if (width >= loadStoreWidth)
return ty;
- // Currently only handle innermost stride being 1, checking
+ // Only handle innermost stride being 1. Rank-0 memrefs have no strides
+ // and trivially satisfy this.
SmallVector<int64_t> strides;
int64_t offset;
if (failed(ty.getStridesAndOffset(strides, offset)))
return nullptr;
- if (!strides.empty() && strides.back() != 1)
+ if (ty.getRank() > 0 && strides.back() != 1)
return nullptr;
auto newElemTy = IntegerType::get(
@@ -866,21 +867,27 @@ void memref::populateMemRefNarrowTypeEmulationConversions(
return nullptr;
StridedLayoutAttr layoutAttr;
- // If the offset is 0, we do not need a strided layout as the stride is
- // 1, so we only use the strided layout if the offset is not 0.
+ // The default layout (no `layoutAttr`) covers the common case of a
+ // zero-offset, unit-innermost-stride memref. We only build a strided
+ // layout when the original offset is non-zero. Rank-0 memrefs carry no
+ // strides, so the layout uses an empty stride list.
if (offset != 0) {
+ SmallVector<int64_t, 1> resultStrides =
+ ty.getRank() == 0 ? SmallVector<int64_t, 1>{}
+ : SmallVector<int64_t, 1>{1};
if (offset == ShapedType::kDynamic) {
- layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
- ArrayRef<int64_t>{1});
+ layoutAttr =
+ StridedLayoutAttr::get(ty.getContext(), offset, resultStrides);
} else {
- // Check if the number of bytes are a multiple of the loadStoreWidth
- // and if so, divide it by the loadStoreWidth to get the offset.
+ // Scale the static offset from emulated-element units to
+ // load-store-element units. Reject offsets that are not a whole
+ // number of load-store elements.
if ((offset * width) % loadStoreWidth != 0)
return std::nullopt;
offset = (offset * width) / loadStoreWidth;
- layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
- ArrayRef<int64_t>{1});
+ layoutAttr =
+ StridedLayoutAttr::get(ty.getContext(), offset, resultStrides);
}
}
diff --git a/mlir/test/Dialect/MemRef/emulate-narrow-type-extract-strided-metadata.mlir b/mlir/test/Dialect/MemRef/emulate-narrow-type-extract-strided-metadata.mlir
index fa79f48f82114..8defc0724bee1 100644
--- a/mlir/test/Dialect/MemRef/emulate-narrow-type-extract-strided-metadata.mlir
+++ b/mlir/test/Dialect/MemRef/emulate-narrow-type-extract-strided-metadata.mlir
@@ -134,6 +134,42 @@ func.func @extract_strided_metadata_i4_base_buffer(%arg: memref<8xi4>)
// -----
+// Rank-0 sub-byte source with zero offset: the converted source is a rank-0
+// i8 memref, and the pattern returns only `{baseBuffer, emulatedOffset}` with
+// no size/stride results (matching `extract_strided_metadata`'s rank-0
+// signature). The emulated offset is the static zero constant.
+
+// CHECK-LABEL: func.func @extract_strided_metadata_i4_rank0
+// CHECK-SAME: %{{.+}}: memref<i8>
+// CHECK-DAG: %[[OFF:.+]] = arith.constant 0 : index
+// CHECK: return %[[OFF]]
+func.func @extract_strided_metadata_i4_rank0(%arg: memref<i4>) -> index {
+ %base, %offset = memref.extract_strided_metadata %arg :
+ memref<i4> -> memref<i4>, index
+ return %offset : index
+}
+
+// -----
+
+// Rank-0 sub-byte source with a dynamic offset: the converter produces a
+// rank-0 i8 memref with a `strided<[], offset: ?>` layout (empty strides),
+// and the pattern scales the runtime i8 offset back to i4 units via muli.
+
+// CHECK-LABEL: func.func @extract_strided_metadata_i4_rank0_dynamic_offset
+// CHECK-SAME: %[[ARG:.+]]: memref<i8, strided<[], offset: ?>>
+// CHECK-DAG: %[[SCALE:.+]] = arith.constant 2 : index
+// CHECK-DAG: %{{.+}}, %[[I8OFF:.+]] = memref.extract_strided_metadata %[[ARG]]
+// CHECK-DAG: %[[OFF:.+]] = arith.muli %[[I8OFF]], %[[SCALE]]
+// CHECK: return %[[OFF]]
+func.func @extract_strided_metadata_i4_rank0_dynamic_offset(
+ %arg: memref<i4, strided<[], offset: ?>>) -> index {
+ %base, %offset = memref.extract_strided_metadata %arg :
+ memref<i4, strided<[], offset: ?>> -> memref<i4>, index
+ return %offset : index
+}
+
+// -----
+
// Non-sub-byte source: pattern must not fire and the op is left for the
// existing `populateResolveExtractStridedMetadataPatterns` pipeline (which
// here resolves it on the unchanged i8 source).
>From 242d71cf33d63c2aff7f6736a93f0bb2fb6d7da4 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Mon, 11 May 2026 13:41:37 -0700
Subject: [PATCH 3/4] [MLIR][MemRef] Drop incomplete cf-conversion test
plumbing
The `enable-cf-conversion` test option and its companion lit file
`emulate-narrow-type-cf.mlir` depended on a 4-arg overload of
`populateFunctionOpInterfaceTypeConversionPattern` that does not exist
upstream. Remove the option, the cf-related includes/dialect registration,
and the lit file so the PR builds standalone. End-to-end cf support will be
re-introduced in a follow-up PR once the upstream helper lands.
---
.../MemRef/emulate-narrow-type-cf.mlir | 69 -------------------
.../Dialect/MemRef/TestEmulateNarrowType.cpp | 40 +----------
2 files changed, 3 insertions(+), 106 deletions(-)
delete mode 100644 mlir/test/Dialect/MemRef/emulate-narrow-type-cf.mlir
diff --git a/mlir/test/Dialect/MemRef/emulate-narrow-type-cf.mlir b/mlir/test/Dialect/MemRef/emulate-narrow-type-cf.mlir
deleted file mode 100644
index a5cebb0110048..0000000000000
--- a/mlir/test/Dialect/MemRef/emulate-narrow-type-cf.mlir
+++ /dev/null
@@ -1,69 +0,0 @@
-// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8 enable-cf-conversion=true" --cse --verify-diagnostics --split-input-file %s | FileCheck %s
-// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8 enable-cf-conversion=true" --split-input-file %s | FileCheck %s --check-prefix=CHECK-NOCSE
-
-// Sub-byte memref type carried through cf.br block args. The
-// BranchOpInterface type-conversion pattern must rewrite both the cf.br
-// operand type and the successor block-arg type to the i8 container, so the
-// downstream uses in the successor block see an i8 source.
-
-// CHECK-LABEL: func.func @cf_br_block_arg_narrow_type
-// CHECK-SAME: %[[ARG:[A-Za-z0-9_]+]]: memref<{{[0-9]+}}xi8>
-// CHECK: cf.br ^[[BB1:.+]](%[[ARG]] : memref<{{[0-9]+}}xi8>)
-// CHECK: ^[[BB1]](%[[BARG:[A-Za-z0-9_]+]]: memref<{{[0-9]+}}xi8>):
-// CHECK: return %[[BARG]]
-// CHECK-NOT: memref<{{[0-9]+}}xi4>
-func.func @cf_br_block_arg_narrow_type(%arg: memref<8xi4>) -> memref<8xi4> {
- cf.br ^bb1(%arg : memref<8xi4>)
-^bb1(%a: memref<8xi4>):
- return %a : memref<8xi4>
-}
-
-// -----
-
-// Sub-byte memref with dynamic offset carried through cf.br block-arg,
-// then loaded via vector.load. After FunctionOpInterfaceAllBlocksSignatureConversion
-// converts the block-arg to i8, an unrealized_conversion_cast is inserted.
-// ConvertVectorLoad calls extract_strided_metadata on the original sub-byte
-// source (op.getBase()), which is illegal without ConvertExtractStridedMetadata.
-
-// CHECK-LABEL: func.func @cf_br_block_arg_vector_load_i4
-// CHECK: vector.load {{.*}} : memref<{{[0-9]+}}xi8, strided<[1], offset: ?>>, vector<{{[0-9]+}}xi8>
-// CHECK-NOT: memref<{{[0-9]+}}xi4>
-// CHECK-NOCSE-LABEL: func.func @cf_br_block_arg_vector_load_i4
-// CHECK-NOCSE-NOT: memref<{{[0-9]+}}xi4>
-func.func @cf_br_block_arg_vector_load_i4(%arg: memref<8xi4, strided<[1], offset: ?>>) -> vector<8xi4> {
- cf.br ^bb1(%arg : memref<8xi4, strided<[1], offset: ?>>)
-^bb1(%a: memref<8xi4, strided<[1], offset: ?>>):
- %c0 = arith.constant 0 : index
- %v = vector.load %a[%c0] : memref<8xi4, strided<[1], offset: ?>>, vector<8xi4>
- return %v : vector<8xi4>
-}
-
-// -----
-
-// Sub-byte memref with static non-zero offset carried through cf.br block-arg,
-// then loaded via vector.load. The static offset (4 i4-elements) is converted
-// to 2 i8-elements in the container type. ConvertExtractStridedMetadata must
-// NOT over-scale the static offset: srcStaticOffset is already in
-// emulated-element (i4) units, so it should be passed through unchanged rather
-// than multiplied by containerBits/emulatedBits (which would give 8, wrong).
-//
-// The CHECK-NOCSE check verifies that the emulated offset constant is 4 (in i4
-// units, as returned by ConvertExtractStridedMetadata) and not 8 (the
-// over-scaled value that would appear with the bug). After CSE the constant
-// is dead and the distinction disappears, so this check requires the no-CSE
-// run line.
-
-// CHECK-LABEL: func.func @cf_br_block_arg_vector_load_i4_static_offset
-// CHECK: vector.load {{.*}} : memref<{{[0-9]+}}xi8, strided<[1], offset: 2>>, vector<{{[0-9]+}}xi8>
-// CHECK-NOT: memref<{{[0-9]+}}xi4>
-// CHECK-NOCSE-LABEL: func.func @cf_br_block_arg_vector_load_i4_static_offset
-// CHECK-NOCSE-NOT: memref<{{[0-9]+}}xi4>
-// CHECK-NOCSE: %c4 = arith.constant 4 : index
-func.func @cf_br_block_arg_vector_load_i4_static_offset(%arg: memref<8xi4, strided<[1], offset: 4>>) -> vector<4xi4> {
- cf.br ^bb1(%arg : memref<8xi4, strided<[1], offset: 4>>)
-^bb1(%a: memref<8xi4, strided<[1], offset: 4>>):
- %c0 = arith.constant 0 : index
- %v = vector.load %a[%c0] : memref<8xi4, strided<[1], offset: 4>>, vector<4xi4>
- return %v : vector<4xi4>
-}
diff --git a/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp b/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
index 3e87cd955a627..bec83a8dcbef9 100644
--- a/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
+++ b/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
@@ -11,14 +11,11 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
-#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
-#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -37,9 +34,9 @@ struct TestEmulateNarrowTypePass
: PassWrapper(pass) {}
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, cf::ControlFlowDialect,
- func::FuncDialect, memref::MemRefDialect,
- vector::VectorDialect, affine::AffineDialect>();
+ registry
+ .insert<arith::ArithDialect, func::FuncDialect, memref::MemRefDialect,
+ vector::VectorDialect, affine::AffineDialect>();
}
StringRef getArgument() const final { return "test-emulate-narrow-int"; }
StringRef getDescription() const final {
@@ -99,19 +96,6 @@ struct TestEmulateNarrowTypePass
arith::ArithDialect, vector::VectorDialect, memref::MemRefDialect,
affine::AffineDialect>(opLegalCallback);
- if (enableCFConversion) {
- target.addDynamicallyLegalDialect<cf::ControlFlowDialect>(
- [&typeConverter](Operation *op) -> bool {
- // Only apply legality check to BranchOpInterface ops; other cf ops
- // (e.g. cf.assert) have no successor-block-arg conversion concern
- // and should remain legal.
- if (!isa<BranchOpInterface>(op))
- return true;
- return isLegalForBranchOpInterfaceTypeConversionPattern(
- op, typeConverter);
- });
- }
-
RewritePatternSet patterns(ctx);
arith::populateArithNarrowTypeEmulationPatterns(typeConverter, patterns);
@@ -120,17 +104,6 @@ struct TestEmulateNarrowTypePass
vector::populateVectorNarrowTypeEmulationPatterns(
typeConverter, patterns, disableAtomicRMW, assumeAligned);
- if (enableCFConversion) {
- populateBranchOpInterfaceTypeConversionPattern(patterns, typeConverter);
- // Opt the FunctionOpInterface signature-conversion pattern into the
- // all-blocks mode so non-entry block-arg types are converted in
- // lockstep with the function signature and branch operands. Use a
- // higher benefit so it supersedes the default-benefit pattern
- // populated by `populateArithNarrowTypeEmulationPatterns`.
- populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(
- patterns, typeConverter, /*benefit=*/2, /*convertAllBlocks=*/true);
- }
-
if (failed(applyPartialConversion(op, target, std::move(patterns))))
signalPassFailure();
}
@@ -160,13 +133,6 @@ struct TestEmulateNarrowTypePass
llvm::cl::desc("assume store offsets are aligned to container element "
"boundaries"),
llvm::cl::init(false)};
-
- Option<bool> enableCFConversion{
- *this, "enable-cf-conversion",
- llvm::cl::desc("register populateBranchOpInterfaceTypeConversionPattern "
- "and mark cf dialect ops dynamically legal based on "
- "isLegalForBranchOpInterfaceTypeConversionPattern"),
- llvm::cl::init(false)};
};
struct TestMemRefFlattenAndVectorNarrowTypeEmulationPass
>From dfca37eedf57473e0266352d765517a60662b6b0 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Thu, 14 May 2026 04:58:32 -0700
Subject: [PATCH 4/4] [MLIR][MemRef] clang-format narrow-type stride init
---
mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
index 5634803645f84..782d56fb0e9a4 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
@@ -872,9 +872,9 @@ void memref::populateMemRefNarrowTypeEmulationConversions(
// layout when the original offset is non-zero. Rank-0 memrefs carry no
// strides, so the layout uses an empty stride list.
if (offset != 0) {
- SmallVector<int64_t, 1> resultStrides =
- ty.getRank() == 0 ? SmallVector<int64_t, 1>{}
- : SmallVector<int64_t, 1>{1};
+ SmallVector<int64_t, 1> resultStrides = ty.getRank() == 0
+ ? SmallVector<int64_t, 1>{}
+ : SmallVector<int64_t, 1>{1};
if (offset == ShapedType::kDynamic) {
layoutAttr =
StridedLayoutAttr::get(ty.getContext(), offset, resultStrides);
More information about the Mlir-commits
mailing list