[mlir][shape] Fold shape.broadcast with scalar operands

[jpienaar]

Nice, thanks

[tpopp]

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

70

If you're willing to add even more changes, this should be commutative.

93

Nit: Include in the description that you added assemblyFormats for these ops.

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

240

Nit: Change commit description to say "Fold shape.broadcast with at least 1 scalar operand" or similar since I read that as both being scalar.

mlir/test/Dialect/Shape/canonicalize.mlir

74

This test would have behaved the same before your changes because both are constant. I think the example above with the order of inputs flipped would currently fail.

[tpopp]

LGTM otherwise

[herhut]

Comments

[herhut]

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

70

See https://reviews.llvm.org/D83856

mlir/test/Dialect/Shape/canonicalize.mlir

74

I added the flipped case.

[herhut]

Add back test and unbreak.

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

def Shape_AddOp : Shape_Op<"add", [SameOperandsAndResultType]> {
   }];
 
   let arguments = (ins Shape_SizeType:$lhs, Shape_SizeType:$rhs);
   let results = (outs Shape_SizeType:$result);
 
   let assemblyFormat = "$lhs `,` $rhs attr-dict";
 }
 
 def Shape_BroadcastOp : Shape_Op<"broadcast", []> {
   let summary = "Returns the broadcasted output shape of two inputs";
   let description = [{
     Computes the broadcasted output shape following:
     1. If any inputs are unranked, output is unranked;
     2. Else the input array with number of dimensions smaller than the max
        input dimension, has 1’s prepended to its shapes and the output shape is
        calculated as follows:
 
            output[i] = lhs[i] if lhs[i] == rhs[i] or rhs[i] is unknown/undefined
                      = rhs[i] if lhs[i] is unknown/undefined
                      = lhs[i] if rhs[i] == 1
                      = rhs[i] if lhs[i] == 1
                      = error  if lhs[i] != rhs[i]
 
     Op has an optional string attribute for the error case where there is no
     broadcastable output shape possible for the given inputs.
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs,
                    OptionalAttr<StrAttr>:$error);
   let results = (outs Shape_ShapeType:$result);
 
+  let assemblyFormat = "$lhs `,` $rhs attr-dict";
+
   let hasFolder = 1;
 }
 
 def Shape_ConstShapeOp : Shape_Op<"const_shape", [ConstantLike, NoSideEffect]> {
   let summary = "Creates a constant of !shape.shape type";
   let description = [{
     Creates a !shape.shape with rank given by the length of `shape` and with
     dimension sizes given by the values of `shape`.
let description = [{
     concat([2,3], [4,5]) -> [2,3,4,5]
     concat([], []) -> []
     concat([], [4,5,6]) -> [4,5,6]
   }];
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
   let results = (outs Shape_ShapeType:$result);
 
+  let assemblyFormat = "$lhs `,` $rhs attr-dict";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Shape constraint related ops.
 //===----------------------------------------------------------------------===//
 
 // TODO: Move the code below and witnesses to a different file.

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

static LogicalResult verify(AssumingAllOp op) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
-  if (!operands[0] || !operands[1])
+  if (!operands[1])
     return nullptr;
-  auto lhsShape = llvm::to_vector<6>(
-      operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
+
   auto rhsShape = llvm::to_vector<6>(
       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  if (rhsShape.empty())
+    return lhs();
+
+  if (!operands[0])
+    return nullptr;
+
+  auto lhsShape = llvm::to_vector<6>(
+      operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
+  if (lhsShape.empty())
+    return rhs();
+
   SmallVector<int64_t, 6> resultShape;
   // If the shapes are not compatible, we can't fold it.
   // TODO: Fold to an "error".
   if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
     return nullptr;
   Builder builder(getContext());
   return builder.getIndexTensorAttr(resultShape);
 }

mlir/test/Dialect/Shape/canonicalize.mlir

 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.const_shape [7, 2]
   %0 = shape.const_shape [1, 2]
   %1 = shape.const_shape [7, 1]
-  %2 = "shape.broadcast"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  %2 = shape.broadcast %0, %1
+  return %2 : !shape.shape
+}
+
+// -----
+
+// Rhs is a scalar.
+// CHECK-LABEL: func @f
+func @f(%arg0 : !shape.shape) -> !shape.shape {
+  // CHECK: return %arg0
+  %0 = shape.const_shape []
+  %1 = shape.broadcast %arg0, %0
+  return %1 : !shape.shape
+}
+
+// -----
+
+// Lhs is a scalar.
+// CHECK-LABEL: func @f
+func @f() -> !shape.shape {
+  // CHECK-NOT: broadcast
+  %0 = shape.const_shape []
+  %1 = shape.const_shape [1,2,3]
+  %2 = shape.broadcast %0, %1
   return %2 : !shape.shape
 }
 
 // -----
 
 // Incompatible shapes. No folding.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.broadcast
   %0 = shape.const_shape [2]
   %1 = shape.const_shape [7]
-  %2 = "shape.broadcast"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  %2 = shape.broadcast %0, %1
   return %2 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> !shape.shape {
   // CHECK: shape.const_shape [0, 1, 2, 3]
   %lhs = shape.const_shape [0, 1]
   %rhs = shape.const_shape [2, 3]
-  %0 = "shape.concat"(%lhs, %rhs) : (!shape.shape, !shape.shape) -> !shape.shape
+  %0 = shape.concat %lhs, %rhs
   return %0 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func @f() -> tensor<2xindex> {

mlir/test/Dialect/Shape/ops.mlir

 
 func @test_shape_num_elements_fixed() {
   %0 = shape.const_shape [1, 57, 92]
   %1 = call @shape_num_elements(%0) : (!shape.shape) -> (!shape.size)
   %3 = "shape.print"(%1) : (!shape.size) -> !shape.size
   return
 }
 
-func @test_broadcastable_fixed() {
+func @test_broadcast_fixed() {
   %0 = shape.const_shape [10, 1, 57, 92]
   %1 = shape.const_shape [4, 57, 92]
-  %2 = "shape.broadcastable"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  %2 = shape.broadcast %0, %1
   %3 = "shape.print"(%2) : (!shape.shape) -> !shape.shape
   return
 }
 
 func @test_shape_any_fixed() {
   %0 = shape.const_shape [4, 57, 92]
   %1 = shape.const_shape [4, 57, 92]
   %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape