[mlir][Analysis] Allow Slice Analysis to work with linalg::LinalgOp

[antiagainst]

Can we add a test for it in slicing-utils.mlir?

[mravishankar]

Adding test for backward slice of Linalg operation.

[mravishankar]

rebase

mlir/lib/Analysis/SliceAnalysis.cpp

 //===- UseDefAnalysis.cpp - Analysis for Transitive UseDef chains ---------===//
 //
 // 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 Analysis functions specific to slicing in Function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SetVector.h"
 
 ///
 /// Implements Analysis functions specific to slicing in Function.

 }
 
 static void getBackwardSliceImpl(Operation *op,
                                  SetVector<Operation *> *backwardSlice,
                                  TransitiveFilter filter) {
   if (!op)
     return;
 
-  assert((op->getNumRegions() == 0 || isa<AffineForOp, scf::ForOp>(op)) &&
+  assert((op->getNumRegions() == 0 ||
+          isa<AffineForOp, scf::ForOp, linalg::LinalgOp>(op)) &&
          "unexpected generic op with regions");
 
   // Evaluate whether we should keep this def.
   // This is useful in particular to implement scoping; i.e. return the
   // transitive forwardSlice in the current scope.
   if (!filter(op)) {
     return;
   }

mlir/test/IR/slice.mlir

+// RUN: mlir-opt -slice-analysis-test %s
+
+func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) {
+  %a = alloc(%arg0, %arg2) : memref<?x?xf32>
+  %b = alloc(%arg2, %arg1) : memref<?x?xf32>
+  %c = alloc(%arg0, %arg1) : memref<?x?xf32>
+  %d = alloc(%arg0, %arg1) : memref<?x?xf32>
+  linalg.matmul %a, %b, %c : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+  linalg.matmul %a, %b, %d : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>)
+  dealloc %c : memref<?x?xf32>
+  dealloc %b : memref<?x?xf32>
+  dealloc %a : memref<?x?xf32>
+  dealloc %d : memref<?x?xf32>
+  return
+}
+
+// CHECK-LABEL: func @slicing_linalg_op__backward_slice__0
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: index
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: index
+//   CHECK-DAG:   %[[A:.+]] = alloc(%[[ARG0]], %[[ARG2]]) : memref<?x?xf32>
+//   CHECK-DAG:   %[[B:.+]] = alloc(%[[ARG2]], %[[ARG1]]) : memref<?x?xf32>
+//   CHECK-DAG:   %[[C:.+]] = alloc(%[[ARG0]], %[[ARG1]]) : memref<?x?xf32>
+//       CHECK:   return
+
+// CHECK-LABEL: func @slicing_linalg_op__backward_slice__1
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: index
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: index
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: index
+//   CHECK-DAG:   %[[A:.+]] = alloc(%[[ARG0]], %[[ARG2]]) : memref<?x?xf32>
+//   CHECK-DAG:   %[[B:.+]] = alloc(%[[ARG2]], %[[ARG1]]) : memref<?x?xf32>
+//   CHECK-DAG:   %[[C:.+]] = alloc(%[[ARG0]], %[[ARG1]]) : memref<?x?xf32>
+//       CHECK:   return

mlir/test/lib/IR/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestIR
   TestFunc.cpp
   TestInterfaces.cpp
   TestMatchers.cpp
   TestSideEffects.cpp
+  TestSlicing.cpp
   TestSymbolUses.cpp
   TestTypes.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRPass
   MLIRTestDialect
   )
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../Dialect/Test)
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/../Dialect/Test)

mlir/test/lib/IR/TestSlicing.cpp

+//===- TestSlicing.cpp - Testing slice functionality ----------------------===//
+//
+// 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 a simple testing pass for slicing.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Module.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+
+using namespace mlir;
+
+/// Create a function with the same signature as the parent function of `op`
+/// with name being the function name and a `suffix`.
+static LogicalResult createBackwardSliceFunction(Operation *op,
+                                                 StringRef suffix) {
+  FuncOp parentFuncOp = op->getParentOfType<FuncOp>();
+  OpBuilder builder(parentFuncOp);
+  Location loc = op->getLoc();
+  std::string clonedFuncOpName = parentFuncOp.getName().str() + suffix.str();
+  FuncOp clonedFuncOp =
+      builder.create<FuncOp>(loc, clonedFuncOpName, parentFuncOp.getType());
+  BlockAndValueMapping mapper;
+  builder.setInsertionPointToEnd(clonedFuncOp.addEntryBlock());
+  for (auto arg : enumerate(parentFuncOp.getArguments()))
+    mapper.map(arg.value(), clonedFuncOp.getArgument(arg.index()));
+  llvm::SetVector<Operation *> slice;
+  getBackwardSlice(op, &slice);
+  for (Operation *slicedOp : slice)
+    builder.clone(*slicedOp, mapper);
+  builder.create<ReturnOp>(loc);
+  return success();
+}
+
+namespace {
+/// Pass to test slice generated from slice analysis.
+struct SliceAnalysisTestPass
+    : public PassWrapper<SliceAnalysisTestPass, OperationPass<ModuleOp>> {
+  void runOnOperation() override;
+  SliceAnalysisTestPass() = default;
+  SliceAnalysisTestPass(const SliceAnalysisTestPass &) {}
+};
+} // namespace
+
+void SliceAnalysisTestPass::runOnOperation() {
+  ModuleOp module = getOperation();
+  auto funcOps = module.getOps<FuncOp>();
+  unsigned opNum = 0;
+  for (auto funcOp : funcOps) {
+    // TODO: For now this is just looking for Linalg ops. It can be generalized
+    // to look for other ops using flags.
+    funcOp.walk([&](Operation *op) {
+      if (!isa<linalg::LinalgOp>(op))
+        return WalkResult::advance();
+      std::string append =
+          std::string("__backward_slice__") + std::to_string(opNum);
+      createBackwardSliceFunction(op, append);
+      opNum++;
+      return WalkResult::advance();
+    });
+  }
+}
+
+namespace mlir {
+void registerSliceAnalysisTestPass() {
+  PassRegistration<SliceAnalysisTestPass> pass(
+      "slice-analysis-test", "Test Slice analysis functionality.");
+}
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerConvertToTargetEnvPass();
 void registerInliner();
 void registerMemRefBoundCheck();
 void registerPassManagerTestPass();
 void registerPatternsTestPass();
 void registerPrintOpAvailabilityPass();
 void registerSideEffectTestPasses();
 void registerSimpleParametricTilingPass();
+void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestBufferPlacementPreparationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
void registerTestPasses() {
   registerConvertToTargetEnvPass();
   registerInliner();
   registerMemRefBoundCheck();
   registerPassManagerTestPass();
   registerPatternsTestPass();
   registerPrintOpAvailabilityPass();
   registerSideEffectTestPasses();
   registerSimpleParametricTilingPass();
+  registerSliceAnalysisTestPass();
   registerSymbolTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAllReduceLoweringPass();
   registerTestAffineLoopUnswitchingPass();
   registerTestLoopPermutationPass();
   registerTestCallGraphPass();
   registerTestConstantFold();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
#endif
   registerAsmPrinterCLOptions();
   registerMLIRContextCLOptions();
   registerPassManagerCLOptions();
   PassPipelineCLParser passPipeline("", "Compiler passes to run");
 
   // Parse pass names in main to ensure static initialization completed.
   cl::ParseCommandLineOptions(argc, argv, "MLIR modular optimizer driver\n");
 
-  if(showDialects) {
+  if (showDialects) {
     llvm::outs() << "Registered Dialects:\n";
     MLIRContext context;
-    for(Dialect *dialect : context.getRegisteredDialects()) {
+    for (Dialect *dialect : context.getRegisteredDialects()) {
       llvm::outs() << dialect->getNamespace() << "\n";
     }
     return 0;
   }
 
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);