[mlir][LinAlg][Transform][GPU] Add GPU memory hierarchy to the transform.promote op

[tavakkoliamirmohammad]

Summary

In this patch we are adding the support of copying a a memref.subview to the shared or private memory in GPU. The global to shared memory copy is adopted from codes implemented in IREE (https://github.com/iree-org/iree), but the private memory copy part has not been implemented in IREE. This patch enables transferring a subview from global->shared, global->private, and shared->private.

Our final aim is to provide a copy layout as an affine map to the transform.promote op to support transpose memory copy. This map is a permutation of the original affine index map. Although this has been implemented and user can copy data to arbitrary layout , this attempt is not included in this patch since we have still problem with linalg.generic operations to change their index map to the transformed index map. You can find more in following links (Initial attempt to support layout map in promote op in transform dialect) (Fix data transpose in shared memory)

[tavakkoliamirmohammad]

Format the file and fix wrong diff file not including the private memory copy

[ftynse]

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

1784

Functions that are only necessary in this translation unit should be declared static.

Top-level entities such as function must be documented.

I would also consider placing that somewhere in Linalg/Transform/ rather than in TransformOps.

1796

Nit: don't specify the number of stack elements for the vector.

1804

It's a good practice to use OpBuilder:InsertionGuard to restore the location before the function returns.

1821

We don't commit commented-out code.

1847

This function only differs from the one above by not using the shared address space here. This is something that can be generalized.

We also have an explicit address space for private memory on GPUs. Likely, it should be used here instead of the default address space.

1880–1894

This would do strange things if both attributes are set. Have you considered having a single attribute that specifies which kind of memory to use?

[tavakkoliamirmohammad]

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

1804

This is done in line 1788. Is there anything needed to reset the insertion point? since by destroying the OpBuilder::InsertionGuard guard the insertion point is restored

[tavakkoliamirmohammad]

Added a new device mapping enum for GPU memory hierarchy, removed comments, refactor allocation, deallocation, and copy function to linalg/Transform/promotion.cpp

[ftynse]

Nice, almost there, thanks you!

mlir/include/mlir/Dialect/GPU/TransformOps/GPUDeviceMappingAttr.td

99

Nit: we tend to prefer OpenCL-style terminology, so s/shared/workgroup (which is the actual name of the attribute). s/thread/workitem, s/block/workgroup.

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

396

Nit: top-level comments in MLIR should be prefixed with ///. Also, terminate sentences in documentation with a period.

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

1804

We don't see this line in the review. Please reupload the patch with more context git diff -U99999 or use Arcanist https://llvm.org/docs/Phabricator.html.

1808

Nit: MLIR uses camelCase for variable names.

1808–1812

Nit: any chance of adding a temporary to make this more readable.

... mapping = getMapping()->getValue()
auto address_space = mapping.take_front()...

I also don't understand why this needs access to raw .data()

1813–1815

It is less expensive to query the details of the attribute than to construct a new one, the latter needs to take a lock. Something like address_space.getAddressSpace() == gpu::GPUDialect::getWorkgroupAddressSpace() is preferable here.

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

406

Ultra-nit: extra whitespace between the function and its documentation.

[tavakkoliamirmohammad]

This comment was removed by tavakkoliamirmohammad.

[tavakkoliamirmohammad]

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

1808–1812

I wanted to get the first element of the mapping so I have to through the chain from mlir:ArrayAttr to llvm::ArrayRef to the actual data. I thought this be more readable but it seems the new approach is more readable.

[tavakkoliamirmohammad]

Add /// comments. Replaced expensive query attribute with a simple code. Change the description of GPUMemorySpaceMappingAttr to reflect OpenCL terminology. Also, uploaded the patch with -U99999 flag

[tavakkoliamirmohammad]

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

[hw-yuexy]

Hi Amir,

I have a question regarding the promote op patch.

When I chain 2 promote op together, both trying to promote operands in a tiled linalg.matmul, with the first one copies from global memory to L1, the second one copies from L1 to L0. I found all memref.copy op are placed in the inner most loop with copy sizes dependent on the inner most loop IV. However, this can lead to bad performance because when moving from GM to L1, it's often more profitable to move larger chunks of data (i.e. we can hoist the GM->L1 memref.copy ops into an outer loop, copying more data). I'm wondering if there are plans to make promote op more flexible, so that it can hoist some memref.copy ops to outer loops when possible?

Thanks.

[tavakkoliamirmohammad]

In D144666#4214523, @hw-yuexy wrote:

Hi Amir,

I have a question regarding the promote op patch.

When I chain 2 promote op together, both trying to promote operands in a tiled linalg.matmul, with the first one copies from global memory to L1, the second one copies from L1 to L0. I found all memref.copy op are placed in the inner most loop with copy sizes dependent on the inner most loop IV. However, this can lead to bad performance because when moving from GM to L1, it's often more profitable to move larger chunks of data (i.e. we can hoist the GM->L1 memref.copy ops into an outer loop, copying more data). I'm wondering if there are plans to make promote op more flexible, so that it can hoist some memref.copy ops to outer loops when possible?

Thanks.

Hi.
Sorry I didn't get the notification for this message. The reason is mostly, the workgroup or private memory cannot hold the entire data due to low capacity. Mostly, the technique is to do tiling and, in a loop, only copy the data needed to workgroup or global memory. If you know your entire data fits in those memories what you can do is to promote the data before tiling for computation.

[kaitingwang]

Hi Amir,

Thanks for the nice patch! I'm trying to use this patch as-is with as little changes as possible, for my accelerator which unfortunately has a few different types of 'private' memories.
Since the 'mapping' attribute is defined in the td file as DeviceMappingArrayAttr, which requires definition of address space (#gpu.memory_space<workgroup> or #gpu.memory_space<private>) inside the td file (GPUDeviceMappingAttr.td) to inherit from GPUMemorySpaceMappingAttr. This means, I would need to modify MLIR (e.g. potentially add my hardware dialect) with my private memory spaces in order to use promote op's mapping. I know it is not recommended to expose the address space as a IntegerAttr or UnitAttr. However, it would be more flexible to be honest since memspace in a memref is an integer...

Any advice would be greatly appreciated! Thank you!

Amy

mlir/include/mlir/Dialect/GPU/TransformOps/GPUDeviceMappingAttr.td

let description = [{
 
     Thread blocks (aka work-group) are grouped into a grid where grid may be
     described by a 1-, 2-, or 3-dimensional rectangle. This attribute indicates
     that thread block parallelism is desired. It can be consumed by lowering to
     generate GPU code.
   }];
 }
 
+
+def GPUMemorySpaceMappingAttr : GPU_Attr<"GPUMemorySpaceMapping", "memory_space", [
+  DeclareAttrInterfaceMethods<DeviceMappingAttrInterface> ] >  {
+  let parameters = (ins
+    EnumParameter<GPU_AddressSpaceEnum>:$address_space
+  );
+  let assemblyFormat = "`<` params `>`";
+  let description = [{
+    An attribute that allows defining memory hierarchy for GPU devices.
+
+    GPU Memory has three memory space, global, workgroup, and private. The global memory
+    is visible to all workitems and workgroups, the workgroup memory is only available for workitems
+    within a workgroup, and private memory is only visible to a single workitem. This attribute indicates
+    that using memory hiearchy is desired. It can be consumed by lowering to
+    move data to a specific address space in GPU code.
+  }];
+}
+
+
 #endif // GPU_DEVICE_MAPPING_ATTR

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

::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::transform::TransformState &state);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // PromoteOp
 //===----------------------------------------------------------------------===//
 
+
 def PromoteOp : Op<Transform_Dialect, "structured.promote",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
     TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Promotes the specified operands of the target into a separate memory buffer.
 
     At this point, this transform does not allow customizing alloc/dealloc
     functions nor the behavior on copy in/out operations.
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,
                        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

 ///   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,
+                                        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,
+                                         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.
+LogicalResult deallocateGPUPrivateMemory(OpBuilder &, Value /*buffer*/);
+
 /// Emit a suitable vector form for a Linalg op. If provided, `inputVectorSizes`
 /// are used to vectorize this operation. `inputVectorSizes` must match the rank
 /// of the iteration space of the operation and the sizes must be smaller or
 /// equal than their counterpart interation space sizes, if static.
 /// `inputVectorShapes` also allows the vectorization of operations with dynamic
 /// shapes.
 LogicalResult vectorize(RewriterBase &rewriter, LinalgOp linalgOp,
                         ArrayRef<int64_t> inputVectorSizes = {},

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 int64_t GPUWarpMappingAttr::getMappingId() const {
   return static_cast<int64_t>(getWarp());
 }
 
 int64_t GPUThreadMappingAttr::getMappingId() const {
   return static_cast<int64_t>(getThread());
 }
 
+int64_t GPUMemorySpaceMappingAttr::getMappingId() const {
+  return static_cast<int64_t>(getAddressSpace());
+}
+
 //===----------------------------------------------------------------------===//
 // MMAMatrixType
 //===----------------------------------------------------------------------===//
 
 MMAMatrixType MMAMatrixType::get(ArrayRef<int64_t> shape, Type elementType,
                                  StringRef operand) {
   return Base::get(elementType.getContext(), shape, elementType, operand);
 }

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

if (!std::is_permutation(sequence.begin(), sequence.end(),
              << attr;
     }
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // PromoteOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::PromoteOp::applyToOne(LinalgOp target,
                                  transform::ApplyToEachResultList &results,
                                  transform::TransformState &state) {
   LinalgPromotionOptions promotionOptions;
   if (!getOperandsToPromote().empty())
     promotionOptions = promotionOptions.setOperandsToPromote(
         extractFromI64ArrayAttr(getOperandsToPromote()));
   if (getUseFullTilesByDefault())
     promotionOptions = promotionOptions.setUseFullTileBuffersByDefault(
         getUseFullTilesByDefault());
   if (getUseAlloca())
     promotionOptions = promotionOptions.setUseAlloca(getUseAlloca());
   if (!getUseFullTileBuffers().empty())
     promotionOptions = promotionOptions.setUseFullTileBuffers(
         llvm::to_vector(getUseFullTileBuffers().getAsValueRange<BoolAttr>()));
   if (getAlignment().has_value())
     promotionOptions = promotionOptions.setAlignment(*getAlignment());
 
+  if (getMapping().has_value()) {
+    // The mapping should only contain an element
+    auto mapping = *getMapping();
+    if (mapping.size() > 1)
+      return emitDefaultDefiniteFailure(target);
+
+    auto addressSpace = mapping[0].cast<gpu::GPUMemorySpaceMappingAttr>();
+
+    if (addressSpace.getAddressSpace() ==
+        gpu::GPUDialect::getWorkgroupAddressSpace()) {
+      promotionOptions =
+          promotionOptions
+              .setAllocationDeallocationFns(allocateWorkgroupMemory,
+                                            deallocateWorkgroupMemory)
+              .setCopyInOutFns(copyToWorkgroupMemory, copyToWorkgroupMemory)
+              .setUseFullTileBuffers({false, false});
+    } else if (addressSpace.getAddressSpace() ==
+               gpu::GPUDialect::getPrivateAddressSpace()) {
+      promotionOptions =
+          promotionOptions
+              .setAllocationDeallocationFns(allocateGPUPrivateMemory,
+                                            deallocateGPUPrivateMemory)
+              .setCopyInOutFns(copyToGPUPrivateMemory, copyToGPUPrivateMemory)
+              .setUseFullTileBuffers({false, false});
+    } else {
+      return emitDefaultDefiniteFailure(target);
+    }
+  }
+
   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);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // ReplaceOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ReplaceOp::apply(TransformResults &transformResults,
                             TransformState &state) {
   ArrayRef<Operation *> payload = state.getPayloadOps(getTarget());
 
   // Check for invalid targets.
for (Operation *target : payload) {
     rewriter.setInsertionPoint(target);
     Operation *replacement = rewriter.clone(*pattern);
     rewriter.replaceOp(target, replacement->getResults());
     replacements.push_back(replacement);
   }
   transformResults.set(getReplacement().cast<OpResult>(), replacements);
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::ReplaceOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   producesHandle(getReplacement(), effects);
   modifiesPayload(effects);
 }
 
 LogicalResult transform::ReplaceOp::verify() {
   if (!getBodyRegion().hasOneBlock())
     return emitOpError() << "expected one block";
   if (std::distance(getBodyRegion().front().begin(),
                     getBodyRegion().front().end()) != 1)
     return emitOpError() << "expected one operation in block";
   Operation *replacement = &getBodyRegion().front().front();
   if (replacement->getNumOperands() > 0)
     return replacement->emitOpError()
            << "expected replacement without operands";
   if (!replacement->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
       replacement->getNumRegions() > 0)
     return replacement->emitOpError()
            << "expect op that is isolated from above";
   return success();

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

 //===- Promotion.cpp - Implementation of linalg Promotion -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the linalg dialect Promotion pass.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
mlir::linalg::promoteSubViews(OpBuilder &builder, LinalgOp linalgOp,
   LinalgOpInstancePromotionOptions linalgOptions(linalgOp, options);
   auto layout = DataLayout::closest(linalgOp);
   ImplicitLocOpBuilder b(linalgOp.getLoc(), builder);
   auto res = ::promoteSubViews(b, linalgOp, linalgOptions, layout);
   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,
+    gpu::AddressSpace addressSpace) {
+  OpBuilder::InsertionGuard guard(builder);
+
+  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{},
+      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,
+    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,
+    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) {
+  Operation *copyOp = 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

 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1 = transform.structured.promote %0
 }
 
 // -----
+func.func @gemm_shared(%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_shared
+// 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<16x16xf32, #gpu.address_space<workgroup>>
+// CHECK: %[[alloc_B:.*]] = memref.alloc() : memref<16x16xf32, #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:         %[[shared_A:.*]] = memref.subview %[[alloc_B]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<16x16xf32, #gpu.address_space<workgroup>> to memref<?x?xf32, strided<[16, 1]>, #gpu.address_space<workgroup>>
+// CHECK:         %[[shared_B:.*]] = memref.subview %[[alloc_A]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<16x16xf32, #gpu.address_space<workgroup>> to memref<?x?xf32, strided<[16, 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<[16, 1]>, #gpu.address_space<workgroup>>
+// CHECK-NEXT:    gpu.barrier
+
+// CHECK-NEXT:    gpu.barrier
+// CHECK-NEXT:    memref.copy %[[subview_B]], %[[shared_B]] :  memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[16, 1]>, #gpu.address_space<workgroup>>
+// CHECK-NEXT:    gpu.barrier
+
+// CHECK:         linalg.matmul 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, 16, 16] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+  %2 = transform.structured.promote %1 { operands_to_promote = [0, 1], mapping = [#gpu.memory_space<workgroup>] }
+}
+
+
+// -----
+
+func.func @gemm_private(%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_private
+// 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.alloca() : memref<16x16xf32, #gpu.address_space<private>>
+// CHECK: %[[alloc_B:.*]] = memref.alloca() : memref<16x16xf32, #gpu.address_space<private>>
+// 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:         %[[private_A:.*]] = memref.subview %[[alloc_B]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<16x16xf32, #gpu.address_space<private>> to memref<?x?xf32, strided<[16, 1]>, #gpu.address_space<private>>
+// CHECK:         %[[private_B:.*]] = memref.subview %[[alloc_A]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<16x16xf32, #gpu.address_space<private>> to memref<?x?xf32, strided<[16, 1]>, #gpu.address_space<private>>
+
+// CHECK-NEXT:    memref.copy %[[subview_A]], %[[private_A]] :  memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[16, 1]>, #gpu.address_space<private>>
+// CHECK-NEXT:    memref.copy %[[subview_B]], %[[private_B]] :  memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[16, 1]>, #gpu.address_space<private>>
+
+// CHECK:         linalg.matmul ins(%[[private_A]], %[[private_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, 16, 16] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+  %2 = transform.structured.promote %1 { operands_to_promote = [0, 1], mapping = [#gpu.memory_space<private>] }
+}
+
+
+// -----
 
 #map6 = affine_map<(d0, d1, d2) -> (d0, d2)>
 #map7 = affine_map<(d0, d1, d2) -> (d1, d2)>
 #map8 = affine_map<(d0, d1, d2) -> (d0, d1)>
 
 // CHECK: promote_rank_reducing_subviews(%[[arg0:.+]]: memref<{{.*}}>, %[[arg1:.+]]: memref<{{.*}}>, %[[arg2:.+]]: memref<{{.*}}>, %[[lb1:.+]]: index, %[[lb2:.+]]: index, %[[lb3:.+]]: index, %[[lb4:.+]]: index, %[[lb5:.+]]: index, %[[lb6:.+]]: index, %[[ub1:.+]]: index, %[[ub2:.+]]: index
 func.func @promote_rank_reducing_subviews(%arg0:  memref<?x?x?x64xf32, strided<[?, ?, ?, ?], offset: ?>>, %arg1: memref<128x3x3x64xf32, strided<[?, ?, ?, ?], offset: ?>>, %arg2: memref<?x?x?x128xf32>,
                                           %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %ub1: index, %ub2: index) {