This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 76e239e2e1 [Unity] Fix Cutlass Codegen for Dense (#16252)
76e239e2e1 is described below
commit 76e239e2e19b6579f19ca588a798526c3bfeaad8
Author: Siyuan Feng <[email protected]>
AuthorDate: Sun Dec 17 07:25:56 2023 +0800
[Unity] Fix Cutlass Codegen for Dense (#16252)
[Unity] Fix Cutlass Codegen for Dense gen_tensor_op
As dense op is not used in `relax`, the cutlass codegen is not tested.
However, the codegen is broken for dense op. This PR fixes the codegen.
---
python/tvm/contrib/cutlass/gen_tensor_op.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py
b/python/tvm/contrib/cutlass/gen_tensor_op.py
index e8f9bc9261..298d789572 100644
--- a/python/tvm/contrib/cutlass/gen_tensor_op.py
+++ b/python/tvm/contrib/cutlass/gen_tensor_op.py
@@ -562,7 +562,8 @@ def instantiate_template(func_name, annotations, func_args):
elif "dense" in func_name or "matmul" in func_name:
batched = "batch" in annotations
- transposed = "transposed" in func_name
+ # dense is equal to transposed_matmul
+ transposed = "transposed" in func_name or "dense" in func_name
lhs_arg_idx = _get_optional_int_annotation(annotations, "lhs_arg_idx",
0)
rhs_arg_idx = _get_optional_int_annotation(annotations, "rhs_arg_idx",
1)
bias_arg_idx = _get_optional_int_annotation(annotations,
"bias_arg_idx", None)
