[mlir][Transform][LinAlg] Add copy permutation for GPU memory hierarchy in the `transform.promote` op

[tavakkoliamirmohammad]

Summary

In this patch, we are adding a copy permutation to the transform.promote op. This option is available when using the new mapping attribute to the transform.promote op. This allows changing the logical view of the data (memref.subview) using memref.transpose without moving the data.

If the copy permutation is given by the user, the result operation will be the generated operation linalg.generic otherwise the original linalg target operation.

[tavakkoliamirmohammad]

Fix merge conflict

[ftynse]

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

87–88

isa followed by cast (let alone dyn_cast) is an anti-pattern. dyn_cast itself does the type check and returns null on failure, so something like if (auto subview = dyn_cast<memref::SubViewOp>(op)) should be used instead. If you want to be fancy, you can do

auto viewType = llvm::TypeSwitch<ShapedType>(op)
  .Case<SubViewOp, TransposeOp>([](auto casted) { return casted.getType(); })
  .Default([](Operation *) { return nullptr; });
if (!viewType) return std::nullopt;

mlir/test/Dialect/Linalg/promote.mlir

287

CHECK-LABEL

299

Do we really need to pattern-match exact types here? We strive to have minimal tests that exercise the relevant behavior.

327

Please add the newline.

[tavakkoliamirmohammad]

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

87–88

Thank you for suggesting this elegant way!

[tavakkoliamirmohammad]

Update the diff based on the comment

[ftynse]

LGTM when the last comment is addressed, but I'd like an opinion from @nicolasvasilache on the promotion logic.

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

478–484

Same as the code above, isa followed by dyn_cast is an anti-pattern.

[tavakkoliamirmohammad]

Fixed anti-pattern and return condition of allocateSubviewGPUMemoryInAddressSpace function

[tavakkoliamirmohammad]

Hi @nicolasvasilache, just wanted to check in and see if you had a chance to review the promotion logic yet. If you have any thoughts or feedback, it would be great to hear them. Thanks!

[tavakkoliamirmohammad]

Thank you! I don't have the commit access to LLVM project for applying this patch.

[ftynse]

Could you please rebase?

[tavakkoliamirmohammad]

rebased

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let description = [{
     was modified inplace.
   }];
 
   let arguments = (ins PDL_Operation:$target,
                        DefaultValuedAttr<I64ArrayAttr, "{}">:$operands_to_promote,
                        DefaultValuedAttr<BoolArrayAttr, "{}">:$use_full_tile_buffers,
                        UnitAttr:$use_full_tiles_by_default,
                        UnitAttr:$use_alloca,
-                       OptionalAttr<DeviceMappingArrayAttr>:$mapping,
+                       ConfinedAttr<OptionalAttr<DeviceMappingArrayAttr>, [ArrayMaxCount<1>]>:$mapping,
+                       OptionalAttr<DenseI64ArrayAttr>:$copy_permutation,
                        OptionalAttr<I64Attr>:$alignment);
   let results = (outs PDL_Operation:$transformed);
 
   let assemblyFormat = "$target attr-dict";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

FailureOr<GenericOp> generalizeNamedOp(RewriterBase &rewriter,
                                        LinalgOp namedOp);
 
 /// Callback function type used to perform the allocation for the promoted
 /// `subView`. In `boundingSubViewsize` a best attempt is made to find the
 /// smallest constant value for the size of the buffer needed for each
 /// dimension. If that is not possible, contains the dynamic size of the
 /// subview. The call back should return the buffer to use.
 using AllocBufferCallbackFn = std::function<std::optional<Value>(
-    OpBuilder &b, memref::SubViewOp subView,
-    ArrayRef<Value> boundingSubViewSize, DataLayout &layout)>;
+    OpBuilder &b, Operation *op, ArrayRef<Value> boundingSubViewSize,
+    DataLayout &layout)>;
 
 /// Callback function type used to deallocate the buffers used to hold the
 /// promoted subview.
 using DeallocBufferCallbackFn =
     std::function<LogicalResult(OpBuilder &b, Value buffer)>;
 
 /// Callback function type used to insert copy from original subview to
 /// subview of the promoted region for the read operands/subview of promoted
struct LinalgPromotionOptions {
   std::optional<CopyCallbackFn> copyInFn;
   std::optional<CopyCallbackFn> copyOutFn;
   LinalgPromotionOptions &setCopyInOutFns(CopyCallbackFn const &copyIn,
                                           CopyCallbackFn const &copyOut) {
     copyInFn = copyIn;
     copyOutFn = copyOut;
     return *this;
   }
+  ArrayRef<int64_t> copyPermutation;
+  LinalgPromotionOptions &
+  setCopyPermutation(ArrayRef<int64_t> permutationArrayRef) {
+    copyPermutation = permutationArrayRef;
+    return *this;
+  }
 };
 
 /// Create a new buffer using the `allocationFn` provided. The size of this
 /// buffer is the smallest constant bounding size along each dimension that
 /// can be computed for the size of the result of `subView`. Returns the
-/// allocated buffer as `fullLocalView` and the view that matches the size of
-/// the result of subview operation as `partialLocalView`.
+/// allocated buffer as `fullLocalView`, the view that matches the size of
+/// the result of subview operation as `partialLocalView`, and the input or
+/// resulted operation from the promotion as`referenceOp` that can be used
+/// subsequent steps.
 struct PromotionInfo {
   Value fullLocalView;
   Value partialLocalView;
+  Value referenceOp;
 };
 FailureOr<PromotionInfo>
 promoteSubviewAsNewBuffer(OpBuilder &b, Location loc, memref::SubViewOp subView,
                           const AllocBufferCallbackFn &allocationFn,
-                          DataLayout &layout);
+                          DataLayout &layout, ArrayRef<int64_t> permutation);
 
 /// Promote the `subViews` into a new buffer allocated at the insertion point
 /// `b`. Promotion occurs in 3 steps:
 ///   1. Create a new buffer for a full tile (i.e. not clipped at the
 ///   boundary).
 ///   2. Take a full view on the buffer.
 ///   3. Take a partial slice of the full view in step 2. and copy into it.
 ///
 /// Return the modified linalg op (the modification happens in place) as well
 /// as all the copy ops created.
 FailureOr<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
                                     const LinalgPromotionOptions &options);
 
 /// Allocate the subview in the GPU workgroup memory.
-Optional<Value> allocateWorkgroupMemory(OpBuilder &builder,
-                                        memref::SubViewOp subview,
+Optional<Value> allocateWorkgroupMemory(OpBuilder &builder, Operation *subview,
                                         ArrayRef<Value> sizeBounds,
                                         DataLayout &);
 
 /// In case of GPU group memory there is no need to deallocate.
 LogicalResult deallocateWorkgroupMemory(OpBuilder &, Value /*buffer*/);
 
 /// Create Memref copy operations and add gpu barrier guards before and after
 /// the copy operation to ensure data integrity.
 LogicalResult copyToWorkgroupMemory(OpBuilder &b, Value src, Value dst);
 
 /// Allocate the subview in the GPU private memory.
-Optional<Value> allocateGPUPrivateMemory(OpBuilder &builder,
-                                         memref::SubViewOp subview,
+Optional<Value> allocateGPUPrivateMemory(OpBuilder &builder, Operation *subview,
                                          ArrayRef<Value> sizeBounds,
                                          DataLayout &);
 
 /// Normal copy to between src and dst.
 LogicalResult copyToGPUPrivateMemory(OpBuilder &b, Value src, Value dst);
 
 /// In case of GPU private memory there is no need to deallocate since the
 /// memory is freed when going outside of the scope.

mlir/include/mlir/IR/OpBase.td

 def IntPositive : AttrConstraint<
     CPred<"$_self.cast<::mlir::IntegerAttr>().getValue().isStrictlyPositive()">,
     "whose value is positive">;
 
 class ArrayMinCount<int n> : AttrConstraint<
     CPred<"$_self.cast<::mlir::ArrayAttr>().size() >= " # n>,
     "with at least " # n # " elements">;
 
+class ArrayMaxCount<int n> : AttrConstraint<
+    CPred<"$_self.cast<::mlir::ArrayAttr>().size() <= " # n>,
+    "with at most " # n # " elements">;
+
 class ArrayCount<int n> : AttrConstraint<
     CPred<"$_self.cast<::mlir::ArrayAttr>().size() == " #n>,
     "with exactly " # n # " elements">;
 
 class DenseArrayCount<int n> : AttrConstraint<
     CPred<"$_self.cast<::mlir::DenseArrayAttr>().size() == " #n>,
     "with exactly " # n # " elements">;
 

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

if (addressSpace.getAddressSpace() ==
           promotionOptions
               .setAllocationDeallocationFns(allocateGPUPrivateMemory,
                                             deallocateGPUPrivateMemory)
               .setCopyInOutFns(copyToGPUPrivateMemory, copyToGPUPrivateMemory)
               .setUseFullTileBuffers({false, false});
     } else {
       return emitDefaultDefiniteFailure(target);
     }
+    if (getCopyPermutation().has_value()) {
+      promotionOptions.setCopyPermutation(*getCopyPermutation());
+    }
   }
 
   if (failed(promoteSubviewsPrecondition(target, promotionOptions)))
     return emitDefaultDefiniteFailure(target);
 
   TrivialPatternRewriter rewriter(target->getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<LinalgOp> res = promoteSubViews(rewriter, target, promotionOptions);
   if (failed(res))
     return emitDefaultDefiniteFailure(target);
-  results.push_back(target);
+  results.push_back(*res);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // ReplaceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure

mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp

static Value allocBuffer(ImplicitLocOpBuilder &b,
   return b.create<memref::AllocOp>(dynamicBufferType, mul, alignmentAttr);
 }
 
 /// Default allocation callback function. This allocates a promoted buffer when
 /// no call back to do so is provided. The default is to allocate a
 /// memref<..xi8> and return a view to get a memref type of shape
 /// boundingSubViewSize.
 static std::optional<Value> defaultAllocBufferCallBack(
-    const LinalgPromotionOptions &options, OpBuilder &builder,
-    memref::SubViewOp subView, ArrayRef<Value> boundingSubViewSize,
-    std::optional<unsigned> alignment, DataLayout &layout) {
-  ShapedType viewType = subView.getType();
-  ImplicitLocOpBuilder b(subView.getLoc(), builder);
+    const LinalgPromotionOptions &options, OpBuilder &builder, Operation *op,
+    ArrayRef<Value> boundingSubViewSize, std::optional<unsigned> alignment,
+    DataLayout &layout) {
+  ShapedType viewType;
+  if (isa<memref::SubViewOp>(op)) {
+    viewType = dyn_cast<memref::SubViewOp>(op).getType();
+  } else if (isa<memref::TransposeOp>(op)) {
+    viewType = dyn_cast<memref::TransposeOp>(op).getType();
+  } else {
+    return std::nullopt;
+  }
+  ImplicitLocOpBuilder b(op->getLoc(), builder);
   auto zero = b.createOrFold<arith::ConstantIndexOp>(0);
   auto one = b.createOrFold<arith::ConstantIndexOp>(1);
 
   Value allocSize = one;
   for (const auto &size : llvm::enumerate(boundingSubViewSize))
     allocSize = b.createOrFold<arith::MulIOp>(allocSize, size.value());
   Value buffer = allocBuffer(b, options, viewType.getElementType(), allocSize,
                              layout, alignment);
struct LinalgOpInstancePromotionOptions {
   /// well as to copy the data into and out of these buffers.
   AllocBufferCallbackFn allocationFn;
   DeallocBufferCallbackFn deallocationFn;
   CopyCallbackFn copyInFn;
   CopyCallbackFn copyOutFn;
 
   /// Alignment of promoted buffer.
   std::optional<unsigned> alignment;
+  ArrayRef<int64_t> copyPermutation;
 };
 } // namespace
 
 LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
     LinalgOp linalgOp, const LinalgPromotionOptions &options)
     : subViews(), alignment(options.alignment) {
   assert(linalgOp.hasBufferSemantics() && "revisit usage of shaped operand");
   auto vUseFullTileBuffers =
if (auto sv = dyn_cast_or_null<memref::SubViewOp>(op)) {
       subViews[operandNumber] = sv;
       useFullTileBuffers[sv] = vUseFullTileBuffers[operandNumber];
     }
   }
 
   if (options.allocationFn) {
     allocationFn = *options.allocationFn;
   } else {
-    allocationFn = [&](OpBuilder &b, memref::SubViewOp subViewOp,
+    allocationFn = [&](OpBuilder &b, Operation *subViewOp,
                        ArrayRef<Value> boundingSubViewSize,
                        DataLayout &layout) -> std::optional<Value> {
       return defaultAllocBufferCallBack(options, b, subViewOp,
                                         boundingSubViewSize, alignment, layout);
     };
   }
 
   if (options.deallocationFn) {
     deallocationFn = *options.deallocationFn;
   } else {
     deallocationFn = [&](OpBuilder &b, Value buffer) {
       return defaultDeallocBufferCallBack(options, b, buffer);
     };
   }
 
   // Save the loc because `linalgOp` goes out of scope.
   Location loc = linalgOp.getLoc();
   auto defaultCopyCallBack = [loc](OpBuilder &b, Value src,
                                    Value dst) -> LogicalResult {
     b.create<memref::CopyOp>(loc, src, dst);
     return success();
   };
   copyInFn = (options.copyInFn ? *(options.copyInFn) : defaultCopyCallBack);
   copyOutFn = (options.copyOutFn ? *(options.copyOutFn) : defaultCopyCallBack);
+  copyPermutation = options.copyPermutation;
 }
 
 // Performs promotion of a `subView` into a local buffer of the size of the
 // *ranges* of the `subView`. This produces a buffer whose size may be bigger
 // than the actual size of the `subView` at the boundaries.
 // This is related to the full/partial tile problem.
 // Returns a PromotionInfo containing a `buffer`, `fullLocalView` and
-// `partialLocalView` such that:
+// `partialLocalView` and `referenceOp` such that:
 //   * `buffer` is always the size of the full tile.
 //   * `fullLocalView` is a dense contiguous view into that buffer.
 //   * `partialLocalView` is a dense non-contiguous slice of `fullLocalView`
 //     that corresponds to the size of `subView` and accounting for boundary
 //     effects.
+//   * `referenceOp` which can be the original subviewOp or created
+//   memref:TransposeOp
+//
 // The point of the full tile buffer is that constant static tile sizes are
 // folded and result in a buffer type with statically known size and alignment
 // properties.
 // To account for general boundary effects, padding must be performed on the
 // boundary tiles. For now this is done with an unconditional `fill` op followed
 // by a partial `copy` op.
 FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
     OpBuilder &b, Location loc, memref::SubViewOp subView,
-    const AllocBufferCallbackFn &allocationFn, DataLayout &layout) {
+    const AllocBufferCallbackFn &allocationFn, DataLayout &layout,
+    ArrayRef<int64_t> permutation) {
   auto viewType = subView.getType();
   auto rank = viewType.getRank();
-  SmallVector<Value, 4> fullSizes;
+  SmallVector<Value> fullSizes;
   SmallVector<OpFoldResult> partialSizes;
   fullSizes.reserve(rank);
   partialSizes.reserve(rank);
+  // Get identity map.
+  AffineMap permutationMap =
+      AffineMap::getMultiDimIdentityMap(rank, b.getContext());
+  // If permutation is given update the permutation map.
+  if (permutation.size() > 0) {
+    permutationMap = AffineMap::getPermutationMap(
+        llvm::to_vector_of<unsigned>(permutation), b.getContext());
+  }
   llvm::SmallBitVector droppedDims = subView.getDroppedDims();
   int64_t resultDimIdx = 0;
   for (const auto &en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) {
     if (droppedDims[en.index()])
       continue;
     auto rangeValue = en.value();
     // Try to extract a tight constant. If the size is known statically, no need
     // to look for the bound.
     LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n");
-    Value size;
-    if (auto attr = rangeValue.size.dyn_cast<Attribute>()) {
-      size = getValueOrCreateConstantIndexOp(b, loc, rangeValue.size);
-    } else {
-      Value materializedSize =
-          getValueOrCreateConstantIndexOp(b, loc, rangeValue.size);
-      FailureOr<int64_t> upperBound =
-          getConstantUpperBoundForIndex(materializedSize);
+    Value size = getValueOrCreateConstantIndexOp(b, loc, rangeValue.size);
+    auto attr = rangeValue.size.dyn_cast<Attribute>();
+    if (!attr) {
+      FailureOr<int64_t> upperBound = getConstantUpperBoundForIndex(size);
       size = failed(upperBound)
-                 ? materializedSize
+                 ? size
                  : b.create<arith::ConstantIndexOp>(loc, *upperBound);
     }
     LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n");
     fullSizes.push_back(size);
     partialSizes.push_back(
         b.createOrFold<memref::DimOp>(loc, subView, resultDimIdx++));
   }
-  SmallVector<int64_t, 4> dynSizes(fullSizes.size(), ShapedType::kDynamic);
   // If a callback is not specified, then use the default implementation for
   // allocating the promoted buffer.
-  std::optional<Value> fullLocalView =
-      allocationFn(b, subView, fullSizes, layout);
+  std::optional<Value> fullLocalView;
+  Value referenceOp = subView;
+  if (permutationMap.isIdentity()) {
+    fullLocalView = allocationFn(b, subView, fullSizes, layout);
+  } else {
+    auto transposeOp = b.create<memref::TransposeOp>(
+        loc, subView, AffineMapAttr::get(permutationMap));
+    referenceOp = transposeOp;
+    fullLocalView = allocationFn(
+        b, transposeOp,
+        applyPermutationMap(permutationMap, ArrayRef(fullSizes)), layout);
+  }
+
   if (!fullLocalView)
     return failure();
   SmallVector<OpFoldResult, 4> zeros(fullSizes.size(), b.getIndexAttr(0));
   SmallVector<OpFoldResult, 4> ones(fullSizes.size(), b.getIndexAttr(1));
   auto partialLocalView = b.createOrFold<memref::SubViewOp>(
-      loc, *fullLocalView, zeros, partialSizes, ones);
-  return PromotionInfo{*fullLocalView, partialLocalView};
+      loc, *fullLocalView, zeros,
+      applyPermutationMap(permutationMap, ArrayRef(partialSizes)), ones);
+  return PromotionInfo{*fullLocalView, partialLocalView, referenceOp};
 }
 
 static FailureOr<MapVector<int64_t, PromotionInfo>>
 promoteSubViews(ImplicitLocOpBuilder &b,
                 LinalgOpInstancePromotionOptions options, DataLayout &layout) {
   if (options.subViews.empty())
     return failure();
 
   MapVector<int64_t, PromotionInfo> promotionInfoMap;
 
   for (auto v : options.subViews) {
     memref::SubViewOp subView =
         cast<memref::SubViewOp>(v.second.getDefiningOp());
-    auto promotionInfo = promoteSubviewAsNewBuffer(
-        b, b.getLoc(), subView, options.allocationFn, layout);
+    auto promotionInfo =
+        promoteSubviewAsNewBuffer(b, b.getLoc(), subView, options.allocationFn,
+                                  layout, options.copyPermutation);
     if (failed(promotionInfo))
       return failure();
     promotionInfoMap[v.first] = *promotionInfo;
 
     // Only fill the buffer if the full local view is used
     if (!options.useFullTileBuffers[v.second])
       continue;
     Type subviewEltType = subView.getType().getElementType();
for (auto v : options.subViews) {
     b.create<linalg::FillOp>(fillVal, promotionInfo->fullLocalView);
   }
 
   // Copy data into the promoted buffers. Use callback if provided.
   for (auto v : options.subViews) {
     auto info = promotionInfoMap.find(v.first);
     if (info == promotionInfoMap.end())
       continue;
-    if (failed(options.copyInFn(
-            b, cast<memref::SubViewOp>(v.second.getDefiningOp()),
+    if (failed(options.copyInFn(b, info->second.referenceOp,
                                 info->second.partialLocalView)))
       return failure();
   }
   return promotionInfoMap;
 }
 
 static FailureOr<LinalgOp>
 promoteSubViews(ImplicitLocOpBuilder &b, LinalgOp op,
                 LinalgOpInstancePromotionOptions options, DataLayout &layout) {
   assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
 
   // 1. Promote the specified views and use them in the new op.
   auto promotedBuffersAndViews = promoteSubViews(b, options, layout);
   if (failed(promotedBuffersAndViews) ||
       promotedBuffersAndViews->size() != options.subViews.size())
     return failure();
 
   // 2. Append all other operands as they appear, this enforces that such
   // operands are not views. This is to support cases such as FillOp taking
   // extra scalars etc.  Keep a reference to output buffers;
   SmallVector<Value, 8> opViews;
   opViews.reserve(op->getNumOperands());
   SmallVector<std::pair<Value, Value>, 8> writebackViews;
   writebackViews.reserve(promotedBuffersAndViews->size());
+  SmallVector<AffineMap> indexingMaps = op.getIndexingMapsArray();
   for (OpOperand &opOperand : op->getOpOperands()) {
     int64_t operandNumber = opOperand.getOperandNumber();
     if (options.subViews.count(operandNumber) != 0) {
       if (options.useFullTileBuffers[opOperand.get()])
         opViews.push_back(
             (*promotedBuffersAndViews)[operandNumber].fullLocalView);
       else
         opViews.push_back(
             (*promotedBuffersAndViews)[operandNumber].partialLocalView);
       if (operandNumber >= op.getNumDpsInputs())
         writebackViews.emplace_back(std::make_pair(
-            opOperand.get(),
+            (*promotedBuffersAndViews)[operandNumber].referenceOp,
             (*promotedBuffersAndViews)[operandNumber].partialLocalView));
+
+      // 2.1 Get the identity affine map.
+      AffineMap permutationMap = AffineMap::getMultiDimIdentityMap(
+          opOperand.get().getType().cast<MemRefType>().getRank(),
+          b.getContext());
+
+      // 2.2 Apply the given copy permutation to the original indexing map.
+      if (options.copyPermutation.size() > 0) {
+        permutationMap = AffineMap::getPermutationMap(
+            llvm::to_vector_of<unsigned>(options.copyPermutation),
+            b.getContext());
+      }
+      AffineMap transposedMap =
+          permutationMap.compose(op.getMatchingIndexingMap(&opOperand));
+      indexingMaps[op.getIndexingMapIndex(&opOperand)] = transposedMap;
     } else {
       opViews.push_back(opOperand.get());
     }
   }
   op->setOperands(0, opViews.size(), opViews);
+  linalg::GenericOp transposedGenericOp;
+  // 2.3 If the copy permutation is given replace the current Linalg op with a
+  // linalg.generic.
+  if (options.copyPermutation.size() > 0) {
+    ValueRange operandsRef(op->getOperands());
+    transposedGenericOp = b.create<linalg::GenericOp>(
+        /*location=*/op->getLoc(),
+        /*inputs=*/operandsRef.take_front(op.getNumDpsInputs()),
+        /*outputs=*/operandsRef.drop_front(op.getNumDpsInputs()),
+        /*indexingMaps=*/indexingMaps,
+        /*iteratorTypes=*/op.getIteratorTypesArray());
+    transposedGenericOp.getRegion().takeBody(op->getRegion(0));
+  }
 
   OpBuilder::InsertionGuard guard(b);
   b.setInsertionPointAfter(op);
   // 3. Emit write-back for the promoted output views: copy the partial view.
   for (auto viewAndPartialLocalView : writebackViews) {
     if (failed(options.copyOutFn(b, viewAndPartialLocalView.second,
                                  viewAndPartialLocalView.first)))
       return failure();
   }
 
   // 4. Dealloc all local buffers.
   for (const auto &pi : *promotedBuffersAndViews)
     (void)options.deallocationFn(b, pi.second.fullLocalView);
+
+  // 5. If the copy permutation is given replace the current Linalg op with
+  // created linalg.generic.
+  if (options.copyPermutation.size() > 0) {
+    IRRewriter rewriter(b);
+    rewriter.replaceOp(op, transposedGenericOp->getResults());
+    return cast<LinalgOp>(*transposedGenericOp);
+  }
   return op;
 }
 
 LogicalResult
 mlir::linalg::promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options) {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   // Transformation applies to buffers only.
mlir::linalg::promoteSubViews(OpBuilder &builder, LinalgOp linalgOp,
   if (failed(res))
     return failure();
   return res;
 }
 
 /// Allocate the given subview to a memory address space in GPU by creating a
 /// allocation operation and setting the memref type address space to desired
 /// address space.
-static Optional<Value> allocateSubviewGPUMemoryInAddressSpace(
-    OpBuilder &builder, memref::SubViewOp subview, ArrayRef<Value> sizeBounds,
+static Optional<Value>
+allocateSubviewGPUMemoryInAddressSpace(OpBuilder &builder, Operation *op,
+                                       ArrayRef<Value> sizeBounds,
                                        gpu::AddressSpace addressSpace) {
   OpBuilder::InsertionGuard guard(builder);
+  func::FuncOp funcOp;
+  MemRefType subview;
+  if (isa<memref::SubViewOp>(op)) {
+    subview = dyn_cast<memref::SubViewOp>(op).getType();
+    funcOp = dyn_cast<memref::SubViewOp>(op)->getParentOfType<func::FuncOp>();
+  } else if (isa<memref::TransposeOp>(op)) {
+    subview = dyn_cast<memref::TransposeOp>(op).getType();
+    funcOp = dyn_cast<memref::TransposeOp>(op)->getParentOfType<func::FuncOp>();
+  } else {
+    return std::nullopt;
+  }
 
-  func::FuncOp funcOp = subview->getParentOfType<func::FuncOp>();
   if (!funcOp)
     return std::nullopt;
 
   // The subview size bounds are expected to be constant; they specify the shape
   // of the allocation.
   SmallVector<int64_t> shape;
   for (Value bound : sizeBounds) {
     APInt value;
     if (!matchPattern(bound, m_ConstantInt(&value)))
       return std::nullopt;
     shape.push_back(value.getSExtValue());
   }
 
   builder.setInsertionPoint(&funcOp.front(), funcOp.front().begin());
   auto type = MemRefType::get(
-      shape, subview.getType().getElementType(), MemRefLayoutAttrInterface{},
+      shape, subview.getElementType(), MemRefLayoutAttrInterface{},
       gpu::AddressSpaceAttr::get(builder.getContext(), addressSpace));
   Value buffer;
   if (addressSpace == gpu::GPUDialect::getWorkgroupAddressSpace()) {
     buffer = builder.create<memref::AllocOp>(funcOp.getLoc(), type);
   } else if (addressSpace == gpu::GPUDialect::getPrivateAddressSpace()) {
     buffer = builder.create<memref::AllocaOp>(funcOp.getLoc(), type);
   } else {
     return std::nullopt;
   }
   return buffer;
 }
 
 /// Allocate the subview in the GPU workgroup memory.
-Optional<Value> mlir::linalg::allocateWorkgroupMemory(
-    OpBuilder &builder, memref::SubViewOp subview, ArrayRef<Value> sizeBounds,
+Optional<Value>
+mlir::linalg::allocateWorkgroupMemory(OpBuilder &builder, Operation *subview,
+                                      ArrayRef<Value> sizeBounds,
                                       DataLayout &) {
   return allocateSubviewGPUMemoryInAddressSpace(
       builder, subview, sizeBounds,
       gpu::GPUDialect::getWorkgroupAddressSpace());
 }
 
 /// In case of GPU group memory there is no need to deallocate.
 LogicalResult mlir::linalg::deallocateWorkgroupMemory(OpBuilder &,
                                                       Value /*buffer*/) {
   return success();
 }
 
 /// Create Memref copy operations and add gpu barrier guards before and after
 /// the copy operation to ensure data integrity.
 LogicalResult mlir::linalg::copyToWorkgroupMemory(OpBuilder &b, Value src,
                                                   Value dst) {
   b.create<gpu::BarrierOp>(src.getLoc());
   Operation *copyOp = b.create<memref::CopyOp>(src.getLoc(), src, dst);
   b.create<gpu::BarrierOp>(copyOp->getLoc());
   return success();
 }
 
 /// Allocate the subview in the GPU private memory.
-Optional<Value> mlir::linalg::allocateGPUPrivateMemory(
-    OpBuilder &builder, memref::SubViewOp subview, ArrayRef<Value> sizeBounds,
+Optional<Value>
+mlir::linalg::allocateGPUPrivateMemory(OpBuilder &builder, Operation *subview,
+                                       ArrayRef<Value> sizeBounds,
                                        DataLayout &) {
   return allocateSubviewGPUMemoryInAddressSpace(
       builder, subview, sizeBounds, gpu::GPUDialect::getPrivateAddressSpace());
 }
 
 /// Normal copy to between src and dst.
 LogicalResult mlir::linalg::copyToGPUPrivateMemory(OpBuilder &b, Value src,
                                                    Value dst) {
   b.create<memref::CopyOp>(src.getLoc(), src, dst);
   return success();
 }
 
 /// In case of GPU private memory there is no need to deallocate since the
 /// memory is freed when going outside of the scope.
 LogicalResult mlir::linalg::deallocateGPUPrivateMemory(OpBuilder &,
                                                        Value /*buffer*/) {
   return success();
 }
 No newline at end of file

mlir/test/Dialect/Linalg/promote.mlir

func.func @promote_rank_reducing_subviews(%arg0:  memref<?x?x?x64xf32, strided<[?, ?, ?, ?], offset: ?>>, %arg1: memref<128x3x3x64xf32, strided<[?, ?, ?, ?], offset: ?>>, %arg2: memref<?x?x?x128xf32>,
   return
 }
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   %0 = transform.structured.match interface{LinalgOp} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = transform.structured.promote %0
 }
+
+// -----
+func.func @gemm_transposed(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
+{
+   linalg.matmul ins(%a, %b: memref<?x?xf32>, memref<?x?xf32>)
+               outs(%c: memref<?x?xf32>)
+   return
+}
+
+//  CHECK: func @gemm_transposed
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32>
+// CHECK: %[[alloc_A:.*]] = memref.alloc() : memref<16x4xf32, #gpu.address_space<workgroup>>
+// CHECK: %[[alloc_B:.*]] = memref.alloc() : memref<8x4xf32, #gpu.address_space<workgroup>>
+// CHECK-DAG: %[[C16:.*]] = arith.constant 16
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1
+// CHECK:   scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK:     scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK:       scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK:         %[[subview_A:.*]] = memref.subview {{.*}} : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+// CHECK:         %[[subview_B:.*]] = memref.subview {{.*}} : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+// CHECK:         %[[subview_C:.*]] = memref.subview {{.*}} : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
+
+// CHECK:         %[[transposed_B:.*]] = memref.transpose %[[subview_B]] (d0, d1) -> (d1, d0) : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[1, ?], offset: ?>> 
+// CHECK:         %[[shared_B:.*]] = memref.subview %[[alloc_B]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<8x4xf32, #gpu.address_space<workgroup>> to memref<?x?xf32, strided<[4, 1]>, #gpu.address_space<workgroup>>
+
+// CHECK-NEXT:    gpu.barrier
+// CHECK-NEXT:    memref.copy %[[transposed_B]], %[[shared_B]] :  memref<?x?xf32, strided<[1, ?], offset: ?>> to memref<?x?xf32, strided<[4, 1]>, #gpu.address_space<workgroup>>
+// CHECK-NEXT:    gpu.barrier
+
+// CHECK:         %[[shared_A:.*]] = memref.subview %[[alloc_A]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<16x4xf32, #gpu.address_space<workgroup>> to memref<?x?xf32, strided<[4, 1]>, #gpu.address_space<workgroup>>
+
+
+// CHECK-NEXT:    gpu.barrier
+// CHECK-NEXT:    memref.copy %[[subview_A]], %[[shared_A]] :  memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[4, 1]>, #gpu.address_space<workgroup>>
+// CHECK-NEXT:    gpu.barrier
+
+// CHECK:         linalg.generic {{.+}} ins(%[[shared_A]], %[[shared_B]]{{.*}} outs(%[[subview_C]]
+
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1, %loops:3 = transform.structured.tile %0 [16, 8, 4] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+  %2 = transform.structured.promote %1 { operands_to_promote = [1], mapping = [#gpu.memory_space<workgroup>], copy_permutation = array<i64: 1, 0> }
+  %3 = transform.structured.promote %2 { operands_to_promote = [0], mapping = [#gpu.memory_space<workgroup>]}
+}
 No newline at end of file