[TE] Support schedulable TIR compute definitions in TOPI#11589

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.

Sign up for GitHub

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

Merged
junrushao merged 7 commits into from
Jun 13, 2022
Merged

[TE] Support schedulable TIR compute definitions in TOPI #11589

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
eb3feba
[TIR] Update domain touched to track touched domains for all buffers
csullivan May 28, 2022
fa888bf
[TE] Refactor logical steps of CreatePrimFunc into separate functions.
csullivan May 30, 2022
f75d82d
[TE] Update CreatePrimFunc to pre-populate buffer binds map with
csullivan May 30, 2022
5efedea
[TE] Introduce te.extern_primfunc which allows for inlining a
csullivan May 31, 2022
659dce6
Add test showing the use of te.extern on a TVMScript PrimFunc
csullivan May 25, 2022
1936ae4
Fix task extraction default filter func skipping TE ExternOpNodes.
csullivan Jun 4, 2022
ad2b285
Address @Hzfengsy's review comments.
csullivan Jun 8, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions python/tvm/te/__init__.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -40,6 +40,7 @@
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc

from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
from .autodiff import gradient
Expand Down
82 changes: 82 additions & 0 deletions python/tvm/te/operation.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -24,6 +24,7 @@
import tvm._ffi
import tvm.tir
import tvm.tir._ffi_api
import tvm.arith._ffi_api
from tvm._ffi.base import string_types
from tvm.ir import Array
from tvm.runtime import convert
Expand DownExpand Up@@ -354,6 +355,87 @@ def extern(
return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
"""Compute tensors via a schedulable TIR PrimFunc

Parameters
----------
input_tensors: list of Tensor
Input tensors that map to the corresponding primfunc input params.

primfunc: PrimFunc
The TIR PrimFunc

Returns
-------
tensor: Tensor or list of Tensors
The created tensor or tuple of tensors if it contains multiple outputs.

Example
-------
In the code below, a TVMScript defined TIR PrimFunc is inlined into
a TE ExternOp. Applying te.create_prim_func on this

.. code-block:: python

A = te.placeholder((128, 128), name="A")
B = te.placeholder((128, 128), name="B")

@T.prim_func
def before_split(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
for i, j in T.grid(128, 128):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = A[vi, vj] * 2.0

C = te.extern_primfunc([A, B], func)
"""
access_map = {
k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
}
in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
assert in_buffers, "PrimFunc has no input buffers"
assert out_buffers, "PrimFunc has no output buffers"

outputs = []
inplace = []
input_buffers = in_buffers
for obuf in out_buffers:
if obuf in in_buffers:
inplace.append(obuf)
else:
outputs.append(obuf)

if not outputs:
iobuf = inplace.pop()
input_buffers.remove(iobuf)
outputs = [iobuf]

assert len(input_buffers) == len(input_tensors), (
"The number of provided input input_tensors does not match the number of ",
"input buffers in the primfunc",
)
for tensor, buffer in zip(input_tensors, input_buffers):
# TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
assert tensor.shape == buffer.shape, (
"The input input_tensors provided do not match the input buffers in the ",
"primfunc. Please check that the order of input te.Input_Tensors and the ",
"order of the primfunc variables in the params list agree.",
)
output = extern(
[buf.shape for buf in outputs],
input_tensors,
lambda ins, outs: primfunc.body,
in_buffers=input_buffers,
out_buffers=outputs,
**kwargs,
)
return output


def var(name="tindex", dtype="int32", span=None):
"""Create a new variable with specified name and dtype

Expand Down
106 changes: 86 additions & 20 deletions src/arith/domain_touched.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <tuple>
#include <unordered_map>
#include <unordered_set>

Expand All@@ -34,18 +35,54 @@ namespace arith {

using namespace tir;

namespace {

using BufferTouches = std::vector<std::vector<IntSet>>;

struct LoadAccess {
Comment thread
Hzfengsy marked this conversation as resolved.
BufferTouches set;
};

struct StoreAccess {
BufferTouches set;
};

struct CombinedAccess {
BufferTouches set;
};

using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;

} // namespace

// Find Read region of the tensor in the stmt.
class BufferTouchedDomain final : public StmtExprVisitor {
public:
BufferTouchedDomain(const Buffer& buffer, bool consider_loads, bool consider_stores)
: buffer_(buffer), consider_loads_(consider_loads), consider_stores_(consider_stores) {}
BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }

std::unordered_map<const BufferNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
return buffer_access_map_;
}

Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) {
auto kv = buffer_access_map_.find(buffer.get());
CHECK(kv != buffer_access_map_.end())
<< "The requested buffer is not contained in the provided stmt body.";

Region Find(const Stmt& stmt) {
operator()(stmt);
Region ret;
Range none;
for (size_t i = 0; i < bounds_.size(); ++i) {
ret.push_back(arith::Union(bounds_[i]).CoverRange(none));
BufferTouches bounds;
if (consider_loads && consider_stores) {
bounds = std::get<CombinedAccess>(kv->second).set;
} else if (consider_loads) {
bounds = std::get<LoadAccess>(kv->second).set;
} else if (consider_stores) {
bounds = std::get<StoreAccess>(kv->second).set;
} else {
CHECK(false) << "Must consider at least on of either loads and stores, but both are false";
}
for (size_t i = 0; i < bounds.size(); ++i) {
ret.push_back(arith::Union(bounds[i]).CoverRange(none));
}
return ret;
}
Expand DownExpand Up@@ -78,41 +115,70 @@ class BufferTouchedDomain final : public StmtExprVisitor {
}

void VisitExpr_(const BufferLoadNode* op) final {
if (consider_loads_ && buffer_.same_as(op->buffer)) {
Touch(op->indices);
}
// Record load-exclusive buffer access
Touch(&std::get<LoadAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
// Record load-store inclusive buffer access
Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
StmtExprVisitor::VisitExpr_(op);
}

void VisitStmt_(const BufferStoreNode* op) final {
if (consider_stores_ && buffer_.same_as(op->buffer)) {
Touch(op->indices);
}
// Record store-exclusive buffer access
Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
// Record load-store inclusive buffer access
Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
StmtExprVisitor::VisitStmt_(op);
}

private:
void Touch(const Array<PrimExpr>& args) {
if (args.size() > bounds_.size()) {
bounds_.resize(args.size());
template <typename ArrayType>
void Touch(BufferTouches* bounds, const ArrayType& args) const {
if (args.size() > bounds->size()) {
bounds->resize(args.size());
}
for (size_t i = 0; i < args.size(); ++i) {
bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
(*bounds)[i].emplace_back(EvalSet(args[i], dom_map_));
}
}

const Buffer& buffer_;
bool consider_loads_, consider_stores_;
std::vector<std::vector<IntSet> > bounds_;
std::unordered_map<const BufferNode*, BufferDomainAccess> buffer_access_map_;
std::unordered_map<const VarNode*, IntSet> dom_map_;
};

Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads,
bool consider_stores) {
return BufferTouchedDomain(buffer, consider_loads, consider_stores).Find(stmt);
return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores);
}

Map<Buffer, runtime::ADT> DomainTouchedAccessMap(const PrimFunc& func) {
auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
Map<Buffer, runtime::ADT> ret;
auto& buffer_map = func->buffer_map;
for (auto& var : func->params) {
auto& buffer = buffer_map[var];
auto& access = buffer_access_map[buffer.get()];
Array<Array<IntSet>> loads, stores, combined;
for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
loads.push_back(Array<IntSet>(touch));
}
for (std::vector<IntSet>& touch : std::get<StoreAccess>(access).set) {
stores.push_back(Array<IntSet>(touch));
}
for (std::vector<IntSet>& touch : std::get<CombinedAccess>(access).set) {
combined.push_back(Array<IntSet>(touch));
}

std::vector<ObjectRef> fields;
fields.push_back(loads);
fields.push_back(stores);
fields.push_back(combined);
ret.Set(buffer, runtime::ADT::Tuple(fields));
}
return ret;
}

TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);

} // namespace arith
} // namespace tvm
2 changes: 1 addition & 1 deletion src/relay/backend/task_extraction.cc
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -52,7 +52,7 @@ bool DefaultTaskFilter(const Array<te::Tensor>& args) {
stack.pop_back();
if (tensor->op->IsInstance<PlaceholderOpNode>()) {
// do nothing
} else if (tensor->op->IsInstance<ComputeOpNode>()) {
} else if (tensor->op->IsInstance<ComputeOpNode>() || tensor->op->IsInstance<ExternOpNode>()) {
Array<Tensor> inputs = tensor->op->InputTensors();
for (const Tensor& v : inputs) {
if (!visited.count(v.get())) {
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
[TE] Support schedulable TIR compute definitions in TOPI#11589

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.

 Sign up for GitHub

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

 Merged

junrushao
merged 7 commits into

from

Jun 13, 2022
 Merged


[TE] Support schedulable TIR compute definitions in TOPI
#11589
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
eb3feba

[TIR] Update domain touched to track touched domains for all buffers


csullivan May 28, 2022
fa888bf

[TE] Refactor logical steps of CreatePrimFunc into separate functions.


csullivan May 30, 2022
f75d82d

[TE] Update CreatePrimFunc to pre-populate buffer binds map with


csullivan May 30, 2022
5efedea

[TE] Introduce te.extern_primfunc which allows for inlining a


csullivan May 31, 2022
659dce6

Add test showing the use of te.extern on a TVMScript PrimFunc


csullivan May 25, 2022
1936ae4

Fix task extraction default filter func skipping TE ExternOpNodes.


csullivan Jun 4, 2022
ad2b285

Address @Hzfengsy's review comments.


csullivan Jun 8, 2022

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
python/tvm/te/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -40,6 +40,7 @@
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc

from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
from .autodiff import gradient
Expand Down

82 changes: 82 additions & 0 deletions
python/tvm/te/operation.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -24,6 +24,7 @@
import tvm._ffi
import tvm.tir
import tvm.tir._ffi_api
import tvm.arith._ffi_api
from tvm._ffi.base import string_types
from tvm.ir import Array
from tvm.runtime import convert
Expand DownExpand Up@@ -354,6 +355,87 @@ def extern(
 return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
 """Compute tensors via a schedulable TIR PrimFunc

 Parameters
 ----------
 input_tensors: list of Tensor
 Input tensors that map to the corresponding primfunc input params.

 primfunc: PrimFunc
 The TIR PrimFunc

 Returns
 -------
 tensor: Tensor or list of Tensors
 The created tensor or tuple of tensors if it contains multiple outputs.

 Example
 -------
 In the code below, a TVMScript defined TIR PrimFunc is inlined into
 a TE ExternOp. Applying te.create_prim_func on this

 .. code-block:: python

 A = te.placeholder((128, 128), name="A")
 B = te.placeholder((128, 128), name="B")

 @T.prim_func
 def before_split(a: T.handle, b: T.handle) -> None:
 A = T.match_buffer(a, (128, 128))
 B = T.match_buffer(b, (128, 128))
 for i, j in T.grid(128, 128):
 with T.block("B"):
 vi, vj = T.axis.remap("SS", [i, j])
 B[vi, vj] = A[vi, vj] * 2.0

 C = te.extern_primfunc([A, B], func)
 """
 access_map = {
 k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
 }
 in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
 out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
 assert in_buffers, "PrimFunc has no input buffers"
 assert out_buffers, "PrimFunc has no output buffers"

 outputs = []
 inplace = []
 input_buffers = in_buffers
 for obuf in out_buffers:
 if obuf in in_buffers:
 inplace.append(obuf)
 else:
 outputs.append(obuf)

 if not outputs:
 iobuf = inplace.pop()
 input_buffers.remove(iobuf)
 outputs = [iobuf]

 assert len(input_buffers) == len(input_tensors), (
 "The number of provided input input_tensors does not match the number of ",
 "input buffers in the primfunc",
 )
 for tensor, buffer in zip(input_tensors, input_buffers):
 # TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
 assert tensor.shape == buffer.shape, (
 "The input input_tensors provided do not match the input buffers in the ",
 "primfunc. Please check that the order of input te.Input_Tensors and the ",
 "order of the primfunc variables in the params list agree.",
 )
 output = extern(
 [buf.shape for buf in outputs],
 input_tensors,
 lambda ins, outs: primfunc.body,
 in_buffers=input_buffers,
 out_buffers=outputs,
 **kwargs,
 )
 return output


def var(name="tindex", dtype="int32", span=None):
 """Create a new variable with specified name and dtype

Expand Down

106 changes: 86 additions & 20 deletions
src/arith/domain_touched.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <tuple>
#include <unordered_map>
#include <unordered_set>

Expand All@@ -34,18 +35,54 @@ namespace arith {

using namespace tir;

namespace {

using BufferTouches = std::vector<std::vector<IntSet>>;

struct LoadAccess {
Comment thread
Hzfengsy marked this conversation as resolved.

 BufferTouches set;
};

struct StoreAccess {
 BufferTouches set;
};

struct CombinedAccess {
 BufferTouches set;
};

using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;

} // namespace

// Find Read region of the tensor in the stmt.
class BufferTouchedDomain final : public StmtExprVisitor {
 public:
 BufferTouchedDomain(const Buffer& buffer, bool consider_loads, bool consider_stores)
 : buffer_(buffer), consider_loads_(consider_loads), consider_stores_(consider_stores) {}
 BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }

 std::unordered_map<const BufferNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
 return buffer_access_map_;
 }

 Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) {
 auto kv = buffer_access_map_.find(buffer.get());
 CHECK(kv != buffer_access_map_.end())
 << "The requested buffer is not contained in the provided stmt body.";

 Region Find(const Stmt& stmt) {
 operator()(stmt);
 Region ret;
 Range none;
 for (size_t i = 0; i < bounds_.size(); ++i) {
 ret.push_back(arith::Union(bounds_[i]).CoverRange(none));
 BufferTouches bounds;
 if (consider_loads && consider_stores) {
 bounds = std::get<CombinedAccess>(kv->second).set;
 } else if (consider_loads) {
 bounds = std::get<LoadAccess>(kv->second).set;
 } else if (consider_stores) {
 bounds = std::get<StoreAccess>(kv->second).set;
 } else {
 CHECK(false) << "Must consider at least on of either loads and stores, but both are false";
 }
 for (size_t i = 0; i < bounds.size(); ++i) {
 ret.push_back(arith::Union(bounds[i]).CoverRange(none));
 }
 return ret;
 }
Expand DownExpand Up@@ -78,41 +115,70 @@ class BufferTouchedDomain final : public StmtExprVisitor {
 }

 void VisitExpr_(const BufferLoadNode* op) final {
 if (consider_loads_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record load-exclusive buffer access
 Touch(&std::get<LoadAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitExpr_(op);
 }

 void VisitStmt_(const BufferStoreNode* op) final {
 if (consider_stores_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record store-exclusive buffer access
 Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitStmt_(op);
 }

 private:
 void Touch(const Array<PrimExpr>& args) {
 if (args.size() > bounds_.size()) {
 bounds_.resize(args.size());
 template <typename ArrayType>
 void Touch(BufferTouches* bounds, const ArrayType& args) const {
 if (args.size() > bounds->size()) {
 bounds->resize(args.size());
 }
 for (size_t i = 0; i < args.size(); ++i) {
bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
(*bounds)[i].emplace_back(EvalSet(args[i], dom_map_));
 }
 }

 const Buffer& buffer_;
 bool consider_loads_, consider_stores_;
 std::vector<std::vector<IntSet> > bounds_;
 std::unordered_map<const BufferNode*, BufferDomainAccess> buffer_access_map_;
 std::unordered_map<const VarNode*, IntSet> dom_map_;
};

Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads,
 bool consider_stores) {
 return BufferTouchedDomain(buffer, consider_loads, consider_stores).Find(stmt);
 return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores);
}

Map<Buffer, runtime::ADT> DomainTouchedAccessMap(const PrimFunc& func) {
 auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
 Map<Buffer, runtime::ADT> ret;
 auto& buffer_map = func->buffer_map;
 for (auto& var : func->params) {
 auto& buffer = buffer_map[var];
 auto& access = buffer_access_map[buffer.get()];
 Array<Array<IntSet>> loads, stores, combined;
 for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
 loads.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<StoreAccess>(access).set) {
 stores.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<CombinedAccess>(access).set) {
 combined.push_back(Array<IntSet>(touch));
 }

 std::vector<ObjectRef> fields;
 fields.push_back(loads);
 fields.push_back(stores);
 fields.push_back(combined);
 ret.Set(buffer, runtime::ADT::Tuple(fields));
 }
 return ret;
}

TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);

} // namespace arith
} // namespace tvm

2 changes: 1 addition & 1 deletion
src/relay/backend/task_extraction.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -52,7 +52,7 @@ bool DefaultTaskFilter(const Array<te::Tensor>& args) {
 stack.pop_back();
 if (tensor->op->IsInstance<PlaceholderOpNode>()) {
 // do nothing
 } else if (tensor->op->IsInstance<ComputeOpNode>()) {
 } else if (tensor->op->IsInstance<ComputeOpNode>() || tensor->op->IsInstance<ExternOpNode>()) {
 Array<Tensor> inputs = tensor->op->InputTensors();
 for (const Tensor& v : inputs) {
 if (!visited.count(v.get())) {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[TE] Support schedulable TIR compute definitions in TOPI#11589

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.

 Sign up for GitHub

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

 Merged

junrushao
merged 7 commits into

from

Jun 13, 2022
 Merged


[TE] Support schedulable TIR compute definitions in TOPI
#11589
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
eb3feba

[TIR] Update domain touched to track touched domains for all buffers


csullivan May 28, 2022
fa888bf

[TE] Refactor logical steps of CreatePrimFunc into separate functions.


csullivan May 30, 2022
f75d82d

[TE] Update CreatePrimFunc to pre-populate buffer binds map with


csullivan May 30, 2022
5efedea

[TE] Introduce te.extern_primfunc which allows for inlining a


csullivan May 31, 2022
659dce6

Add test showing the use of te.extern on a TVMScript PrimFunc


csullivan May 25, 2022
1936ae4

Fix task extraction default filter func skipping TE ExternOpNodes.


csullivan Jun 4, 2022
ad2b285

Address @Hzfengsy's review comments.


csullivan Jun 8, 2022

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
python/tvm/te/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -40,6 +40,7 @@
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc

from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
from .autodiff import gradient
Expand Down

82 changes: 82 additions & 0 deletions
python/tvm/te/operation.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -24,6 +24,7 @@
import tvm._ffi
import tvm.tir
import tvm.tir._ffi_api
import tvm.arith._ffi_api
from tvm._ffi.base import string_types
from tvm.ir import Array
from tvm.runtime import convert
Expand DownExpand Up@@ -354,6 +355,87 @@ def extern(
 return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
 """Compute tensors via a schedulable TIR PrimFunc

 Parameters
 ----------
 input_tensors: list of Tensor
 Input tensors that map to the corresponding primfunc input params.

 primfunc: PrimFunc
 The TIR PrimFunc

 Returns
 -------
 tensor: Tensor or list of Tensors
 The created tensor or tuple of tensors if it contains multiple outputs.

 Example
 -------
 In the code below, a TVMScript defined TIR PrimFunc is inlined into
 a TE ExternOp. Applying te.create_prim_func on this

 .. code-block:: python

 A = te.placeholder((128, 128), name="A")
 B = te.placeholder((128, 128), name="B")

 @T.prim_func
 def before_split(a: T.handle, b: T.handle) -> None:
 A = T.match_buffer(a, (128, 128))
 B = T.match_buffer(b, (128, 128))
 for i, j in T.grid(128, 128):
 with T.block("B"):
 vi, vj = T.axis.remap("SS", [i, j])
 B[vi, vj] = A[vi, vj] * 2.0

 C = te.extern_primfunc([A, B], func)
 """
 access_map = {
 k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
 }
 in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
 out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
 assert in_buffers, "PrimFunc has no input buffers"
 assert out_buffers, "PrimFunc has no output buffers"

 outputs = []
 inplace = []
 input_buffers = in_buffers
 for obuf in out_buffers:
 if obuf in in_buffers:
 inplace.append(obuf)
 else:
 outputs.append(obuf)

 if not outputs:
 iobuf = inplace.pop()
 input_buffers.remove(iobuf)
 outputs = [iobuf]

 assert len(input_buffers) == len(input_tensors), (
 "The number of provided input input_tensors does not match the number of ",
 "input buffers in the primfunc",
 )
 for tensor, buffer in zip(input_tensors, input_buffers):
 # TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
 assert tensor.shape == buffer.shape, (
 "The input input_tensors provided do not match the input buffers in the ",
 "primfunc. Please check that the order of input te.Input_Tensors and the ",
 "order of the primfunc variables in the params list agree.",
 )
 output = extern(
 [buf.shape for buf in outputs],
 input_tensors,
 lambda ins, outs: primfunc.body,
 in_buffers=input_buffers,
 out_buffers=outputs,
 **kwargs,
 )
 return output


def var(name="tindex", dtype="int32", span=None):
 """Create a new variable with specified name and dtype

Expand Down

106 changes: 86 additions & 20 deletions
src/arith/domain_touched.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <tuple>
#include <unordered_map>
#include <unordered_set>

Expand All@@ -34,18 +35,54 @@ namespace arith {

using namespace tir;

namespace {

using BufferTouches = std::vector<std::vector<IntSet>>;

struct LoadAccess {
Comment thread
Hzfengsy marked this conversation as resolved.

 BufferTouches set;
};

struct StoreAccess {
 BufferTouches set;
};

struct CombinedAccess {
 BufferTouches set;
};

using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;

} // namespace

// Find Read region of the tensor in the stmt.
class BufferTouchedDomain final : public StmtExprVisitor {
 public:
 BufferTouchedDomain(const Buffer& buffer, bool consider_loads, bool consider_stores)
 : buffer_(buffer), consider_loads_(consider_loads), consider_stores_(consider_stores) {}
 BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }

 std::unordered_map<const BufferNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
 return buffer_access_map_;
 }

 Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) {
 auto kv = buffer_access_map_.find(buffer.get());
 CHECK(kv != buffer_access_map_.end())
 << "The requested buffer is not contained in the provided stmt body.";

 Region Find(const Stmt& stmt) {
 operator()(stmt);
 Region ret;
 Range none;
 for (size_t i = 0; i < bounds_.size(); ++i) {
 ret.push_back(arith::Union(bounds_[i]).CoverRange(none));
 BufferTouches bounds;
 if (consider_loads && consider_stores) {
 bounds = std::get<CombinedAccess>(kv->second).set;
 } else if (consider_loads) {
 bounds = std::get<LoadAccess>(kv->second).set;
 } else if (consider_stores) {
 bounds = std::get<StoreAccess>(kv->second).set;
 } else {
 CHECK(false) << "Must consider at least on of either loads and stores, but both are false";
 }
 for (size_t i = 0; i < bounds.size(); ++i) {
 ret.push_back(arith::Union(bounds[i]).CoverRange(none));
 }
 return ret;
 }
Expand DownExpand Up@@ -78,41 +115,70 @@ class BufferTouchedDomain final : public StmtExprVisitor {
 }

 void VisitExpr_(const BufferLoadNode* op) final {
 if (consider_loads_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record load-exclusive buffer access
 Touch(&std::get<LoadAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitExpr_(op);
 }

 void VisitStmt_(const BufferStoreNode* op) final {
 if (consider_stores_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record store-exclusive buffer access
 Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitStmt_(op);
 }

 private:
 void Touch(const Array<PrimExpr>& args) {
 if (args.size() > bounds_.size()) {
 bounds_.resize(args.size());
 template <typename ArrayType>
 void Touch(BufferTouches* bounds, const ArrayType& args) const {
 if (args.size() > bounds->size()) {
 bounds->resize(args.size());
 }
 for (size_t i = 0; i < args.size(); ++i) {
bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
(*bounds)[i].emplace_back(EvalSet(args[i], dom_map_));
 }
 }

 const Buffer& buffer_;
 bool consider_loads_, consider_stores_;
 std::vector<std::vector<IntSet> > bounds_;
 std::unordered_map<const BufferNode*, BufferDomainAccess> buffer_access_map_;
 std::unordered_map<const VarNode*, IntSet> dom_map_;
};

Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads,
 bool consider_stores) {
 return BufferTouchedDomain(buffer, consider_loads, consider_stores).Find(stmt);
 return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores);
}

Map<Buffer, runtime::ADT> DomainTouchedAccessMap(const PrimFunc& func) {
 auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
 Map<Buffer, runtime::ADT> ret;
 auto& buffer_map = func->buffer_map;
 for (auto& var : func->params) {
 auto& buffer = buffer_map[var];
 auto& access = buffer_access_map[buffer.get()];
 Array<Array<IntSet>> loads, stores, combined;
 for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
 loads.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<StoreAccess>(access).set) {
 stores.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<CombinedAccess>(access).set) {
 combined.push_back(Array<IntSet>(touch));
 }

 std::vector<ObjectRef> fields;
 fields.push_back(loads);
 fields.push_back(stores);
 fields.push_back(combined);
 ret.Set(buffer, runtime::ADT::Tuple(fields));
 }
 return ret;
}

TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);

} // namespace arith
} // namespace tvm

2 changes: 1 addition & 1 deletion
src/relay/backend/task_extraction.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -52,7 +52,7 @@ bool DefaultTaskFilter(const Array<te::Tensor>& args) {
 stack.pop_back();
 if (tensor->op->IsInstance<PlaceholderOpNode>()) {
 // do nothing
 } else if (tensor->op->IsInstance<ComputeOpNode>()) {
 } else if (tensor->op->IsInstance<ComputeOpNode>() || tensor->op->IsInstance<ExternOpNode>()) {
 Array<Tensor> inputs = tensor->op->InputTensors();
 for (const Tensor& v : inputs) {
 if (!visited.count(v.get())) {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[TE] Support schedulable TIR compute definitions in TOPI#11589

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.

 Sign up for GitHub

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

 Merged

junrushao
merged 7 commits into

from

Jun 13, 2022
 Merged


[TE] Support schedulable TIR compute definitions in TOPI
#11589
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
eb3feba

[TIR] Update domain touched to track touched domains for all buffers


csullivan May 28, 2022
fa888bf

[TE] Refactor logical steps of CreatePrimFunc into separate functions.


csullivan May 30, 2022
f75d82d

[TE] Update CreatePrimFunc to pre-populate buffer binds map with


csullivan May 30, 2022
5efedea

[TE] Introduce te.extern_primfunc which allows for inlining a


csullivan May 31, 2022
659dce6

Add test showing the use of te.extern on a TVMScript PrimFunc


csullivan May 25, 2022
1936ae4

Fix task extraction default filter func skipping TE ExternOpNodes.


csullivan Jun 4, 2022
ad2b285

Address @Hzfengsy's review comments.


csullivan Jun 8, 2022

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
python/tvm/te/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -40,6 +40,7 @@
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc

from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
from .autodiff import gradient
Expand Down

82 changes: 82 additions & 0 deletions
python/tvm/te/operation.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -24,6 +24,7 @@
import tvm._ffi
import tvm.tir
import tvm.tir._ffi_api
import tvm.arith._ffi_api
from tvm._ffi.base import string_types
from tvm.ir import Array
from tvm.runtime import convert
Expand DownExpand Up@@ -354,6 +355,87 @@ def extern(
 return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
 """Compute tensors via a schedulable TIR PrimFunc

 Parameters
 ----------
 input_tensors: list of Tensor
 Input tensors that map to the corresponding primfunc input params.

 primfunc: PrimFunc
 The TIR PrimFunc

 Returns
 -------
 tensor: Tensor or list of Tensors
 The created tensor or tuple of tensors if it contains multiple outputs.

 Example
 -------
 In the code below, a TVMScript defined TIR PrimFunc is inlined into
 a TE ExternOp. Applying te.create_prim_func on this

 .. code-block:: python

 A = te.placeholder((128, 128), name="A")
 B = te.placeholder((128, 128), name="B")

 @T.prim_func
 def before_split(a: T.handle, b: T.handle) -> None:
 A = T.match_buffer(a, (128, 128))
 B = T.match_buffer(b, (128, 128))
 for i, j in T.grid(128, 128):
 with T.block("B"):
 vi, vj = T.axis.remap("SS", [i, j])
 B[vi, vj] = A[vi, vj] * 2.0

 C = te.extern_primfunc([A, B], func)
 """
 access_map = {
 k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
 }
 in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
 out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
 assert in_buffers, "PrimFunc has no input buffers"
 assert out_buffers, "PrimFunc has no output buffers"

 outputs = []
 inplace = []
 input_buffers = in_buffers
 for obuf in out_buffers:
 if obuf in in_buffers:
 inplace.append(obuf)
 else:
 outputs.append(obuf)

 if not outputs:
 iobuf = inplace.pop()
 input_buffers.remove(iobuf)
 outputs = [iobuf]

 assert len(input_buffers) == len(input_tensors), (
 "The number of provided input input_tensors does not match the number of ",
 "input buffers in the primfunc",
 )
 for tensor, buffer in zip(input_tensors, input_buffers):
 # TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
 assert tensor.shape == buffer.shape, (
 "The input input_tensors provided do not match the input buffers in the ",
 "primfunc. Please check that the order of input te.Input_Tensors and the ",
 "order of the primfunc variables in the params list agree.",
 )
 output = extern(
 [buf.shape for buf in outputs],
 input_tensors,
 lambda ins, outs: primfunc.body,
 in_buffers=input_buffers,
 out_buffers=outputs,
 **kwargs,
 )
 return output


def var(name="tindex", dtype="int32", span=None):
 """Create a new variable with specified name and dtype

Expand Down

106 changes: 86 additions & 20 deletions
src/arith/domain_touched.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <tuple>
#include <unordered_map>
#include <unordered_set>

Expand All@@ -34,18 +35,54 @@ namespace arith {

using namespace tir;

namespace {

using BufferTouches = std::vector<std::vector<IntSet>>;

struct LoadAccess {
Comment thread
Hzfengsy marked this conversation as resolved.

 BufferTouches set;
};

struct StoreAccess {
 BufferTouches set;
};

struct CombinedAccess {
 BufferTouches set;
};

using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;

} // namespace

// Find Read region of the tensor in the stmt.
class BufferTouchedDomain final : public StmtExprVisitor {
 public:
 BufferTouchedDomain(const Buffer& buffer, bool consider_loads, bool consider_stores)
 : buffer_(buffer), consider_loads_(consider_loads), consider_stores_(consider_stores) {}
 BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }

 std::unordered_map<const BufferNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
 return buffer_access_map_;
 }

 Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) {
 auto kv = buffer_access_map_.find(buffer.get());
 CHECK(kv != buffer_access_map_.end())
 << "The requested buffer is not contained in the provided stmt body.";

 Region Find(const Stmt& stmt) {
 operator()(stmt);
 Region ret;
 Range none;
 for (size_t i = 0; i < bounds_.size(); ++i) {
 ret.push_back(arith::Union(bounds_[i]).CoverRange(none));
 BufferTouches bounds;
 if (consider_loads && consider_stores) {
 bounds = std::get<CombinedAccess>(kv->second).set;
 } else if (consider_loads) {
 bounds = std::get<LoadAccess>(kv->second).set;
 } else if (consider_stores) {
 bounds = std::get<StoreAccess>(kv->second).set;
 } else {
 CHECK(false) << "Must consider at least on of either loads and stores, but both are false";
 }
 for (size_t i = 0; i < bounds.size(); ++i) {
 ret.push_back(arith::Union(bounds[i]).CoverRange(none));
 }
 return ret;
 }
Expand DownExpand Up@@ -78,41 +115,70 @@ class BufferTouchedDomain final : public StmtExprVisitor {
 }

 void VisitExpr_(const BufferLoadNode* op) final {
 if (consider_loads_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record load-exclusive buffer access
 Touch(&std::get<LoadAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitExpr_(op);
 }

 void VisitStmt_(const BufferStoreNode* op) final {
 if (consider_stores_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record store-exclusive buffer access
 Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitStmt_(op);
 }

 private:
 void Touch(const Array<PrimExpr>& args) {
 if (args.size() > bounds_.size()) {
 bounds_.resize(args.size());
 template <typename ArrayType>
 void Touch(BufferTouches* bounds, const ArrayType& args) const {
 if (args.size() > bounds->size()) {
 bounds->resize(args.size());
 }
 for (size_t i = 0; i < args.size(); ++i) {
bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
(*bounds)[i].emplace_back(EvalSet(args[i], dom_map_));
 }
 }

 const Buffer& buffer_;
 bool consider_loads_, consider_stores_;
 std::vector<std::vector<IntSet> > bounds_;
 std::unordered_map<const BufferNode*, BufferDomainAccess> buffer_access_map_;
 std::unordered_map<const VarNode*, IntSet> dom_map_;
};

Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads,
 bool consider_stores) {
 return BufferTouchedDomain(buffer, consider_loads, consider_stores).Find(stmt);
 return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores);
}

Map<Buffer, runtime::ADT> DomainTouchedAccessMap(const PrimFunc& func) {
 auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
 Map<Buffer, runtime::ADT> ret;
 auto& buffer_map = func->buffer_map;
 for (auto& var : func->params) {
 auto& buffer = buffer_map[var];
 auto& access = buffer_access_map[buffer.get()];
 Array<Array<IntSet>> loads, stores, combined;
 for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
 loads.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<StoreAccess>(access).set) {
 stores.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<CombinedAccess>(access).set) {
 combined.push_back(Array<IntSet>(touch));
 }

 std::vector<ObjectRef> fields;
 fields.push_back(loads);
 fields.push_back(stores);
 fields.push_back(combined);
 ret.Set(buffer, runtime::ADT::Tuple(fields));
 }
 return ret;
}

TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);

} // namespace arith
} // namespace tvm

2 changes: 1 addition & 1 deletion
src/relay/backend/task_extraction.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -52,7 +52,7 @@ bool DefaultTaskFilter(const Array<te::Tensor>& args) {
 stack.pop_back();
 if (tensor->op->IsInstance<PlaceholderOpNode>()) {
 // do nothing
 } else if (tensor->op->IsInstance<ComputeOpNode>()) {
 } else if (tensor->op->IsInstance<ComputeOpNode>() || tensor->op->IsInstance<ExternOpNode>()) {
 Array<Tensor> inputs = tensor->op->InputTensors();
 for (const Tensor& v : inputs) {
 if (!visited.count(v.get())) {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
[TE] Support schedulable TIR compute definitions in TOPI#11589

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.

 Sign up for GitHub

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

 Merged

junrushao
merged 7 commits into

from

Jun 13, 2022
 Merged


[TE] Support schedulable TIR compute definitions in TOPI
#11589
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
eb3feba

[TIR] Update domain touched to track touched domains for all buffers


csullivan May 28, 2022
fa888bf

[TE] Refactor logical steps of CreatePrimFunc into separate functions.


csullivan May 30, 2022
f75d82d

[TE] Update CreatePrimFunc to pre-populate buffer binds map with


csullivan May 30, 2022
5efedea

[TE] Introduce te.extern_primfunc which allows for inlining a


csullivan May 31, 2022
659dce6

Add test showing the use of te.extern on a TVMScript PrimFunc


csullivan May 25, 2022
1936ae4

Fix task extraction default filter func skipping TE ExternOpNodes.


csullivan Jun 4, 2022
ad2b285

Address @Hzfengsy's review comments.


csullivan Jun 8, 2022

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
python/tvm/te/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -40,6 +40,7 @@
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc

from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
from .autodiff import gradient
Expand Down

82 changes: 82 additions & 0 deletions
python/tvm/te/operation.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -24,6 +24,7 @@
import tvm._ffi
import tvm.tir
import tvm.tir._ffi_api
import tvm.arith._ffi_api
from tvm._ffi.base import string_types
from tvm.ir import Array
from tvm.runtime import convert
Expand DownExpand Up@@ -354,6 +355,87 @@ def extern(
 return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
 """Compute tensors via a schedulable TIR PrimFunc

 Parameters
 ----------
 input_tensors: list of Tensor
 Input tensors that map to the corresponding primfunc input params.

 primfunc: PrimFunc
 The TIR PrimFunc

 Returns
 -------
 tensor: Tensor or list of Tensors
 The created tensor or tuple of tensors if it contains multiple outputs.

 Example
 -------
 In the code below, a TVMScript defined TIR PrimFunc is inlined into
 a TE ExternOp. Applying te.create_prim_func on this

 .. code-block:: python

 A = te.placeholder((128, 128), name="A")
 B = te.placeholder((128, 128), name="B")

 @T.prim_func
 def before_split(a: T.handle, b: T.handle) -> None:
 A = T.match_buffer(a, (128, 128))
 B = T.match_buffer(b, (128, 128))
 for i, j in T.grid(128, 128):
 with T.block("B"):
 vi, vj = T.axis.remap("SS", [i, j])
 B[vi, vj] = A[vi, vj] * 2.0

 C = te.extern_primfunc([A, B], func)
 """
 access_map = {
 k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
 }
 in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
 out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
 assert in_buffers, "PrimFunc has no input buffers"
 assert out_buffers, "PrimFunc has no output buffers"

 outputs = []
 inplace = []
 input_buffers = in_buffers
 for obuf in out_buffers:
 if obuf in in_buffers:
 inplace.append(obuf)
 else:
 outputs.append(obuf)

 if not outputs:
 iobuf = inplace.pop()
 input_buffers.remove(iobuf)
 outputs = [iobuf]

 assert len(input_buffers) == len(input_tensors), (
 "The number of provided input input_tensors does not match the number of ",
 "input buffers in the primfunc",
 )
 for tensor, buffer in zip(input_tensors, input_buffers):
 # TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
 assert tensor.shape == buffer.shape, (
 "The input input_tensors provided do not match the input buffers in the ",
 "primfunc. Please check that the order of input te.Input_Tensors and the ",
 "order of the primfunc variables in the params list agree.",
 )
 output = extern(
 [buf.shape for buf in outputs],
 input_tensors,
 lambda ins, outs: primfunc.body,
 in_buffers=input_buffers,
 out_buffers=outputs,
 **kwargs,
 )
 return output


def var(name="tindex", dtype="int32", span=None):
 """Create a new variable with specified name and dtype

Expand Down

106 changes: 86 additions & 20 deletions
src/arith/domain_touched.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <tuple>
#include <unordered_map>
#include <unordered_set>

Expand All@@ -34,18 +35,54 @@ namespace arith {

using namespace tir;

namespace {

using BufferTouches = std::vector<std::vector<IntSet>>;

struct LoadAccess {
Comment thread
Hzfengsy marked this conversation as resolved.

 BufferTouches set;
};

struct StoreAccess {
 BufferTouches set;
};

struct CombinedAccess {
 BufferTouches set;
};

using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;

} // namespace

// Find Read region of the tensor in the stmt.
class BufferTouchedDomain final : public StmtExprVisitor {
 public:
 BufferTouchedDomain(const Buffer& buffer, bool consider_loads, bool consider_stores)
 : buffer_(buffer), consider_loads_(consider_loads), consider_stores_(consider_stores) {}
 BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }

 std::unordered_map<const BufferNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
 return buffer_access_map_;
 }

 Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) {
 auto kv = buffer_access_map_.find(buffer.get());
 CHECK(kv != buffer_access_map_.end())
 << "The requested buffer is not contained in the provided stmt body.";

 Region Find(const Stmt& stmt) {
 operator()(stmt);
 Region ret;
 Range none;
 for (size_t i = 0; i < bounds_.size(); ++i) {
 ret.push_back(arith::Union(bounds_[i]).CoverRange(none));
 BufferTouches bounds;
 if (consider_loads && consider_stores) {
 bounds = std::get<CombinedAccess>(kv->second).set;
 } else if (consider_loads) {
 bounds = std::get<LoadAccess>(kv->second).set;
 } else if (consider_stores) {
 bounds = std::get<StoreAccess>(kv->second).set;
 } else {
 CHECK(false) << "Must consider at least on of either loads and stores, but both are false";
 }
 for (size_t i = 0; i < bounds.size(); ++i) {
 ret.push_back(arith::Union(bounds[i]).CoverRange(none));
 }
 return ret;
 }
Expand DownExpand Up@@ -78,41 +115,70 @@ class BufferTouchedDomain final : public StmtExprVisitor {
 }

 void VisitExpr_(const BufferLoadNode* op) final {
 if (consider_loads_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record load-exclusive buffer access
 Touch(&std::get<LoadAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitExpr_(op);
 }

 void VisitStmt_(const BufferStoreNode* op) final {
 if (consider_stores_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record store-exclusive buffer access
 Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitStmt_(op);
 }

 private:
 void Touch(const Array<PrimExpr>& args) {
 if (args.size() > bounds_.size()) {
 bounds_.resize(args.size());
 template <typename ArrayType>
 void Touch(BufferTouches* bounds, const ArrayType& args) const {
 if (args.size() > bounds->size()) {
 bounds->resize(args.size());
 }
 for (size_t i = 0; i < args.size(); ++i) {
bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
(*bounds)[i].emplace_back(EvalSet(args[i], dom_map_));
 }
 }

 const Buffer& buffer_;
 bool consider_loads_, consider_stores_;
 std::vector<std::vector<IntSet> > bounds_;
 std::unordered_map<const BufferNode*, BufferDomainAccess> buffer_access_map_;
 std::unordered_map<const VarNode*, IntSet> dom_map_;
};

Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads,
 bool consider_stores) {
 return BufferTouchedDomain(buffer, consider_loads, consider_stores).Find(stmt);
 return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores);
}

Map<Buffer, runtime::ADT> DomainTouchedAccessMap(const PrimFunc& func) {
 auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
 Map<Buffer, runtime::ADT> ret;
 auto& buffer_map = func->buffer_map;
 for (auto& var : func->params) {
 auto& buffer = buffer_map[var];
 auto& access = buffer_access_map[buffer.get()];
 Array<Array<IntSet>> loads, stores, combined;
 for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
 loads.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<StoreAccess>(access).set) {
 stores.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<CombinedAccess>(access).set) {
 combined.push_back(Array<IntSet>(touch));
 }

 std::vector<ObjectRef> fields;
 fields.push_back(loads);
 fields.push_back(stores);
 fields.push_back(combined);
 ret.Set(buffer, runtime::ADT::Tuple(fields));
 }
 return ret;
}

TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);

} // namespace arith
} // namespace tvm

2 changes: 1 addition & 1 deletion
src/relay/backend/task_extraction.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -52,7 +52,7 @@ bool DefaultTaskFilter(const Array<te::Tensor>& args) {
 stack.pop_back();
 if (tensor->op->IsInstance<PlaceholderOpNode>()) {
 // do nothing
 } else if (tensor->op->IsInstance<ComputeOpNode>()) {
 } else if (tensor->op->IsInstance<ComputeOpNode>() || tensor->op->IsInstance<ExternOpNode>()) {
 Array<Tensor> inputs = tensor->op->InputTensors();
 for (const Tensor& v : inputs) {
 if (!visited.count(v.get())) {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[TE] Support schedulable TIR compute definitions in TOPI#11589

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.

 Sign up for GitHub

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

 Merged

junrushao
merged 7 commits into

from

Jun 13, 2022
 Merged


[TE] Support schedulable TIR compute definitions in TOPI
#11589
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
eb3feba

[TIR] Update domain touched to track touched domains for all buffers


csullivan May 28, 2022
fa888bf

[TE] Refactor logical steps of CreatePrimFunc into separate functions.


csullivan May 30, 2022
f75d82d

[TE] Update CreatePrimFunc to pre-populate buffer binds map with


csullivan May 30, 2022
5efedea

[TE] Introduce te.extern_primfunc which allows for inlining a


csullivan May 31, 2022
659dce6

Add test showing the use of te.extern on a TVMScript PrimFunc


csullivan May 25, 2022
1936ae4

Fix task extraction default filter func skipping TE ExternOpNodes.


csullivan Jun 4, 2022
ad2b285

Address @Hzfengsy's review comments.


csullivan Jun 8, 2022

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
python/tvm/te/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -40,6 +40,7 @@
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc

from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
from .autodiff import gradient
Expand Down

82 changes: 82 additions & 0 deletions
python/tvm/te/operation.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -24,6 +24,7 @@
import tvm._ffi
import tvm.tir
import tvm.tir._ffi_api
import tvm.arith._ffi_api
from tvm._ffi.base import string_types
from tvm.ir import Array
from tvm.runtime import convert
Expand DownExpand Up@@ -354,6 +355,87 @@ def extern(
 return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
 """Compute tensors via a schedulable TIR PrimFunc

 Parameters
 ----------
 input_tensors: list of Tensor
 Input tensors that map to the corresponding primfunc input params.

 primfunc: PrimFunc
 The TIR PrimFunc

 Returns
 -------
 tensor: Tensor or list of Tensors
 The created tensor or tuple of tensors if it contains multiple outputs.

 Example
 -------
 In the code below, a TVMScript defined TIR PrimFunc is inlined into
 a TE ExternOp. Applying te.create_prim_func on this

 .. code-block:: python

 A = te.placeholder((128, 128), name="A")
 B = te.placeholder((128, 128), name="B")

 @T.prim_func
 def before_split(a: T.handle, b: T.handle) -> None:
 A = T.match_buffer(a, (128, 128))
 B = T.match_buffer(b, (128, 128))
 for i, j in T.grid(128, 128):
 with T.block("B"):
 vi, vj = T.axis.remap("SS", [i, j])
 B[vi, vj] = A[vi, vj] * 2.0

 C = te.extern_primfunc([A, B], func)
 """
 access_map = {
 k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
 }
 in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
 out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
 assert in_buffers, "PrimFunc has no input buffers"
 assert out_buffers, "PrimFunc has no output buffers"

 outputs = []
 inplace = []
 input_buffers = in_buffers
 for obuf in out_buffers:
 if obuf in in_buffers:
 inplace.append(obuf)
 else:
 outputs.append(obuf)

 if not outputs:
 iobuf = inplace.pop()
 input_buffers.remove(iobuf)
 outputs = [iobuf]

 assert len(input_buffers) == len(input_tensors), (
 "The number of provided input input_tensors does not match the number of ",
 "input buffers in the primfunc",
 )
 for tensor, buffer in zip(input_tensors, input_buffers):
 # TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
 assert tensor.shape == buffer.shape, (
 "The input input_tensors provided do not match the input buffers in the ",
 "primfunc. Please check that the order of input te.Input_Tensors and the ",
 "order of the primfunc variables in the params list agree.",
 )
 output = extern(
 [buf.shape for buf in outputs],
 input_tensors,
 lambda ins, outs: primfunc.body,
 in_buffers=input_buffers,
 out_buffers=outputs,
 **kwargs,
 )
 return output


def var(name="tindex", dtype="int32", span=None):
 """Create a new variable with specified name and dtype

Expand Down

106 changes: 86 additions & 20 deletions
src/arith/domain_touched.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <tuple>
#include <unordered_map>
#include <unordered_set>

Expand All@@ -34,18 +35,54 @@ namespace arith {

using namespace tir;

namespace {

using BufferTouches = std::vector<std::vector<IntSet>>;

struct LoadAccess {
Comment thread
Hzfengsy marked this conversation as resolved.

 BufferTouches set;
};

struct StoreAccess {
 BufferTouches set;
};

struct CombinedAccess {
 BufferTouches set;
};

using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;

} // namespace

// Find Read region of the tensor in the stmt.
class BufferTouchedDomain final : public StmtExprVisitor {
 public:
 BufferTouchedDomain(const Buffer& buffer, bool consider_loads, bool consider_stores)
 : buffer_(buffer), consider_loads_(consider_loads), consider_stores_(consider_stores) {}
 BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }

 std::unordered_map<const BufferNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
 return buffer_access_map_;
 }

 Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) {
 auto kv = buffer_access_map_.find(buffer.get());
 CHECK(kv != buffer_access_map_.end())
 << "The requested buffer is not contained in the provided stmt body.";

 Region Find(const Stmt& stmt) {
 operator()(stmt);
 Region ret;
 Range none;
 for (size_t i = 0; i < bounds_.size(); ++i) {
 ret.push_back(arith::Union(bounds_[i]).CoverRange(none));
 BufferTouches bounds;
 if (consider_loads && consider_stores) {
 bounds = std::get<CombinedAccess>(kv->second).set;
 } else if (consider_loads) {
 bounds = std::get<LoadAccess>(kv->second).set;
 } else if (consider_stores) {
 bounds = std::get<StoreAccess>(kv->second).set;
 } else {
 CHECK(false) << "Must consider at least on of either loads and stores, but both are false";
 }
 for (size_t i = 0; i < bounds.size(); ++i) {
 ret.push_back(arith::Union(bounds[i]).CoverRange(none));
 }
 return ret;
 }
Expand DownExpand Up@@ -78,41 +115,70 @@ class BufferTouchedDomain final : public StmtExprVisitor {
 }

 void VisitExpr_(const BufferLoadNode* op) final {
 if (consider_loads_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record load-exclusive buffer access
 Touch(&std::get<LoadAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitExpr_(op);
 }

 void VisitStmt_(const BufferStoreNode* op) final {
 if (consider_stores_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record store-exclusive buffer access
 Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitStmt_(op);
 }

 private:
 void Touch(const Array<PrimExpr>& args) {
 if (args.size() > bounds_.size()) {
 bounds_.resize(args.size());
 template <typename ArrayType>
 void Touch(BufferTouches* bounds, const ArrayType& args) const {
 if (args.size() > bounds->size()) {
 bounds->resize(args.size());
 }
 for (size_t i = 0; i < args.size(); ++i) {
bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
(*bounds)[i].emplace_back(EvalSet(args[i], dom_map_));
 }
 }

 const Buffer& buffer_;
 bool consider_loads_, consider_stores_;
 std::vector<std::vector<IntSet> > bounds_;
 std::unordered_map<const BufferNode*, BufferDomainAccess> buffer_access_map_;
 std::unordered_map<const VarNode*, IntSet> dom_map_;
};

Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads,
 bool consider_stores) {
 return BufferTouchedDomain(buffer, consider_loads, consider_stores).Find(stmt);
 return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores);
}

Map<Buffer, runtime::ADT> DomainTouchedAccessMap(const PrimFunc& func) {
 auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
 Map<Buffer, runtime::ADT> ret;
 auto& buffer_map = func->buffer_map;
 for (auto& var : func->params) {
 auto& buffer = buffer_map[var];
 auto& access = buffer_access_map[buffer.get()];
 Array<Array<IntSet>> loads, stores, combined;
 for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
 loads.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<StoreAccess>(access).set) {
 stores.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<CombinedAccess>(access).set) {
 combined.push_back(Array<IntSet>(touch));
 }

 std::vector<ObjectRef> fields;
 fields.push_back(loads);
 fields.push_back(stores);
 fields.push_back(combined);
 ret.Set(buffer, runtime::ADT::Tuple(fields));
 }
 return ret;
}

TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);

} // namespace arith
} // namespace tvm

2 changes: 1 addition & 1 deletion
src/relay/backend/task_extraction.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -52,7 +52,7 @@ bool DefaultTaskFilter(const Array<te::Tensor>& args) {
 stack.pop_back();
 if (tensor->op->IsInstance<PlaceholderOpNode>()) {
 // do nothing
 } else if (tensor->op->IsInstance<ComputeOpNode>()) {
 } else if (tensor->op->IsInstance<ComputeOpNode>() || tensor->op->IsInstance<ExternOpNode>()) {
 Array<Tensor> inputs = tensor->op->InputTensors();
 for (const Tensor& v : inputs) {
 if (!visited.count(v.get())) {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[TE] Support schedulable TIR compute definitions in TOPI#11589

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.

 Sign up for GitHub

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

 Merged

junrushao
merged 7 commits into

from

Jun 13, 2022
 Merged


[TE] Support schedulable TIR compute definitions in TOPI
#11589
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
eb3feba

[TIR] Update domain touched to track touched domains for all buffers


csullivan May 28, 2022
fa888bf

[TE] Refactor logical steps of CreatePrimFunc into separate functions.


csullivan May 30, 2022
f75d82d

[TE] Update CreatePrimFunc to pre-populate buffer binds map with


csullivan May 30, 2022
5efedea

[TE] Introduce te.extern_primfunc which allows for inlining a


csullivan May 31, 2022
659dce6

Add test showing the use of te.extern on a TVMScript PrimFunc


csullivan May 25, 2022
1936ae4

Fix task extraction default filter func skipping TE ExternOpNodes.


csullivan Jun 4, 2022
ad2b285

Address @Hzfengsy's review comments.


csullivan Jun 8, 2022

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
python/tvm/te/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -40,6 +40,7 @@
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc

from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
from .autodiff import gradient
Expand Down

82 changes: 82 additions & 0 deletions
python/tvm/te/operation.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -24,6 +24,7 @@
import tvm._ffi
import tvm.tir
import tvm.tir._ffi_api
import tvm.arith._ffi_api
from tvm._ffi.base import string_types
from tvm.ir import Array
from tvm.runtime import convert
Expand DownExpand Up@@ -354,6 +355,87 @@ def extern(
 return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
 """Compute tensors via a schedulable TIR PrimFunc

 Parameters
 ----------
 input_tensors: list of Tensor
 Input tensors that map to the corresponding primfunc input params.

 primfunc: PrimFunc
 The TIR PrimFunc

 Returns
 -------
 tensor: Tensor or list of Tensors
 The created tensor or tuple of tensors if it contains multiple outputs.

 Example
 -------
 In the code below, a TVMScript defined TIR PrimFunc is inlined into
 a TE ExternOp. Applying te.create_prim_func on this

 .. code-block:: python

 A = te.placeholder((128, 128), name="A")
 B = te.placeholder((128, 128), name="B")

 @T.prim_func
 def before_split(a: T.handle, b: T.handle) -> None:
 A = T.match_buffer(a, (128, 128))
 B = T.match_buffer(b, (128, 128))
 for i, j in T.grid(128, 128):
 with T.block("B"):
 vi, vj = T.axis.remap("SS", [i, j])
 B[vi, vj] = A[vi, vj] * 2.0

 C = te.extern_primfunc([A, B], func)
 """
 access_map = {
 k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
 }
 in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
 out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
 assert in_buffers, "PrimFunc has no input buffers"
 assert out_buffers, "PrimFunc has no output buffers"

 outputs = []
 inplace = []
 input_buffers = in_buffers
 for obuf in out_buffers:
 if obuf in in_buffers:
 inplace.append(obuf)
 else:
 outputs.append(obuf)

 if not outputs:
 iobuf = inplace.pop()
 input_buffers.remove(iobuf)
 outputs = [iobuf]

 assert len(input_buffers) == len(input_tensors), (
 "The number of provided input input_tensors does not match the number of ",
 "input buffers in the primfunc",
 )
 for tensor, buffer in zip(input_tensors, input_buffers):
 # TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
 assert tensor.shape == buffer.shape, (
 "The input input_tensors provided do not match the input buffers in the ",
 "primfunc. Please check that the order of input te.Input_Tensors and the ",
 "order of the primfunc variables in the params list agree.",
 )
 output = extern(
 [buf.shape for buf in outputs],
 input_tensors,
 lambda ins, outs: primfunc.body,
 in_buffers=input_buffers,
 out_buffers=outputs,
 **kwargs,
 )
 return output


def var(name="tindex", dtype="int32", span=None):
 """Create a new variable with specified name and dtype

Expand Down

106 changes: 86 additions & 20 deletions
src/arith/domain_touched.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <tuple>
#include <unordered_map>
#include <unordered_set>

Expand All@@ -34,18 +35,54 @@ namespace arith {

using namespace tir;

namespace {

using BufferTouches = std::vector<std::vector<IntSet>>;

struct LoadAccess {
Comment thread
Hzfengsy marked this conversation as resolved.

 BufferTouches set;
};

struct StoreAccess {
 BufferTouches set;
};

struct CombinedAccess {
 BufferTouches set;
};

using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;

} // namespace

// Find Read region of the tensor in the stmt.
class BufferTouchedDomain final : public StmtExprVisitor {
 public:
 BufferTouchedDomain(const Buffer& buffer, bool consider_loads, bool consider_stores)
 : buffer_(buffer), consider_loads_(consider_loads), consider_stores_(consider_stores) {}
 BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }

 std::unordered_map<const BufferNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
 return buffer_access_map_;
 }

 Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) {
 auto kv = buffer_access_map_.find(buffer.get());
 CHECK(kv != buffer_access_map_.end())
 << "The requested buffer is not contained in the provided stmt body.";

 Region Find(const Stmt& stmt) {
 operator()(stmt);
 Region ret;
 Range none;
 for (size_t i = 0; i < bounds_.size(); ++i) {
 ret.push_back(arith::Union(bounds_[i]).CoverRange(none));
 BufferTouches bounds;
 if (consider_loads && consider_stores) {
 bounds = std::get<CombinedAccess>(kv->second).set;
 } else if (consider_loads) {
 bounds = std::get<LoadAccess>(kv->second).set;
 } else if (consider_stores) {
 bounds = std::get<StoreAccess>(kv->second).set;
 } else {
 CHECK(false) << "Must consider at least on of either loads and stores, but both are false";
 }
 for (size_t i = 0; i < bounds.size(); ++i) {
 ret.push_back(arith::Union(bounds[i]).CoverRange(none));
 }
 return ret;
 }
Expand DownExpand Up@@ -78,41 +115,70 @@ class BufferTouchedDomain final : public StmtExprVisitor {
 }

 void VisitExpr_(const BufferLoadNode* op) final {
 if (consider_loads_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record load-exclusive buffer access
 Touch(&std::get<LoadAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitExpr_(op);
 }

 void VisitStmt_(const BufferStoreNode* op) final {
 if (consider_stores_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record store-exclusive buffer access
 Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitStmt_(op);
 }

 private:
 void Touch(const Array<PrimExpr>& args) {
 if (args.size() > bounds_.size()) {
 bounds_.resize(args.size());
 template <typename ArrayType>
 void Touch(BufferTouches* bounds, const ArrayType& args) const {
 if (args.size() > bounds->size()) {
 bounds->resize(args.size());
 }
 for (size_t i = 0; i < args.size(); ++i) {
bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
(*bounds)[i].emplace_back(EvalSet(args[i], dom_map_));
 }
 }

 const Buffer& buffer_;
 bool consider_loads_, consider_stores_;
 std::vector<std::vector<IntSet> > bounds_;
 std::unordered_map<const BufferNode*, BufferDomainAccess> buffer_access_map_;
 std::unordered_map<const VarNode*, IntSet> dom_map_;
};

Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads,
 bool consider_stores) {
 return BufferTouchedDomain(buffer, consider_loads, consider_stores).Find(stmt);
 return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores);
}

Map<Buffer, runtime::ADT> DomainTouchedAccessMap(const PrimFunc& func) {
 auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
 Map<Buffer, runtime::ADT> ret;
 auto& buffer_map = func->buffer_map;
 for (auto& var : func->params) {
 auto& buffer = buffer_map[var];
 auto& access = buffer_access_map[buffer.get()];
 Array<Array<IntSet>> loads, stores, combined;
 for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
 loads.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<StoreAccess>(access).set) {
 stores.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<CombinedAccess>(access).set) {
 combined.push_back(Array<IntSet>(touch));
 }

 std::vector<ObjectRef> fields;
 fields.push_back(loads);
 fields.push_back(stores);
 fields.push_back(combined);
 ret.Set(buffer, runtime::ADT::Tuple(fields));
 }
 return ret;
}

TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);

} // namespace arith
} // namespace tvm

2 changes: 1 addition & 1 deletion
src/relay/backend/task_extraction.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -52,7 +52,7 @@ bool DefaultTaskFilter(const Array<te::Tensor>& args) {
 stack.pop_back();
 if (tensor->op->IsInstance<PlaceholderOpNode>()) {
 // do nothing
 } else if (tensor->op->IsInstance<ComputeOpNode>()) {
 } else if (tensor->op->IsInstance<ComputeOpNode>() || tensor->op->IsInstance<ExternOpNode>()) {
 Array<Tensor> inputs = tensor->op->InputTensors();
 for (const Tensor& v : inputs) {
 if (!visited.count(v.get())) {
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
[TE] Support schedulable TIR compute definitions in TOPI#11589

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.

 Sign up for GitHub

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

 Merged

junrushao
merged 7 commits into

from

Jun 13, 2022
 Merged


[TE] Support schedulable TIR compute definitions in TOPI
#11589
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


7 commits


Select commit
Hold shift + click to select a range
eb3feba

[TIR] Update domain touched to track touched domains for all buffers


csullivan May 28, 2022
fa888bf

[TE] Refactor logical steps of CreatePrimFunc into separate functions.


csullivan May 30, 2022
f75d82d

[TE] Update CreatePrimFunc to pre-populate buffer binds map with


csullivan May 30, 2022
5efedea

[TE] Introduce te.extern_primfunc which allows for inlining a


csullivan May 31, 2022
659dce6

Add test showing the use of te.extern on a TVMScript PrimFunc


csullivan May 25, 2022
1936ae4

Fix task extraction default filter func skipping TE ExternOpNodes.


csullivan Jun 4, 2022
ad2b285

Address @Hzfengsy's review comments.


csullivan Jun 8, 2022

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


1 change: 1 addition & 0 deletions
python/tvm/te/__init__.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -40,6 +40,7 @@
from .operation import placeholder, compute, scan, extern, var, size_var, const
from .operation import thread_axis, reduce_axis
from .operation import create_prim_func
from .operation import extern_primfunc

from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp, HybridOp
from .autodiff import gradient
Expand Down

82 changes: 82 additions & 0 deletions
python/tvm/te/operation.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -24,6 +24,7 @@
import tvm._ffi
import tvm.tir
import tvm.tir._ffi_api
import tvm.arith._ffi_api
from tvm._ffi.base import string_types
from tvm.ir import Array
from tvm.runtime import convert
Expand DownExpand Up@@ -354,6 +355,87 @@ def extern(
 return res[0] if len(res) == 1 else res


def extern_primfunc(input_tensors: List[_tensor.Tensor], primfunc: tvm.tir.PrimFunc, **kwargs):
 """Compute tensors via a schedulable TIR PrimFunc

 Parameters
 ----------
 input_tensors: list of Tensor
 Input tensors that map to the corresponding primfunc input params.

 primfunc: PrimFunc
 The TIR PrimFunc

 Returns
 -------
 tensor: Tensor or list of Tensors
 The created tensor or tuple of tensors if it contains multiple outputs.

 Example
 -------
 In the code below, a TVMScript defined TIR PrimFunc is inlined into
 a TE ExternOp. Applying te.create_prim_func on this

 .. code-block:: python

 A = te.placeholder((128, 128), name="A")
 B = te.placeholder((128, 128), name="B")

 @T.prim_func
 def before_split(a: T.handle, b: T.handle) -> None:
 A = T.match_buffer(a, (128, 128))
 B = T.match_buffer(b, (128, 128))
 for i, j in T.grid(128, 128):
 with T.block("B"):
 vi, vj = T.axis.remap("SS", [i, j])
 B[vi, vj] = A[vi, vj] * 2.0

 C = te.extern_primfunc([A, B], func)
 """
 access_map = {
 k: tuple(v) for k, v in tvm.arith._ffi_api.DomainTouchedAccessMap(primfunc).items()
 }
 in_buffers = [buf for buf, access in access_map.items() if len(access[0])]
 out_buffers = [buf for buf, access in access_map.items() if len(access[1])]
 assert in_buffers, "PrimFunc has no input buffers"
 assert out_buffers, "PrimFunc has no output buffers"

 outputs = []
 inplace = []
 input_buffers = in_buffers
 for obuf in out_buffers:
 if obuf in in_buffers:
 inplace.append(obuf)
 else:
 outputs.append(obuf)

 if not outputs:
 iobuf = inplace.pop()
 input_buffers.remove(iobuf)
 outputs = [iobuf]

 assert len(input_buffers) == len(input_tensors), (
 "The number of provided input input_tensors does not match the number of ",
 "input buffers in the primfunc",
 )
 for tensor, buffer in zip(input_tensors, input_buffers):
 # TODO(csullivan): Can a stronger comparison between Tensor<>Buffer be made?
 assert tensor.shape == buffer.shape, (
 "The input input_tensors provided do not match the input buffers in the ",
 "primfunc. Please check that the order of input te.Input_Tensors and the ",
 "order of the primfunc variables in the params list agree.",
 )
 output = extern(
 [buf.shape for buf in outputs],
 input_tensors,
 lambda ins, outs: primfunc.body,
 in_buffers=input_buffers,
 out_buffers=outputs,
 **kwargs,
 )
 return output


def var(name="tindex", dtype="int32", span=None):
 """Create a new variable with specified name and dtype

Expand Down

106 changes: 86 additions & 20 deletions
src/arith/domain_touched.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -26,6 +26,7 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <tuple>
#include <unordered_map>
#include <unordered_set>

Expand All@@ -34,18 +35,54 @@ namespace arith {

using namespace tir;

namespace {

using BufferTouches = std::vector<std::vector<IntSet>>;

struct LoadAccess {
Comment thread
Hzfengsy marked this conversation as resolved.

 BufferTouches set;
};

struct StoreAccess {
 BufferTouches set;
};

struct CombinedAccess {
 BufferTouches set;
};

using BufferDomainAccess = std::tuple<LoadAccess, StoreAccess, CombinedAccess>;

} // namespace

// Find Read region of the tensor in the stmt.
class BufferTouchedDomain final : public StmtExprVisitor {
 public:
 BufferTouchedDomain(const Buffer& buffer, bool consider_loads, bool consider_stores)
 : buffer_(buffer), consider_loads_(consider_loads), consider_stores_(consider_stores) {}
 BufferTouchedDomain(const Stmt& stmt) { operator()(stmt); }

 std::unordered_map<const BufferNode*, BufferDomainAccess>& GetAccessedBufferRegions() {
 return buffer_access_map_;
 }

 Region FindUnion(const Buffer& buffer, bool consider_loads, bool consider_stores) {
 auto kv = buffer_access_map_.find(buffer.get());
 CHECK(kv != buffer_access_map_.end())
 << "The requested buffer is not contained in the provided stmt body.";

 Region Find(const Stmt& stmt) {
 operator()(stmt);
 Region ret;
 Range none;
 for (size_t i = 0; i < bounds_.size(); ++i) {
 ret.push_back(arith::Union(bounds_[i]).CoverRange(none));
 BufferTouches bounds;
 if (consider_loads && consider_stores) {
 bounds = std::get<CombinedAccess>(kv->second).set;
 } else if (consider_loads) {
 bounds = std::get<LoadAccess>(kv->second).set;
 } else if (consider_stores) {
 bounds = std::get<StoreAccess>(kv->second).set;
 } else {
 CHECK(false) << "Must consider at least on of either loads and stores, but both are false";
 }
 for (size_t i = 0; i < bounds.size(); ++i) {
 ret.push_back(arith::Union(bounds[i]).CoverRange(none));
 }
 return ret;
 }
Expand DownExpand Up@@ -78,41 +115,70 @@ class BufferTouchedDomain final : public StmtExprVisitor {
 }

 void VisitExpr_(const BufferLoadNode* op) final {
 if (consider_loads_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record load-exclusive buffer access
 Touch(&std::get<LoadAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitExpr_(op);
 }

 void VisitStmt_(const BufferStoreNode* op) final {
 if (consider_stores_ && buffer_.same_as(op->buffer)) {
 Touch(op->indices);
 }
 // Record store-exclusive buffer access
 Touch(&std::get<StoreAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 // Record load-store inclusive buffer access
 Touch(&std::get<CombinedAccess>(buffer_access_map_[op->buffer.get()]).set, op->indices);
 StmtExprVisitor::VisitStmt_(op);
 }

 private:
 void Touch(const Array<PrimExpr>& args) {
 if (args.size() > bounds_.size()) {
 bounds_.resize(args.size());
 template <typename ArrayType>
 void Touch(BufferTouches* bounds, const ArrayType& args) const {
 if (args.size() > bounds->size()) {
 bounds->resize(args.size());
 }
 for (size_t i = 0; i < args.size(); ++i) {
bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
(*bounds)[i].emplace_back(EvalSet(args[i], dom_map_));
 }
 }

 const Buffer& buffer_;
 bool consider_loads_, consider_stores_;
 std::vector<std::vector<IntSet> > bounds_;
 std::unordered_map<const BufferNode*, BufferDomainAccess> buffer_access_map_;
 std::unordered_map<const VarNode*, IntSet> dom_map_;
};

Region DomainTouched(const Stmt& stmt, const Buffer& buffer, bool consider_loads,
 bool consider_stores) {
 return BufferTouchedDomain(buffer, consider_loads, consider_stores).Find(stmt);
 return BufferTouchedDomain(stmt).FindUnion(buffer, consider_loads, consider_stores);
}

Map<Buffer, runtime::ADT> DomainTouchedAccessMap(const PrimFunc& func) {
 auto buffer_access_map = BufferTouchedDomain(func->body).GetAccessedBufferRegions();
 Map<Buffer, runtime::ADT> ret;
 auto& buffer_map = func->buffer_map;
 for (auto& var : func->params) {
 auto& buffer = buffer_map[var];
 auto& access = buffer_access_map[buffer.get()];
 Array<Array<IntSet>> loads, stores, combined;
 for (std::vector<IntSet>& touch : std::get<LoadAccess>(access).set) {
 loads.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<StoreAccess>(access).set) {
 stores.push_back(Array<IntSet>(touch));
 }
 for (std::vector<IntSet>& touch : std::get<CombinedAccess>(access).set) {
 combined.push_back(Array<IntSet>(touch));
 }

 std::vector<ObjectRef> fields;
 fields.push_back(loads);
 fields.push_back(stores);
 fields.push_back(combined);
 ret.Set(buffer, runtime::ADT::Tuple(fields));
 }
 return ret;
}

TVM_REGISTER_GLOBAL("arith.DomainTouched").set_body_typed(DomainTouched);
TVM_REGISTER_GLOBAL("arith.DomainTouchedAccessMap").set_body_typed(DomainTouchedAccessMap);

} // namespace arith
} // namespace tvm

2 changes: 1 addition & 1 deletion
src/relay/backend/task_extraction.cc

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -52,7 +52,7 @@ bool DefaultTaskFilter(const Array<te::Tensor>& args) {
 stack.pop_back();
 if (tensor->op->IsInstance<PlaceholderOpNode>()) {
 // do nothing
 } else if (tensor->op->IsInstance<ComputeOpNode>()) {
 } else if (tensor->op->IsInstance<ComputeOpNode>() || tensor->op->IsInstance<ExternOpNode>()) {
 Array<Tensor> inputs = tensor->op->InputTensors();
 for (const Tensor& v : inputs) {
 if (!visited.count(v.get())) {
Expand Down
Loading