-
Notifications
You must be signed in to change notification settings - Fork 0
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[TIR, Relay] improve bfloat16 support #2
Changes from 10 commits
539e5e0
9fd4dd9
8bf6144
a786001
c959702
7e77f56
8e0766c
12fb7b1
9e32fde
6c29073
07ea8de
5583b86
a526fee
a1a3c31
689eead
09a66f4
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -835,10 +835,18 @@ TVM_DLL PrimExpr q_multiply_shift(PrimExpr x, PrimExpr y, PrimExpr q, PrimExpr s | |
Span span = Span()); | ||
|
||
// Intrinsic operators | ||
#define TVM_DECLARE_INTRIN_UNARY(OpName) \ | ||
inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \ | ||
static const Op& op = Op::Get("tir." #OpName); \ | ||
return tir::Call(x.dtype(), op, {x}, span); \ | ||
#define TVM_DECLARE_INTRIN_UNARY(OpName) \ | ||
inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \ | ||
static const Op& op = Op::Get("tir." #OpName); \ | ||
if (x.dtype().is_bfloat16()) { \ | ||
DataType srcType = x.dtype(); \ | ||
DataType dstType(kDLFloat, 32, srcType.lanes()); \ | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Make those There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ok |
||
PrimExpr castX = tir::Cast(dstType, {x}, span); \ | ||
PrimExpr result = tir::Call(dstType, op, {castX}, span); \ | ||
return tir::Cast(srcType, {result}, span); \ | ||
} else { \ | ||
return tir::Call(x.dtype(), op, {x}, span); \ | ||
} \ | ||
} | ||
|
||
TVM_DECLARE_INTRIN_UNARY(exp); | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -40,6 +40,8 @@ | |
"nn.conv3d_transpose", | ||
"nn.dense", | ||
"nn.batch_matmul", | ||
"nn.bias_add", | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Not sure if we can change this default list. Better to have another CPU list, otherwise you need to evaluate the impact to NV hardware. |
||
"nn.batch_norm", | ||
] | ||
DEFAULT_FOLLOW_LIST = [ | ||
# These ops add new data or change shape | ||
|
@@ -80,8 +82,6 @@ | |
"subtract", | ||
"multiply", | ||
"divide", | ||
"nn.bias_add", | ||
"nn.batch_norm", | ||
"sqrt", | ||
"shape_of", | ||
# Simple activations | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please don't change this. You can change it locally, but don't upsteam.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ok, I'll fix this.