Skip to content

Commit

Permalink
Implementation of CallDPS (apache#3)
Browse files Browse the repository at this point in the history
* Implementation of call_dps.

* Implementation of PackedFuncExpr.

* Test CallDPS for TIR function.

* Rename.

* Add header and comments.

* Update.

* Address comments.
  • Loading branch information
ZihengJiang authored and YuchenJin committed Jan 14, 2023
1 parent 03fd11e commit 9a6de0e
Showing 10 changed files with 247 additions and 8 deletions.
36 changes: 35 additions & 1 deletion include/tvm/relax/expr.h
Original file line number Diff line number Diff line change
@@ -32,6 +32,9 @@ namespace tvm {
namespace relax {

using relay::Id;
using relay::Call;
using relay::Tuple;
using relay::TupleGetItem;
using ExprNode = RelayExprNode;
using Expr = RelayExpr;

@@ -388,7 +391,7 @@ class FunctionNode : public BaseFuncNode {
static constexpr const char* _type_key = "relax.expr.Function";
static constexpr const bool _type_has_method_sequal_reduce = true;
static constexpr const bool _type_has_method_shash_reduce = true;
TVM_DECLARE_FINAL_OBJECT_INFO(FunctionNode, ExprNode);
TVM_DECLARE_FINAL_OBJECT_INFO(FunctionNode, BaseFuncNode);
};

class Function : public Expr {
@@ -398,6 +401,37 @@ class Function : public Expr {
TVM_DEFINE_OBJECT_REF_METHODS(Function, Expr, FunctionNode);
};


/*! \brief The extern function, which can represent packed function. */
class ExternFuncNode : public BaseFuncNode {
public:
/*! \brief The name of global symbol. */
String global_symbol;

void VisitAttrs(AttrVisitor* v) {
v->Visit("global_symbol", &global_symbol);
}

bool SEqualReduce(const ExternFuncNode* other, SEqualReducer equal) const {
return equal(global_symbol, other->global_symbol);
}

void SHashReduce(SHashReducer hash_reduce) const {
hash_reduce(global_symbol);
}

static constexpr const char* _type_key = "relax.expr.ExternFunc";
static constexpr const bool _type_has_method_sequal_reduce = true;
static constexpr const bool _type_has_method_shash_reduce = true;
TVM_DECLARE_FINAL_OBJECT_INFO(ExternFuncNode, BaseFuncNode);
};

class ExternFunc : public Expr {
public:
TVM_DLL ExternFunc(String global_symbol);
TVM_DEFINE_OBJECT_REF_METHODS(ExternFunc, Expr, ExternFuncNode);
};

} // namespace relax
} // namespace tvm

7 changes: 7 additions & 0 deletions python/tvm/relax/__init__.py
Original file line number Diff line number Diff line change
@@ -19,6 +19,7 @@
from . import expr
from . import ty
from . import vm
from . import op


# Expr
@@ -36,10 +37,13 @@
DataflowBlock = expr.DataflowBlock
SeqExpr = expr.SeqExpr
ShapeExpr = expr.ShapeExpr
Tuple = expr.Tuple
Function = expr.Function
ExternFunc = expr.ExternFunc

# helper functions
const = expr.const
extern = expr.extern

# Type
ShapeType = ty.ShapeType
@@ -49,3 +53,6 @@
ExecBuilder = exec_builder.ExecBuilder
VirtualMachine = vm.VirtualMachine
load_exec_from_file = vm.load_exec_from_file

# Operator
from .op.base import call_dps
21 changes: 17 additions & 4 deletions python/tvm/relax/expr.py
Original file line number Diff line number Diff line change
@@ -16,15 +16,17 @@
# under the License.
from typing import List, Optional, Union, Dict
import tvm._ffi
from ..ir.base import Node, Span, SourceName
from ..relay.base import Id
from ..ir import Node, Span, SourceName, BaseFunc
from ..runtime import String
from ..relay import Id, Tuple, TupleGetItem
from ..tir import PrimExpr
from . import _ffi_api
from .. import relay

GlobalVar = relay.GlobalVar
Expr = relay.Expr
Type = relay.Type
GlobalVar = relay.GlobalVar
Call = relay.Call
const = relay.const


@@ -106,7 +108,7 @@ def __init__(self, blocks: List[BindingBlock], body: Expr) -> None:


@tvm._ffi.register_object("relax.expr.Function")
class Function(Expr):
class Function(BaseFunc):
name: Optional[GlobalVar]
params: List[Var]
body: Expr
@@ -116,3 +118,14 @@ def __init__(self, params: List[Var], body: Expr,
ret_type: Type, name: Optional[GlobalVar] = None) -> None:
self.__init_handle_by_constructor__(_ffi_api.Function, name, params,
body, ret_type)


@tvm._ffi.register_object("relax.expr.ExternFunc")
class ExternFunc(BaseFunc):
global_symbol: String

def __init__(self, global_symbol: String) -> None:
self.__init_handle_by_constructor__(_ffi_api.ExternFunc, global_symbol)

def extern(name):
return ExternFunc(name)
21 changes: 21 additions & 0 deletions python/tvm/relax/op/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=wildcard-import, redefined-builtin
"""Relax core operators."""

# Operators
from .base import *
18 changes: 18 additions & 0 deletions python/tvm/relax/op/_ffi_api.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
import tvm._ffi

tvm._ffi._init_api("relax.op", __name__)
47 changes: 47 additions & 0 deletions python/tvm/relax/op/base.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
from ...ir import BaseFunc
from ..expr import Expr, ShapeExpr, Tuple, Call
from . import _ffi_api
from typing import Union, List

def call_dps(shape: Union[ShapeExpr, List[int]],
func: Expr,
args: Union[Tuple, List[Expr]]) -> Call:
"""
Call a destination-passing-style function and return the output.
Parameters
----------
shape: ShapeExpr
The output shape.
func : ExternFunc or PrimFunc
The destination-passing-style function.
args : Tuple[Expr]
The input arguments.
Returns
-------
ret: Call
A call node for the call_dps operator.
"""
if isinstance(shape, (list, tuple)):
shape = ShapeExpr(shape)
if isinstance(args, (list, tuple)):
args = Tuple(args)
return _ffi_api.call_dps(shape, func, args)
1 change: 1 addition & 0 deletions python/tvm/relay/__init__.py
Original file line number Diff line number Diff line change
@@ -70,6 +70,7 @@
# Span
Span = base.Span
SourceName = base.SourceName
Id = base.Id

# Type
Type = ty.Type
18 changes: 15 additions & 3 deletions src/relax/expr.cc
Original file line number Diff line number Diff line change
@@ -16,7 +16,6 @@
* specific language governing permissions and limitations
* under the License.
*/

#include <tvm/relax/expr.h>

namespace tvm {
@@ -155,6 +154,8 @@ TVM_REGISTER_GLOBAL("relax.SeqExpr")
});


TVM_REGISTER_NODE_TYPE(FunctionNode);

Function::Function(runtime::Optional<GlobalVar> name,
Array<Var> params,
Expr body,
@@ -167,8 +168,6 @@ Function::Function(runtime::Optional<GlobalVar> name,
data_ = std::move(n);
}

TVM_REGISTER_NODE_TYPE(FunctionNode);

TVM_REGISTER_GLOBAL("relax.Function")
.set_body_typed([](runtime::Optional<GlobalVar> name,
Array<Var> params,
@@ -177,5 +176,18 @@ TVM_REGISTER_GLOBAL("relax.Function")
return Function(name, params, body, ret_type);
});

TVM_REGISTER_NODE_TYPE(ExternFuncNode);

ExternFunc::ExternFunc(String global_symbol) {
ObjectPtr<ExternFuncNode> n = make_object<ExternFuncNode>();
n->global_symbol = std::move(global_symbol);
data_ = std::move(n);
}

TVM_REGISTER_GLOBAL("relax.ExternFunc")
.set_body_typed([](String global_symbol) {
return ExternFunc(global_symbol);
});

} // namespace relax
} // namespace tvm
40 changes: 40 additions & 0 deletions src/relax/op.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
#include <tvm/relay/op.h>
#include <tvm/relax/expr.h>

namespace tvm {
namespace relax {

Expr MakeCallDPS(ShapeExpr shape, Expr func, Tuple args) {
static const Op& op = Op::Get("call_dps");
return Call(op, {shape, func, args}, {}, {});
}

TVM_REGISTER_GLOBAL("relax.op.call_dps")
.set_body_typed(MakeCallDPS);

RELAY_REGISTER_OP("call_dps")
.set_num_inputs(3)
.add_argument("shape", "ShapeExpr", "The output shape.")
.add_argument("func", "Expr", "The destination-passing-style function.")
.add_argument("args", "Tuple", "The input arguments.");

} // namespace relax
} // namespace tvm
46 changes: 46 additions & 0 deletions tests/python/relax/test_op.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import tvm
from tvm import tir
from tvm import relax as rx
from tvm.script import ty
from tvm.ir import TensorType
import numpy as np

@tvm.register_func("test.op.identity")
def identity_packed(a):
return tvm.nd.array(a.asnumpy)

@tvm.script.tir
def identity_tir(a: ty.handle, b: ty.handle) -> None:
A = tir.match_buffer(a, [54, 96])
B = tir.match_buffer(b, [54, 96])

with tir.block([54, 96], "compute") as [vi, vj]:
B[vi, vj] = A[vi, vj]


def test_call_dps() -> None:
shape_anno = [54, 96]
type_anno = rx.DynTensorType(2, "float32")
v0 = rx.Var("v0", shape_anno, type_anno)
v1 = rx.call_dps([54, 96], rx.extern("test.op.identity"), [v0])
v1 = rx.call_dps([54, 96], identity_tir, [v0])


if __name__ == "__main__":
test_call_dps()

0 comments on commit 9a6de0e

Please sign in to comment.