Skip to content

[mlir][NFC] update mlir/Dialect create APIs (27/n) #150638

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 25, 2025
Merged

[mlir][NFC] update mlir/Dialect create APIs (27/n) #150638

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,16 +481,16 @@ struct MemoryCounterWaitOpLowering
if (chipset.majorVersion >= 12) {
Location loc = op.getLoc();
if (std::optional<int> ds = adaptor.getDs())
rewriter.create<ROCDL::WaitDscntOp>(loc, *ds);
ROCDL::WaitDscntOp::create(rewriter, loc, *ds);

if (std::optional<int> load = adaptor.getLoad())
rewriter.create<ROCDL::WaitLoadcntOp>(loc, *load);
ROCDL::WaitLoadcntOp::create(rewriter, loc, *load);

if (std::optional<int> store = adaptor.getStore())
rewriter.create<ROCDL::WaitStorecntOp>(loc, *store);
ROCDL::WaitStorecntOp::create(rewriter, loc, *store);

if (std::optional<int> exp = adaptor.getExp())
rewriter.create<ROCDL::WaitExpcntOp>(loc, *exp);
ROCDL::WaitExpcntOp::create(rewriter, loc, *exp);

rewriter.eraseOp(op);
return success();
Expand Down
22 changes: 12 additions & 10 deletions mlir/lib/Dialect/ArmNeon/Transforms/LowerContractToNeonPatterns.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,8 +145,8 @@ class VectorContractRewriter {
return rewriter.createOrFold<arm_neon::UsmmlaOp>(loc, acc.getType(), acc,
lhs, rhs);
case MMLA::Bfloat:
return rewriter.create<arm_neon::BfmmlaOp>(loc, acc.getType(), acc, lhs,
rhs);
return arm_neon::BfmmlaOp::create(rewriter, loc, acc.getType(), acc, lhs,
rhs);
case MMLA::Nop:
llvm_unreachable("Uninitialized operation type");
}
Expand Down Expand Up @@ -226,8 +226,9 @@ class VectorContractRewriter {

// Initial accumulator for the final result. This is the un-tiled result if
// tiling is done.
Value result = rewriter.create<arith::ConstantOp>(
loc, op.getResultType(), rewriter.getZeroAttr(op.getResultType()));
Value result =
arith::ConstantOp::create(rewriter, loc, op.getResultType(),
rewriter.getZeroAttr(op.getResultType()));

SmallVector<int64_t, 3> loopOrder = {0, 1};
if (iterationBounds.size() == 3)
Expand Down Expand Up @@ -263,8 +264,9 @@ class VectorContractRewriter {
if (dimM == 1) {
auto expandRowVector = [&](Value tiledOperand,
VectorType expandedTypeType) {
auto emptyOperand = rewriter.create<arith::ConstantOp>(
loc, expandedTypeType, rewriter.getZeroAttr(expandedTypeType));
auto emptyOperand =
arith::ConstantOp::create(rewriter, loc, expandedTypeType,
rewriter.getZeroAttr(expandedTypeType));
SmallVector<int64_t> offsets(
cast<ShapedType>(emptyOperand.getType()).getRank(), 0);
SmallVector<int64_t> strides(
Expand All @@ -280,8 +282,8 @@ class VectorContractRewriter {
// using the instruction for unsigned by signed multiplication with
// reversed operands.
if (swapOperands)
tiledAcc = rewriter.create<vector::TransposeOp>(
loc, tiledAcc, ArrayRef<int64_t>({1, 0}));
tiledAcc = vector::TransposeOp::create(rewriter, loc, tiledAcc,
ArrayRef<int64_t>({1, 0}));

// Collapse tiled operands to 1D vectors required by the ArmNeon ops
auto collapsedLhs = rewriter.createOrFold<vector::ShapeCastOp>(
Expand Down Expand Up @@ -309,8 +311,8 @@ class VectorContractRewriter {
// Because of the reversed operands the result is obtained transposed.
// Transpose it back,
if (swapOperands)
tiledRes = rewriter.create<vector::TransposeOp>(
loc, tiledRes, ArrayRef<int64_t>({1, 0}));
tiledRes = vector::TransposeOp::create(rewriter, loc, tiledRes,
ArrayRef<int64_t>({1, 0}));

// With vecmat, only one row of tiled ACC can be inserted into the final
// result
Expand Down
77 changes: 39 additions & 38 deletions mlir/lib/Dialect/ArmSVE/Transforms/LowerContractToSVEPatterns.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,13 +214,13 @@ Value VectorContractRewriter::createMMLA(PatternRewriter &rewriter,

switch (mmlaOp) {
case MMLA::SignedInt:
return rewriter.create<arm_sve::SmmlaOp>(loc, resTy, acc, lhs, rhs);
return arm_sve::SmmlaOp::create(rewriter, loc, resTy, acc, lhs, rhs);
case MMLA::UnsignedInt:
return rewriter.create<arm_sve::UmmlaOp>(loc, resTy, acc, lhs, rhs);
return arm_sve::UmmlaOp::create(rewriter, loc, resTy, acc, lhs, rhs);
case MMLA::MixedInt:
return rewriter.create<arm_sve::UsmmlaOp>(loc, resTy, acc, lhs, rhs);
return arm_sve::UsmmlaOp::create(rewriter, loc, resTy, acc, lhs, rhs);
case MMLA::Bfloat:
return rewriter.create<arm_sve::BfmmlaOp>(loc, resTy, acc, lhs, rhs);
return arm_sve::BfmmlaOp::create(rewriter, loc, resTy, acc, lhs, rhs);
default:
llvm_unreachable("Uninitialized operation kind");
}
Expand Down Expand Up @@ -316,62 +316,63 @@ Value VectorContractRewriter::lower(vector::ContractionOp op,
for (int64_t i = 0; i < M; i += 2) {
// Extract two consecutive rows of the LHS tile.
auto r0 =
rewriter.create<vector::ExtractOp>(loc, lhs, ArrayRef<int64_t>{i});
vector::ExtractOp::create(rewriter, loc, lhs, ArrayRef<int64_t>{i});
auto r1 =
rewriter.create<vector::ExtractOp>(loc, lhs, ArrayRef<int64_t>{i + 1});
vector::ExtractOp::create(rewriter, loc, lhs, ArrayRef<int64_t>{i + 1});
// Concatenate to obtain a 2 x K x <input-type> flattened sub-tile.
SmallVector<int64_t> shuffleIdx(2 * K);
std::iota(shuffleIdx.begin(), shuffleIdx.end(), 0);
auto t = rewriter.create<vector::ShuffleOp>(loc, r0, r1, shuffleIdx);
auto t = vector::ShuffleOp::create(rewriter, loc, r0, r1, shuffleIdx);
// Turn it into a scalable vector.
auto s = rewriter.create<vector::ScalableInsertOp>(
loc, t, rewriter.create<ub::PoisonOp>(loc, flatLhsType), 0);
auto s = vector::ScalableInsertOp::create(
rewriter, loc, t, ub::PoisonOp::create(rewriter, loc, flatLhsType), 0);
// Replicate the sub-tile VSCALE times to fill the entire vector.
auto r = rewriter.create<arm_sve::DupQLaneOp>(loc, s, 0);
auto r = arm_sve::DupQLaneOp::create(rewriter, loc, s, 0);
lhsTile.push_back(r);
}

// "Flatten" the RHS tile from <[N]xK> to <[N*K]>.
auto rhs = rewriter.create<vector::ShapeCastOp>(this->rhs.getLoc(),
flatRhsTileType, this->rhs);
auto rhs = vector::ShapeCastOp::create(rewriter, this->rhs.getLoc(),
flatRhsTileType, this->rhs);

// Extract the RHS sub-tiles with logical shape <Kx[2]>.
SmallVector<Value> rhsTile;
for (int64_t j = 0; j < N; j += 2)
rhsTile.push_back(rewriter.create<vector::ScalableExtractOp>(
loc, flatRhsType, rhs, j * K));
rhsTile.push_back(vector::ScalableExtractOp::create(
rewriter, loc, flatRhsType, rhs, j * K));

// Extract and pack the ACC sub-tiles.
SmallVector<Value> accTile;
for (int64_t i = 0; i < M; i += 2) {
// Extract two consecutive rows of the accumulator tile.
auto r0 = rewriter.create<vector::ExtractOp>(loc, op.getAcc(),
ArrayRef<int64_t>{i});
auto r1 = rewriter.create<vector::ExtractOp>(loc, op.getAcc(),
ArrayRef<int64_t>{i + 1});
auto r0 = vector::ExtractOp::create(rewriter, loc, op.getAcc(),
ArrayRef<int64_t>{i});
auto r1 = vector::ExtractOp::create(rewriter, loc, op.getAcc(),
ArrayRef<int64_t>{i + 1});
Value accTileVec;
if (swapOperands) {
// We are performing the operation with swapped LHS and RHS we need to
// transpose each individual 2x2 tile of the accumulator and (later) the
// final result.
accTileVec = rewriter.create<vector::InterleaveOp>(loc, r0, r1);
accTileVec = vector::InterleaveOp::create(rewriter, loc, r0, r1);
} else {
// Bitcast accumulator rows to double-width integer elements, so
// subsequent interleave/deinterleave work on pairs of elements.
auto r0I64 = rewriter.create<vector::BitCastOp>(loc, accRow64Ty, r0);
auto r1I64 = rewriter.create<vector::BitCastOp>(loc, accRow64Ty, r1);
auto r0I64 = vector::BitCastOp::create(rewriter, loc, accRow64Ty, r0);
auto r1I64 = vector::BitCastOp::create(rewriter, loc, accRow64Ty, r1);

// Interleave the rows, effectively flattening each 2x2 tile into 4
// consecutive elements.
auto intrI64 = rewriter.create<vector::InterleaveOp>(loc, r0I64, r1I64);
auto intrI64 = vector::InterleaveOp::create(rewriter, loc, r0I64, r1I64);

// Bitcast back to original element type.
accTileVec = rewriter.create<vector::BitCastOp>(loc, accRowX2Ty, intrI64);
accTileVec =
vector::BitCastOp::create(rewriter, loc, accRowX2Ty, intrI64);
}
// Extract ACC sub-tiles.
for (int64_t j = 0; j < N; j += 2)
accTile.push_back(rewriter.create<vector::ScalableExtractOp>(
loc, flatAccType, accTileVec, j * 2));
accTile.push_back(vector::ScalableExtractOp::create(
rewriter, loc, flatAccType, accTileVec, j * 2));
}

// Emit sub-tile matrix multiplications.
Expand All @@ -384,36 +385,36 @@ Value VectorContractRewriter::lower(vector::ContractionOp op,
}

// Unpack the OUT sub-tiles and insert into the result.
Value result = rewriter.create<ub::PoisonOp>(loc, op.getResultType());
Value result = ub::PoisonOp::create(rewriter, loc, op.getResultType());
for (int64_t i = 0; i < M / 2; ++i) {
// Collect a number of sub-tiles in a row.
Value row = rewriter.create<ub::PoisonOp>(loc, accRowX2Ty);
Value row = ub::PoisonOp::create(rewriter, loc, accRowX2Ty);
for (int64_t j = 0; j < N / 2; ++j)
row = rewriter.create<vector::ScalableInsertOp>(
loc, outTile[i * N / 2 + j], row, j * 4);
row = vector::ScalableInsertOp::create(
rewriter, loc, outTile[i * N / 2 + j], row, j * 4);

// Unpack the row to obtain two rows of the output. If we have the out
// sub-tiles transposed we obtain two consecutive output rows by
// separating even and odd elements, i.e. a simple deinterleave.
// Otherwise, the interleave is by pairs.
Value out0, out1;
if (swapOperands) {
auto tmp = rewriter.create<vector::DeinterleaveOp>(loc, row);
auto tmp = vector::DeinterleaveOp::create(rewriter, loc, row);
out0 = tmp.getRes1();
out1 = tmp.getRes2();
} else {
// Deinterleave by pairs.
auto row64 = rewriter.create<vector::BitCastOp>(loc, accRowX264Ty, row);
auto deintr64 = rewriter.create<vector::DeinterleaveOp>(loc, row64);
auto row64 = vector::BitCastOp::create(rewriter, loc, accRowX264Ty, row);
auto deintr64 = vector::DeinterleaveOp::create(rewriter, loc, row64);

// Bitcast back into original element type and insert into the result.
out0 =
rewriter.create<vector::BitCastOp>(loc, accRowTy, deintr64.getRes1());
out1 =
rewriter.create<vector::BitCastOp>(loc, accRowTy, deintr64.getRes2());
out0 = vector::BitCastOp::create(rewriter, loc, accRowTy,
deintr64.getRes1());
out1 = vector::BitCastOp::create(rewriter, loc, accRowTy,
deintr64.getRes2());
}
result = rewriter.create<vector::InsertOp>(loc, out0, result, i * 2);
result = rewriter.create<vector::InsertOp>(loc, out1, result, i * 2 + 1);
result = vector::InsertOp::create(rewriter, loc, out0, result, i * 2);
result = vector::InsertOp::create(rewriter, loc, out1, result, i * 2 + 1);
}

return result;
Expand Down
4 changes: 2 additions & 2 deletions mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2198,8 +2198,8 @@ vectorizeAsLinalgContraction(RewriterBase &rewriter, VectorizationState &state,
}

// Create contraction.
Operation *contractOp = rewriter.create<vector::ContractionOp>(
loc, /*lhs=*/vecOperands[0],
Operation *contractOp = vector::ContractionOp::create(
rewriter, loc, /*lhs=*/vecOperands[0],
/*rhs=*/vecOperands[1], /*acc=*/vecOperands[2],
linalgOp.getIndexingMaps(), rewriter.getArrayAttr(iterAttrs), *maybeKind);
contractOp = state.maskOperation(rewriter, contractOp, linalgOp);
Expand Down
12 changes: 6 additions & 6 deletions mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -237,8 +237,8 @@ LogicalResult resolveSourceIndicesExpandShape(
llvm::map_to_vector(group, [&](int64_t d) { return destShape[d]; });
SmallVector<Value> groupIndices =
llvm::map_to_vector(group, [&](int64_t d) { return indices[d]; });
Value collapsedIndex = rewriter.create<affine::AffineLinearizeIndexOp>(
loc, groupIndices, groupBasis, /*disjoint=*/startsInbounds);
Value collapsedIndex = affine::AffineLinearizeIndexOp::create(
rewriter, loc, groupIndices, groupBasis, /*disjoint=*/startsInbounds);
sourceIndices.push_back(collapsedIndex);
}
return success();
Expand All @@ -250,8 +250,8 @@ resolveSourceIndicesCollapseShape(Location loc, PatternRewriter &rewriter,
ValueRange indices,
SmallVectorImpl<Value> &sourceIndices) {
// Note: collapse_shape requires a strided memref, we can do this.
auto metadata = rewriter.create<memref::ExtractStridedMetadataOp>(
loc, collapseShapeOp.getSrc());
auto metadata = memref::ExtractStridedMetadataOp::create(
rewriter, loc, collapseShapeOp.getSrc());
SmallVector<OpFoldResult> sourceSizes = metadata.getConstifiedMixedSizes();
for (auto [index, group] :
llvm::zip(indices, collapseShapeOp.getReassociationIndices())) {
Expand All @@ -265,8 +265,8 @@ resolveSourceIndicesCollapseShape(Location loc, PatternRewriter &rewriter,

SmallVector<OpFoldResult> basis =
llvm::map_to_vector(group, [&](int64_t d) { return sourceSizes[d]; });
auto delinearize = rewriter.create<affine::AffineDelinearizeIndexOp>(
loc, index, basis, /*hasOuterBound=*/true);
auto delinearize = affine::AffineDelinearizeIndexOp::create(
rewriter, loc, index, basis, /*hasOuterBound=*/true);
llvm::append_range(sourceIndices, delinearize.getResults());
}
if (collapseShapeOp.getReassociationIndices().empty()) {
Expand Down
10 changes: 5 additions & 5 deletions mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,7 +207,7 @@ struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {
// Subtract startOfRange from the original subgroup id to get the adjusted
// sg id
Value startOfRangeVal =
rewriter.create<arith::ConstantIndexOp>(loc, startOfRange);
arith::ConstantIndexOp::create(rewriter, loc, startOfRange);
adjustedSgId =
rewriter.createOrFold<index::SubOp>(loc, linearSgId, startOfRangeVal);
}
Expand Down Expand Up @@ -431,8 +431,8 @@ struct WgToSgVectorBroadcastOp

SmallVector<Value> newBroadcastOps;
for (auto operand : adaptor.getOperands().front()) {
auto newBroadcast = rewriter.create<vector::BroadcastOp>(
op.getLoc(), newResultType, operand);
auto newBroadcast = vector::BroadcastOp::create(rewriter, op.getLoc(),
newResultType, operand);
xegpu::setLayoutAttr(newBroadcast->getResult(0),
layout.dropSgLayoutAndData());
newBroadcastOps.push_back(newBroadcast.getResult());
Expand Down Expand Up @@ -563,8 +563,8 @@ struct WgToSgConvertLayoutOp
if (input && target) {
// keep the ConvertLayoutOp for rest fields, e.g., inst_data.
for (auto [i, src] : llvm::enumerate(adaptor.getSource())) {
auto newOp = rewriter.create<xegpu::ConvertLayoutOp>(
op.getLoc(), src.getType(), src, input, target);
auto newOp = xegpu::ConvertLayoutOp::create(
rewriter, op.getLoc(), src.getType(), src, input, target);
newOps[i] = newOp;
}
}
Expand Down
Loading