[TIR] [bfloat16] add bfloat16 promotion for CallNode (#12370)

* add bfloat16 promotion for CallNode

* add softmax to bfloat16 build test

Author: yangulei

src/tir/transforms/bf16_legalize.cc

@@ -55,6 +55,7 @@ class BF16PromoteRewriter : public StmtExprMutator {
   PrimExpr VisitExpr_(const LENode* op) final;
   PrimExpr VisitExpr_(const GTNode* op) final;
   PrimExpr VisitExpr_(const GENode* op) final;
+  PrimExpr VisitExpr_(const CallNode* op) final;
 };
 
 #define DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC, NEEDCAST)                \
@@ -88,6 +89,26 @@ DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LENode, operator<=, false)
 DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GTNode, operator>, false)
 DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GENode, operator>=, false)
 
+PrimExpr BF16PromoteRewriter::VisitExpr_(const CallNode* op) {
+  Array<PrimExpr> args;
+  for (auto& arg : op->args) {
+    PrimExpr x = this->VisitExpr(arg);
+    if (x.dtype().is_bfloat16()) {
+      DataType fp32_dtype(kDLFloat, 32, x.dtype().lanes());
+      args.push_back(Cast(fp32_dtype, {x}, op->span));
+    } else {
+      args.push_back(x);
+    }
+  }
+  if (op->dtype.is_bfloat16()) {
+    DataType fp32_dtype(kDLFloat, 32, op->dtype.lanes());
+    PrimExpr result_fp32 = Call(fp32_dtype, op->op, args, op->span);
+    return Cast(op->dtype, {result_fp32}, op->span);
+  } else {
+    return Call(op->dtype, op->op, args, op->span);
+  }
+}
+
 /*
  * Eliminate verbose casting between fp32 and bf16
  * Checks if the AST has the pattern:

tests/python/relay/test_cpp_build_module.py

@@ -127,7 +127,9 @@ def test_bf16_build():
     relu_bf16 = relay.nn.relu(bn_bf16[0])
     maxpool_bf16 = relay.nn.max_pool2d(relu_bf16, pool_size=(2, 2), strides=(2, 2))
     avgpool_bf16 = relay.nn.avg_pool2d(maxpool_bf16, pool_size=(2, 2), strides=(2, 2))
-    mod_bf16 = tvm.IRModule.from_expr(avgpool_bf16)
+    flattened_bf16 = relay.nn.batch_flatten(avgpool_bf16)
+    softmax_bf16 = relay.nn.softmax(flattened_bf16)
+    mod_bf16 = tvm.IRModule.from_expr(softmax_bf16)
     with tvm.transform.PassContext(opt_level=3):
         relay.build(mod_bf16, target="llvm", params=params)