[MLIR][GPU] Simplify memcpy of cast

[wsmoses]

Summary

Introduce a simplification that allows memcpy of a cast to simply use the underlying op

[mehdi_amini]

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

1073

I'd rather not use LogicalResult for something that isn't "success" or "failure", this seems like a bool to me here.

Also, isn't this method duplicated in many places now? Can it be made inlined in a header in the memref dialect somehow?

[rriddle]

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

1076

Drop the mlir:: here.

mlir/include/mlir/Dialect/GPU/GPUOps.td

let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    Arg<AnyMemRef, "", [MemRead]>:$src);
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $dst`,` $src `:` type($dst)`,` type($src) attr-dict
   }];
   let verifier = [{ return ::verify(*this); }];
+  let hasFolder = 1;
 }
 
 def GPU_SubgroupMmaLoadMatrixOp : GPU_Op<"subgroup_mma_load_matrix",
     [MemoryEffects<[MemRead]>]>{
 
   let summary = "GPU warp synchronous matrix load";
 
   let description = [{

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

 //===- GPUDialect.cpp - MLIR Dialect for GPU Kernels implementation -------===//
 //
 // 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 GPU kernel-related dialect and its operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/GPUDialect.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/OpImplementation.h"
static LogicalResult verify(SubgroupMmaComputeOp op) {
 
   if (aShape[1] != bShape[0] || aShape[0] != cShape[0] ||
       bShape[1] != cShape[1])
     return op.emitError("operand shapes do not satisfy matmul constraints");
 
   return success();
 }
 
+/// This is a common class used for patterns of the form
+/// "someop(memrefcast) -> someop".  It folds the source of any memref.cast
+/// into the root operation directly.
+static LogicalResult foldMemRefCast(Operation *op) {
+  bool folded = false;
+  for (OpOperand &operand : op->getOpOperands()) {
+    auto cast = operand.get().getDefiningOp<mlir::memref::CastOp>();
+    if (cast) {
+      operand.set(cast.getOperand());
+      folded = true;
+    }
+  }
+  return success(folded);
+}
+
+LogicalResult MemcpyOp::fold(ArrayRef<Attribute> operands,
+                             SmallVectorImpl<::mlir::OpFoldResult> &results) {
+  return foldMemRefCast(*this);
+}
+
 #include "mlir/Dialect/GPU/GPUOpInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"

mlir/test/Dialect/GPU/canonicalize.mlir

+// RUN: mlir-opt %s -canonicalize --split-input-file -allow-unregistered-dialect | FileCheck %s
+
+// CHECK-LABEL: @memcpy_after_cast
+func @memcpy_after_cast(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
+  // CHECK-NOT: memref.cast
+  // CHECK: gpu.memcpy
+  %0 = memref.cast %arg0 : memref<10xf32> to memref<?xf32>
+  %1 = memref.cast %arg1 : memref<10xf32> to memref<?xf32>
+  gpu.memcpy %0,%1 : memref<?xf32>, memref<?xf32>
+  return
+}