[mlir] Add a ViewLikeOpInterface

This can help provide a common interface for view-like
ops so that for example Linalg's dependency analysis
can avoid relying on concrete ops.

Differential Revision: https://reviews.llvm.org/D78645

Author: antiagainst

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h

@@ -23,6 +23,7 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/SideEffects.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

@@ -16,6 +16,7 @@
 include "mlir/Dialect/Affine/IR/AffineOpsBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Interfaces/SideEffects.td"
+include "mlir/Interfaces/ViewLikeInterface.td"
 
 // Base class for Linalg dialect ops that do not correspond to library calls.
 class Linalg_Op<string mnemonic, list<OpTrait> traits = []> :
@@ -179,7 +180,8 @@ def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<"tensor_reshape">,
   }];
 }
 
-def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
+def Linalg_SliceOp : Linalg_Op<"slice", [
+      DeclareOpInterfaceMethods<ViewLikeOpInterface>, NoSideEffect]>,
     Arguments<(ins AnyStridedMemRef:$view,
                    Variadic<AnyTypeOf<[Range, Index]>>:$indexings)>,
     Results<(outs AnyStridedMemRef)> {

mlir/include/mlir/Dialect/StandardOps/IR/Ops.h

@@ -21,6 +21,7 @@
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffects.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/StandardOps/IR/OpsEnums.h.inc"

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

@@ -17,6 +17,7 @@ include "mlir/IR/OpAsmInterface.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffects.td"
+include "mlir/Interfaces/ViewLikeInterface.td"
 
 def StandardOps_Dialect : Dialect {
   let name = "std";
@@ -2315,7 +2316,11 @@ def SubIOp : IntArithmeticOp<"subi"> {
 // SubViewOp
 //===----------------------------------------------------------------------===//
 
-def SubViewOp : Std_Op<"subview", [AttrSizedOperandSegments, NoSideEffect]> {
+def SubViewOp : Std_Op<"subview", [
+    AttrSizedOperandSegments, 
+    DeclareOpInterfaceMethods<ViewLikeOpInterface>,
+    NoSideEffect,
+  ]> {
   let summary = "memref subview operation";
   let description = [{
     The "subview" operation converts a memref type to another memref type
@@ -2785,7 +2790,8 @@ def UnsignedShiftRightOp : IntArithmeticOp<"shift_right_unsigned"> {
 // ViewOp
 //===----------------------------------------------------------------------===//
 
-def ViewOp : Std_Op<"view", [NoSideEffect]> {
+def ViewOp : Std_Op<"view", [
+    DeclareOpInterfaceMethods<ViewLikeOpInterface>, NoSideEffect]> {
   let summary = "memref view operation";
   let description = [{
     The "view" operation converts a 1-D memref with i8 element type,

mlir/include/mlir/Interfaces/CMakeLists.txt

@@ -27,3 +27,8 @@ set(LLVM_TARGET_DEFINITIONS SideEffects.td)
 mlir_tablegen(SideEffectInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(SideEffectInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRSideEffectOpInterfacesIncGen)
+
+set(LLVM_TARGET_DEFINITIONS ViewLikeInterface.td)
+mlir_tablegen(ViewLikeInterface.h.inc -gen-op-interface-decls)
+mlir_tablegen(ViewLikeInterface.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRViewLikeInterfaceIncGen)

mlir/include/mlir/Interfaces/ViewLikeInterface.h

@@ -0,0 +1,24 @@
+//===- ViewLikeInterface.h - View-like operations interface ---------------===//
+//
+// 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 operation interface for view-like operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_VIEWLIKEINTERFACE_H_
+#define MLIR_INTERFACES_VIEWLIKEINTERFACE_H_
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+
+#include "mlir/Interfaces/ViewLikeInterface.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_INTERFACES_VIEWLIKEINTERFACE_H_

mlir/include/mlir/Interfaces/ViewLikeInterface.td

@@ -0,0 +1,32 @@
+//===- ViewLikeInterface.td - ViewLike interface -----------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface for view-like operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_VIEWLIKEINTERFACE
+#define MLIR_INTERFACES_VIEWLIKEINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def ViewLikeOpInterface : OpInterface<"ViewLikeOpInterface"> {
+  let description = [{
+    A view-like operation "views" a buffer in a potentially different way. It
+    takes in a (view of) buffer (and potentially some other operands) and returns
+    another view of buffer.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      "Returns the source buffer from which the view is created.",
+      "Value", "getViewSource">
+  ];
+}
+
+#endif // MLIR_INTERFACES_VIEWLIKEINTERFACE

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

@@ -37,22 +37,17 @@ Value Aliases::find(Value v) {
   while (true) {
     if (v.isa<BlockArgument>())
       return v;
-    if (auto alloc = dyn_cast_or_null<AllocOp>(v.getDefiningOp())) {
+    Operation *defOp = v.getDefiningOp();
+    if (auto alloc = dyn_cast_or_null<AllocOp>(defOp)) {
       if (isStrided(alloc.getType()))
         return alloc.getResult();
     }
-    if (auto slice = dyn_cast_or_null<SliceOp>(v.getDefiningOp())) {
-      auto it = aliases.insert(std::make_pair(v, find(slice.view())));
+    if (auto viewLikeOp = dyn_cast_or_null<ViewLikeOpInterface>(defOp)) {
+      auto it =
+          aliases.insert(std::make_pair(v, find(viewLikeOp.getViewSource())));
       return it.first->second;
     }
-    if (auto view = dyn_cast_or_null<ViewOp>(v.getDefiningOp())) {
-      auto it = aliases.insert(std::make_pair(v, view.source()));
-      return it.first->second;
-    }
-    if (auto view = dyn_cast_or_null<SubViewOp>(v.getDefiningOp())) {
-      v = view.source();
-      continue;
-    }
+
     llvm::errs() << "View alias analysis reduces to: " << v << "\n";
     llvm_unreachable("unsupported view alias case");
   }

mlir/lib/Dialect/Linalg/IR/CMakeLists.txt

@@ -17,5 +17,6 @@ target_link_libraries(MLIRLinalgOps
   PUBLIC
   MLIRIR
   MLIRSideEffects
+  MLIRViewLikeInterface
   MLIRStandardOps
   )

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -680,6 +680,8 @@ static LogicalResult verify(SliceOp op) {
   return success();
 }
 
+Value SliceOp::getViewSource() { return view(); }
+
 //===----------------------------------------------------------------------===//
 // TransposeOp
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/StandardOps/CMakeLists.txt

@@ -16,5 +16,6 @@ target_link_libraries(MLIRStandardOps
   MLIREDSC
   MLIRIR
   MLIRSideEffects
+  MLIRViewLikeInterface
   LLVMSupport
   )

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

@@ -2324,6 +2324,8 @@ SubViewOp::getStaticStrides(SmallVectorImpl<int64_t> &staticStrides) {
   return success();
 }
 
+Value SubViewOp::getViewSource() { return source(); }
+
 namespace {
 
 /// Pattern to rewrite a subview op with constant size arguments.
@@ -2669,6 +2671,8 @@ static LogicalResult verify(ViewOp op) {
   return success();
 }
 
+Value ViewOp::getViewSource() { return source(); }
+
 namespace {
 
 struct ViewOpShapeFolder : public OpRewritePattern<ViewOp> {

mlir/lib/Interfaces/CMakeLists.txt

@@ -5,6 +5,7 @@ set(LLVM_OPTIONAL_SOURCES
   InferTypeOpInterface.cpp
   LoopLikeInterface.cpp
   SideEffects.cpp
+  ViewLikeInterface.cpp
   )
 
 add_mlir_library(MLIRCallInterfaces
@@ -90,3 +91,17 @@ target_link_libraries(MLIRSideEffects
   PUBLIC
   MLIRIR
   )
+
+add_mlir_library(MLIRViewLikeInterface
+  ViewLikeInterface.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces
+
+  DEPENDS
+  MLIRViewLikeInterfaceIncGen
+  )
+target_link_libraries(MLIRLoopLikeInterface
+  PUBLIC
+  MLIRIR
+  )

mlir/lib/Interfaces/ViewLikeInterface.cpp

@@ -0,0 +1,18 @@
+//===- ViewLikeInterface.cpp - View-like operations in MLIR ---------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/ViewLikeInterface.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// ViewLike Interfaces
+//===----------------------------------------------------------------------===//
+
+/// Include the definitions of the loop-like interfaces.
+#include "mlir/Interfaces/ViewLikeInterface.cpp.inc"