[Mlir-commits] [mlir] [MLIR][LLVM] Use ViewLikeOpInterface (PR #111663)

[Tobias Gysi]

https://github.com/gysit created https://github.com/llvm/llvm-project/pull/111663

This commit adds the ViewLikeOpInterface to the GEP and AddrSpaceCast operations. This allows us to simplify the inliner interface. At the same time, the change also makes the inliner interface more extensible for downstream users that have custom view-like operations.

>From 1ca6a1f9751bdcf9969f7ba3eb87275f7ee53468 Mon Sep 17 00:00:00 2001
From: Tobias Gysi <tobias.gysi at nextsilicon.com>
Date: Wed, 9 Oct 2024 11:43:05 +0000
Subject: [PATCH] [MLIR][LLVM] Use ViewLikeOpInterface

This commit adds the ViewLikeOpInterface to the GEP and AddrSpaceCast
operations. This allows us to simplify the inliner interface.
At the same time, the change also makes the inliner interface more
extensible for downstream users that have custom view-like operations.
---
 mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h         |  1 +
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td            |  8 ++++++--
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp             |  4 ++++
 .../Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp | 10 +++++-----
 4 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h
index 9341a5a11cd629..d236cae0d80882 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h
@@ -29,6 +29,7 @@
 #include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/ThreadLocalCache.h"
 #include "llvm/ADT/PointerEmbeddedInt.h"
 #include "llvm/IR/DerivedTypes.h"
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 88e82ce48959b0..825b48fb3d4a34 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -24,6 +24,7 @@ include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/MemorySlotInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/ViewLikeInterface.td"
 
 class LLVM_Builder<string builder> {
   string llvmBuilder = builder;
@@ -246,7 +247,9 @@ def LLVM_AllocaOp : LLVM_Op<"alloca",
 def LLVM_GEPOp : LLVM_Op<"getelementptr", [Pure,
     DeclareOpInterfaceMethods<PromotableOpInterface>,
     DeclareOpInterfaceMethods<SafeMemorySlotAccessOpInterface>,
-    DeclareOpInterfaceMethods<DestructurableAccessorOpInterface>]> {
+    DeclareOpInterfaceMethods<DestructurableAccessorOpInterface>,
+    DeclareOpInterfaceMethods<ViewLikeOpInterface>
+    ]> {
   let arguments = (ins LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$base,
                    Variadic<LLVM_ScalarOrVectorOf<AnySignlessInteger>>:$dynamicIndices,
                    DenseI32ArrayAttr:$rawConstantIndices,
@@ -495,7 +498,8 @@ def LLVM_BitcastOp : LLVM_CastOp<"bitcast", "BitCast", LLVM_AnyNonAggregate,
 def LLVM_AddrSpaceCastOp : LLVM_CastOp<"addrspacecast", "AddrSpaceCast",
     LLVM_ScalarOrVectorOf<LLVM_AnyPointer>,
     LLVM_ScalarOrVectorOf<LLVM_AnyPointer>,
-    [DeclareOpInterfaceMethods<PromotableOpInterface>]> {
+    [DeclareOpInterfaceMethods<PromotableOpInterface>,
+     DeclareOpInterfaceMethods<ViewLikeOpInterface>]> {
   let hasFolder = 1;
 }
 def LLVM_IntToPtrOp : LLVM_CastOp<"inttoptr", "IntToPtr",
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index fb7024a14f8d4e..9e361848f8c0bf 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -3226,6 +3226,8 @@ OpFoldResult LLVM::AddrSpaceCastOp::fold(FoldAdaptor adaptor) {
   return foldChainableCast(*this, adaptor);
 }
 
+Value LLVM::AddrSpaceCastOp::getViewSource() { return getArg(); }
+
 //===----------------------------------------------------------------------===//
 // Folder for LLVM::GEPOp
 //===----------------------------------------------------------------------===//
@@ -3276,6 +3278,8 @@ OpFoldResult LLVM::GEPOp::fold(FoldAdaptor adaptor) {
   return {};
 }
 
+Value LLVM::GEPOp::getViewSource() { return getBase(); }
+
 //===----------------------------------------------------------------------===//
 // ShlOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
index 031930dcfc2131..cb79083846c742 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.cpp
@@ -16,6 +16,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
@@ -229,11 +230,10 @@ static FailureOr<SmallVector<Value>>
 getUnderlyingObjectSet(Value pointerValue) {
   SmallVector<Value> result;
   WalkContinuation walkResult = walkSlice(pointerValue, [&](Value val) {
-    if (auto gepOp = val.getDefiningOp<LLVM::GEPOp>())
-      return WalkContinuation::advanceTo(gepOp.getBase());
-
-    if (auto addrCast = val.getDefiningOp<LLVM::AddrSpaceCastOp>())
-      return WalkContinuation::advanceTo(addrCast.getOperand());
+    // Attempt to advance to the underlying object of a view-like operation such
+    // as a GEPOp or an AddrSpaceCastOp.
+    if (auto viewOp = val.getDefiningOp<ViewLikeOpInterface>())
+      return WalkContinuation::advanceTo(viewOp.getViewSource());
 
     // Attempt to advance to control flow predecessors.
     std::optional<SmallVector<Value>> controlFlowPredecessors =