[MLIR][Tensor] Add ValueBoundsOpInterface for tensor::CollapseShapeOp

[Groverkss]

Summary

This patch teaches ValueBoundOpInterface how to handle tensor::CollapseShapeOp.
This cannot reason about two dynamic dimensions being collapsed into
a single dimension, as Presburger lib cannot represent multiplication of
variables.

[springerm]

mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp

129

If I remember correctly, multiplication of two symbols is currently not supported in ValueBoundsOpInterface because there is no flat representation in the FlatAffineValueConstraints, so cstr.bound() may not do anything. The implementation here is good, but maybe we should add a note/TODO as a comment.

mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir

242

You could add this:

// TODO: Bounds that are a multiplication of two symbols are not supported.
// expected-error @below{{could not reify bound}}
%0 = "test.reify_bound"(%out) {dim = 0} : (tensor<?x?xf32>) -> (index)

mlir/lib/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.cpp

void populateBoundsForIndexValue(Operation *op, Value value,
     auto tensorType =
         llvm::dyn_cast<RankedTensorType>(rankOp.getTensor().getType());
     if (!tensorType)
       return;
     cstr.bound(value) == tensorType.getRank();
   }
 };
 
+struct CollapseShapeOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<CollapseShapeOpInterface,
+                                                   CollapseShapeOp> {
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto collapseShapeOp = cast<CollapseShapeOp>(op);
+    assert(value == collapseShapeOp.getResult() && "invalid value");
+
+    auto tensorType = dyn_cast<RankedTensorType>(collapseShapeOp.getSrcType());
+    if (!tensorType)
+      return;
+
+    ArrayRef<int64_t> shape = tensorType.getShape();
+    ReassociationIndices dimReassoc =
+        collapseShapeOp.getReassociationIndices()[dim];
+    AffineExpr expr = getAffineConstantExpr(1, op->getContext());
+    for (int64_t index : dimReassoc) {
+      if (shape[index] == ShapedType::kDynamic) {
+        expr = expr * cstr.getExpr(collapseShapeOp.getSrc(), index);
+      } else {
+        expr = expr * cstr.getExpr(shape[index]);
+      }
+    }
+
+    cstr.bound(value)[dim] == expr;
+  }
+};
+
 } // namespace
 } // namespace tensor
 } // namespace mlir
 
 void mlir::tensor::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
     tensor::CastOp::attachInterface<tensor::CastOpInterface>(*ctx);
     tensor::DimOp::attachInterface<tensor::DimOpInterface>(*ctx);
     tensor::EmptyOp::attachInterface<tensor::EmptyOpInterface>(*ctx);
     tensor::ExtractSliceOp::attachInterface<tensor::ExtractSliceOpInterface>(
         *ctx);
     tensor::InsertOp::attachInterface<
         DstValueBoundsOpInterfaceExternalModel<tensor::InsertOp>>(*ctx);
     tensor::InsertSliceOp::attachInterface<
         DstValueBoundsOpInterfaceExternalModel<tensor::InsertSliceOp>>(*ctx);
     tensor::PadOp::attachInterface<tensor::PadOpInterface>(*ctx);
     tensor::RankOp::attachInterface<tensor::RankOpInterface>(*ctx);
+    tensor::CollapseShapeOp::attachInterface<tensor::CollapseShapeOpInterface>(
+        *ctx);
   });
 }

mlir/test/Dialect/Tensor/value-bounds-op-interface-impl.mlir

func.func @dynamic_dims_are_maybe_equal_2(%t: tensor<?x?xf32>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %dim0 = tensor.dim %t, %c0 : tensor<?x?xf32>
   %dim1 = tensor.dim %t, %c1 : tensor<?x?xf32>
   // expected-error @below {{could not determine equality}}
   "test.are_equal"(%dim0, %dim1) : (index, index) -> ()
   return
 }
+
+// -----
+
+func.func @collapse_shape_1(%t: tensor<1x?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %dim1 = tensor.dim %t, %c1 : tensor<1x?x?xf32>
+  %out = tensor.collapse_shape %t [[0, 1], [2]] : tensor<1x?x?xf32> into tensor<?x?xf32>
+  %dim0 = tensor.dim %out, %c0 : tensor<?x?xf32>
+  // expected-remark @below {{equal}}
+  "test.are_equal"(%dim0, %dim1) : (index, index) -> ()
+  return
+}
+
+// -----
+
+func.func @collapse_shape_2(%t: tensor<?x?x?xf32>) {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %dim1 = tensor.dim %t, %c2 : tensor<?x?x?xf32>
+  %out = tensor.collapse_shape %t [[0, 1], [2]] : tensor<?x?x?xf32> into tensor<?x?xf32>
+  %dim0 = tensor.dim %out, %c1 : tensor<?x?xf32>
+  // expected-remark @below {{equal}}
+  "test.are_equal"(%dim0, %dim1) : (index, index) -> ()
+  return
+}
+
+// -----
+
+func.func @collapse_shape_3(%t: tensor<?x?x?xf32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %dim1 = tensor.dim %t, %c1 : tensor<?x?x?xf32>
+  %out = tensor.collapse_shape %t [[0, 1], [2]] : tensor<?x?x?xf32> into tensor<?x?xf32>
+  %dim0 = tensor.dim %out, %c0 : tensor<?x?xf32>
+  // expected-error @below {{could not determine equality}}
+  "test.are_equal"(%dim0, %dim1) : (index, index) -> ()
+  return
+}