@@ -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