[MLIR][Shape] Allow `shape.any` to operate on extent tensors

[frgossen]

Summary

Depends On D84431

[pifon2a]

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

517–519

nit: " " (two spaces) -> " "

[frgossen]

Address nit

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

def Shape_ConcatOp : Shape_Op<"concat", []> {
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Shape constraint related ops.
 //===----------------------------------------------------------------------===//
 
 // TODO: Move the code below and witnesses to a different file.
-def Shape_AnyOp : Shape_Op<"any", [Commutative, NoSideEffect]> {
+def Shape_AnyOp : Shape_Op<"any", [Commutative,
+                                   NoSideEffect,
+                                   SameOperandsAndResultType]> {
   let summary = "Return any combination of the input shapes";
   let description = [{
-    This operation takes multiple input shapes and returns some combination of
-    their dimensions. This can be best seen with examples below.
+    This operation takes multiple input shapes or extent tensors and returns
+    some combination of their dimensions. This can be best seen with examples
+    below.
 
     The result is undefined, but still side-effect free, in cases where the
     inputs have differing ranks or differ in extents of shared dimensions.
 
     Example:
     ```mlir
     %s0 = shape.any [2,?], [?,3] // [2,3]
     %s1 = shape.any [?,?], [1,2] // [1,2]
     ```
   }];
 
-  let arguments = (ins Variadic<Shape_ShapeType>:$inputs);
-  let results = (outs Shape_ShapeType:$result);
-
-  let assemblyFormat = "$inputs attr-dict";
+  let arguments = (ins Variadic<Shape_ShapeOrExtentTensorType>:$inputs);
+  let results = (outs Shape_ShapeOrExtentTensorType:$result);
 
+  let assemblyFormat = "$inputs `:` type($result)  attr-dict";
   let hasFolder = 1;
 }
 
 def Shape_AssumingAllOp : Shape_Op<"assuming_all", [Commutative, NoSideEffect]> {
   let summary = "Return a logical AND of all witnesses";
   let description = [{
     Used to simplify constraints as any single failing precondition is enough
     to prevent execution.

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

func @get_extent_from_extent_tensor(%extents : tensor<?xindex>,
   return %result : !shape.size
 }
 
 // -----
 
 // Lower `any` to its first operand.
 // CHECK-LABEL: @any_of_three
 // CHECK-SAME:  (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<?xindex>) -> tensor<?xindex>
-func @any_of_three(%a : !shape.shape, %b : !shape.shape, %c : !shape.shape)
-    -> !shape.shape {
+func @any_of_three(%a : tensor<?xindex>,
+                   %b : tensor<?xindex>,
+                   %c : tensor<?xindex>) -> tensor<?xindex> {
   // CHECK: return %[[A]] : tensor<?xindex>
-  %result = shape.any %a, %b, %c
-  return %result : !shape.shape
+  %result = shape.any %a, %b, %c : tensor<?xindex>
+  return %result : tensor<?xindex>
 }
 
 // -----
 
 // Lower `any` to its first operand.
 // CHECK-LABEL: @any_of_one
 // CHECK-SAME:  (%[[A:.*]]: tensor<?xindex>) -> tensor<?xindex>
-func @any_of_one(%a : !shape.shape) -> !shape.shape {
+func @any_of_one(%a : tensor<?xindex>) -> tensor<?xindex> {
   // CHECK: return %[[A]] : tensor<?xindex>
-  %result = shape.any %a
-  return %result : !shape.shape
+  %result = shape.any %a : tensor<?xindex>
+  return %result : tensor<?xindex>
 }
 

mlir/test/Dialect/Shape/canonicalize.mlir

func @f() {
   "consume.witness"(%2) : (!shape.witness) -> ()
   return
 }
 
 // -----
 
 // any can be replaced with a constant input if it has one.
 // CHECK-LABEL: func @f
-func @f(%arg0 : !shape.shape) -> !shape.shape {
+func @f(%arg : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: %[[CS:.*]] = shape.const_shape
   // CHECK-NEXT: return %[[CS]]
   %0 = shape.const_shape [2, 3, 4] : !shape.shape
-  %1 = shape.any %0, %arg0
+  %1 = shape.any %0, %arg : !shape.shape
   return %1 : !shape.shape
 }
 
+// -----
+
+// any can be replaced with a constant input if it has one.
+// CHECK-LABEL: func @f
+func @f(%arg : tensor<?xindex>) -> tensor<?xindex> {
+  // CHECK-NEXT: %[[CS:.*]] = shape.const_shape [2, 3, 4] : tensor<?xindex>
+  // CHECK-NEXT: return %[[CS]] : tensor<?xindex>
+  %0 = shape.const_shape [2, 3, 4] : tensor<?xindex>
+  %1 = shape.any %0, %arg : tensor<?xindex>
+  return %1 : tensor<?xindex>
+}
 
 // -----
 
 // Folding of any with partially constant operands is not yet implemented.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape, %arg1 : !shape.shape) -> !shape.shape {
   // CHECK-NEXT: %[[CS:.*]] = shape.any
   // CHECK-NEXT: return %[[CS]]
-  %1 = shape.any %arg0, %arg1
+  %1 = shape.any %arg0, %arg1 : !shape.shape
   return %1 : !shape.shape
 }
 
 // -----
 
 // assuming with a known passing witness can be removed
 // CHECK-LABEL: func @f
 func @f() {

mlir/test/Dialect/Shape/ops.mlir

-// RUN: mlir-opt -split-input-file %s | mlir-opt | FileCheck %s
 // Verify the printed output can be parsed.
 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 // Verify the generic form can be parsed.
 // RUN: mlir-opt -mlir-print-op-generic %s | mlir-opt | FileCheck %s
 
 // CHECK-LABEL: shape_num_elements
 func @shape_num_elements(%shape : !shape.shape) -> !shape.size {
   %init = shape.const_size 1
func @test_constraints() {
   %0 = shape.const_shape [] : !shape.shape
   %1 = shape.const_shape [1, 2, 3] : !shape.shape
   %w0 = shape.cstr_broadcastable %0, %1 : !shape.shape, !shape.shape
   %w1 = shape.cstr_eq %0, %1
   %w2 = shape.const_witness true
   %w3 = shape.const_witness false
   %w4 = shape.assuming_all %w0, %w1, %w2, %w3
   shape.assuming %w4 -> !shape.shape {
-    %2 = shape.any %0, %1
+    %2 = shape.any %0, %1 : !shape.shape
     shape.assuming_yield %2 : !shape.shape
   }
   return
 }
 
 func @broadcastable_on_extent_tensors(%lhs : tensor<?xindex>,
                                       %rhs : tensor<?xindex>) {
   %w0 = shape.cstr_broadcastable %lhs, %rhs : tensor<?xindex>, tensor<?xindex>
func @get_extent_on_shape(%arg : !shape.shape) -> !shape.size {
   return %result : !shape.size
 }
 
 func @get_extent_on_extent_tensor(%arg : tensor<?xindex>) -> !shape.size {
   %c0 = shape.const_size 0
   %result = shape.get_extent %arg, %c0 : tensor<?xindex>
   return %result : !shape.size
 }
+
+func @any() {
+  %0 = shape.const_shape [1, 2, 3] : !shape.shape
+  %1 = shape.const_shape [4, 5, 6] : !shape.shape
+  %2 = shape.any %0, %1 : !shape.shape
+  %3 = shape.const_shape [1, 2, 3] : tensor<?xindex>
+  %4 = shape.const_shape [4, 5, 6] : tensor<?xindex>
+  %5 = shape.any %3, %4 : tensor<?xindex>
+  return
+}
+