Add the support of getting the length of tensor list

This patch added both the op definition and the Python translation to support
getting the length of tensor list and comparing the results with a constant.

PiperOrigin-RevId: 346157648
Change-Id: If7fcf42a1632fd08292426551e191fd9919d67a4

Author: liufengdb

tensorflow/compiler/mlir/tfr/ir/tfr_ops.cc

@@ -437,6 +437,23 @@ struct RemoveRedundantGetElement : public OpRewritePattern<GetElementOp> {
   }
 };
 
+struct RemoveRedundantGetLength : public OpRewritePattern<GetLengthOp> {
+  using OpRewritePattern<GetLengthOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GetLengthOp gl_op,
+                                PatternRewriter &rewriter) const override {
+    auto preceding_build_list = llvm::dyn_cast_or_null<BuildListOp>(
+        gl_op.tensor_list().getDefiningOp());
+    if (!preceding_build_list) {
+      return failure();
+    }
+    int64_t num_tensors = preceding_build_list.getNumOperands();
+    rewriter.replaceOpWithNewOp<ConstantOp>(gl_op,
+                                            rewriter.getIndexAttr(num_tensors));
+    return success();
+  }
+};
+
 struct BuildConstantListAsAttr : public OpRewritePattern<BuildListOp> {
   using OpRewritePattern<BuildListOp>::OpRewritePattern;
 
@@ -477,6 +494,11 @@ void GetElementOp::getCanonicalizationPatterns(
   results.insert<RemoveRedundantGetElement>(context);
 }
 
+void GetLengthOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                              MLIRContext *context) {
+  results.insert<RemoveRedundantGetLength>(context);
+}
+
 void BuildListOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                               MLIRContext *context) {
   results.insert<BuildConstantListAsAttr>(context);

tensorflow/compiler/mlir/tfr/ir/tfr_ops.td

@@ -349,6 +349,30 @@ def TFR_BuildListOp : TFR_Op<"build_list", [NoSideEffect]> {
   let hasCanonicalizer = 1;
 }
 
+def TFR_GetLengthOp : TFR_Op<"get_length", [NoSideEffect]> {
+  let description = [{
+    The `get_length` operation returns the number of tensors for a
+    tfr.tensor_list.
+
+    Example:
+
+    ```mlir
+    %2 = tfr.get_length(%1) : tfr.tensor -> index
+    %2 = tfr.get_length %1 -> index
+    ```
+  }];
+
+  let arguments = (ins TFR_TensorListType:$tensor_list);
+
+  let results = (outs Index:$out);
+
+  let hasCanonicalizer = 1;
+
+  let assemblyFormat = [{
+    $tensor_list attr-dict `->` type($out)
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Function related classes
 //===----------------------------------------------------------------------===//

tensorflow/compiler/mlir/tfr/python/tfr_gen.py

@@ -643,6 +643,15 @@ def _index_to_I64(self, value, ty):
     else:
       return value, ty
 
+  def _i64_to_index(self, value, ty):
+    if ty == TFRTypes.I64:
+      casted = self._ssa_name('casted')
+      self._emit_with_loc('\n{} = index_cast {} : i64 to index'.format(
+          casted, value))
+      return casted, TFRTypes.INDEX
+    else:
+      return value, ty
+
   def _value_to_tensor(self, value, ty, node):
     value, ty = self._index_to_I64(value, ty)
     cst_tensor = self._ssa_name('cst')
@@ -828,15 +837,19 @@ def visit_Call(self, node):
       if func_name == 'len':
         arg, ty = self.visit(node.args[0])
         ty = self._get_inferred_type(node.args[0], ty)
-        assert ty == TFRTypes.TF_TENSOR_SHAPE_LIST, ty
-        len_value = self._ssa_name('len')
-        self._emit_with_loc(
-            '\n{} = shape.rank {} : !shape.shape -> !shape.size'.format(
-                len_value, arg), node)
-        size_value = self._ssa_name('len_size')
-        self._emit_with_loc(
-            '\n{} = shape.size_to_index {} : !shape.size'.format(
-                size_value, len_value), node)
+        if ty == TFRTypes.TF_TENSOR_SHAPE_LIST:
+          len_value = self._ssa_name('len')
+          self._emit_with_loc(
+              '\n{} = shape.rank {} : !shape.shape -> !shape.size'.format(
+                  len_value, arg), node)
+          size_value = self._ssa_name('len_size')
+          self._emit_with_loc(
+              '\n{} = shape.size_to_index {} : !shape.size'.format(
+                  size_value, len_value), node)
+        elif ty == TFRTypes.TENSOR_LIST:
+          size_value = self._ssa_name('len')
+          self._emit_with_loc(
+              '\n{} = tfr.get_length {} -> index'.format(size_value, arg), node)
         return (size_value, TFRTypes.INDEX)
 
     raise NotImplementedError('call operator not recognized: {} {}'.format(
@@ -845,7 +858,7 @@ def visit_Call(self, node):
   def visit_Compare(self, node):
     lhs, lhs_ty = self.visit(node.left)
     for op, right in zip(node.ops, node.comparators):
-      rhs, _ = self.visit(right)
+      rhs, rhs_ty = self.visit(right)
       if isinstance(op, ast.Eq):
         pred = 'eq'
       elif isinstance(op, ast.Lt):
@@ -870,6 +883,10 @@ def visit_Compare(self, node):
           code = 'cmpi'
         elif lhs_ty == TFRTypes.F32:
           code = 'cmpf'
+        elif lhs_ty == TFRTypes.INDEX:
+          code = 'cmpi'
+          # TODO(fengliuai): the reverse type inference should solve the issue.
+          rhs, _ = self._i64_to_index(rhs, rhs_ty)
         else:
           raise NotImplementedError('Compare operand type not recognized')
         self._emit_with_loc(

tensorflow/compiler/mlir/tfr/python/tfr_gen_test.py

@@ -126,6 +126,15 @@ def _tfr_control_flow_range_for(x):
   return x_sum
 
 
+@composite.Composite('TestInputNOp')
+def _tfr_control_flow_tensor_list_size(ins):
+  n = len(ins)
+  if n == 1:
+    return ins[0]
+  else:
+    return math_ops.AddN(ins)
+
+
 #--- test fn for tf ops ---
 
 
@@ -403,6 +412,24 @@ def test_tfr_control_flow(self):
       CHECK-NEXT:   %{{.*}} = constant true
       CHECK-NEXT:   tfr.return %[[for_stmt]] : !tfr.tensor
       CHECK-NEXT: }
+
+      CHECK-LABEL: tfr.func @tf__test_input_n_op(%ins: !tfr.tensor_list) -> (!tfr.tensor) {
+      CHECK-NEXT:   %[[len:.*]] = tfr.get_length %ins -> index
+      CHECK-NEXT:   %[[cst:.*]] = constant 1 : i64
+      CHECK-NEXT:   %[[casted:.*]] = index_cast %[[cst]] : i64 to index
+      CHECK-NEXT:   %[[eq:.*]] = cmpi "eq", %[[len]], %[[casted]] : index
+      CHECK-NEXT:   %[[if:.*]] = scf.if %[[eq]] -> (!tfr.tensor) {
+      CHECK-NEXT:     %{{.*}} = constant true
+      CHECK-NEXT:     %{{.*}} = constant 0 : index
+      CHECK-NEXT:     %[[elt:.*]] = tfr.get_element %ins[%cst_2] : (!tfr.tensor_list, index) -> !tfr.tensor
+      CHECK-NEXT:     scf.yield %[[elt]] : !tfr.tensor
+      CHECK-NEXT:   } else {
+      CHECK-NEXT:     %{{.*}} = constant true
+      CHECK-NEXT:     %[[AddN:.*]] = tfr.call @tf__add_n(%ins) : (!tfr.tensor_list) -> (!tfr.tensor)
+      CHECK-NEXT:     scf.yield %[[AddN]] : !tfr.tensor
+      CHECK-NEXT:   }
+      CHECK-NEXT:   tfr.return %[[if_stmt]] : !tfr.tensor
+      CHECK-NEXT:  }
     """
     self._check_code(mlir_code, mlir_code_exp)
 

tensorflow/compiler/mlir/tfr/tests/ops.mlir

@@ -260,6 +260,19 @@ func @build_const_list() -> !tfr.attr {
 
 // -----
 
+// CHECK-LABEL: get_length
+// CANON-LABEL: get_length
+func @get_length(%arg0: !tfr.tensor<A>, %arg1: !tfr.tensor<B>) -> index {
+  %0 = "tfr.build_list"(%arg0, %arg1) : (!tfr.tensor<A>, !tfr.tensor<B>) -> !tfr.tensor_list
+  %1 = "tfr.get_length"(%0) : (!tfr.tensor_list) -> index
+  return %1 : index
+
+// CANON-NEXT: %[[c:.*]] = constant 2 : index
+// CANON-NEXT: return %[[c]] : index
+}
+
+// -----
+
 // CHECK-LABEL: tfr.func
 tfr.func @External(%arg0: !tfr.tensor<A>,
               %arg1: !tfr.tensor_list<C>,