[mlir][wip] Start Shape dialect

[jpienaar]

Summary

• Add basic skeleton for Shape dialect;
• Add description of types and ops to be used;

Naming still open and pretty form would enhance this. As mentioned on
discourse, will do this piece wise for discussion/easier feedback cycles. Tried
to reduce parts while allowing enough to get idea/write tests with.

[rriddle]

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

3

MLIR -> LLVM

9

nit: & -> and

10

typo: relations?

[rriddle]

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

51

///?

55

///

68

Can we get comments on each of these types?

96

///

115

weird spacing here.

117

This header comment looks off.

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

3

MLIR -> LLVM

71

Are these wrapped at 80 characters? They look like they are overflowing.

mlir/include/mlir/IR/DialectSymbolRegistry.def

27

Not this revision, but we should sort these.

mlir/lib/Dialect/Shape/DialectRegistration.cpp

3

MLIR -> LLVM

[merge_guards_bot]

Unit tests: fail. 62150 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: fail. clang-tidy found 2 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[stellaraccident]

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

88

Just "Value" now, right?

[stellaraccident]

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

181

I might just be mis-understanding, but to my reading, I'm not sure this formulation is right. I believe this would yield "left aligned" broadcasting, whereas the traditional rules for a ranked broadcast are "right aligned":

OUTPUT DIM: 0 1 2 
       LHS: Z Y X
       RHS:     X

I think there is a step before where you want to align the lower ranked input dimensions to fill towards the right, with respect to the output dims. For full generality, you may also need a dimension map (similar to xhl_hlo broadcast_dimensions attributes).

[jpienaar]

Addresses review comments.

[jpienaar]

Updated, thanks

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

68

This made me think about having ODS generate these from the description :)

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

88

Ha, good point I had written this before that change :)

181

Good catch, I forgot to add that step. In the code I had used zip_longest and reverse iterator, but that doesn't lend itself to words that well :). Here I'm pursuing broadcasting as defined in numpy and used by ML frameworks (https://docs.scipy.org/doc/numpy/reference/ufuncs.html#broadcasting) but I could be wrong so please check.

Re: broadcast_dimensions attribute, while that is more general that is not the same semantics as what folks think of when broadcasting is mentioned. That is a combination of reshape & broadcast IMHO.

mlir/include/mlir/IR/DialectSymbolRegistry.def

27

Yeah I was very tempted ... I'll do in NFC follow up.

[jpienaar]

Ran clang-format

[merge_guards_bot]

Unit tests: fail. 62150 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: fail. clang-tidy found 2 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: fail. 62150 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: fail. clang-tidy found 2 errors and 1 warnings. 0 of them are added as review comments below (why?).

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

mlir/include/mlir/Dialect/CMakeLists.txt

9

Seems like this should come before SPIRV?

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

19

Why do you need Function.h here?

37

Any reason not to sort these?

68

Realistically, we should be able to auto-generate simple types pretty trivially.

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

152

Can we sort these operations alphabetically?

[antiagainst]

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

182

I'd typically think an op means some action so the name is typically a verb but here we have broadcastable which sounds like a trait. What about just name it as broadcast?

208

My 2c: join sounds better for me as it indicates some operation behind. any_of gives me the feeling that the result will be the same as one of the input, which is not always true. :)

281

Would it make sense to name it as debug_print to be more clear?

[herhut]

Great to see progress on this! Some comments.

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

47

Nit: It is not used for dimensions, only, but generally as a positive integer value with support for being unknown and invalid. In that sense it is more like a size_t than a dim. Maybe shape.size?

56

I find the name type somewhat misleading. I think Shape_Shape would describe better what it is (a shape). But I get the desire to name it type. Should they all be Shape_DimType and Shape_ShapeType?

101

DimOrShape?

111

+1 to the constant op. Also, why expose a -1 here? Why not use

shape.constant unknown : !shape.size
shape.constant invalid : !shape.size
shape.constant <positive integer> : !shape.size

and then parse that internally into some special integer values? This would leave more flexibility to change the backing type to something other than int later.

135

Is there a reason for this? Also, thinking of library calls that return ranked tensors of statically unknown rank, this function would not allow me to create a shape for them. Should we also have a version that takes just the rank?

182

I agree. Broadcast seems a better name, as it broadcasts two shapes into a resulting broadcasted shape.

208

any_of is less constrained in that it does not have to pick the join of the two values, whereas join has rigid semantics. On the other hand, the description below describes the join of the two values.

222

I would expect

join ([], []) -> []
join([], [*]) -> []
join([], [?,...,?]) -> [invalid]

223

join([1,?], [2, ?, ?]) -> [invalid] ?

243

"for each every" is one too many.

264

Should this be shape.mul if it existed?

272

Does this mean the rank is unknown? Or also if parts of the shape are unknown?

If known rank was allowed, one could still implement computing the rank by counting the elements of the vector (as a somewhat useless example).

[nicolasvasilache]

I question the relevance of carrying implicit broadcasting concepts from HLO and other frontends into MLIR core.
I see this as a red flag that the required separation of concerns has not happened.
Anything related to broadcast should be on the XLA / frontend side IMO.

I do not see a mention of broadcast in the discourse discussion: https://llvm.discourse.group/t/shape-type-dialect/423/4.

[nicolasvasilache]

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

191

This is frontend logic that I do not think should leak into MLIR.

191

I meant into MLIR *core*, sorry.

[stellaraccident]

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

191

I tend to agree. We do need this to lower frontend dialects that choose to implement numpy-style broadcasting, but that does not mean it needs to be in core (or at least in this level of core -- I could see having a numpy_broadcast_shape op somewhere).

I suspect that there may have been some thought of needing this in some follow-on analysis algorithms. It would probably be better to discuss more in that context.

[jpienaar]

Updating to address review comments (sorting ops, types, making type names uniform, renaming broadcastable op, renaming dim type, fixing example).

[jpienaar]

Updated, thanks

In D73944#1864082, @nicolasvasilache wrote:

I question the relevance of carrying implicit broadcasting concepts from HLO and other frontends into MLIR core.

I disagree. This is about enabling frontend dialects to express their shape constraints. And need not even be implicit (e.g., explicit broadcast operator has broadcast shape op in its shape function). And I see this similar to signed integer discussion. While we have preferences and would prefer to never use implicit broadcast, many frontend dialects do and this is about enabling their modelling. Enabling frontend dialects to express their constraints is very useful. This is similar to broadcastable trait that we have at the moment.

I see this as a red flag that the required separation of concerns has not happened.
Anything related to broadcast should be on the XLA / frontend side IMO.

Are you saying we should not have any frontend dialects? This moves responsibility into op construction of the frontend time, which is often not where we wanted & optimizing frontend dialect is definitely something that we want, which might require legalizing from one frontend dialect to another.

mlir/include/mlir/Dialect/CMakeLists.txt

9

Depends on if you follow sort order case sensitive or insensitive :) I normally default to case sensitive, but can do insensitive.

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

19

Ah this was left over from previous testing, thanks

37

It was a semi hierarchy in my mind, but that doesn't really matter, so sorted.

68

Exactly. Should be easy to address.

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

47

Good point and I like that, renaming.

56

Yeah, Shape_Shape was just weird, changing to uniformly suffixing with Type.

111

Mostly as it made writing tests easier while allowing reusing the integer attribute parser/dense storage and so switch later. But I can start with an ArrayAttr and then do the above.

For invalid case I feel constant doesn't highlight the error state well enough and would rather that be a separate op that takes a string attribute. WDYT?

135

The one below does. The main reason was that then you need to special case value interpretation (where I was using only integers, e.g., require rank followed by dims so that you could do [?] for unranked [1, ?] for 1D ranked but unknown etc.) but with expansion to using strings, one can use * now and so rolled it into constant.

191

Correct, we need this for frontend dialects (TF, TFLite, XLA [client] HLO, etc) to be able to express their shape constraints. it need not be used by standard etc. and we could add specialization that only works for ranked types. But I see unranked inputs and broadcast behavior as common for many frontend dialects and this similar to broadcastable trait or signed integers. We can move the op into a different file.

Now, of course nobody is forced to use this and standard ops will never generate nor need to solve shape equations using this.

223

I changed notations and didn't do it correctly.

243

It was just so good it happened every every time 😁

264

D'oh yes. This was computing sum of dimensions :)

272

Unraked, updated to clarify. Correct. well and I was using this for op where one computed the product of known indices, so inside the expression it would conditionally reduce 🙂

281

Sounds good. I quite like it, made writing tests easier (although I should refactor those into non-executable tests ...).

[jpienaar]

Reduce with mul should be init with 1 instead of 0

[mehdi_amini]

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

245

I think si.dim should be shape.size?

[herhut]

Thanks for starting this. I'd be happy to see this land.

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

111

I like the idea. Invalid shapes should always have some reason attached to them to ease debugging whereas normal constants probably don't. Something like shape.error.

191

Ultimately, we could think of making the shape dialect and the machinery built on top of it extensible so that other front-ends can add their own high-level shape operations to model their constraints. In such a world, I would agree with @nicolasvasilache that the numpy-style broadcast operation should live with the corresponding front-end dialect.
However, as a first step, we should build a closed system that is powerful enough to model constraints for the input languages we care about (which includes HLO). This requires computing shapes for broadcast, so we should have the operation.

[jpienaar]

Fix incorrectly named si.dim to shape.size

[jpienaar]

Updated typo, thanks

[stephenneuendorffer]

It seems like this landed without enough code to compile? There's no .cpp file which includes the implementation .cpp.inc generated from tablegen?
As a result BUILD_SHARED_LIBS=on fails.

Steve

mlir/include/mlir/Dialect/CMakeLists.txt

 add_subdirectory(AffineOps)
 add_subdirectory(FxpMathOps)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(LoopOps)
 add_subdirectory(QuantOps)
 add_subdirectory(SPIRV)
+add_subdirectory(Shape)
 add_subdirectory(StandardOps)
 add_subdirectory(VectorOps)

mlir/include/mlir/Dialect/Shape/CMakeLists.txt

+add_subdirectory(IR)

mlir/include/mlir/Dialect/Shape/IR/CMakeLists.txt

+add_mlir_dialect(ShapeOps ShapeOps)

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

+//===- Shape.h - MLIR Shape dialect -----------------------------*- C++ -*-===//
+//
+// Part of the MLIR 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 defines the shape dialect that is used to describe & solve shape
+// relantions of MLIR operations using ShapedType.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SHAPE_SHAPE_H
+#define MLIR_SHAPE_SHAPE_H
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/FunctionSupport.h"
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+namespace shape {
+
+/// This dialect contains shape inference related operations and facilities.
+class ShapeDialect : public Dialect {
+public:
+  /// Create the dialect in the given `context`.
+  explicit ShapeDialect(MLIRContext *context);
+};
+
+namespace ShapeTypes {
+enum Kind {
+  Dim = Type::FIRST_SHAPE_TYPE,
+  Shape,
+  Element,
+  Component,
+  ValueShape,
+};
+} // namespace ShapeTypes
+
+// The type of a single dimension.
+class DimType : public Type::TypeBase<DimType, Type> {
+ public:
+  using Base::Base;
+
+  static DimType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::Dim);
+  }
+
+  // Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) { return kind == ShapeTypes::Kind::Dim; }
+};
+
+// The shape descriptor type represents rank and dimension sizes.
+class ShapeType : public Type::TypeBase<ShapeType, Type> {
+ public:
+  using Base::Base;
+
+  static ShapeType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::Shape);
+  }
+
+  // Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) { return kind == ShapeTypes::Kind::Shape; }
+};
+
+class ElementType : public Type::TypeBase<ElementType, Type> {
+public:
+  using Base::Base;
+
+  static ElementType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::Element);
+  }
+
+  // Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) {
+    return kind == ShapeTypes::Kind::Element;
+  }
+};
+
+class ComponentType : public Type::TypeBase<ComponentType, Type> {
+public:
+  using Base::Base;
+
+  static ComponentType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::Component);
+  }
+
+  // Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) {
+    return kind == ShapeTypes::Kind::Component;
+  }
+};
+
+// The ValueShape represents a (potentially unknown) runtime value and shape.
+class ValueShapeType : public Type::TypeBase<ValueShapeType, Type> {
+public:
+  using Base::Base;
+
+  static ValueShapeType get(MLIRContext *context) {
+    return Base::get(context, ShapeTypes::Kind::ValueShape);
+  }
+
+  // Support method to enable LLVM-style type casting.
+  static bool kindof(unsigned kind) {
+    return kind == ShapeTypes::Kind::ValueShape;
+  }
+};
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Shape/IR/ShapeOps.h.inc"
+
+} // namespace shape
+}  // namespace mlir
+
+#endif // MLIR_SHAPE_SHAPE_H

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

+//===- Shape.td - Shape operations definition --------------*- tablegen -*-===//
+//
+// Part of the MLIR 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 is the operation definition file for Shape dialect operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SHAPE_OPS
+#define SHAPE_OPS
+
+include "mlir/IR/OpBase.td"
+
+// TODO(jpienaar): Move to base.
+def AnyShaped: ShapedContainerType<[AnyType], IsShapedTypePred, "shaped">;
+
+//===----------------------------------------------------------------------===//
+// Shape Inference dialect definitions
+//===----------------------------------------------------------------------===//
+
+def ShapeDialect : Dialect {
+  let name = "shape";
+
+  let summary = "Types and operations for shape dialect";
+  let description = [{
+    This dialect contains operations for shape inference.
+
+    Note: Unless explicitly stated, all functions that return a shape and take
+    shapes as input, return the invalid shape if one of its operands is an
+    invalid shape. This avoids flagging multiple errors for one verification
+    failure. The dialect itself does not specify how errors should be combined
+    (there are multiple different options, from always chosing first operand,
+    concatting etc. on how to combine them).
+  }];
+
+  let cppNamespace = "shape";
+}
+
+def Shape_Dim : DialectType<ShapeDialect, CPred<"$_self.isa<DimType>()">, "dim"> {
+  let typeDescription = [{
+    `shape.dim` represents a dimension type which could be either a non-negative
+    integer, unknown or error.
+
+    Operations on `shape.dim` types are specialized to handle
+    unknown/dynamic value. So, for example, `<unknown> + x == <unknown>` for all
+    non-error `x : !shape.dim` (e.g., an unknown value does not become known due to
+    addition).
+  }];
+}
+
+def Shape_Type : DialectType<ShapeDialect, CPred<"$_self.isa<ShapeType>()">, "shape"> {
+  let typeDescription = [{
+    `shape.type` represents either an unranked shape, a ranked shape with possibly
+    unknown dimensions or an invalid shape. The rank is of type `shape.dim` and, if
+    rank is known, the extent is a 1D tensor of type `shape.dim`.
+
+    Shape is printed:
+    * `[*]` if it is an unranked shape
+    * `[?, 2]` if a rank 2 tensor with one unknown dimension
+    * `[3, 4]` is a rank 2 static tensor
+    * `[]` is a scalar
+    * `[1]` is a rank 1 tensor with 1 element
+    * `[invalid]` for an invalid shape
+  }];
+}
+
+def Shape_ElementType : DialectType<ShapeDialect, CPred<"$_self.isa<ElementType>()">, "element type"> {
+  let typeDescription = [{
+    `shape.element_type` represents the element type of the ShapedType. It may
+    be unknown, error or regular element type supported by ShapedType.
+  }];
+}
+
+def Shape_ComponentType : DialectType<ShapeDialect, CPred<"$_self.isa<ComponentType>()">, "component type"> {
+  let typeDescription = [{
+    `shape.element_type` represents the element type of the ShapedType. It may
+    be unknown, error or regular element type supported by ShapedType.
+  }];
+}
+
+def Shape_ValueShape : DialectType<ShapeDialect, CPred<"$_self.isa<ValueShapeType>()">, "value shape"> {
+  let typeDescription = [{
+    `shape.value_shape` represents the value produced by an operation (this corresponds to
+    `Value*` in the compiler) and a shape. Conceptually this is a tuple of a
+    value (potentially unknown) and `shape.type`. The value and shape can either
+    or both be unknown. If both the `value` and `shape` are known, then the
+    shape of `value` is conformant with `shape`.
+  }];
+}
+
+def Shape_DimOrType: AnyTypeOf<[Shape_Dim, Shape_Type], "dim or type">;
+
+//===----------------------------------------------------------------------===//
+// Shape op definitions
+//===----------------------------------------------------------------------===//
+
+// Base class for the operation in this dialect
+class Shape_Op<string mnemonic, list<OpTrait> traits = []> :
+    Op<ShapeDialect, mnemonic, traits>;
+
+// TODO: consider just making a constant op instead.
+def Shape_ConstantDimOp : Shape_Op<"constant_dim", []> {
+  let summary = "Creates a constant dim";
+  let description = [{
+    An operation that builds a dim from integer attribute. -1 is treated as
+    unknown, no other negative values are supported.
+
+    ```mlir
+    %x = shape.constant_dim 10 : !shape.dim
+    ```
+  }];
+
+  let arguments = (ins I32Attr:$value);
+  let results = (outs Shape_Dim:$result);
+}
+
+def Shape_ConstantShapeOp : Shape_Op<"constant_shape", []> {
+  let summary = "Create a constant shape descripton";
+  let description = [{
+    An operation that builds a Shape of a given rank and extents. The value is
+    treated like an array of dimensions (converted as if fed to
+    `shape.constant_dim`) with the rank of the shape equal to the number of
+    elements in the array.
+
+    Note: it is not possible to create an unranked shape using the op.
+  }];
+
+  // TODO: ArrayAttr is not specific enough.
+  let arguments = (ins ArrayAttr:$value);
+  let results = (outs Shape_Type:$output);
+}
+
+def Shape_CreateUnrankedShapeOp : Shape_Op<"unranked_shape", []> {
+  let summary = "Creates an unranked Shape descriptor";
+  let description = [{
+    An operation that builds a shape without known rank.
+  }];
+
+  let results = (outs Shape_Type:$output);
+}
+
+def Shape_CreateShapeOp : Shape_Op<"create_shape", []> {
+  let summary = "Creates a shape descriptor from a tensor";
+  let description = [{
+    Creates a shape from a 1D integral tensor. The rank equals the number of
+    elements in the tensor, and extent matches the values of the elements.
+  }]
+
+  let arguments = (ins I32Tensor:$input);
+  let results = (outs Shape_Type:$result);
+}
+
+def Shape_ShapeOfOp : Shape_Op<"shape_of", []> {
+  let summary = "Returns shape of a value or shaped type operand";
+
+  let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueType]>:$arg);
+  let results = (outs Shape_Type:$result);
+}
+
+def Shape_AddOp : Shape_Op<"add", [SameOperandsAndResultType]> {
+  let summary = "Adds two dimensions";
+  let description = [{
+    Adds two valid dimensions as follows:
+    * lhs + rhs = unknown if either lhs or rhs unknown;
+    * lhs + rhs = (int)lhs + (int)rhs if known;
+  }];
+
+  let arguments = (ins Shape_DimOrType:$lhs, Shape_DimOrType:$rhs);
+  let results = (outs Shape_DimOrType:$result);
+}
+
+def Shape_BroadcastableOp : Shape_Op<"broadcastable", []> {
+  let summary = "Returns the broadcastable output shape";
+  let description = [{
+    Computes the broadcastable output shape following:
+    1. If any inputs are unranked, output is unranked;
+    2. Then for all 0 <= i < max rank of input shapes
+
+           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 possible for the given inputs.
+  }];
+
+  let arguments = (ins Shape_Type:$lhs, Shape_Shape:$rhs,
+                   OptionalAttr<StringAttr>:$error);
+  let results = (outs Shape_Type:$result);
+}
+
+// Previously called any_of and I do like that name, not sure which one is
+// easier for folks.
+def Shape_JoinOp : Shape_Op<"join", []> {
+  let summary = "Returns the least general shape/dim of its operands";
+  let description = [{
+    An operation that computes the least general shape of input operands. This
+    effectively asserts that corresponding static dimensions are equal. The
+    behavior is to match each element of the `shape.type` and propagate the most
+    restrictive information, returning an invalid shape if there are
+    contradictory requirements. E.g., using pseudo code
+
+    ```
+    shape.join([*], [*]) -> [*]
+    shape.join([*], [1, ?]) -> [1, ?]
+    shape.join([1, 2], [1, ?]) -> [1, 2]
+    shape.join([*], [1, 2]) -> [1, 2]
+    shape.join([], X) -> []
+    shape.join([1, U], [2, U, U]) -> []
+    ```
+
+    `shape.join` also allows specifying an optional error string, that may be
+    used to return an error to the user upon mismatch of dimensions.
+
+    ```mlir
+    %c = shape.join %a, %b, error="<reason>" : !shape.type
+    ```
+  }];
+
+  let arguments = (ins Shape_DimOrType:$arg0, Shape_DimOrType:$arg1,
+                   OptionalAttr<StrAttr>:$error);
+  let results = (outs Shape_DimOrType:$result);
+}
+
+def Shape_ReduceOp : Shape_Op<"reduce", []> {
+  let summary = "Returns an expression reduced over shape";
+  let description = [{
+    An operation that takes as input a shape, number of initial values and has a
+    region/function that is applied repeatedly for each every dimension of the
+    shape.
+
+    Conceptually this op performs the following reduction:
+
+    ```
+    res[] = init;
+    for (int i = 0, e = shape.rank(); i != e; ++i) {
+      res = fn(i, shape[i], res[0], ..., res[n]);
+    }
+    ```
+
+    Where fn is provided by the user and the result of the reduce op is the
+    last computed output of the reduce function. As an example, computing the
+    number of elements
+
+    ```mlir
+    func @shape_num_elements(%shape : !shape.type) -> !shape.dim {
+      %0 = "shape.constant_dim"() {value = 0 : i32} : () -> !shape.dim
+      %1 = "shape.reduce"(%shape, %0) ( {
+        ^bb0(%index: i32, %dim: !si.dim, %lci: !shape.dim):
+          %acc = "shape.add"(%lci, %dim) : (!shape.dim, !shape.dim) -> !shape.dim
+          "shape.return"(%acc) : (!shape.dim) -> ()
+        }) : (!shape.type, !shape.dim) -> (!shape.dim)
+      return %1 : !shape.dim
+    }
+    ```
+
+    If the shape is unknown, then the results of the op is also unknown.
+  }];
+
+  let arguments = (ins Shape_Type:$shape, Variadic<AnyType>:$args);
+  let results = (outs Variadic<AnyType>:$result);
+
+  let regions = (region SizedRegion<1>:$body);
+}
+
+def Shape_PrintOp : Shape_Op<"print", []> {
+  let summary = "Prints the input dim or shape";
+  let description = [{
+    Prints the input dim or shape and passes through input.
+
+    Note: This is simply a debugging tool and may be removed.
+  }];
+
+  let arguments = (ins Shape_DimOrType:$input);
+  let results =  (outs Shape_DimOrType:$output);
+}
+
+// TODO: Add Ops: if_static, if_ranked
+
+#endif // SHAPE_OPS

mlir/include/mlir/IR/DialectSymbolRegistry.def

 DEFINE_SYM_KIND_RANGE(QUANTIZATION)
 DEFINE_SYM_KIND_RANGE(IREE) // IREE stands for IR Execution Engine
 DEFINE_SYM_KIND_RANGE(LINALG) // Linear Algebra Dialect
 DEFINE_SYM_KIND_RANGE(FIR) // Flang Fortran IR Dialect
 DEFINE_SYM_KIND_RANGE(OPENMP) // OpenMP IR Dialect
 DEFINE_SYM_KIND_RANGE(TOY) // Toy language (tutorial) Dialect
 DEFINE_SYM_KIND_RANGE(SPIRV) // SPIR-V dialect
 DEFINE_SYM_KIND_RANGE(XLA_HLO) // XLA HLO dialect
+DEFINE_SYM_KIND_RANGE(SHAPE) // Shape dialect
 
 // The following ranges are reserved for experimenting with MLIR dialects in a
 // private context without having to register them here.
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_0)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_1)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_2)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_3)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_4)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_5)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_6)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_7)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_8)
 DEFINE_SYM_KIND_RANGE(PRIVATE_EXPERIMENTAL_9)
 
 #undef DEFINE_SYM_KIND_RANGE

mlir/lib/Dialect/CMakeLists.txt

 add_subdirectory(AffineOps)
 add_subdirectory(FxpMathOps)
 add_subdirectory(GPU)
 add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(LoopOps)
 add_subdirectory(QuantOps)
 add_subdirectory(SDBM)
 add_subdirectory(SPIRV)
+add_subdirectory(Shape)
 add_subdirectory(StandardOps)
 add_subdirectory(VectorOps)
 
 add_llvm_library(MLIRDialect
   Traits.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect
   )
 target_link_libraries(MLIRDialect MLIRIR)

mlir/lib/Dialect/Shape/CMakeLists.txt

+file(GLOB globbed *.c *.cpp)
+add_llvm_library(MLIRShape
+  ${globbed}
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
+  )
+add_dependencies(MLIRShape MLIRShapeOpsIncGen LLVMSupport)
+target_link_libraries(MLIRShape LLVMSupport)

mlir/lib/Dialect/Shape/DialectRegistration.cpp

+//===- DialectRegistration.cpp - Register shape dialect -------------------===//
+//
+// Part of the MLIR 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/Dialect/Shape/Shape.h"
+using namespace mlir;
+
+// Static initialization for shape dialect registration.
+static DialectRegistration<shape::ShapeDialect> ShapeOps;