You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "yelite (via GitHub)" <gi...@apache.org> on 2023/03/01 18:49:54 UTC
[GitHub] [tvm] yelite opened a new pull request, #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
yelite opened a new pull request, #14166:
URL: https://github.com/apache/tvm/pull/14166
This PR adds the support of batch matmul to Relax CUTLASS BYOC. This includes batch matmul with shape 2D x ND, ND x 2D and ND x ND. Additionally, this supports bias, transposed y and all supported epilogues from regular matmul.
This PR also includes:
1. Add check function to all APIs related to PatternRegistryEntry (a follow up of #14109)
2. Add check function to CUTLASS pattern table, blocking unsupported matmul from being partitioned.
3. Reorganize the tests related to CUTLASS matmul.
cc @vinx13 @masahi
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122288676
##########
tests/python/relax/test_codegen_cutlass.py:
##########
@@ -255,13 +177,48 @@ def test_conv2d_offload():
tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+def test_kernel_sharing():
+ data_np = np.random.randn(16, 32, 32, 16).astype("float16")
+ weight1_np = np.random.randn(16, 3, 3, 16).astype("float16")
+ weight2_np = np.random.randn(16, 3, 3, 16).astype("float16")
+
+ out = get_result_with_relax_cutlass_offload(Conv2dx2, data_np, weight1_np, weight2_np)
+
+ relay_expr = get_relay_conv2d_relu_x2(data_np.shape, weight1_np.shape)
+ ref = get_relay_ref(relay_expr, data_np, weight1_np, weight2_np)
+
+ tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+
+
+def get_reference_matmul_result(x, y, bias=None, transposed_y=False, activation=None):
Review Comment:
Nice, I changed the matmul tests to use integer inputs in https://github.com/apache/tvm/pull/14162, that might not be needed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122270289
##########
src/relax/transform/fuse_ops.cc:
##########
@@ -900,7 +900,7 @@ class PatternBasedPartitioner : ExprVisitor {
using Group = GraphPartitioner::Group;
using GroupMap = OperatorFusor::GroupMap;
using ExprVisitor::VisitExpr_;
- using FCheckMatch = backend::PatternRegistryEntryNode::FCheckMatch;
+ using FCheckMatch = runtime::TypedPackedFunc<bool(const Map<DFPattern, Expr>&, const Expr&)>;
Review Comment:
we can remove https://github.com/apache/tvm/blob/03c5f998b37afe41257e2c37b1e95e0ef5357ce8/src/relax/transform/fuse_ops.cc#L43 if it's no longer needed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122422145
##########
python/tvm/relax/backend/pattern_registry.py:
##########
@@ -46,25 +47,59 @@ class PatternRegistryEntry(Object):
arg_patterns: Mapping[str, DFPattern]
The mapping from arg name to its pattern. It can be used to extract arg expression
from match result. All DFPattern in this map should be part of the `pattern`.
+
+ check: Callable[[Mapping[DFPattern, Expr], Expr], bool]
+ The function to check whether the match result is accepted.
"""
name: str
pattern: DFPattern
arg_patterns: Mapping[str, DFPattern]
-
- def __init__(self, name: str, pattern: DFPattern, arg_patterns: Mapping[str, DFPattern]):
+ check: Callable[[Mapping[DFPattern, Expr], Expr], bool]
+
+ def __init__(
+ self,
+ name: str,
+ pattern: DFPattern,
+ arg_patterns: Mapping[str, DFPattern],
+ check: Callable[[Mapping[DFPattern, Expr], Expr], bool],
+ ):
self.__init_handle_by_constructor__(
- _ffi_api.PatternRegistryEntry, name, pattern, arg_patterns # type: ignore
+ _ffi_api.PatternRegistryEntry, name, pattern, arg_patterns, check # type: ignore
)
+_REGISTERED_PATTERN_NAMES: Set[str] = set()
+
+
+def _cleanup_registered_patterns():
+ _ffi_api.RemovePatterns(list(_REGISTERED_PATTERN_NAMES)) # type: ignore # pylint: disable=no-member
+
+
+_CLEANUP_REGISTERED = False
+
+
+def _ensure_cleanup_function_registered():
+ """
+ Add a cleanup function to be called on interpreter termination, to remove all
+ patterns registered on the Python side. Without cleaning up those patterns,
+ program will segfault on termination. It's because the check functiosn of pattern
Review Comment:
this seems a general issue of ffi
cc @tqchen @junrushao
##########
python/tvm/relax/backend/contrib/cutlass.py:
##########
@@ -116,7 +173,6 @@ def partition_for_cutlass(mod):
compiled by the CUTLASS backend.
"""
- cutlass_patterns = get_patterns_with_prefix("cutlass")
- return transform.FuseOpsByPattern(cutlass_patterns, bind_constants=True, annotate_codegen=True)(
- mod
- )
+ cutlass_pattern_entries = get_patterns_with_prefix("cutlass")
+ patterns = [(e.name, e.pattern, e.check) for e in cutlass_pattern_entries]
+ return transform.FuseOpsByPattern(patterns, bind_constants=True, annotate_codegen=True)(mod)
Review Comment:
```suggestion
return transform.FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
```
I made a typo in the previous PR, it should be False
##########
python/tvm/relax/backend/pattern_registry.py:
##########
@@ -46,25 +47,59 @@ class PatternRegistryEntry(Object):
arg_patterns: Mapping[str, DFPattern]
The mapping from arg name to its pattern. It can be used to extract arg expression
from match result. All DFPattern in this map should be part of the `pattern`.
+
+ check: Callable[[Mapping[DFPattern, Expr], Expr], bool]
+ The function to check whether the match result is accepted.
"""
name: str
pattern: DFPattern
arg_patterns: Mapping[str, DFPattern]
-
- def __init__(self, name: str, pattern: DFPattern, arg_patterns: Mapping[str, DFPattern]):
+ check: Callable[[Mapping[DFPattern, Expr], Expr], bool]
+
+ def __init__(
+ self,
+ name: str,
+ pattern: DFPattern,
+ arg_patterns: Mapping[str, DFPattern],
+ check: Callable[[Mapping[DFPattern, Expr], Expr], bool],
+ ):
self.__init_handle_by_constructor__(
- _ffi_api.PatternRegistryEntry, name, pattern, arg_patterns # type: ignore
+ _ffi_api.PatternRegistryEntry, name, pattern, arg_patterns, check # type: ignore
)
+_REGISTERED_PATTERN_NAMES: Set[str] = set()
+
+
+def _cleanup_registered_patterns():
+ _ffi_api.RemovePatterns(list(_REGISTERED_PATTERN_NAMES)) # type: ignore # pylint: disable=no-member
+
+
+_CLEANUP_REGISTERED = False
+
+
+def _ensure_cleanup_function_registered():
+ """
+ Add a cleanup function to be called on interpreter termination, to remove all
+ patterns registered on the Python side. Without cleaning up those patterns,
+ program will segfault on termination. It's because the check functiosn of pattern
Review Comment:
```suggestion
program will segfault on termination. It's because the check functions of pattern
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tvm-bot commented on pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "tvm-bot (via GitHub)" <gi...@apache.org>.
tvm-bot commented on PR #14166:
URL: https://github.com/apache/tvm/pull/14166#issuecomment-1450685847
<!---bot-comment-->
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @-ing them in a comment.
<!--bot-comment-ccs-start-->
* cc @billishyahao, @quic-sanirudh <sub>See [#10317](https://github.com/apache/tvm/issues/10317) for details</sub><!--bot-comment-ccs-end-->
<sub>Generated by [tvm-bot](https://github.com/apache/tvm/blob/main/ci/README.md#github-actions)</sub>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122547524
##########
python/tvm/relax/backend/contrib/cutlass.py:
##########
@@ -17,84 +17,141 @@
"""Pattern table for CUTLASS backend"""
-from tvm.relax import transform
+from typing import Mapping, Optional, Tuple
+
+import tvm
+from tvm.contrib.cutlass.build import is_valid_for_cutlass_matmul
+from tvm.relax import Call, Expr, ShapeExpr, transform
+from tvm.relax.dpl import DFPattern
from ..pattern_registry import get_patterns_with_prefix, register_patterns
from ..patterns import make_fused_bias_activation_pattern, make_matmul_pattern
+
+def _get_static_shape(shape: ShapeExpr) -> Optional[Tuple[int]]:
+ result = []
+ for dim in shape.values:
+ if isinstance(dim, tvm.tir.expr.IntImm):
+ result.append(int(dim))
+ else:
+ return None
+ return result
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype):
+ """Check if dtypes in the given workload are supported by CUTLASS."""
+ return (
+ (lhs_dtype == "float16" and rhs_dtype == "float16")
+ or (lhs_dtype == "float32" and rhs_dtype == "float32")
+ or (lhs_dtype in ("int8", "uint8") and rhs_dtype in ("int8", "uint8"))
+ )
+
+
+def _check_matmul(
+ match_result: Mapping[DFPattern, Expr],
+ _: Expr,
+) -> bool:
+ matmul_call: Call = None
+ for _, expr in match_result.items():
+ if isinstance(expr, Call) and expr.op.name == "relax.matmul":
Review Comment:
The problem here is it will iterate over all nodes in the pattern, including wildcard that matches anything. This can actually be a problem when there are two consecutive matmuls, if the first matmul is not offloaded (e.g doesn’t pass check func), and when we are checking the match for the second matmul, the match result will have two matmuls because the wildcard matches the first matmul. One possible solution is to also check pattern is instance of CallPattern
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] yelite commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "yelite (via GitHub)" <gi...@apache.org>.
yelite commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122174636
##########
tests/python/relax/test_codegen_cutlass.py:
##########
@@ -255,13 +177,48 @@ def test_conv2d_offload():
tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+def test_kernel_sharing():
+ data_np = np.random.randn(16, 32, 32, 16).astype("float16")
+ weight1_np = np.random.randn(16, 3, 3, 16).astype("float16")
+ weight2_np = np.random.randn(16, 3, 3, 16).astype("float16")
+
+ out = get_result_with_relax_cutlass_offload(Conv2dx2, data_np, weight1_np, weight2_np)
+
+ relay_expr = get_relay_conv2d_relu_x2(data_np.shape, weight1_np.shape)
+ ref = get_relay_ref(relay_expr, data_np, weight1_np, weight2_np)
+
+ tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+
+
+def get_reference_matmul_result(x, y, bias=None, transposed_y=False, activation=None):
Review Comment:
Will update this after #14162 gets merged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] yelite commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "yelite (via GitHub)" <gi...@apache.org>.
yelite commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122283144
##########
tests/python/relax/test_codegen_cutlass.py:
##########
@@ -255,13 +177,48 @@ def test_conv2d_offload():
tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+def test_kernel_sharing():
+ data_np = np.random.randn(16, 32, 32, 16).astype("float16")
+ weight1_np = np.random.randn(16, 3, 3, 16).astype("float16")
+ weight2_np = np.random.randn(16, 3, 3, 16).astype("float16")
+
+ out = get_result_with_relax_cutlass_offload(Conv2dx2, data_np, weight1_np, weight2_np)
+
+ relay_expr = get_relay_conv2d_relu_x2(data_np.shape, weight1_np.shape)
+ ref = get_relay_ref(relay_expr, data_np, weight1_np, weight2_np)
+
+ tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+
+
+def get_reference_matmul_result(x, y, bias=None, transposed_y=False, activation=None):
Review Comment:
Yes, all tests in this file pass. The tolerance it uses is 'rtol=1e-2, atol=1e-3'.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122547524
##########
python/tvm/relax/backend/contrib/cutlass.py:
##########
@@ -17,84 +17,141 @@
"""Pattern table for CUTLASS backend"""
-from tvm.relax import transform
+from typing import Mapping, Optional, Tuple
+
+import tvm
+from tvm.contrib.cutlass.build import is_valid_for_cutlass_matmul
+from tvm.relax import Call, Expr, ShapeExpr, transform
+from tvm.relax.dpl import DFPattern
from ..pattern_registry import get_patterns_with_prefix, register_patterns
from ..patterns import make_fused_bias_activation_pattern, make_matmul_pattern
+
+def _get_static_shape(shape: ShapeExpr) -> Optional[Tuple[int]]:
+ result = []
+ for dim in shape.values:
+ if isinstance(dim, tvm.tir.expr.IntImm):
+ result.append(int(dim))
+ else:
+ return None
+ return result
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype):
+ """Check if dtypes in the given workload are supported by CUTLASS."""
+ return (
+ (lhs_dtype == "float16" and rhs_dtype == "float16")
+ or (lhs_dtype == "float32" and rhs_dtype == "float32")
+ or (lhs_dtype in ("int8", "uint8") and rhs_dtype in ("int8", "uint8"))
+ )
+
+
+def _check_matmul(
+ match_result: Mapping[DFPattern, Expr],
+ _: Expr,
+) -> bool:
+ matmul_call: Call = None
+ for _, expr in match_result.items():
+ if isinstance(expr, Call) and expr.op.name == "relax.matmul":
Review Comment:
The problem here is it will iterate over all nodes in the pattern, including wildcard that matches anything. This can actually be a problem when there are two consecutive matmuls, if the first matmul is not offloaded (e.g doesn’t pass check func), and when we are checking the match for the second matmul, the match result will have two matmuls because the wildcard matches the first matmul
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122267174
##########
tests/python/relax/test_codegen_cutlass.py:
##########
@@ -255,13 +177,48 @@ def test_conv2d_offload():
tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+def test_kernel_sharing():
+ data_np = np.random.randn(16, 32, 32, 16).astype("float16")
+ weight1_np = np.random.randn(16, 3, 3, 16).astype("float16")
+ weight2_np = np.random.randn(16, 3, 3, 16).astype("float16")
+
+ out = get_result_with_relax_cutlass_offload(Conv2dx2, data_np, weight1_np, weight2_np)
+
+ relay_expr = get_relay_conv2d_relu_x2(data_np.shape, weight1_np.shape)
+ ref = get_relay_ref(relay_expr, data_np, weight1_np, weight2_np)
+
+ tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+
+
+def get_reference_matmul_result(x, y, bias=None, transposed_y=False, activation=None):
Review Comment:
Were you able to pass accuracy check with this function?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] yelite commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "yelite (via GitHub)" <gi...@apache.org>.
yelite commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122377609
##########
python/tvm/relax/backend/contrib/cutlass.py:
##########
@@ -17,84 +17,132 @@
"""Pattern table for CUTLASS backend"""
-from tvm.relax import transform
+from typing import Mapping, Optional, Tuple
+
+import tvm
+from tvm.contrib.cutlass.build import is_valid_for_cutlass_matmul
+from tvm.relax import Call, Expr, ShapeExpr, transform
+from tvm.relax.dpl import DFPattern
from ..pattern_registry import get_patterns_with_prefix, register_patterns
from ..patterns import make_fused_bias_activation_pattern, make_matmul_pattern
+
+def _get_static_shape(shape: ShapeExpr) -> Optional[Tuple[int]]:
+ result = []
+ for dim in shape.values:
+ if isinstance(dim, tvm.tir.expr.IntImm):
+ result.append(int(dim))
+ else:
+ return None
+ return result
+
+
+def _check_matmul(
+ match_result: Mapping[DFPattern, Expr],
+ _: Expr,
+) -> bool:
+ matmul_call: Call = None
+ for _, expr in match_result.items():
+ if isinstance(expr, Call) and expr.op.name == "relax.matmul":
+ matmul_call = expr
+ if matmul_call is None:
+ raise ValueError("Cannot find call to matmul from match_result.")
+
+ lhs_shape = _get_static_shape(matmul_call.args[0].struct_info.shape)
+ rhs_shape = _get_static_shape(matmul_call.args[1].struct_info.shape)
+ if len(lhs_shape) < 2 or len(rhs_shape) < 2:
+ return False
+
+ lhs_dtype = matmul_call.args[0].struct_info.dtype
+ rhs_dtype = matmul_call.args[1].struct_info.dtype
+ if lhs_dtype != "float16" or rhs_dtype != "float16":
Review Comment:
I updated the dtype check to match what we have in Relay CUTLASS BYOC.
```
(lhs_dtype == "float16" and rhs_dtype == "float16")
or (lhs_dtype == "float32" and rhs_dtype == "float32")
or (lhs_dtype in ("int8", "uint8") and rhs_dtype in ("int8", "uint8"))
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122385560
##########
python/tvm/contrib/cutlass/gen_tensor_op.py:
##########
@@ -548,17 +549,20 @@ def get_batch_stride(stride_annot, arg0_idx, arg1_idx, arg0_axis_idx, arg1_axis_
attrs["ElementInputB"] = DataTypeTag[dtype_map[annotations[f"arg{rhs_arg_idx}_dtype"]]]
attrs["ElementOutput"] = DataTypeTag[dtype_map[annotations["ret_dtype"]]]
- attrs["K"] = str(int(lhs_shape[batched_offset + 1]))
- attrs["M"] = get_dim(lhs_shape[batched_offset], lhs_arg, 0, batched_offset)
+ attrs["K"] = str(int(lhs_shape[lhs_batched_offset + 1]))
+ attrs["M"] = get_dim(lhs_shape[lhs_batched_offset], lhs_arg, 0, lhs_batched_offset)
if transposed:
- attrs["N"] = get_dim(rhs_shape[batched_offset], rhs_arg, 0, batched_offset)
+ attrs["N"] = get_dim(rhs_shape[rhs_batched_offset], rhs_arg, 0, rhs_batched_offset)
else:
- attrs["N"] = get_dim(rhs_shape[batched_offset + 1], rhs_arg, 1, batched_offset)
+ attrs["N"] = get_dim(rhs_shape[rhs_batched_offset + 1], rhs_arg, 1, rhs_batched_offset)
if batched:
headers.append("cutlass/gemm/device/gemm_batched.h")
- attrs["batch"] = get_dim(lhs_shape[0], lhs_arg, 0)
+ # TODO: Support dynamic shape
+ # With the support of more general broadcasting in batch matmul,
+ # it needs more sophisticated code to handle dynamic shape
+ attrs["batch"] = str(int(annotations["batch"]))
Review Comment:
What do you mean by `more general broadcasting`? Like the conceptual "batch" dimension spanning multiple axes?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122267174
##########
tests/python/relax/test_codegen_cutlass.py:
##########
@@ -255,13 +177,48 @@ def test_conv2d_offload():
tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+def test_kernel_sharing():
+ data_np = np.random.randn(16, 32, 32, 16).astype("float16")
+ weight1_np = np.random.randn(16, 3, 3, 16).astype("float16")
+ weight2_np = np.random.randn(16, 3, 3, 16).astype("float16")
+
+ out = get_result_with_relax_cutlass_offload(Conv2dx2, data_np, weight1_np, weight2_np)
+
+ relay_expr = get_relay_conv2d_relu_x2(data_np.shape, weight1_np.shape)
+ ref = get_relay_ref(relay_expr, data_np, weight1_np, weight2_np)
+
+ tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)
+
+
+def get_reference_matmul_result(x, y, bias=None, transposed_y=False, activation=None):
Review Comment:
Were you able to pass accuracy check with this function and float inputs?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] yelite commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "yelite (via GitHub)" <gi...@apache.org>.
yelite commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122616165
##########
tests/python/relax/test_codegen_cutlass.py:
##########
@@ -120,12 +119,27 @@ def test_conv2d_offload():
np.testing.assert_equal(out, ref)
+def test_kernel_sharing():
+ data_np = np.random.randn(16, 32, 32, 16).astype("float16")
+ weight1_np = np.random.randn(16, 3, 3, 16).astype("float16")
+ weight2_np = np.random.randn(16, 3, 3, 16).astype("float16")
+
+ out = get_result_with_relax_cutlass_offload(Conv2dx2, data_np, weight1_np, weight2_np)
+
+ relay_expr = get_relay_conv2d_relu_x2(data_np.shape, weight1_np.shape)
+ ref = get_relay_ref(relay_expr, data_np, weight1_np, weight2_np)
Review Comment:
Fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi merged pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi merged PR #14166:
URL: https://github.com/apache/tvm/pull/14166
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] yelite commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "yelite (via GitHub)" <gi...@apache.org>.
yelite commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122615865
##########
python/tvm/contrib/cutlass/gen_tensor_op.py:
##########
@@ -548,17 +549,20 @@ def get_batch_stride(stride_annot, arg0_idx, arg1_idx, arg0_axis_idx, arg1_axis_
attrs["ElementInputB"] = DataTypeTag[dtype_map[annotations[f"arg{rhs_arg_idx}_dtype"]]]
attrs["ElementOutput"] = DataTypeTag[dtype_map[annotations["ret_dtype"]]]
- attrs["K"] = str(int(lhs_shape[batched_offset + 1]))
- attrs["M"] = get_dim(lhs_shape[batched_offset], lhs_arg, 0, batched_offset)
+ attrs["K"] = str(int(lhs_shape[lhs_batched_offset + 1]))
+ attrs["M"] = get_dim(lhs_shape[lhs_batched_offset], lhs_arg, 0, lhs_batched_offset)
if transposed:
- attrs["N"] = get_dim(rhs_shape[batched_offset], rhs_arg, 0, batched_offset)
+ attrs["N"] = get_dim(rhs_shape[rhs_batched_offset], rhs_arg, 0, rhs_batched_offset)
else:
- attrs["N"] = get_dim(rhs_shape[batched_offset + 1], rhs_arg, 1, batched_offset)
+ attrs["N"] = get_dim(rhs_shape[rhs_batched_offset + 1], rhs_arg, 1, rhs_batched_offset)
if batched:
headers.append("cutlass/gemm/device/gemm_batched.h")
- attrs["batch"] = get_dim(lhs_shape[0], lhs_arg, 0)
+ # TODO: Support dynamic shape
+ # With the support of more general broadcasting in batch matmul,
+ # it needs more sophisticated code to handle dynamic shape
+ attrs["batch"] = str(int(annotations["batch"]))
Review Comment:
Yes, for example, matmul with shape (a, b, c, 5, 5) x (5, 5) has batch `a * b * c`. I reverted the change on this line. Right now the check function in the cutlass pattern table requires static shape anyway. I will create a followup PR for formal dynamic shape support in cutlass matmul.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] yelite commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "yelite (via GitHub)" <gi...@apache.org>.
yelite commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1123535730
##########
python/tvm/relax/backend/contrib/cutlass.py:
##########
@@ -17,84 +17,141 @@
"""Pattern table for CUTLASS backend"""
-from tvm.relax import transform
+from typing import Mapping, Optional, Tuple
+
+import tvm
+from tvm.contrib.cutlass.build import is_valid_for_cutlass_matmul
+from tvm.relax import Call, Expr, ShapeExpr, transform
+from tvm.relax.dpl import DFPattern
from ..pattern_registry import get_patterns_with_prefix, register_patterns
from ..patterns import make_fused_bias_activation_pattern, make_matmul_pattern
+
+def _get_static_shape(shape: ShapeExpr) -> Optional[Tuple[int]]:
+ result = []
+ for dim in shape.values:
+ if isinstance(dim, tvm.tir.expr.IntImm):
+ result.append(int(dim))
+ else:
+ return None
+ return result
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype):
+ """Check if dtypes in the given workload are supported by CUTLASS."""
+ return (
+ (lhs_dtype == "float16" and rhs_dtype == "float16")
+ or (lhs_dtype == "float32" and rhs_dtype == "float32")
+ or (lhs_dtype in ("int8", "uint8") and rhs_dtype in ("int8", "uint8"))
+ )
+
+
+def _check_matmul(
+ match_result: Mapping[DFPattern, Expr],
+ _: Expr,
+) -> bool:
+ matmul_call: Call = None
+ for _, expr in match_result.items():
+ if isinstance(expr, Call) and expr.op.name == "relax.matmul":
Review Comment:
I will fix this in my follow up PR about dynamic shape support.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122386539
##########
tests/python/relax/test_codegen_cutlass.py:
##########
@@ -120,12 +119,27 @@ def test_conv2d_offload():
np.testing.assert_equal(out, ref)
+def test_kernel_sharing():
+ data_np = np.random.randn(16, 32, 32, 16).astype("float16")
+ weight1_np = np.random.randn(16, 3, 3, 16).astype("float16")
+ weight2_np = np.random.randn(16, 3, 3, 16).astype("float16")
+
+ out = get_result_with_relax_cutlass_offload(Conv2dx2, data_np, weight1_np, weight2_np)
+
+ relay_expr = get_relay_conv2d_relu_x2(data_np.shape, weight1_np.shape)
+ ref = get_relay_ref(relay_expr, data_np, weight1_np, weight2_np)
Review Comment:
This diff seems like a merge bug
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "masahi (via GitHub)" <gi...@apache.org>.
masahi commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122386539
##########
tests/python/relax/test_codegen_cutlass.py:
##########
@@ -120,12 +119,27 @@ def test_conv2d_offload():
np.testing.assert_equal(out, ref)
+def test_kernel_sharing():
+ data_np = np.random.randn(16, 32, 32, 16).astype("float16")
+ weight1_np = np.random.randn(16, 3, 3, 16).astype("float16")
+ weight2_np = np.random.randn(16, 3, 3, 16).astype("float16")
+
+ out = get_result_with_relax_cutlass_offload(Conv2dx2, data_np, weight1_np, weight2_np)
+
+ relay_expr = get_relay_conv2d_relu_x2(data_np.shape, weight1_np.shape)
+ ref = get_relay_ref(relay_expr, data_np, weight1_np, weight2_np)
Review Comment:
This diff seems like a merge bug. There should be no use of relay. (Also remove `import relay`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122276298
##########
python/tvm/relax/backend/contrib/cutlass.py:
##########
@@ -17,84 +17,132 @@
"""Pattern table for CUTLASS backend"""
-from tvm.relax import transform
+from typing import Mapping, Optional, Tuple
+
+import tvm
+from tvm.contrib.cutlass.build import is_valid_for_cutlass_matmul
+from tvm.relax import Call, Expr, ShapeExpr, transform
+from tvm.relax.dpl import DFPattern
from ..pattern_registry import get_patterns_with_prefix, register_patterns
from ..patterns import make_fused_bias_activation_pattern, make_matmul_pattern
+
+def _get_static_shape(shape: ShapeExpr) -> Optional[Tuple[int]]:
+ result = []
+ for dim in shape.values:
+ if isinstance(dim, tvm.tir.expr.IntImm):
+ result.append(int(dim))
+ else:
+ return None
+ return result
+
+
+def _check_matmul(
+ match_result: Mapping[DFPattern, Expr],
+ _: Expr,
+) -> bool:
+ matmul_call: Call = None
+ for _, expr in match_result.items():
+ if isinstance(expr, Call) and expr.op.name == "relax.matmul":
+ matmul_call = expr
+ if matmul_call is None:
+ raise ValueError("Cannot find call to matmul from match_result.")
+
+ lhs_shape = _get_static_shape(matmul_call.args[0].struct_info.shape)
+ rhs_shape = _get_static_shape(matmul_call.args[1].struct_info.shape)
+ if len(lhs_shape) < 2 or len(rhs_shape) < 2:
+ return False
+
+ lhs_dtype = matmul_call.args[0].struct_info.dtype
+ rhs_dtype = matmul_call.args[1].struct_info.dtype
+ if lhs_dtype != "float16" or rhs_dtype != "float16":
Review Comment:
fp32 is also supported (if not, let's revisit them later)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] vinx13 commented on a diff in pull request #14166: [Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC
Posted by "vinx13 (via GitHub)" <gi...@apache.org>.
vinx13 commented on code in PR #14166:
URL: https://github.com/apache/tvm/pull/14166#discussion_r1122437329
##########
python/tvm/relax/backend/contrib/cutlass.py:
##########
@@ -17,84 +17,141 @@
"""Pattern table for CUTLASS backend"""
-from tvm.relax import transform
+from typing import Mapping, Optional, Tuple
+
+import tvm
+from tvm.contrib.cutlass.build import is_valid_for_cutlass_matmul
+from tvm.relax import Call, Expr, ShapeExpr, transform
+from tvm.relax.dpl import DFPattern
from ..pattern_registry import get_patterns_with_prefix, register_patterns
from ..patterns import make_fused_bias_activation_pattern, make_matmul_pattern
+
+def _get_static_shape(shape: ShapeExpr) -> Optional[Tuple[int]]:
+ result = []
+ for dim in shape.values:
+ if isinstance(dim, tvm.tir.expr.IntImm):
+ result.append(int(dim))
+ else:
+ return None
+ return result
+
+
+def _is_supported_dtype(lhs_dtype, rhs_dtype):
+ """Check if dtypes in the given workload are supported by CUTLASS."""
+ return (
+ (lhs_dtype == "float16" and rhs_dtype == "float16")
+ or (lhs_dtype == "float32" and rhs_dtype == "float32")
+ or (lhs_dtype in ("int8", "uint8") and rhs_dtype in ("int8", "uint8"))
+ )
+
+
+def _check_matmul(
+ match_result: Mapping[DFPattern, Expr],
+ _: Expr,
+) -> bool:
+ matmul_call: Call = None
+ for _, expr in match_result.items():
+ if isinstance(expr, Call) and expr.op.name == "relax.matmul":
Review Comment:
```suggestion
if isinstance(expr, Call) and isinstance(expr.op, tvm.ir.Op) and expr.op.name == "relax.matmul":
```
it's possible it's calling a PrimFunc, it will be a global var
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org