-
Notifications
You must be signed in to change notification settings - Fork 361
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(aten::std|aten::masked_fill): Implement masked_fill, aten::std
works for non bias corrected cases Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
- Loading branch information
1 parent
fa7d6d9
commit a086a5b
Showing
16 changed files
with
565 additions
and
35 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,30 @@ | ||
| #include "torch/csrc/jit/passes/subgraph_rewrite.h" | ||
|
|
||
| #include "core/util/prelude.h" | ||
|
|
||
| namespace trtorch { | ||
| namespace core { | ||
| namespace lowering { | ||
| namespace passes { | ||
|
|
||
| void UnpackStd(std::shared_ptr<torch::jit::Graph>& graph) { | ||
| std::string std_pattern = R"IR( | ||
| graph(%1, %dim, %unbiased, %keepdim): | ||
| %out: Tensor = aten::std(%1, %dim, %unbiased, %keepdim) | ||
| return (%out))IR"; | ||
| std::string unpacked_pattern = R"IR( | ||
| graph(%1, %dim, %unbiased, %keepdim): | ||
| %z: Tensor = aten::var(%1, %dim, %unbiased, %keepdim) | ||
| %out: Tensor = aten::sqrt(%z) | ||
| return (%out))IR"; | ||
|
|
||
| torch::jit::SubgraphRewriter std_rewriter; | ||
| std_rewriter.RegisterRewritePattern(std_pattern, unpacked_pattern); | ||
| std_rewriter.runOnGraph(graph); | ||
| LOG_GRAPH("Post unpack std: " << *graph); | ||
| } | ||
|
|
||
| } // namespace passes | ||
| } // namespace lowering | ||
| } // namespace core | ||
| } // namespace trtorch |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| #include "torch/csrc/jit/passes/subgraph_rewrite.h" | ||
|
|
||
| #include "core/util/prelude.h" | ||
|
|
||
| namespace trtorch { | ||
| namespace core { | ||
| namespace lowering { | ||
| namespace passes { | ||
|
|
||
| void UnpackVar(std::shared_ptr<torch::jit::Graph>& graph) { | ||
| std::string var_pattern = R"IR( | ||
| graph(%input, %dim, %unbiased, %keepdim): | ||
| %out: Tensor = aten::var(%input, %dim, %unbiased, %keepdim) | ||
| return (%out))IR"; | ||
| std::string unpacked_pattern = R"IR( | ||
| graph(%input, %dims, %unbiased, %keepdim): | ||
| %none: None = prim::Constant() | ||
| %false: bool = prim::Constant[value=0]() | ||
| %0: int = prim::Constant[value=0]() | ||
| %1: int = prim::Constant[value=1]() | ||
| %sqrd: Tensor = aten::mul(%input, %input) | ||
| %sqrdmean: Tensor = aten::mean(%sqrd, %dims, %keepdim, %none) | ||
| %mean: Tensor = aten::mean(%input, %dims, %keepdim, %none) | ||
| %meansqrd: Tensor = aten::mul(%mean, %mean) | ||
| %var: Tensor = aten::sub(%sqrdmean, %meansqrd, %1) | ||
| %varout : Tensor = prim::If(%unbiased) | ||
| block0(): | ||
| %shape: int[] = aten::size(%input) | ||
| %shapet: Tensor = aten::tensor(%shape, %0, %none, %false) | ||
| %dim: int = prim::ListUnpack(%dims) | ||
| %reduceddims: Tensor = aten::select(%shapet, %0, %dim) | ||
| %numel: Tensor = aten::prod(%reduceddims, %dim, %keepdim, %none) | ||
| %mul: Tensor = aten::mul(%var, %numel) | ||
| %sub: Tensor = aten::sub(%numel, %1, %1) | ||
| %v: Tensor = aten::div(%mul, %sub) | ||
| -> (%v) | ||
| block1(): | ||
| -> (%var) | ||
| return(%varout))IR"; | ||
|
|
||
| torch::jit::SubgraphRewriter var_rewriter; | ||
| var_rewriter.RegisterRewritePattern(var_pattern, unpacked_pattern); | ||
| var_rewriter.runOnGraph(graph); | ||
| LOG_DEBUG("Post unpack var: " << *graph); | ||
|
|
||
| } | ||
|
|
||
| } // namespace passes | ||
| } // namespace lowering | ||
| } // namespace core | ||
| } // namespace trtorch |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.