【PIR OpTest Fix No.5】 fix test_tril_triu_op (#59734)

xingmingyyj · web-flow · commit 644cbcc920da · 2023-12-25T15:37:47.000+08:00
* fix

* fix

* Update pir_op_test_white_list
diff --git a/paddle/fluid/ir_adaptor/translator/op_translator.cc b/paddle/fluid/ir_adaptor/translator/op_translator.cc
@@ -1343,6 +1343,28 @@ struct TrilAndTriuOpTranscriber : public OpTranscriber {
   }
 };
 
+struct TrilAndTriuGradOpTranscriber : public OpTranscriber {
+  pir::OpInfo LoopkUpOpInfo(pir::IrContext* ctx,
+                            const OpDesc& op_desc) override {
+    bool lower = PADDLE_GET_CONST(bool, op_desc.GetAttr("lower"));
+    std::string target_op_name = "";
+    if (lower) {
+      target_op_name = "pd_op.tril_grad";
+    } else {
+      target_op_name = "pd_op.triu_grad";
+    }
+    const auto& op_info = ctx->GetRegisteredOpInfo(target_op_name);
+    if (!op_info) {
+      IR_THROW(
+          "Op tril_triu_grad should have corresponding OpInfo pd_op.tril_grad "
+          "or "
+          "pd_op.triu_grad.");
+    }
+
+    return op_info;
+  }
+};
+
 using ValueInfo =
     std::tuple<std::vector<int64_t>, dialect::DenseTensorType, pir::OpResult>;
 
@@ -2988,6 +3010,7 @@ OpTranslator::OpTranslator() {
   special_handlers["split"] = SplitOpTranscriber();
   special_handlers["sum"] = AddNOpTranscriber();
   special_handlers["tril_triu"] = TrilAndTriuOpTranscriber();
+  special_handlers["tril_triu_grad"] = TrilAndTriuGradOpTranscriber();
   special_handlers["matrix_rank"] = MatrixRankOpTranscriber();
   special_handlers["mul"] = MulOpTranscriber();
   special_handlers["mul_grad"] = MulGradOpTranscriber();
diff --git a/test/white_list/pir_op_test_white_list b/test/white_list/pir_op_test_white_list
@@ -301,6 +301,7 @@ test_transpose_int8_mkldnn_op
 test_transpose_op
 test_triangular_solve_op
 test_tril_indices_op
+test_tril_triu_op
 test_trilinear_interp_v2_op
 test_triu_indices_op
 test_trunc_op
