[MLIR][Shape] Generalize `shape.concat` to extent tensors

[qingyunqu]

Summary

The operation shape.concat was used for type shape only.
We now enable it for extent tensors.

[jpienaar]

LG, thanks

[mehdi_amini]

Do you have commit access or do you need someone to push this for you?

[qingyunqu]

In D127321#3567440, @jpienaar wrote:

LG, thanks

Should I change the $lhs and $rhs to Variadic<Shape_ShapeOrExtentTensorType>? @jpienaar

[qingyunqu]

In D127321#3567615, @mehdi_amini wrote:

Do you have commit access or do you need someone to push this for you?

I don't really get the point. What should I do between these two methods.

[mehdi_amini]

In D127321#3567648, @qingyunqu wrote:

In D127321#3567615, @mehdi_amini wrote:

Do you have commit access or do you need someone to push this for you?

I don't really get the point. What should I do between these two methods.

Regular contributors have direct write access to the repository and take care themselves to git push their changes to the main branch. On the other hand new contributors don't have write access and rely on the reviewer to push the change for them after approving.

[qingyunqu]

In D127321#3568353, @mehdi_amini wrote:

In D127321#3567648, @qingyunqu wrote:

In D127321#3567615, @mehdi_amini wrote:

Do you have commit access or do you need someone to push this for you?

I don't really get the point. What should I do between these two methods.

Regular contributors have direct write access to the repository and take care themselves to git push their changes to the main branch. On the other hand new contributors don't have write access and rely on the reviewer to push the change for them after approving.

I think I'm the new contributor.

[jpienaar]

In D127321#3567629, @qingyunqu wrote:

In D127321#3567440, @jpienaar wrote:

LG, thanks

Should I change the $lhs and $rhs to Variadic<Shape_ShapeOrExtentTensorType>? @jpienaar

No, these are currently intended to be explicit and fixed.

If you have a shape function with a lot concats where this would simplify things, I'd be interested to know about

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

let description = [{
     followed by the dimensions of `rhs`.
 
     Example:
     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 arguments = (ins Shape_ShapeOrExtentTensorType:$lhs, Shape_ShapeOrExtentTensorType:$rhs);
+  let results = (outs Shape_ShapeOrExtentTensorType:$result);
+
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($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/test/Dialect/Shape/canonicalize.mlir

 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
 func.func @f() -> !shape.shape {
   // CHECK: shape.const_shape [0, 1, 2, 3] : !shape.shape
   %lhs = shape.const_shape [0, 1] : !shape.shape
   %rhs = shape.const_shape [2, 3] : !shape.shape
-  %0 = shape.concat %lhs, %rhs
+  %0 = shape.concat %lhs, %rhs : !shape.shape , !shape.shape -> !shape.shape
   return %0 : !shape.shape
 }
 
 // -----
 
 // Basic case.
 // CHECK-LABEL: func @f
+func.func @f() -> tensor<4xindex> {
+  // CHECK: shape.const_shape [0, 1, 2, 3] : tensor<4xindex>
+  %lhs = shape.const_shape [0, 1] : tensor<2xindex>
+  %rhs = shape.const_shape [2, 3] : tensor<2xindex>
+  %0 = shape.concat %lhs, %rhs : tensor<2xindex>, tensor<2xindex> -> tensor<4xindex>
+  return %0 : tensor<4xindex>
+}
+
+// -----
+
+// Basic case.
+// CHECK-LABEL: func @f
 func.func @f() -> tensor<2xindex> {
   // CHECK: shape.const_shape [0, 1] : tensor<2xindex>
   %cs = shape.const_shape [0, 1] : !shape.shape
   %0 = shape.to_extent_tensor %cs : !shape.shape -> tensor<2xindex>
   return %0 : tensor<2xindex>
 }
 
 // -----