Handle sequence_lens for GRU on CPU

src/Conversion/ONNXToKrnl/RNN/GRU.cpp

@@ -389,7 +389,8 @@ template <>
 void calculateState<GruState, GruActivationPack, GruWeightPack, GruBiasPack>(
     ConversionPatternRewriter &rewriter, Location loc, Value Xt, GruState state,
     GruActivationPack activationPack, GruWeightPack weightPack,
-    GruBiasPack biasPack, Value sequenceIV, Value directionIV, bool isForward) {
+    GruBiasPack biasPack, Value sequenceIV, Value directionIV,
+    Value sequenceLens, Value initialH, bool isForward) {
   // Equations (Default: f=Sigmoid, g=Tanh):"
   // zt = f(Xt*(Wz^T) + Ht-1*(Rz^T) + Wbz + Rbz)"
   // rt = f(Xt*(Wr^T) + Ht-1*(Rr^T) + Wbr + Rbr)"
@@ -498,7 +499,10 @@ void calculateState<GruState, GruActivationPack, GruWeightPack, GruBiasPack>(
           Value nextHt = createMath.add(ztht, ztHt);
 
           // Store the intermediate Ht.
-          createKrnl.store(nextHt, Ht, indices);
+          // Handle sequence_lens
+          nextHt = handleSequenceLens(createKrnl, createMath, sequenceLens,
+              initialH, nextHt, sequenceIV, directionIV, bs, hs, Ht);
+
           if (!isNoneValue(state.allH))
             createKrnl.store(
                 nextHt, state.allH, {sequenceIV, directionIV, bs, hs});
@@ -602,7 +606,10 @@ void calculateState<GruState, GruActivationPack, GruWeightPack, GruBiasPack>(
           Value nextHt = createMath.add(ztht, ztHt);
 
           // Store the intermediate Ht.
-          createKrnl.store(nextHt, Ht, indices);
+          // Handle sequence_lens
+          nextHt = handleSequenceLens(createKrnl, createMath, sequenceLens,
+              initialH, nextHt, sequenceIV, directionIV, bs, hs, Ht);
+
           if (!isNoneValue(state.allH))
             createKrnl.store(
                 nextHt, state.allH, {sequenceIV, directionIV, bs, hs});

src/Conversion/ONNXToKrnl/RNN/LSTM.cpp

@@ -443,7 +443,7 @@ void calculateState<LstmState, LstmActivationPack, LstmWeightPack,
     LstmBiasPack>(ConversionPatternRewriter &rewriter, Location loc, Value Xt,
     LstmState state, LstmActivationPack activationPack,
     LstmWeightPack weightPack, LstmBiasPack biasPack, Value sequenceIV,
-    Value directionIV, bool isForward) {
+    Value directionIV, Value sequenceLens, Value initialH, bool isForward) {
   // Equations for LSTM.
   // it = f(Xt*(Wi^T) + Ht-1*(Ri^T) + Pi (.) Ct-1 + Wbi + Rbi)
   // ft = f(Xt*(Wf^T) + Ht-1*(Rf^T) + Pf (.) Ct-1 + Wbf + Rbf)
@@ -452,6 +452,13 @@ void calculateState<LstmState, LstmActivationPack, LstmWeightPack,
   // ot = f(Xt*(Wo^T) + Ht-1*(Ro^T) + Po (.) Ct + Wbo + Rbo)
   // Ht = ot (.) h(Ct)
 
+  // ToFix: add support of sequence lens for LSTM
+  // The assert will fail the test_lstm_with_peephole.
+  // In that test case, the length of the input is used as sequence_lens.
+  // Therefore, onnx-mlir can pass the test by ignoring the sequence_lens
+  // paramenter.
+  // assert(isNoneValue(sequenceLens) && "not implemented yet");
+
   // TODO remove scope
   MultiDialectBuilder<KrnlBuilder, MathBuilder, MemRefBuilder, OnnxBuilder>
       create(rewriter, loc);

src/Conversion/ONNXToKrnl/RNN/RNN.cpp

@@ -300,7 +300,8 @@ template <>
 void calculateState<RnnState, RnnActivationPack, RnnWeightPack, RnnBiasPack>(
     ConversionPatternRewriter &rewriter, Location loc, Value Xt, RnnState state,
     RnnActivationPack activationPack, RnnWeightPack weightPack,
-    RnnBiasPack biasPack, Value sequenceIV, Value directionIV, bool isForward) {
+    RnnBiasPack biasPack, Value sequenceIV, Value directionIV,
+    Value sequenceLens, Value initialH, bool isForward) {
   // Equations for RNN.
   // Ht = f(Xt*(Wi^T) + Ht-1*(Ri^T) + Wbi + Rbi)
   // Shape information:
@@ -311,6 +312,9 @@ void calculateState<RnnState, RnnActivationPack, RnnWeightPack, RnnBiasPack>(
   // Wbi: [hidden_size]
   // Rbi: [hidden_size]
 
+  // ToFix: add support of sequenceLens for RNN
+  assert(isNoneValue(sequenceLens) && "not implemented yet");
+
   MultiDialectBuilder<KrnlBuilder, MathBuilder, MemRefBuilder, OnnxBuilder>
       create(rewriter, loc);
 

src/Conversion/ONNXToKrnl/RNN/RNNBase.cpp

@@ -336,4 +336,33 @@ Value emitXSliceAt(ConversionPatternRewriter &rewriter, Location loc, Value X,
   return sliceX;
 }
 
+// Change the nextHt and Ht value if sequenceLens is defined.
+// When a sample reachs the limit of its sequence len, nextHt will be padded
+// with 0 (or initialH), and Ht will keep the last value at the sequence end
+// so that the final value Ht is the last value at their sequence len.
+Value handleSequenceLens(KrnlBuilder &createKrnl, MathBuilder &createMath,
+    Value sequenceLens, Value initialH, Value nextHt, Value sequenceIV,
+    Value directionIV, Value bs, Value hs, Value Ht) {
+  if (!isNoneValue(sequenceLens)) {
+    Value sequenceUB = createKrnl.load(sequenceLens, {bs});
+    Value initial;
+    if (isNoneValue(initialH)) {
+      initial = createMath.constant(nextHt.getType(), 0.);
+    } else {
+      initial = createKrnl.load(initialH, {directionIV, bs, hs});
+    }
+    Value cond = createMath.sge(
+        createMath.cast(sequenceUB.getType(), sequenceIV), sequenceUB);
+    nextHt = createMath.select(cond, /*padding*/ initial, nextHt);
+
+    // Last HT should be the last in sequenceLens or the current result
+    Value lastHt =
+        createMath.select(cond, createKrnl.load(Ht, {bs, hs}), nextHt);
+    createKrnl.store(lastHt, Ht, {bs, hs});
+  } else {
+    createKrnl.store(nextHt, Ht, {bs, hs});
+  }
+  return nextHt;
+}
+
 } // namespace onnx_mlir

src/Conversion/ONNXToKrnl/RNN/RNNBase.hpp

@@ -75,6 +75,15 @@ mlir::Value applyActivation(mlir::OpBuilder &rewriter, mlir::Location loc,
 mlir::Value emitXSliceAt(mlir::ConversionPatternRewriter &rewriter,
     mlir::Location loc, mlir::Value X, mlir::Value timestep);
 
+// Change the nextHt and Ht value if sequenceLens is defined.
+// When a sample reachs the limit of its sequence len, nextHt will be padded
+// with 0 (or initialH), and Ht will keep the last value at the sequence end
+// so that the final value Ht is the last value at their sequence len.
+mlir::Value handleSequenceLens(KrnlBuilder &createKrnl, MathBuilder &createMath,
+    mlir::Value sequenceLens, mlir::Value initialH, mlir::Value nextHt,
+    mlir::Value sequenceIV, mlir::Value directionIV, mlir::Value bs,
+    mlir::Value hs, mlir::Value Ht);
+
 // Override the following methods when lowering an RNN operation:
 // - hasAllNoneOutput
 // - getActivationPack
@@ -116,7 +125,8 @@ S allocAndInitializeStates(mlir::ConversionPatternRewriter &rewriter,
 template <typename S, typename A, typename W, typename B>
 void calculateState(mlir::ConversionPatternRewriter &rewriter,
     mlir::Location loc, mlir::Value Xt, S state, A activationSet, W weight,
-    B bias, mlir::Value sequenceIV, mlir::Value directionIV, bool isForward);
+    B bias, mlir::Value sequenceIV, mlir::Value directionIV,
+    mlir::Value sequenceLens, mlir::Value initialH, bool isForward);
 
 // Write states to the RNN's outputs.
 template <typename RNNOp, typename S>
@@ -136,6 +146,8 @@ struct ONNXRNNOpLowering : public mlir::OpConversionPattern<RNNOp> {
     mlir::Operation *op = rnnOp.getOperation();
     mlir::Location loc = ONNXLoc<RNNOp>(op);
     mlir::Value X = adaptor.getX();
+    mlir::Value sequenceLens = adaptor.getSequenceLens();
+    mlir::Value initialH = adaptor.getInitialH();
 
     if (hasAllNoneOutput<RNNOp>(&rnnOp)) {
       rewriter.eraseOp(op);
@@ -188,7 +200,7 @@ struct ONNXRNNOpLowering : public mlir::OpConversionPattern<RNNOp> {
             // Emit calculation for one RNN step.
             calculateState<S, A, W, B>(rewriter, loc, Xt, state,
                 activationForward, weightForward, biasForward, sequenceIV,
-                directionIV,
+                directionIV, sequenceLens, initialH,
                 /*isForward=*/true);
           });
     }
@@ -226,7 +238,7 @@ struct ONNXRNNOpLowering : public mlir::OpConversionPattern<RNNOp> {
             // Emit calculation for one RNN step.
             calculateState<S, A, W, B>(rewriter, loc, Xt, state,
                 activationReverse, weightReverse, biasReverse,
-                reverseSequenceIV, directionIV,
+                reverseSequenceIV, directionIV, sequenceLens, initialH,
                 /*isForward=*/false);
           });
     }