Skip to content

Commit 6e2f420

Browse files
charithaintccharithaintc
committed
fix func naming
1 parent 0550d4b commit 6e2f420

File tree

8 files changed

+81
-134
lines changed

8 files changed

+81
-134
lines changed

‎mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td‎

Lines changed: 29 additions & 70 deletions
Original file line numberDiff line numberDiff line change
@@ -194,26 +194,29 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
194194
InterfaceMethod<"Get the num of effective subgroups",
195195
"int64_t",
196196
"getNumSubgroups", (ins), [{
197-
std::optional<SmallVector<int64_t>> sgLayout = llvm::cast<ConcreteAttr>(tablegen_opaque_val).getSgLayoutAsInt();
197+
std::optional<SmallVector<int64_t>> sgLayout = llvm::cast<ConcreteAttr>(tablegen_opaque_val).getEffectiveSgLayoutAsInt();
198198
if (sgLayout.has_value())
199199
return computeProduct(*sgLayout);
200200
return 0;
201201
}], [{}]>,
202-
InterfaceMethod<"Get the SgLayout field of the attribute as integer array",
202+
InterfaceMethod<"Get the order of the layout attribute",
203+
"DenseI32ArrayAttr",
204+
"getOrder">,
205+
InterfaceMethod<"Get the effective SgLayout of the layout attribute as integer array",
203206
"SmallVector<int64_t>",
204-
"getSgLayoutAsInt">,
205-
InterfaceMethod<"Get the SgData field of the attribute as integer array",
207+
"getEffectiveSgLayoutAsInt">,
208+
InterfaceMethod<"Get the effective SgData of the layout attribute as integer array",
206209
"SmallVector<int64_t>",
207-
"getSgDataAsInt">,
208-
InterfaceMethod<"Get the InstData field of the attribute as integer array",
210+
"getEffectiveSgDataAsInt">,
211+
InterfaceMethod<"Get the effective InstData of the layout attribute as integer array",
209212
"SmallVector<int64_t>",
210-
"getInstDataAsInt">,
211-
InterfaceMethod<"Get the LaneLayout field of the attribute as integer array",
213+
"getEffectiveInstDataAsInt">,
214+
InterfaceMethod<"Get the effective LaneLayout of the layout attribute as integer array",
212215
"SmallVector<int64_t>",
213-
"getLaneLayoutAsInt">,
214-
InterfaceMethod<"Get the LaneData field of the attribute as integer array",
216+
"getEffectiveLaneLayoutAsInt">,
217+
InterfaceMethod<"Get the effective LaneData of the layout attribute as integer array",
215218
"SmallVector<int64_t>",
216-
"getLaneDataAsInt">,
219+
"getEffectiveLaneDataAsInt">,
217220
InterfaceMethod<"Derive a new layout by dropping sgLayout and sgData",
218221
"xegpu::DistributeLayoutAttr",
219222
"dropSgLayoutAndData">,
@@ -232,50 +235,6 @@ def DistributeLayoutAttr: AttrInterface<"DistributeLayoutAttr"> {
232235
"FailureOr<SmallVector<SmallVector<Value>>>",
233236
"getOffsets",
234237
(ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId, "ArrayRef<int64_t>":$shape)>,
235-
InterfaceMethod</*desc=*/[{Check if this layout can be achieved by applying a transpose
236-
to some other layout according to given permutation of (0...n-1).}],
237-
/*retTy=*/"bool",
238-
/*methodName=*/"isTransposeOf",
239-
/*args=*/(ins "const xegpu::DistributeLayoutAttr&": $other, "ArrayRef<int64_t>": $perm),
240-
/*methodBody=*/[{
241-
if (!other)
242-
return false;
243-
if ($_self.getRank() != other.getRank() || perm.size() != static_cast<size_t>($_self.getRank()))
244-
return false;
245-
// check if the permutation is valid
246-
int64_t rank = $_self.getRank();
247-
SmallVector<bool, 8> seen(rank, false);
248-
for (const auto &ta : llvm::enumerate(perm)) {
249-
if (ta.value() < 0 || ta.value() >= rank)
250-
return false;
251-
if (seen[ta.value()])
252-
return false;
253-
seen[ta.value()] = true;
254-
}
255-
auto checkTranspose = [](ArrayRef<int64_t> dst, ArrayRef<int64_t> src, ArrayRef<int64_t> perm) {
256-
for (const auto &ta : llvm::enumerate(perm)) {
257-
if (src[ta.index()] != dst[ta.value()])
258-
return false;
259-
}
260-
return true;
261-
};
262-
// check sgLayout
263-
if (!checkTranspose($_self.getSgLayoutAsInt(), other.getSgLayoutAsInt(), perm))
264-
return false;
265-
// check sgData
266-
if (!checkTranspose($_self.getSgDataAsInt(), other.getSgDataAsInt(), perm))
267-
return false;
268-
// check instData
269-
if (!checkTranspose($_self.getInstDataAsInt(), other.getInstDataAsInt(), perm))
270-
return false;
271-
// check laneLayout
272-
if (!checkTranspose($_self.getLaneLayoutAsInt(), other.getLaneLayoutAsInt(), perm))
273-
return false;
274-
// check laneData
275-
if (!checkTranspose($_self.getLaneDataAsInt(), other.getLaneDataAsInt(), perm))
276-
return false;
277-
return true;
278-
}]>,
279238
InterfaceMethod</*desc=*/[{Check if this layout is a slice of some other layout.}],
280239
/*retTy=*/"bool",
281240
/*methodName=*/"isSliceOf",
@@ -439,31 +398,31 @@ def XeGPU_LayoutAttr : XeGPUAttr<"Layout", "layout", [DistributeLayoutAttr]> {
439398
getLaneLayout(), getLaneData(), getOrder());
440399
}
441400

442-
SmallVector<int64_t> getSgLayoutAsInt() const {
401+
SmallVector<int64_t> getEffectiveSgLayoutAsInt() const {
443402
if (DenseI32ArrayAttr layout = getSgLayout())
444403
return llvm::to_vector_of<int64_t>(layout.asArrayRef());
445404
return {};
446405
}
447406

448-
SmallVector<int64_t> getSgDataAsInt() const {
407+
SmallVector<int64_t> getEffectiveSgDataAsInt() const {
449408
if (DenseI32ArrayAttr data = getSgData())
450409
return llvm::to_vector_of<int64_t>(data.asArrayRef());
451410
return {};
452411
}
453412

454-
SmallVector<int64_t> getInstDataAsInt() const {
413+
SmallVector<int64_t> getEffectiveInstDataAsInt() const {
455414
if (DenseI32ArrayAttr inst = getInstData())
456415
return llvm::to_vector_of<int64_t>(inst.asArrayRef());
457416
return {};
458417
}
459418

460-
SmallVector<int64_t> getLaneLayoutAsInt() const {
419+
SmallVector<int64_t> getEffectiveLaneLayoutAsInt() const {
461420
if (DenseI32ArrayAttr layout = getLaneLayout())
462421
return llvm::to_vector_of<int64_t>(layout.asArrayRef());
463422
return {};
464423
}
465424

466-
SmallVector<int64_t> getLaneDataAsInt() const {
425+
SmallVector<int64_t> getEffectiveLaneDataAsInt() const {
467426
if (DenseI32ArrayAttr data = getLaneData())
468427
return llvm::to_vector_of<int64_t>(data.asArrayRef());
469428
return {};
@@ -550,10 +509,10 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
550509

551510
/// Returns the SgLayout of the attribute, computed by applying
552511
/// the slice dimensions to the underlying LayoutAttr.
553-
SmallVector<int64_t> getSgLayoutAsInt() const {
512+
SmallVector<int64_t> getEffectiveSgLayoutAsInt() const {
554513
SliceAttr attr = flatten();
555514
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
556-
auto layout = parent.getSgLayoutAsInt();
515+
auto layout = parent.getEffectiveSgLayoutAsInt();
557516
if (layout.size()) {
558517
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
559518
return XeGPUDialect::slice(ArrayRef<int64_t>(layout), dims);
@@ -563,10 +522,10 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
563522

564523
/// Returns the SgData of the attribute, computed by applying
565524
/// the slice dimensions to the underlying LayoutAttr.
566-
SmallVector<int64_t> getSgDataAsInt() const {
525+
SmallVector<int64_t> getEffectiveSgDataAsInt() const {
567526
SliceAttr attr = flatten();
568527
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
569-
auto data = parent.getSgDataAsInt();
528+
auto data = parent.getEffectiveSgDataAsInt();
570529
if (data.size()) {
571530
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
572531
return XeGPUDialect::slice(ArrayRef<int64_t>(data), dims);
@@ -576,10 +535,10 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
576535

577536
/// Returns the InstData of the attribute, computed by applying
578537
/// the slice dimensions to the underlying LayoutAttr.
579-
SmallVector<int64_t> getInstDataAsInt() const {
538+
SmallVector<int64_t> getEffectiveInstDataAsInt() const {
580539
SliceAttr attr = flatten();
581540
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
582-
auto inst = parent.getInstDataAsInt();
541+
auto inst = parent.getEffectiveInstDataAsInt();
583542
if (inst.size()) {
584543
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
585544
return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(inst), dims);
@@ -589,10 +548,10 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
589548

590549
/// Returns the LaneLayout of the attribute, computed by applying
591550
/// the slice dimensions to the underlying LayoutAttr.
592-
SmallVector<int64_t> getLaneLayoutAsInt() const {
551+
SmallVector<int64_t> getEffectiveLaneLayoutAsInt() const {
593552
SliceAttr attr = flatten();
594553
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
595-
auto layout = parent.getLaneLayoutAsInt();
554+
auto layout = parent.getEffectiveLaneLayoutAsInt();
596555
if (layout.size()) {
597556
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
598557
return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(layout), dims);
@@ -602,10 +561,10 @@ def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [DistributeLayoutAttr]> {
602561

603562
/// Returns the LaneData of the attribute, computed by applying
604563
/// the slice dimensions to the underlying LayoutAttr.
605-
SmallVector<int64_t> getLaneDataAsInt() const {
564+
SmallVector<int64_t> getEffectiveLaneDataAsInt() const {
606565
SliceAttr attr = flatten();
607566
auto parent = dyn_cast<LayoutAttr>(attr.getParent());
608-
auto data = parent.getLaneDataAsInt();
567+
auto data = parent.getEffectiveLaneDataAsInt();
609568
if (data.size()) {
610569
ArrayRef<int64_t> dims = attr.getDims().asArrayRef();
611570
return XeGPUDialect::slice(llvm::ArrayRef<int64_t>(data), dims);

‎mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -163,14 +163,6 @@ SmallVector<OpFoldResult> addWithRightAligned(OpBuilder &builder, Location loc,
163163
ArrayRef<OpFoldResult> lhs,
164164
ArrayRef<OpFoldResult> rhs);
165165

166-
/// Helper function to compute the effective lane layout from a
167-
/// DistributeLayoutAttr which can be either a LayoutAttr or a SliceAttr. For
168-
/// LayoutAttr, this will simply return the lane layout. For SliceAttr, it will
169-
/// compute the effective lane layout by removing the sliced dimensions from the
170-
/// parent lane layout.
171-
SmallVector<int64_t>
172-
computeEffectiveLaneLayout(xegpu::DistributeLayoutAttr layout);
173-
174166
} // namespace xegpu
175167

176168
} // namespace mlir

‎mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -134,22 +134,23 @@ bool XeGPUDialect::isEvenlyDistributable(llvm::ArrayRef<int64_t> shape,
134134
};
135135

136136
// check the sgLayout and sgData
137-
auto maybeSgShape =
138-
tryDistribute(shape, attr.getSgLayoutAsInt(), attr.getSgDataAsInt());
137+
auto maybeSgShape = tryDistribute(shape, attr.getEffectiveSgLayoutAsInt(),
138+
attr.getEffectiveSgDataAsInt());
139139
if (!maybeSgShape)
140140
return false;
141141
auto sgShape = maybeSgShape.value();
142142

143143
// check InstData, it neither have layout nor need round-robin
144144
auto maybeInstShape =
145-
tryDistribute(sgShape, {}, attr.getInstDataAsInt(), false);
145+
tryDistribute(sgShape, {}, attr.getEffectiveInstDataAsInt(), false);
146146
if (!maybeInstShape)
147147
return false;
148148
auto instShape = maybeInstShape.value();
149149

150150
// check LaneLayout and LaneData
151-
auto maybeLaneShape = tryDistribute(instShape, attr.getLaneLayoutAsInt(),
152-
attr.getLaneDataAsInt(), false);
151+
auto maybeLaneShape =
152+
tryDistribute(instShape, attr.getEffectiveLaneLayoutAsInt(),
153+
attr.getEffectiveLaneDataAsInt(), false);
153154
return maybeLaneShape.has_value();
154155
}
155156

@@ -283,9 +284,10 @@ LayoutAttr::delinearizeSubgroupId(OpBuilder &builder, Location loc,
283284
if (!hasDefaultOrder())
284285
return mlir::emitError(loc, "order attribute is currently not supported.");
285286

286-
auto dims = llvm::map_to_vector(getSgLayoutAsInt(), [&](int64_t d) -> Value {
287-
return builder.createOrFold<arith::ConstantIndexOp>(loc, d);
288-
});
287+
auto dims =
288+
llvm::map_to_vector(getEffectiveSgLayoutAsInt(), [&](int64_t d) -> Value {
289+
return builder.createOrFold<arith::ConstantIndexOp>(loc, d);
290+
});
289291

290292
return affine::delinearizeIndex(builder, loc, linearId, dims);
291293
}
@@ -299,8 +301,8 @@ LayoutAttr::getOffsets(OpBuilder &builder, Location loc, Value linearId,
299301
if (!isForWorkgroup())
300302
return failure();
301303

302-
SmallVector<int64_t> sgLayout = getSgLayoutAsInt();
303-
SmallVector<int64_t> sgShape = getSgDataAsInt();
304+
SmallVector<int64_t> sgLayout = getEffectiveSgLayoutAsInt();
305+
SmallVector<int64_t> sgShape = getEffectiveSgDataAsInt();
304306
if (sgShape.empty()) {
305307
if (auto derivedShape = computeShapeRatio(shape, sgLayout))
306308
sgShape = derivedShape.value();
@@ -386,8 +388,8 @@ SliceAttr::getOffsets(OpBuilder &builder, Location loc, Value linearId,
386388
if (!isForWorkgroup())
387389
return failure();
388390

389-
SmallVector<int64_t> sgLayout = getSgLayoutAsInt();
390-
SmallVector<int64_t> sgShape = getSgDataAsInt();
391+
SmallVector<int64_t> sgLayout = getEffectiveSgLayoutAsInt();
392+
SmallVector<int64_t> sgShape = getEffectiveSgDataAsInt();
391393
if (sgShape.empty()) {
392394
if (auto derivedShape = computeShapeRatio(shape, sgLayout))
393395
sgShape = derivedShape.value();

‎mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -85,16 +85,16 @@ struct ConvertLayoutOpPattern
8585
using OpRewritePattern::OpRewritePattern;
8686
LogicalResult matchAndRewrite(xegpu::ConvertLayoutOp op,
8787
PatternRewriter &rewriter) const override {
88-
xegpu::DistributeLayoutAttr input_layout = op.getInputLayoutAttr();
89-
xegpu::DistributeLayoutAttr target_layout = op.getTargetLayoutAttr();
90-
if (input_layout.getInstDataAsInt().empty() ||
91-
target_layout.getInstDataAsInt().empty())
88+
xegpu::DistributeLayoutAttr inputLayout = op.getInputLayoutAttr();
89+
xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
90+
if (inputLayout.getEffectiveInstDataAsInt().empty() ||
91+
targetLayout.getEffectiveInstDataAsInt().empty())
9292
return rewriter.notifyMatchFailure(op, "Not a target ConvertLayoutOp.");
9393

94-
input_layout = input_layout.dropInstData();
95-
target_layout = target_layout.dropInstData();
94+
inputLayout = inputLayout.dropInstData();
95+
targetLayout = targetLayout.dropInstData();
9696
auto newOp = rewriter.createOrFold<xegpu::ConvertLayoutOp>(
97-
op.getLoc(), op.getType(), op.getSource(), input_layout, target_layout);
97+
op.getLoc(), op.getType(), op.getSource(), inputLayout, targetLayout);
9898
rewriter.replaceOp(op, newOp);
9999
return success();
100100
}
@@ -145,8 +145,8 @@ XeGPUBlockingPass::getTileShape(const T &operandOrResult) const {
145145
xegpu::DistributeLayoutAttr layout =
146146
xegpu::getDistributeLayoutAttr(operandOrResult);
147147
if (layout && layout.isForSubgroup()) {
148-
if (!layout.getInstDataAsInt().empty())
149-
return layout.getInstDataAsInt();
148+
if (!layout.getEffectiveInstDataAsInt().empty())
149+
return layout.getEffectiveInstDataAsInt();
150150

151151
if (auto type = dyn_cast<ShapedType>(value.getType()))
152152
return llvm::to_vector(type.getShape());
@@ -226,7 +226,7 @@ bool XeGPUBlockingPass::needsUnroll(Operation *op) const {
226226
Type valTy = value.getType();
227227
if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(valTy)) {
228228
xegpu::DistributeLayoutAttr layout = tdescTy.getLayoutAttr();
229-
return layout && !layout.getInstDataAsInt().empty();
229+
return layout && !layout.getEffectiveInstDataAsInt().empty();
230230
}
231231
auto shapedType = dyn_cast<ShapedType>(valTy);
232232
return shapedType && !llvm::equal(tileShape, shapedType.getShape());

‎mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,7 @@ getDistVecTypeBasedOnLaneLayout(xegpu::DistributeLayoutAttr layout,
8585
assert((isa<xegpu::LayoutAttr>(layout) || isa<xegpu::SliceAttr>(layout)) &&
8686
"Expecting a valid layout.");
8787
SmallVector<int64_t> effectiveLaneLayout =
88-
xegpu::computeEffectiveLaneLayout(layout);
89-
88+
layout.getEffectiveLaneLayoutAsInt();
9089
assert(static_cast<size_t>(originalType.getRank()) >=
9190
effectiveLaneLayout.size() &&
9291
"Rank of the original vector type should be greater or equal to the "
@@ -1234,9 +1233,26 @@ struct VectorShapeCastDistribution : public gpu::WarpDistributionPattern {
12341233
cast<VectorType>(warpOp.getResult(operandNumber).getType());
12351234
xegpu::DistributeLayoutAttr sourceLayout =
12361235
xegpu::getDistributeLayoutAttr(shapeCastOp.getSource());
1237-
if (!sourceLayout)
1236+
xegpu::DistributeLayoutAttr resultLayout =
1237+
xegpu::getDistributeLayoutAttr(shapeCastOp.getResult());
1238+
if (!sourceLayout || !resultLayout)
1239+
return rewriter.notifyMatchFailure(
1240+
warpOp,
1241+
"the source or result of shape_cast op lacks distribution layout");
1242+
1243+
// For rank reducing or increasing shape_cast ops, the lower rank layout
1244+
// must be a slice of higher rank layout.
1245+
int64_t sourceRank = shapeCastOp.getSourceVectorType().getRank();
1246+
int64_t resultRank = shapeCastOp.getResultVectorType().getRank();
1247+
if (sourceRank < resultRank && !sourceLayout.isSliceOf(resultLayout))
1248+
return rewriter.notifyMatchFailure(
1249+
warpOp, "shape_cast is rank reducing but source layout is not a "
1250+
"slice of result layout");
1251+
if (sourceRank > resultRank && !resultLayout.isSliceOf(sourceLayout))
12381252
return rewriter.notifyMatchFailure(
1239-
warpOp, "the source of shape_cast op lacks distribution layout");
1253+
warpOp, "shape_cast is rank increasing but result layout is not a "
1254+
"slice of source layout");
1255+
12401256
FailureOr<VectorType> sourceDistTypeOrFailure =
12411257
getDistVecTypeBasedOnLaneLayout(sourceLayout,
12421258
shapeCastOp.getSourceVectorType());
@@ -1349,10 +1365,8 @@ void XeGPUSubgroupDistributePass::runOnOperation() {
13491365
return AffineMap::getMultiDimMapWithTargets(
13501366
vecRank, {static_cast<unsigned int>(vecRank - 1)}, val.getContext());
13511367
SmallVector<unsigned int> distributedDims;
1352-
// Get the distributed dimensions based on the layout.
1353-
SmallVector<int64_t> laneLayout = xegpu::computeEffectiveLaneLayout(layout);
1354-
for (unsigned i = 0; i < laneLayout.size(); ++i) {
1355-
if (laneLayout[i] > 1)
1368+
for (auto [i, v] : llvm::enumerate(layout.getEffectiveLaneLayoutAsInt())) {
1369+
if (v > 1)
13561370
distributedDims.push_back(i);
13571371
}
13581372
return AffineMap::getMultiDimMapWithTargets(vecRank, distributedDims,

‎mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,9 +52,9 @@ getSgShapeAndCount(ArrayRef<int64_t> shape,
5252
int count = 1;
5353
SmallVector<int64_t> sgShape(shape);
5454
if (layout && layout.isForWorkgroup()) {
55-
SmallVector<int64_t> sgLayout = layout.getSgLayoutAsInt();
56-
if (!layout.getSgDataAsInt().empty())
57-
sgShape = layout.getSgDataAsInt();
55+
SmallVector<int64_t> sgLayout = layout.getEffectiveSgLayoutAsInt();
56+
if (!layout.getEffectiveSgDataAsInt().empty())
57+
sgShape = layout.getEffectiveSgDataAsInt();
5858
else if (auto maybeDerivedSgData = computeShapeRatio(shape, sgLayout))
5959
sgShape = *maybeDerivedSgData;
6060
SmallVector<int64_t> distUnit = computeElementwiseMul(sgLayout, sgShape);
@@ -488,7 +488,7 @@ struct WgToSgVectorBroadcastOp
488488
VectorType::get(sgShape, resultType.getElementType());
489489

490490
// Check if the output layout is distributable
491-
SmallVector<int64_t> sgLayout = layout.getSgLayoutAsInt();
491+
SmallVector<int64_t> sgLayout = layout.getEffectiveSgLayoutAsInt();
492492
if (sgLayout.empty())
493493
return failure();
494494

0 commit comments

Comments
 (0)