[mlir] Move ComposeSubView+ExpandOps from Standard to MemRef

[rriddle]

Summary

These transformations already operate on memref operations (as part of
splitting up the standard dialect). Now that the operations have moved,
it's time for these transformations to move as well.

Depends on D118280

[bondhugula]

@rriddle You've broken the integration tests. Please test with MLIR_INCLUDE_INTEGRATION_TESTS:BOOL=ON or wait for the phab CI.

[rriddle]

In D118285#3274864, @bondhugula wrote:

@rriddle You've broken the integration tests. Please test with MLIR_INCLUDE_INTEGRATION_TESTS:BOOL=ON or wait for the phab CI.

Ooops, let me fix that.

[bondhugula]

In D118285#3274868, @rriddle wrote:

In D118285#3274864, @bondhugula wrote:

@rriddle You've broken the integration tests. Please test with MLIR_INCLUDE_INTEGRATION_TESTS:BOOL=ON or wait for the phab CI.

Ooops, let me fix that.

I've fixed it.

[rriddle]

In D118285#3274873, @bondhugula wrote:

In D118285#3274868, @rriddle wrote:

In D118285#3274864, @bondhugula wrote:

@rriddle You've broken the integration tests. Please test with MLIR_INCLUDE_INTEGRATION_TESTS:BOOL=ON or wait for the phab CI.

Ooops, let me fix that.

I've fixed it.

Thanks! Sorry about that, added integrations to my local build (lost it during a recent move it seems).

mlir/include/mlir/Dialect/MemRef/Transforms/ComposeSubView.h

-//===- ComposeSubView.h - Combining composed subview ops --------*- C++ -*-===//
+//===- ComposeSubView.h - Combining composed memref ops ---------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Patterns for combining composed subview ops.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_DIALECT_STANDARDOPS_TRANSFORMS_COMPOSESUBVIEW_H_
-#define MLIR_DIALECT_STANDARDOPS_TRANSFORMS_COMPOSESUBVIEW_H_
+#ifndef MLIR_DIALECT_MEMREF_TRANSFORMS_COMPOSESUBVIEW_H_
+#define MLIR_DIALECT_MEMREF_TRANSFORMS_COMPOSESUBVIEW_H_
 
 namespace mlir {
-
-// Forward declarations.
 class MLIRContext;
 class RewritePatternSet;
 using OwningRewritePatternList = RewritePatternSet;
 
+namespace memref {
+
 void populateComposeSubViewPatterns(OwningRewritePatternList &patterns,
                                     MLIRContext *context);
 
+} // namespace memref
 } // namespace mlir
 
-#endif // MLIR_DIALECT_STANDARDOPS_TRANSFORMS_COMPOSESUBVIEW_H_
+#endif // MLIR_DIALECT_MEMREF_TRANSFORMS_COMPOSESUBVIEW_H_

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

 #ifndef MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES_H
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 class AffineDialect;
+class StandardOpsDialect;
 namespace tensor {
 class TensorDialect;
 } // namespace tensor
 namespace vector {
 class VectorDialect;
 } // namespace vector
 
 namespace memref {
 
 //===----------------------------------------------------------------------===//
 // Patterns
 //===----------------------------------------------------------------------===//
 
+/// Collects a set of patterns to rewrite ops within the memref dialect.
+void populateExpandOpsPatterns(RewritePatternSet &patterns);
+
 /// Appends patterns for folding memref.subview ops into consumer load/store ops
 /// into `patterns`.
 void populateFoldSubViewOpPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the
 /// `ReifyRankedShapeTypeShapeOpInterface`, in terms of shapes of its input
 /// operands.
 void populateResolveRankedShapeTypeResultDimsPatterns(
     RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the `InferShapedTypeOpInterface`, in
 /// terms of shapes of its input operands.
 void populateResolveShapedTypeResultDimsPatterns(RewritePatternSet &patterns);
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
+/// Creates an instance of the ExpandOps pass that legalizes memref dialect ops
+/// to be convertible to LLVM. For example, `memref.reshape` gets converted to
+/// `memref_reinterpret_cast`.
+std::unique_ptr<Pass> createExpandOpsPass();
+
 /// Creates an operation pass to fold memref.subview ops into consumer
 /// load/store ops into `patterns`.
 std::unique_ptr<Pass> createFoldSubViewOpsPass();
 
 /// Creates an interprocedural pass to normalize memrefs to have a trivial
 /// (identity) layout map.
 std::unique_ptr<OperationPass<ModuleOp>> createNormalizeMemRefsPass();
 

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

 //===-- Passes.td - MemRef transformation definition file --*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES
 #define MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
 
+def ExpandOps : Pass<"memref-expand", "FuncOp"> {
+  let summary = "Legalize memref operations to be convertible to LLVM.";
+  let constructor = "mlir::memref::createExpandOpsPass()";
+  let dependentDialects = ["StandardOpsDialect"];
+}
+
 def FoldSubViewOps : Pass<"fold-memref-subview-ops"> {
   let summary = "Fold memref.subview ops into consumer load/store ops";
   let description = [{
     The pass folds loading/storing from/to subview ops to loading/storing
     from/to the original memref.
   }];
   let constructor = "mlir::memref::createFoldSubViewOpsPass()";
   let dependentDialects = [

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.h

 void populateTensorConstantBufferizePatterns(
     bufferization::GlobalCreator &globalCreator,
     bufferization::BufferizeTypeConverter &typeConverter,
     RewritePatternSet &patterns);
 
 /// Creates an instance of tensor constant bufferization pass.
 std::unique_ptr<Pass> createTensorConstantBufferizePass(unsigned alignment = 0);
 
-/// Creates an instance of the StdExpand pass that legalizes Std
-/// dialect ops to be convertible to LLVM. For example,
-/// `std.arith.ceildivsi` gets transformed to a number of std operations,
-/// which can be lowered to LLVM; `memref.reshape` gets converted to
-/// `memref_reinterpret_cast`.
-std::unique_ptr<Pass> createStdExpandOpsPass();
-
-/// Collects a set of patterns to rewrite ops within the Std dialect.
-void populateStdExpandOpsPatterns(RewritePatternSet &patterns);
-
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_TRANSFORMS_PASSES_H

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

 
 def StdBufferize : Pass<"std-bufferize", "FuncOp"> {
   let summary = "Bufferize the std dialect";
   let constructor = "mlir::createStdBufferizePass()";
   let dependentDialects = ["bufferization::BufferizationDialect",
                            "memref::MemRefDialect", "scf::SCFDialect"];
 }
 
-def StdExpandOps : Pass<"std-expand", "FuncOp"> {
-  let summary = "Legalize std operations to be convertible to LLVM.";
-  let constructor = "mlir::createStdExpandOpsPass()";
-}
-
 def FuncBufferize : Pass<"func-bufferize", "ModuleOp"> {
   let summary = "Bufferize func/call/return ops";
   let description = [{
     A bufferize pass that bufferizes std.func and std.call ops.
 
     Because this pass updates std.func ops, it must be a module pass. It is
     useful to keep this pass separate from other bufferizations so that the
     other ones can be run at function-level in parallel.

mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefTransforms
+  ComposeSubView.cpp
+  ExpandOps.cpp
   FoldSubViewOps.cpp
   NormalizeMemRefs.cpp
   ResolveShapedTypeResultDims.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
 
   DEPENDS
   MLIRMemRefPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRArithmetic
   MLIRInferTypeOpInterface
   MLIRMemRef
   MLIRPass
+  MLIRStandard
   MLIRTensor
   MLIRTransforms
   MLIRVector
 )
 

mlir/lib/Dialect/MemRef/Transforms/ComposeSubView.cpp

 //===- ComposeSubView.cpp - Combining composed subview ops ----------------===//
 //
 // 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 contains patterns for combining composed subview ops (i.e. subview
 // of a subview becomes a single subview).
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/StandardOps/Transforms/ComposeSubView.h"
-
+#include "mlir/Dialect/MemRef/Transforms/ComposeSubView.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
-namespace mlir {
+using namespace mlir;
 
 namespace {
 
 // Replaces a subview of a subview with a single subview. Only supports subview
 // ops with static sizes and static strides of 1 (both static and dynamic
 // offsets are supported).
 struct ComposeSubViewOpPattern : public OpRewritePattern<memref::SubViewOp> {
   using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(memref::SubViewOp op,
     rewriter.replaceOpWithNewOp<memref::SubViewOp>(op, sourceOp.source(),
                                                    offsets, sizes, strides);
     return success();
   }
 };
 
 } // namespace
 
-void populateComposeSubViewPatterns(OwningRewritePatternList &patterns,
-                                    MLIRContext *context) {
+void mlir::memref::populateComposeSubViewPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<ComposeSubViewOpPattern>(context);
 }
-
-} // namespace mlir

mlir/lib/Dialect/MemRef/Transforms/ExpandOps.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
-#include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
 
 /// Converts `atomic_rmw` that cannot be lowered to a simple atomic op with
LogicalResult matchAndRewrite(memref::ReshapeOp op,
     }
     rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(
         op, op.getType(), op.source(), /*offset=*/rewriter.getIndexAttr(0),
         sizes, strides);
     return success();
   }
 };
 
-struct StdExpandOpsPass : public StdExpandOpsBase<StdExpandOpsPass> {
+struct ExpandOpsPass : public ExpandOpsBase<ExpandOpsPass> {
   void runOnOperation() override {
     MLIRContext &ctx = getContext();
 
     RewritePatternSet patterns(&ctx);
-    populateStdExpandOpsPatterns(patterns);
-    ConversionTarget target(getContext());
+    memref::populateExpandOpsPatterns(patterns);
+    ConversionTarget target(ctx);
 
     target.addLegalDialect<arith::ArithmeticDialect, memref::MemRefDialect,
                            StandardOpsDialect>();
     target.addDynamicallyLegalOp<memref::AtomicRMWOp>(
         [](memref::AtomicRMWOp op) {
           return op.kind() != arith::AtomicRMWKind::maxf &&
                  op.kind() != arith::AtomicRMWKind::minf;
         });
     target.addDynamicallyLegalOp<memref::ReshapeOp>([](memref::ReshapeOp op) {
       return !op.shape().getType().cast<MemRefType>().hasStaticShape();
     });
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 
 } // namespace
 
-void mlir::populateStdExpandOpsPatterns(RewritePatternSet &patterns) {
+void mlir::memref::populateExpandOpsPatterns(RewritePatternSet &patterns) {
   patterns.add<AtomicRMWOpConverter, MemRefReshapeOpConverter>(
       patterns.getContext());
 }
 
-std::unique_ptr<Pass> mlir::createStdExpandOpsPass() {
-  return std::make_unique<StdExpandOpsPass>();
+std::unique_ptr<Pass> mlir::memref::createExpandOpsPass() {
+  return std::make_unique<ExpandOpsPass>();
 }

mlir/lib/Dialect/MemRef/Transforms/PassDetail.h

 //===- PassDetail.h - MemRef Pass class details -----------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_MEMREF_TRANSFORMS_PASSDETAIL_H_
 #define DIALECT_MEMREF_TRANSFORMS_PASSDETAIL_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 
 class AffineDialect;
+class StandardOpsDialect;
 
 // Forward declaration from Dialect.h
 template <typename ConcreteDialect>
 void registerDialect(DialectRegistry &registry);
 
 namespace arith {
 class ArithmeticDialect;
 } // namespace arith

mlir/lib/Dialect/StandardOps/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRStandardOpsTransforms
   Bufferize.cpp
-  ComposeSubView.cpp
   DecomposeCallGraphTypes.cpp
-  ExpandOps.cpp
   FuncBufferize.cpp
   FuncConversions.cpp
   TensorConstantBufferize.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/StandardOps/Transforms
 
   DEPENDS

mlir/test/Dialect/MemRef/expand-ops.mlir

-// RUN: mlir-opt -std-expand %s -split-input-file | FileCheck %s
+// RUN: mlir-opt -memref-expand %s -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @atomic_rmw_to_generic
 // CHECK-SAME: ([[F:%.*]]: memref<10xf32>, [[f:%.*]]: f32, [[i:%.*]]: index)
 func @atomic_rmw_to_generic(%F: memref<10xf32>, %f: f32, %i: index) -> f32 {
   %x = memref.atomic_rmw maxf %f, %F[%i] : (f32, memref<10xf32>) -> f32
   return %x : f32
 }
 // CHECK: %0 = memref.generic_atomic_rmw %arg0[%arg2] : memref<10xf32> {

mlir/test/lib/Dialect/CMakeLists.txt

 add_subdirectory(Affine)
 add_subdirectory(DLTI)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(Math)
+add_subdirectory(MemRef)
 add_subdirectory(SCF)
 add_subdirectory(Shape)
 add_subdirectory(SPIRV)
 add_subdirectory(StandardOps)
 add_subdirectory(Test)
 add_subdirectory(Tosa)
 add_subdirectory(Vector)

mlir/test/lib/Dialect/MemRef/CMakeLists.txt

 # Exclude tests from libMLIR.so
-add_mlir_library(MLIRStandardOpsTestPasses
-  TestDecomposeCallGraphTypes.cpp
+add_mlir_library(MLIRMemRefTestPasses
   TestComposeSubView.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
-  MLIRAffine
   MLIRPass
-  MLIRStandardOpsTransforms
+  MLIRMemRefTransforms
   MLIRTestDialect
-  MLIRTransformUtils
   )
 
-target_include_directories(MLIRStandardOpsTestPasses
+target_include_directories(MLIRMemRefTestPasses
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Test
   )

mlir/test/lib/Dialect/MemRef/TestComposeSubView.cpp

 //===- TestComposeSubView.cpp - Test composed subviews --------------------===//
 //
 // 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 pass to test the composed subview patterns.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/StandardOps/Transforms/ComposeSubView.h"
+#include "mlir/Dialect/MemRef/Transforms/ComposeSubView.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestComposeSubViewPass
     : public PassWrapper<TestComposeSubViewPass, OperationPass<FuncOp>> {
   StringRef getArgument() const final { return "test-compose-subview"; }
   StringRef getDescription() const final {
     return "Test combining composed subviews";
   }
   void runOnOperation() override;
   void getDependentDialects(DialectRegistry &registry) const override;
 };
 
 void TestComposeSubViewPass::getDependentDialects(
     DialectRegistry &registry) const {
   registry.insert<AffineDialect>();
 }
 
 void TestComposeSubViewPass::runOnOperation() {
   OwningRewritePatternList patterns(&getContext());
-  populateComposeSubViewPatterns(patterns, &getContext());
+  memref::populateComposeSubViewPatterns(patterns, &getContext());
   (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestComposeSubView() {
   PassRegistration<TestComposeSubViewPass>();
 }
 } // namespace test
 } // namespace mlir

mlir/test/lib/Dialect/StandardOps/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRStandardOpsTestPasses
   TestDecomposeCallGraphTypes.cpp
-  TestComposeSubView.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRPass
   MLIRStandardOpsTransforms
   MLIRTestDialect
   MLIRTransformUtils
   )
 
 target_include_directories(MLIRStandardOpsTestPasses
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/../Test
   ${CMAKE_CURRENT_BINARY_DIR}/../Test
   )

mlir/test/mlir-cpu-runner/memref-reshape.mlir

-// RUN: mlir-opt %s -convert-scf-to-std -std-expand -convert-arith-to-llvm -convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts \
+// RUN: mlir-opt %s -convert-scf-to-std -memref-expand -convert-arith-to-llvm -convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts \
 // RUN: | mlir-cpu-runner -e main -entry-point-result=void \
 // RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext,%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 
 func private @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }
 
 func @main() -> () {

mlir/tools/mlir-opt/CMakeLists.txt

 if(MLIR_INCLUDE_TESTS)
   set(test_libs
     MLIRTestStandardToLLVM
     MLIRAffineTransformsTestPasses
     MLIRDLTITestPasses
     MLIRGPUTestPasses
     MLIRLinalgTestPasses
     MLIRMathTestPasses
+    MLIRMemRefTestPasses
     MLIRSCFTestPasses
     MLIRShapeTestPasses
     MLIRSPIRVTestPasses
     MLIRStandardOpsTestPasses
     MLIRVectorTestPasses
     MLIRTestAnalysis
     MLIRTestDialect
     MLIRTestIR