[PATCH] D73944: [mlir][wip] Start Shape dialect

[Jacques Pienaar via Phabricator via llvm-commits]

jpienaar added a comment.

Updated, thanks

In D73944#1864082 <https://reviews.llvm.org/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.



================
Comment at: mlir/include/mlir/Dialect/CMakeLists.txt:9
 add_subdirectory(SPIRV)
+add_subdirectory(Shape)
 add_subdirectory(StandardOps)
----------------
rriddle wrote:
> Seems like this should come before SPIRV?
Depends on if you follow sort order case sensitive or insensitive :) I normally default to case sensitive, but can do insensitive.


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/Shape.h:68
+
+class ElementType : public Type::TypeBase<ElementType, Type> {
+public:
----------------
rriddle wrote:
> jpienaar wrote:
> > rriddle wrote:
> > > Can we get comments on each of these types?
> > This made me think about having ODS generate these from the description :)
> Realistically, we should be able to auto-generate simple types pretty trivially.
Exactly. Should be easy to address.


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/Shape.h:18
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/FunctionSupport.h"
----------------
rriddle wrote:
> Why do you need Function.h here?
Ah this was left over from previous testing, thanks


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/Shape.h:36
+  Shape,
+  Element,
+  Component,
----------------
rriddle wrote:
> Any reason not to sort these?
It was a semi hierarchy in my mind, but that doesn't really matter, so sorted.


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:46
+  let typeDescription = [{
+    `shape.dim` represents a dimension type which could be either a non-negative
+    integer, unknown or error.
----------------
herhut wrote:
> 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`? 
Good point and I like that, renaming.


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:55
+
+def Shape_Type : DialectType<ShapeDialect,
+    CPred<"$_self.isa<ShapeType>()">, "shape"> {
----------------
herhut wrote:
> 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`?
Yeah, Shape_Shape was just weird, changing to uniformly suffixing with Type.


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:110
+
+// TODO: consider just making a constant op instead.
+def Shape_ConstantDimOp : Shape_Op<"constant_dim", []> {
----------------
herhut wrote:
> +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.
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?


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:134
+
+    Note: it is not possible to create an unranked shape using the op.
+  }];
----------------
herhut wrote:
> 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?
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.


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:190
+
+           output[i] = lhs[i] if lhs[i] == rhs[i] or rhs[i] is unknown/undefined
+                     = rhs[i] if lhs[i] is unknown/undefined
----------------
stellaraccident wrote:
> nicolasvasilache wrote:
> > nicolasvasilache wrote:
> > > This is frontend logic that I do not think should leak into MLIR.
> > I meant into MLIR *core*, sorry.
> 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.
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. 


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:222
+    shape.join([], X) -> []
+    shape.join([1, U], [2, U, U]) -> []
+    ```
----------------
herhut wrote:
> join([1,?], [2, ?, ?]) -> [invalid] ?
I changed notations and didn't do it correctly.


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:242
+    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.
----------------
herhut wrote:
> "for each every" is one too many.
It was just so good it happened every every time 😁 


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:263
+        ^bb0(%index: i32, %dim: !si.dim, %lci: !shape.dim):
+          %acc = "shape.add"(%lci, %dim) :
+            (!shape.dim, !shape.dim) -> !shape.dim
----------------
herhut wrote:
> Should this be `shape.mul` if it existed?
D'oh yes. This was computing sum of dimensions :)


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:271
+
+    If the shape is unknown, then the results of the op is also unknown.
+  }];
----------------
herhut wrote:
> 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). 
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 🙂 


================
Comment at: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td:280
+
+def Shape_PrintOp : Shape_Op<"print", []> {
+  let summary = "Prints the input dim or shape";
----------------
antiagainst wrote:
> Would it make sense to name it as `debug_print` to be more clear?
Sounds good. I quite like it, made writing tests easier (although I should refactor those into non-executable tests ...).


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D73944/new/

https://reviews.llvm.org/D73944