Skip to content

Commit 43f486a

Browse files
committed
[XPU][CVTToLLVM] Adapt CVT to LLVM converstion
After relanding 15c5e55, update our pass to adapt to the introduced changes. Signed-off-by: victor-eds <[email protected]>
1 parent 1c3c3c3 commit 43f486a

File tree

1 file changed

+13
-8
lines changed

1 file changed

+13
-8
lines changed

third_party/intel/lib/TritonIntelGPUToLLVM/ConvertLayoutOpToLLVM.cpp

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -488,17 +488,20 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
488488
assert(to_vector(conversion->getInDimNames()) ==
489489
to_vector(conversion->getOutDimNames()));
490490
auto dims = conversion->getInDimNames();
491-
if (llvm::is_contained(dims, str_attr("block"))) {
491+
if (llvm::is_contained(dims, kBlock)) {
492492
// Case 1: Transfer between values in different CTAs.
493493
// This requires moving values through distributed shared memory.
494494
return rewriter.notifyMatchFailure(
495495
op, "NYI: Transfer between different CTAs");
496-
} else if (llvm::is_contained(dims, str_attr("warp"))) {
497-
return rewriter.notifyMatchFailure(
498-
op, "NYI: Transfer between different warps");
499-
} else if (llvm::is_contained(dims, str_attr("lane"))) {
496+
} else if (llvm::is_contained(dims, kWarp)) {
500497
// Case 2: Transfer between values in the same CTA, in which case we move
501498
// values through shared memory.
499+
// TODO: Implement
500+
return failure();
501+
} else if (llvm::is_contained(dims, kLane)) {
502+
// Case 3. Transfer between values in the same warp, in which case we try
503+
// to move values using warp shuffles, though if the pattern is
504+
// complicated enough we may fall back to using shared memory
502505
// If the operation is a supported sub-group shuffle, perform via shuffle
503506
// operations.
504507
if (intel::cvtIsSubGroupShuffle(srcTy, dstTy)) {
@@ -513,15 +516,17 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
513516
}
514517
// TODO(jlebar): Implement me.
515518
return failure();
516-
} else if (llvm::is_contained(dims, str_attr("register"))) {
519+
} else if (llvm::is_contained(dims, kRegister) ||
520+
dstLayout.getInDimSize(kRegister) !=
521+
srcLayout.getInDimSize(kRegister)) {
517522
// Case 4. Transfer between values in the same thread, in which case we
518523
// simply reorder the elements of adaptor.getSrc().
519524
return transferWithinThread(
520525
op, dstLayout.getFreeVariableMasks()[kRegister],
521526
dstLayout.getInDimSize(kRegister), *conversion, adaptor, rewriter);
522527
} else {
523-
// The two layouts are equivalent. We should probably remove these in
524-
// RemoveLayoutConversion.
528+
// Cast 5. The two layouts are equivalent. We should probably remove
529+
// these in RemoveLayoutConversion.
525530
rewriter.replaceOp(op, adaptor.getSrc());
526531
return success();
527532
}

0 commit comments

Comments
 (0)