[Mlir-commits] [mlir] [mlir] Implement memory-space cast operand fusion into consumers (PR #159454)

[Fabian Mora]

================
@@ -1645,6 +1722,33 @@ OpFoldResult MemorySpaceCastOp::fold(FoldAdaptor adaptor) {
   return Value{};
 }
 
+TypedValue<PtrLikeTypeInterface> MemorySpaceCastOp::getSourcePtr() {
+  return cast<TypedValue<PtrLikeTypeInterface>>(getSource());
+}
+
+TypedValue<PtrLikeTypeInterface> MemorySpaceCastOp::getTargetPtr() {
+  return cast<TypedValue<PtrLikeTypeInterface>>(getDest());
+}
+
+bool MemorySpaceCastOp::isValidMemorySpaceCast(PtrLikeTypeInterface tgt,
+                                               PtrLikeTypeInterface src) {
+  return isa<MemRefType>(tgt) &&
+         tgt.clonePtrWith(src.getMemorySpace(), std::nullopt) == src;
+}
+
+MemorySpaceCastOpInterface
+MemorySpaceCastOp::cloneMemorySpaceCastOp(OpBuilder &b, Type tgt, Value src) {
+  assert(isValidMemorySpaceCast(cast<PtrLikeTypeInterface>(tgt),
+                                cast<PtrLikeTypeInterface>(src.getType())) &&
+         "invalid arguments");
+  return MemorySpaceCastOp::create(b, getLoc(), tgt, src);
+}
+
+bool MemorySpaceCastOp::isFusableMemorySpaceCast() {
+  // Only allow fusion when this is discarding information.
+  return getDest().getType().getMemorySpace() == nullptr;
----------------
fabianmcg wrote:

More than shifting the cost, this is about discarding information/metadata. The generic/default address space doesn't carry any extra data, non-trivial address spaces might carry such information and thus they cannot be elided .

I'll add a comment.

https://github.com/llvm/llvm-project/pull/159454