[mlir][shape] Start a pass that lowers shape constraints.

[silvas]

Summary

This pass converts shape.cstr_* ops to eager (side-effecting)
error-handling code. After that conversion is done, the witnesses are
trivially satisfied and are replaced with shape.const_witness true.

We also throw in a targeted set of canonicalization patterns to clean up
trivial shape.assuming (we can add more as needed; we don't want users
of this pass to always have to run a full canonicalization after this
pass).

[silvas]

Fix variable name.

[herhut]

Nice!

mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp

47

Why not use SelectOp?

124

Is this meant as an optimization to avoid running all canonicalization patterns? Otherwise, it is not clear to me what the benefits over just canonicalizing is.

[silvas]

mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp

47

I was being consistent with ConvertShapeToStandard's lowering of broadcast, which also uses an IfOp

124

Correct. Canonicalizing the whole program is overkill and will lead to fragile interaction of this pass with other passes.

[herhut]

mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp

47

I think both should use SelectOp :)

124

I am worried about the _fragile interaction_ part. Canonicalizing should always be ok and dropping it an canonicalizer pass should not break other passes. Hiding that here seems wrong to me.

Also, users of this now have no choice, as they always get the whole package. I'd much rather just get the constraint lowering patterns. Could this optimization be moved to the use site?

For test, just using --canonicalize won't hurt.

[tpopp]

mlir/include/mlir/Conversion/ShapeToStandard/ShapeToStandard.h

18

Alphabetical order.

[jpienaar]

mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp

2

This doesn't seem to be standard form here

66

Isn't this more that this generates code to ... ? Meaning, if we just wanted to compare we'd use the helper function here

111

Let's add a comment here describing why this works

[silvas]

Address comments

mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp

47

will convert both in a follow-up.

66

Done.

111

Good catch. There's already a comment to that effect in the other pattern. Centralized them to the pass-level comment.

124

As discussed offline, the "fragile interaction" is when debugging miscompiles or other debugging scenarios where one wants to reduce the number of transformations applied to see which one is at fault. Agreed that in a typical bug-free flow canonicalization should always be ok to apply indiscriminately.

Are you ok moving forward with this for now?

[silvas]

Remove shape.assuming canonicalization.

Per offline discussion with Stephan, it LGTM once I remove that.

mlir/include/mlir/Conversion/Passes.td

 
 def ConvertShapeToStandard : Pass<"convert-shape-to-std", "ModuleOp"> {
   let summary = "Convert operations from the shape dialect into the standard "
                 "dialect";
   let constructor = "mlir::createConvertShapeToStandardPass()";
   let dependentDialects = ["StandardOpsDialect", "scf::SCFDialect"];
 }
 
+def ConvertShapeConstraints: Pass<"convert-shape-constraints", "FuncOp"> {
+  let summary = "Convert shape constraint operations to the standard dialect";
+  let description = [{
+    This pass eliminates shape constraints from the program, converting them to
+    eager (side-effecting) error handling code.
+
+    This pass is separate from the regular convert-shape-to-standard, despite
+    converting between the same dialects, because converting shape constraints
+    can happen at a different part of the program than general shape
+    computation lowering.
+  }];
+  let constructor = "mlir::createConvertShapeConstraintsPass()";
+  let dependentDialects = ["StandardOpsDialect", "scf::SCFDialect"];
+}
+
 //===----------------------------------------------------------------------===//
 // SPIRVToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertSPIRVToLLVM : Pass<"convert-spirv-to-llvm", "ModuleOp"> {
   let summary = "Convert SPIR-V dialect to LLVM dialect";
   let constructor = "mlir::createConvertSPIRVToLLVMPass()";
   let dependentDialects = ["LLVM::LLVMDialect"];

mlir/include/mlir/Conversion/ShapeToStandard/ShapeToStandard.h

 #define MLIR_CONVERSION_SHAPETOSTANDARD_SHAPETOSTANDARD_H_
 
 #include <memory>
 
 namespace mlir {
 
 class MLIRContext;
 class ModuleOp;
+class FuncOp;
 template <typename T>
 class OperationPass;
 class OwningRewritePatternList;
 
 void populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx);
 
 std::unique_ptr<OperationPass<ModuleOp>> createConvertShapeToStandardPass();
 
+void populateConvertShapeConstraintsConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx);
+
+std::unique_ptr<OperationPass<FuncOp>> createConvertShapeConstraintsPass();
+
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_SHAPETOSTANDARD_SHAPETOSTANDARD_H_

mlir/lib/Conversion/ShapeToStandard/CMakeLists.txt

 add_mlir_conversion_library(MLIRShapeToStandard
+  ConvertShapeConstraints.cpp
   ShapeToStandard.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ShapeToStandard
 
   DEPENDS
   MLIRConversionPassIncGen
 

mlir/lib/Conversion/ShapeToStandard/ConvertShapeConstraints.cpp

+//===----------------------------------------------------------------------===//
+//
+// 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/Conversion/ShapeToStandard/ShapeToStandard.h"
+
+#include "../PassDetail.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+
+using namespace mlir;
+
+namespace {
+class ConvertCstrBroadcastableOp
+    : public OpRewritePattern<shape::CstrBroadcastableOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(shape::CstrBroadcastableOp op,
+                                PatternRewriter &rewriter) const override {
+    if (op.getType().isa<shape::ShapeType>() ||
+        op.lhs().getType().isa<shape::ShapeType>() ||
+        op.rhs().getType().isa<shape::ShapeType>()) {
+      return rewriter.notifyMatchFailure(
+          op, "cannot convert error-propagating shapes");
+    }
+
+    auto loc = op.getLoc();
+    Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
+    Value one = rewriter.create<ConstantIndexOp>(loc, 1);
+
+    // Find smaller and greater rank and extent tensor.
+    Value lhsRank = rewriter.create<DimOp>(loc, op.lhs(), zero);
+    Value rhsRank = rewriter.create<DimOp>(loc, op.rhs(), zero);
+    Value lhsSmaller =
+        rewriter.create<CmpIOp>(loc, CmpIPredicate::ule, lhsRank, rhsRank);
+    Type indexTy = rewriter.getIndexType();
+    Type extentTensorTy = op.lhs().getType();
+    auto ifOp = rewriter.create<scf::IfOp>(
+        loc, TypeRange{indexTy, extentTensorTy, indexTy, extentTensorTy},
+        lhsSmaller,
+        [&](OpBuilder &b, Location loc) {
+          b.create<scf::YieldOp>(
+              loc, ValueRange{lhsRank, op.lhs(), rhsRank, op.rhs()});
+        },
+        [&](OpBuilder &b, Location loc) {
+          b.create<scf::YieldOp>(
+              loc, ValueRange{rhsRank, op.rhs(), lhsRank, op.lhs()});
+        });
+    Value lesserRank = ifOp.getResult(0);
+    Value lesserRankOperand = ifOp.getResult(1);
+    Value greaterRank = ifOp.getResult(2);
+    Value greaterRankOperand = ifOp.getResult(3);
+
+    Value rankDiff =
+        rewriter.create<SubIOp>(loc, indexTy, greaterRank, lesserRank);
+
+    // Compare the shapes extent by extent, and emit errors for
+    // non-broadcast-compatible shapes.
+    // Two extents are broadcast-compatible if
+    // 1. they are both equal, or
+    // 2. at least one of them is 1.
+
+    rewriter.create<scf::ForOp>(
+        loc, rankDiff, greaterRank, one, llvm::None,
+        [&](OpBuilder &b, Location loc, Value iv, ValueRange) {
+          Value greaterRankOperandExtent = b.create<ExtractElementOp>(
+              loc, greaterRankOperand, ValueRange{iv});
+          Value ivShifted = b.create<SubIOp>(loc, indexTy, iv, rankDiff);
+          Value lesserRankOperandExtent = b.create<ExtractElementOp>(
+              loc, lesserRankOperand, ValueRange{ivShifted});
+
+          Value greaterRankOperandExtentIsOne = b.create<CmpIOp>(
+              loc, CmpIPredicate::eq, greaterRankOperandExtent, one);
+          Value lesserRankOperandExtentIsOne = b.create<CmpIOp>(
+              loc, CmpIPredicate::eq, lesserRankOperandExtent, one);
+          Value extentsAgree =
+              b.create<CmpIOp>(loc, CmpIPredicate::eq, greaterRankOperandExtent,
+                               lesserRankOperandExtent);
+          auto broadcastIsValid =
+              b.create<OrOp>(loc, b.getI1Type(), extentsAgree,
+                             b.create<OrOp>(loc, greaterRankOperandExtentIsOne,
+                                            lesserRankOperandExtentIsOne));
+          b.create<AssertOp>(loc, broadcastIsValid, "invalid broadcast");
+          b.create<scf::YieldOp>(loc);
+        });
+
+    // Now that we have emitted all the assertions, the witness is trivially
+    // satisfied.
+    rewriter.replaceOpWithNewOp<shape::ConstWitnessOp>(op, true);
+    return success();
+  }
+};
+} // namespace
+
+namespace {
+class ConvertCstrRequireOp : public OpRewritePattern<shape::CstrRequireOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(shape::CstrRequireOp op,
+                                PatternRewriter &rewriter) const override {
+    rewriter.create<AssertOp>(op.getLoc(), op.pred(), op.msgAttr());
+    rewriter.replaceOpWithNewOp<shape::ConstWitnessOp>(op, true);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateConvertShapeConstraintsConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  patterns.insert<ConvertCstrBroadcastableOp>(ctx);
+  patterns.insert<ConvertCstrRequireOp>(ctx);
+  // Add in the canonicalization patterns for shape.assuming so that it gets
+  // inlined when its witness becomes a true constant witness.
+  // TODO: Add further targeted canonicalization patterns as needed.
+  shape::AssumingOp::getCanonicalizationPatterns(patterns, ctx);
+}
+
+namespace {
+// This pass eliminates shape constraints from the program, converting them to
+// eager error handling code.
+class ConvertShapeConstraints
+    : public ConvertShapeConstraintsBase<ConvertShapeConstraints> {
+  void runOnOperation() {
+    auto func = getOperation();
+    auto *context = &getContext();
+
+    OwningRewritePatternList patterns;
+    populateConvertShapeConstraintsConversionPatterns(patterns, context);
+
+    if (failed(applyPatternsAndFoldGreedily(func, patterns)))
+      return signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<OperationPass<FuncOp>>
+mlir::createConvertShapeConstraintsPass() {
+  return std::make_unique<ConvertShapeConstraints>();
+}

mlir/test/Conversion/ShapeToStandard/convert-shape-constraints.mlir

+// RUN: mlir-opt -convert-shape-constraints <%s | FileCheck %s
+
+// There's not very much useful to check here other than pasting the output.
+// CHECK-LABEL:   func @cstr_broadcastable(
+// CHECK-SAME:                             %[[LHS:.*]]: tensor<?xindex>,
+// CHECK-SAME:                             %[[RHS:.*]]: tensor<?xindex>) -> !shape.witness {
+// CHECK:           %[[C0:.*]] = constant 0 : index
+// CHECK:           %[[C1:.*]] = constant 1 : index
+// CHECK:           %[[RET:.*]] = shape.const_witness true
+// CHECK:           %[[LHSEXTENT:.*]] = dim %[[LHS]], %[[C0]] : tensor<?xindex>
+// CHECK:           %[[RHSEXTENT:.*]] = dim %[[RHS]], %[[C0]] : tensor<?xindex>
+// CHECK:           %[[LESSEQUAL:.*]] = cmpi "ule", %[[LHSEXTENT]], %[[RHSEXTENT]] : index
+// CHECK:           %[[IFRESULTS:.*]]:4 = scf.if %[[LESSEQUAL]] -> (index, tensor<?xindex>, index, tensor<?xindex>) {
+// CHECK:             scf.yield %[[LHSEXTENT]], %[[LHS]], %[[RHSEXTENT]], %[[RHS]] : index, tensor<?xindex>, index, tensor<?xindex>
+// CHECK:           } else {
+// CHECK:             scf.yield %[[RHSEXTENT]], %[[RHS]], %[[LHSEXTENT]], %[[LHS]] : index, tensor<?xindex>, index, tensor<?xindex>
+// CHECK:           }
+// CHECK:           %[[RANKDIFF:.*]] = subi %[[IFRESULTS:.*]]#2, %[[IFRESULTS]]#0 : index
+// CHECK:           scf.for %[[IV:.*]] = %[[RANKDIFF]] to %[[IFRESULTS]]#2 step %[[C1]] {
+// CHECK:             %[[GREATERRANKOPERANDEXTENT:.*]] = extract_element %[[IFRESULTS]]#3{{\[}}%[[IV]]] : tensor<?xindex>
+// CHECK:             %[[IVSHIFTED:.*]] = subi %[[IV]], %[[RANKDIFF]] : index
+// CHECK:             %[[LESSERRANKOPERANDEXTENT:.*]] = extract_element %[[IFRESULTS]]#1{{\[}}%[[IVSHIFTED]]] : tensor<?xindex>
+// CHECK:             %[[GREATERRANKOPERANDEXTENTISONE:.*]] = cmpi "eq", %[[GREATERRANKOPERANDEXTENT]], %[[C1]] : index
+// CHECK:             %[[LESSERRANKOPERANDEXTENTISONE:.*]] = cmpi "eq", %[[LESSERRANKOPERANDEXTENT]], %[[C1]] : index
+// CHECK:             %[[EXTENTSAGREE:.*]] = cmpi "eq", %[[GREATERRANKOPERANDEXTENT]], %[[LESSERRANKOPERANDEXTENT]] : index
+// CHECK:             %[[OR_TMP:.*]] = or %[[GREATERRANKOPERANDEXTENTISONE]], %[[LESSERRANKOPERANDEXTENTISONE]] : i1
+// CHECK:             %[[BROADCASTISVALID:.*]] = or %[[EXTENTSAGREE]], %[[OR_TMP]] : i1
+// CHECK:             assert %[[BROADCASTISVALID]], "invalid broadcast"
+// CHECK:           }
+// CHECK:           return %[[RET]] : !shape.witness
+// CHECK:         }
+func @cstr_broadcastable(%arg0: tensor<?xindex>, %arg1: tensor<?xindex>) -> !shape.witness {
+  %witness = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  return %witness : !shape.witness
+}
+
+// Check that `shape.assuming` is eliminated after we create the error handling code.
+// CHECK-LABEL: func @assuming
+func @assuming(%arg0: tensor<?xindex>, %arg1: tensor<?xindex>) -> tensor<2xf32> {
+  %witness = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  // CHECK-NOT: shape.assuming
+  // CHECK: %[[CST:.*]] = constant dense<0.000000e+00> : tensor<2xf32>
+  %0 = shape.assuming %witness -> tensor<2xf32> {
+    %c = constant dense<0.0> : tensor<2xf32>
+    shape.assuming_yield %c : tensor<2xf32>
+  }
+  // CHECK: return %[[CST]]
+  return %0 : tensor<2xf32>
+}
+
+// CHECK-LABEL: func @cstr_require
+func @cstr_require(%arg0: i1) -> !shape.witness {
+  // CHECK: %[[RET:.*]] = shape.const_witness true
+  // CHECK: assert %arg0, "msg"
+  // CHECK: return %[[RET]]
+  %witness = shape.cstr_require %arg0, "msg"
+  return %witness : !shape.witness
+}