-
Notifications
You must be signed in to change notification settings - Fork 3.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Hexagon] Generalize builtin for Nd memory alloc with storage scope a…
…nd add lowering for VTCM / Hexagon (#10558) * repurpose texture flatten for vtcm; TIR lowering correct * clean up remaining code in texture flatten pass * add Alloc and FreeTexture, but failing to run over rpc * test passing with malloc in the device api * cleanup * fails in very reliable way with memory corruption * working with non-HexagonBuffer vtcm alloc * cleanup * do not pass scope through mem_copy api * [Hexagon] Resolve breakage in test_hexagon/test_cache_read_write Breakage was caused by #9727, which didn't account for the new `builtin::mem_copy()` when computing the stack size in `StackSizeChecker`. * use HexagonBuffer in Alloc and Free packed funcs * Added comment indicating need for StackSizeChecker::MakeMemCopy. * add AllocVtcmWorkspace and FreeVtcmWorkspace * cleanup * Updated unittests to run all contrib/test_hexagon at CI. * create separate vtcm alloc lowering pass and transform * reset texture_flatten.cc * comments * CI bump * Fix lint formatting error. * Updated fix to remove StackSizeChecker entirely. * pass device and type to device api * Bugfix, verify the precheck's allocations, not own. * Bugfix, pass context information to the precheck. * pass order and shape to device api * working * fix up types and arg passing * pass scope to device api * common builtin for texture / vtcm * add scope to freend api * format and lint * fixed missed format error * restart ci * fix test random value issue + code review feedback * fix test hang * restructure lower vtcm pass per code review feedback (option a) * format error * global.vtcm + tvm_stack_make_shape Co-authored-by: Eric Lunderberg <[email protected]>
- Loading branch information
1 parent
ff54011
commit 2f7bb58
Showing
12 changed files
with
235 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
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,80 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
|
|
||
| #include <tvm/tir/builtin.h> | ||
| #include <tvm/tir/stmt.h> | ||
| #include <tvm/tir/transform.h> | ||
|
|
||
| #include "../../arith/ir_visitor_with_analyzer.h" | ||
|
|
||
| namespace tvm { | ||
| namespace tir { | ||
|
|
||
| inline bool IsVtcmStorage(std::string scope) { | ||
| return scope.find("global.vtcm") != std::string::npos; | ||
| } | ||
|
|
||
| class VtcmAllocator : public StmtExprMutator { | ||
| public: | ||
| using StmtExprMutator::VisitStmt_; | ||
| VtcmAllocator() {} | ||
|
|
||
| Stmt VisitStmt_(const AllocateNode* op) final { | ||
| std::string storage_scope = GetStorageScope(op->buffer_var); | ||
| if (IsVtcmStorage(storage_scope)) { | ||
| Stmt body = this->VisitStmt(op->body); | ||
| Array<PrimExpr> args; | ||
| args.push_back(StringImm(storage_scope)); | ||
| args.push_back(IntImm(DataType::Int(64), op->extents.size())); | ||
| args.push_back(Call(DataType::Handle(), builtin::tvm_stack_make_shape(), op->extents)); | ||
| return LetStmt(op->buffer_var, | ||
| Call(op->buffer_var.dtype(), builtin::nd_mem_alloc_with_scope(), args), body); | ||
| } | ||
| return StmtExprMutator::VisitStmt_(op); | ||
| } | ||
|
|
||
| protected: | ||
| std::string GetStorageScope(const Var& var) { | ||
| auto* ptr = var->type_annotation.as<PointerTypeNode>(); | ||
| ICHECK(ptr) << "Buffer Var's type annotation must be of PointerType"; | ||
| return ptr->storage_scope; | ||
| } | ||
| }; | ||
|
|
||
| PrimFunc LowerVtcmAlloc(PrimFunc func) { | ||
| auto fptr = func.CopyOnWrite(); | ||
| fptr->body = VtcmAllocator()(std::move(fptr->body)); | ||
| return func; | ||
| } | ||
|
|
||
| namespace transform { | ||
|
|
||
| Pass LowerVtcmAlloc() { | ||
| auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) { | ||
| return LowerVtcmAlloc(std::move(f)); | ||
| }; | ||
| return CreatePrimFuncPass(pass_func, 0, "tir.LowerVtcmAlloc", {}); | ||
| } | ||
|
|
||
| TVM_REGISTER_GLOBAL("tir.transform.LowerVtcmAlloc").set_body_typed(LowerVtcmAlloc); | ||
|
|
||
| } // namespace transform | ||
|
|
||
| } // namespace tir | ||
| } // namespace tvm |
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.