[mlir] Add a pass to remove all shape.cstr_ and assuming_ ops.

[tpopp]

Summary

This is to provide a utility to remove unsupported constraints or for
pipelines that happen to receive these but cannot lower them due to not
supporting assertions.

[herhut]

mlir/lib/Dialect/Shape/IR/Shape.cpp

179

Nit: Missing .

mlir/lib/Dialect/Shape/Transforms/RemoveShapeConstraints.cpp

22

Could the ones that just erase the op be a template? Also, you can only do this if there are no more users, right? An alternative would be to replace the constraint operations by a true witness and then normal canonicalization should do the rest.

[silvas]

What is the invariant established by this pass that other passes can rely on? It seems that this pass is needed for correctness in many cases, not just optimization. So it needs to have a well-specified contract for later passes to rely on.

With that said, I think this pass probably needs to signalPassFailure() in the cases it cannot establish that invariant.

[rriddle]

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

541

Why a static method? Also, can you change the name to reflect that it is being inline into the parent? inlineRegionIntoParent?

mlir/lib/Dialect/Shape/Transforms/RemoveShapeConstraints.cpp

15

This is not MLIR style, please fix.

Should be using namespace mlir.

80

These should be mlir::

mlir/test/Dialect/Shape/remove-shape-constraints.mlir

21

nit: Remove the extra newline.

[tpopp]

Match MLIR style.

[tpopp]

I haven't finished handling Sean's comment, but wanted to share my thoughts and pieces so far.

Regarding Sean's comment, I think the guarantee of this pass should be that a success indicates all cstr_ and assuming_ ops are gone.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

541

Changed to be non-static. I wasn't sure if it would be safe to delete the op also with a method on the object.

mlir/lib/Dialect/Shape/Transforms/RemoveShapeConstraints.cpp

22

Random errors started occurring when I tried templating this, and I'd rather avoid digging through templates right now.

Related to Sean's point. I think this pass should only succeed if all cstr_ or assuming_ ops are deleted (and have no users) and otherwise be a failure. This is really just to remove the information, as otherwise the user would have lowerings.

I tried to make this contained and remove them fully as I thought it's nicer to see if a pass works rather than relying on needing canonicalization after as a correctness mechanism.

[rriddle]

mlir/lib/Dialect/Shape/IR/Shape.cpp

193

Oops, missed this. Yeah, can you make this static again?

[herhut]

mlir/lib/Dialect/Shape/Transforms/RemoveShapeConstraints.cpp

22

I am not fighting for one way or another of implementing this. However, if the pass would remove all constraint operations then it also has well defined semantics (removes all constraints). The assuming just would be optimized as it always is and relying on canonicalization to clean up is a common pattern.

[tpopp]

Change constraint remove to replace cstr_ ops with passing witnesses.

[tpopp]

mlir/lib/Dialect/Shape/Transforms/RemoveShapeConstraints.cpp

22

I changed my mind on this and changed it to Stephan's way because it will allow more flexibility in which programs this works on in case people start to implement more complicated ideas with Witnesses such as returning them from blocks or using them in their own specialized ops.

[silvas]

mlir/lib/Dialect/Shape/Transforms/RemoveShapeConstraints.cpp

22

+1, I like that approach.

mlir/test/Dialect/Shape/remove-shape-constraints.mlir

4

can we add a simple test for each op that we remove?

[herhut]

Nice!

mlir/test/Dialect/Shape/remove-shape-constraints.mlir

4

Isn't this what this test does? If we do not remove all, the test will fail?

[tpopp]

Add some tests for individual ops.

[tpopp]

I'm going to land this, as I handled Sean's comment, so there doesn't seem to be anything contentious left.

mlir/include/mlir/Dialect/Shape/IR/Shape.h

 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
+class PatternRewriter;
+
 namespace shape {
 
 namespace ShapeTypes {
 enum Kind {
   Component = Type::FIRST_SHAPE_TYPE,
   Element,
   Shape,
   Size,

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   let arguments = (ins Shape_SizeType:$arg);
   let results = (outs Index:$result);
 
   let assemblyFormat = "$arg attr-dict";
 
   let hasFolder = 1;
 }
 
-def Shape_YieldOp : Shape_Op<"yield", 
+def Shape_YieldOp : Shape_Op<"yield",
     [HasParent<"ReduceOp">,
      NoSideEffect,
      ReturnLike,
      Terminator]> {
   let summary = "Returns the value to parent op";
 
   let arguments = (ins Variadic<AnyType>:$operands);
 
def Shape_AssumingOp : Shape_Op<"assuming",
   }];
   let arguments = (ins Shape_WitnessType:$witness);
   let regions = (region SizedRegion<1>:$doRegion);
   let results = (outs Variadic<AnyType>:$results);
 
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 
+  let extraClassDeclaration = [{
+    // Inline the region into the region containing the AssumingOp and delete
+    // the AssumingOp.
+    //
+    // This does no checks on the inputs to the AssumingOp.
+    static void inlineRegionIntoParent(AssumingOp &op, PatternRewriter &rewriter);
+  }];
+
   let hasCanonicalizer = 1;
 }
 
 def Shape_AssumingYieldOp : Shape_Op<"assuming_yield",
                                      [NoSideEffect, ReturnLike, Terminator]> {
   let summary = "Yield operation";
   let description = [{
     This yield operation represents a return operation within the assert_and_exec

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

 
 #ifndef MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_
 
 #include <memory>
 
 namespace mlir {
 
+class FunctionPass;
 class MLIRContext;
 class OwningRewritePatternList;
 class Pass;
 
 /// Creates an instance of the ShapeToShapeLowering pass that legalizes Shape
 /// dialect to be convertible to Standard. For example, `shape.num_elements` get
 /// transformed to `shape.reduce`, which can be lowered to SCF and Standard.
 std::unique_ptr<Pass> createShapeToShapeLowering();
 
 /// Collects a set of patterns to rewrite ops within the Shape dialect.
 void populateShapeRewritePatterns(MLIRContext *context,
                                   OwningRewritePatternList &patterns);
+
+// Collects a set of patterns to replace all constraints with passing witnesses.
+// This is intended to then allow all ShapeConstraint related ops and data to
+// have no effects and allow them to be freely removed such as through
+// canonicalization and dead code elimination.
+//
+// After this pass, no cstr_ operations exist.
+void populateRemoveShapeConstraintsPatterns(OwningRewritePatternList &patterns,
+                                            MLIRContext *ctx);
+std::unique_ptr<FunctionPass> createRemoveShapeConstraintsPass();
+
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES_H_

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

 //===-- Passes.td - ShapeOps pass 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_SHAPE_TRANSFORMS_PASSES
 #define MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
 
+def RemoveShapeConstraints : FunctionPass<"remove-shape-constraints"> {
+  let summary = "Replace all cstr_ ops with a true witness";
+  let constructor = "mlir::createRemoveShapeConstraintsPass()";
+}
+
 def ShapeToShapeLowering : FunctionPass<"shape-to-shape-lowering"> {
   let summary = "Legalize Shape dialect to be convertible to Standard";
   let constructor = "mlir::createShapeToShapeLowering()";
 }
 
 #endif // MLIR_DIALECT_SHAPE_TRANSFORMS_PASSES

mlir/lib/Dialect/Shape/IR/Shape.cpp

struct AssumingWithTrue : public OpRewritePattern<AssumingOp> {
   using OpRewritePattern<AssumingOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AssumingOp op,
                                 PatternRewriter &rewriter) const override {
     auto witness = op.witness().getDefiningOp<ConstWitnessOp>();
     if (!witness || !witness.passingAttr())
       return failure();
 
-    auto *blockBeforeAssuming = rewriter.getInsertionBlock();
-    auto *assumingBlock = op.getBody();
-    auto initPosition = rewriter.getInsertionPoint();
-    auto *blockAfterAssuming =
-        rewriter.splitBlock(blockBeforeAssuming, initPosition);
+    AssumingOp::inlineRegionIntoParent(op, rewriter);
+    return success();
+  }
+};
+} // namespace
+
+void AssumingOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
+                                             MLIRContext *context) {
+  // If taking a passing witness, inline region.
+  patterns.insert<AssumingWithTrue>(context);
+}
+
+void AssumingOp::inlineRegionIntoParent(AssumingOp &op,
+                                        PatternRewriter &rewriter) {
 
-    // Remove the AssumingOp and AssumingYieldOp.
-    auto &yieldOp = assumingBlock->back();
-    rewriter.inlineRegionBefore(op.doRegion(), blockAfterAssuming);
-    rewriter.replaceOp(op, yieldOp.getOperands());
-    rewriter.eraseOp(&yieldOp);
 
-    // Merge blocks together as there was no branching behavior from the
-    // AssumingOp.
-    rewriter.mergeBlocks(assumingBlock, blockBeforeAssuming);
-    rewriter.mergeBlocks(blockAfterAssuming, blockBeforeAssuming);
-    return success();
-  }
-};
-} // namespace
-
-void AssumingOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
-                                             MLIRContext *context) {
-  // If taking a passing witness, inline region
-  patterns.insert<AssumingWithTrue>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // AssumingAllOp
 //===----------------------------------------------------------------------===//
 OpFoldResult AssumingAllOp::fold(ArrayRef<Attribute> operands) {
   // Iterate in reverse to first handle all constant operands. They are
   // guaranteed to be the tail of the inputs because this is commutative.

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

 add_mlir_dialect_library(MLIRShapeOpsTransforms
+  RemoveShapeConstraints.cpp
   ShapeToShapeLowering.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ShapeOps/Transforms
 
   DEPENDS
   MLIRShapeTransformsIncGen
   )

mlir/lib/Dialect/Shape/Transforms/RemoveShapeConstraints.cpp

+//===-- RemoveShapeConstraints.cpp - Remove Shape Cstr and Assuming 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Shape/Transforms/Passes.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+namespace {
+/// Removal patterns.
+class RemoveCstrBroadcastableOp
+    : public OpRewritePattern<shape::CstrBroadcastableOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(shape::CstrBroadcastableOp op,
+                                PatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<shape::ConstWitnessOp>(op.getOperation(), true);
+    return success();
+  }
+};
+
+class RemoveCstrEqOp : public OpRewritePattern<shape::CstrEqOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(shape::CstrEqOp op,
+                                PatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<shape::ConstWitnessOp>(op.getOperation(), true);
+    return success();
+  }
+};
+
+/// Removal pass.
+class RemoveShapeConstraintsPass
+    : public RemoveShapeConstraintsBase<RemoveShapeConstraintsPass> {
+
+  void runOnFunction() override {
+    MLIRContext &ctx = getContext();
+
+    OwningRewritePatternList patterns;
+    populateRemoveShapeConstraintsPatterns(patterns, &ctx);
+
+    applyPatternsAndFoldGreedily(getFunction(), patterns);
+  }
+};
+
+} // namespace
+
+void mlir::populateRemoveShapeConstraintsPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  patterns.insert<RemoveCstrBroadcastableOp, RemoveCstrEqOp>(ctx);
+}
+
+std::unique_ptr<FunctionPass> mlir::createRemoveShapeConstraintsPass() {
+  return std::make_unique<RemoveShapeConstraintsPass>();
+}

mlir/test/Dialect/Shape/remove-shape-constraints.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints -canonicalize <%s | FileCheck %s --dump-input=fail --check-prefixes=CANON,CHECK-BOTH
+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file -remove-shape-constraints <%s | FileCheck %s --dump-input=fail --check-prefixes=REPLACE,CHECK-BOTH
+
+// -----
+// Check that cstr_broadcastable is removed.
+//
+// CHECK-BOTH: func @f
+func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
+  // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true
+  // REPLACE-NOT: shape.cstr_eq
+  // REPLACE: shape.assuming %[[WITNESS]]
+  // CANON-NEXT: test.source
+  // CANON-NEXT: return
+  %0 = shape.cstr_broadcastable %arg0, %arg1
+  %1 = shape.assuming %0 -> index {
+    %2 = "test.source"() : () -> (index)
+    shape.assuming_yield %2 : index
+  }
+  return %1 : index
+}
+
+// -----
+// Check that cstr_eq is removed.
+//
+// CHECK-BOTH: func @f
+func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
+  // REPLACE-NEXT: %[[WITNESS:.+]] = shape.const_witness true
+  // REPLACE-NOT: shape.cstr_eq
+  // REPLACE: shape.assuming %[[WITNESS]]
+  // CANON-NEXT: test.source
+  // CANON-NEXT: return
+  %0 = shape.cstr_eq %arg0, %arg1
+  %1 = shape.assuming %0 -> index {
+    %2 = "test.source"() : () -> (index)
+    shape.assuming_yield %2 : index
+  }
+  return %1 : index
+}
+
+// -----
+// With a non-const value, we cannot fold away the code, but all constraints
+// should be removed still.
+//
+// CHECK-BOTH: func @f
+func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> index {
+  // CANON-NEXT: test.source
+  // CANON-NEXT: return
+  %0 = shape.cstr_broadcastable %arg0, %arg1
+  %1 = shape.cstr_eq %arg0, %arg1
+  %2 = shape.assuming_all %0, %1
+  %3 = shape.assuming %0 -> index {
+    %4 = "test.source"() : () -> (index)
+    shape.assuming_yield %4 : index
+  }
+  return %3 : index
+}