Fixed merge conflict

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

 };
 
 /// 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 *subViewOp, 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;
+  }
 };
 
 /// Split Reduction options.
 struct SplitReductionOptions {
   // Ratio used to split the reduction dimension.  If the ratio is <= 1,
   // nothing will be done.
   int64_t ratio = 0;
   // Index where the extra dimension is added to the intermediate tensor

 /// namedOp.
 /// Return failure if `namedOp` is a GenericOp or misses a region builder.
 FailureOr<GenericOp> generalizeNamedOp(RewriterBase &rewriter,
                                        LinalgOp namedOp);
 
 /// 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.
-std::optional<Value> allocateWorkgroupMemory(OpBuilder &builder,
-                                             memref::SubViewOp subview,
-                                             ArrayRef<Value> sizeBounds,
-                                             DataLayout &);
+Optional<Value> allocateWorkgroupMemory(OpBuilder &builder, Operation *subview,
 
 /// 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.
-std::optional<Value> allocateGPUPrivateMemory(OpBuilder &builder,
-                                              memref::SubViewOp subview,
-                                              ArrayRef<Value> sizeBounds,
-                                              DataLayout &);
+Optional<Value> allocateGPUPrivateMemory(OpBuilder &builder, Operation *subview,
 
 /// 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.
 LogicalResult deallocateGPUPrivateMemory(OpBuilder &, Value /*buffer*/);
 

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);
 
   IRRewriter 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) {
+  auto viewType = llvm::TypeSwitch<Operation *, ShapedType>(op)
+                      .Case<memref::SubViewOp, memref::TransposeOp>(
+                          [](auto casted) { return casted.getType(); })
+                      .Default([](Operation *) { return nullptr; });
+  if (!viewType)
+    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 std::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);
+  auto subview = llvm::TypeSwitch<Operation *, MemRefType>(op)
+                     .Case<memref::SubViewOp, memref::TransposeOp>(
+                         [](auto casted) { return casted.getType(); })
+                     .Default([](Operation *) { return nullptr; });
+
+  auto funcOp =
+      llvm::TypeSwitch<Operation *, func::FuncOp>(op)
+          .Case<memref::SubViewOp, memref::TransposeOp>([](auto casted) {
+            return casted->template getParentOfType<func::FuncOp>();
+          })
+          .Default([](Operation *) { return nullptr; });
 
-  func::FuncOp funcOp = subview->getParentOfType<func::FuncOp>();
-  if (!funcOp)
+  if (!funcOp || !subview)
     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.
-std::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.
-std::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);

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-LABEL: 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 {{.*}}
+// CHECK:         %[[subview_B:.*]] = memref.subview {{.*}}
+// CHECK:         %[[subview_C:.*]] = memref.subview {{.*}}
+
+// 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>]}
+}