Update comment + address comments + add dest test

Author: qedawkins

mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp

@@ -551,10 +551,9 @@ LogicalResult GatherToLDSOp::verify() {
 }
 
 namespace {
-/// If the source/target of a CopyOp is a CastOp that does not modify the shape
-/// and element type, the cast can be skipped. Such CastOps only cast the layout
-/// of the type.
-struct FoldGatherToLDSOfCast : public OpRewritePattern<GatherToLDSOp> {
+/// If the source/target of a GatherToLDSOp is a CastOp that only removes static
+/// information or changes layout, the cast can be skipped.
+struct FoldGatherToLDSOfCast final : OpRewritePattern<GatherToLDSOp> {
   using OpRewritePattern<GatherToLDSOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(GatherToLDSOp gatherOp,
@@ -563,10 +562,10 @@ struct FoldGatherToLDSOfCast : public OpRewritePattern<GatherToLDSOp> {
 
     // Check source.
     if (auto castOp = gatherOp.getSrc().getDefiningOp<memref::CastOp>()) {
-      auto fromType = llvm::dyn_cast<MemRefType>(castOp.getSource().getType());
-      auto toType = llvm::dyn_cast<MemRefType>(castOp.getSource().getType());
+      auto fromType = dyn_cast<MemRefType>(castOp.getSource().getType());
+      auto toType = dyn_cast<MemRefType>(castOp.getSource().getType());
 
-      if (fromType && toType &&
+      if (memref::CastOp::canFoldIntoConsumerOp(castOp) && fromType && toType &&
           fromType.getElementType() == toType.getElementType()) {
         rewriter.modifyOpInPlace(gatherOp, [&] {
           gatherOp.getSrcMutable().assign(castOp.getSource());
@@ -577,10 +576,10 @@ struct FoldGatherToLDSOfCast : public OpRewritePattern<GatherToLDSOp> {
 
     // Check target.
     if (auto castOp = gatherOp.getDst().getDefiningOp<memref::CastOp>()) {
-      auto fromType = llvm::dyn_cast<MemRefType>(castOp.getSource().getType());
-      auto toType = llvm::dyn_cast<MemRefType>(castOp.getSource().getType());
+      auto fromType = dyn_cast<MemRefType>(castOp.getSource().getType());
+      auto toType = dyn_cast<MemRefType>(castOp.getSource().getType());
 
-      if (fromType && toType &&
+      if (memref::CastOp::canFoldIntoConsumerOp(castOp) && fromType && toType &&
           fromType.getElementType() == toType.getElementType()) {
         rewriter.modifyOpInPlace(gatherOp, [&] {
           gatherOp.getDstMutable().assign(castOp.getSource());

mlir/test/Dialect/AMDGPU/canonicalize.mlir

@@ -144,3 +144,18 @@ func.func @fold_gather_to_lds_of_cast(%global: memref<128x72xf32, 1>, %lds: memr
     : f32, memref<?x?xf32, 1>, memref<64x64xf32, 3>
   func.return
 }
+
+// -----
+
+// CHECK-LABEL: func @fold_gather_to_lds_of_cast_dest
+func.func @fold_gather_to_lds_of_cast_dest(%global: memref<128x72xf32, 1>, %lds: memref<64x64xf32, 3>) {
+// CHECK-SAME: %[[GLOBAL:[A-Za-z0-9]+]]: memref<128x72xf32, 1>
+// CHECK-SAME: %[[LDS:[A-Za-z0-9]+]]: memref<64x64xf32, 3>
+  %c0 = arith.constant 0 : index
+  %0 = memref.cast %lds : memref<64x64xf32, 3> to memref<?x?xf32, 3>
+  // CHECK: amdgpu.gather_to_lds %[[GLOBAL]][{{.*}}], %[[LDS]]
+  // CHECK-SAME: : f32, memref<128x72xf32, 1>, memref<64x64xf32, 3>
+  amdgpu.gather_to_lds %global[%c0, %c0], %0[%c0, %c0]
+    : f32, memref<128x72xf32, 1>, memref<?x?xf32, 3>
+  func.return
+}