[BYOC] Enable mcBlas/mcDnn#14
Conversation
Change-Id: I1cae7e03deb667afb2564cd4c54c195c45f28038
Summary of ChangesHello @JoieAli, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly expands TVM's capabilities by integrating the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request enables mcBlas and mcDnn as BYOC backends for the MACA target. The changes are extensive, touching the CMake build system, TOPI operator library, Relay and Relax BYOC infrastructure, and runtime implementations. The overall approach is sound and follows existing patterns in TVM. However, I've found a few critical issues that need to be addressed, such as undefined variables in Python scripts which will lead to runtime errors. There are also some minor inconsistencies and what appears to be accidentally included files for a 'mctlass' backend. Please see my detailed comments for suggestions.
| return mcdnn.conv_forward( | ||
| data, | ||
| kernel, | ||
| [pad_d, pad_h, pad_w], | ||
| [stride_d, stride_h, stride_w], | ||
| [dilation_d, dilation_h, dilation_w], | ||
| conv_mode=1, | ||
| tensor_format=tensor_format, | ||
| algo=cfg["algo"].val, | ||
| conv_dtype=dtype, | ||
| ) |
There was a problem hiding this comment.
The variable dtype is used on line 122 but it is not defined within the conv3d_mcdnn function. This will cause a NameError. It seems you intended to use out_dtype for the conv_dtype parameter.
| return mcdnn.conv_forward( | |
| data, | |
| kernel, | |
| [pad_d, pad_h, pad_w], | |
| [stride_d, stride_h, stride_w], | |
| [dilation_d, dilation_h, dilation_w], | |
| conv_mode=1, | |
| tensor_format=tensor_format, | |
| algo=cfg["algo"].val, | |
| conv_dtype=dtype, | |
| ) | |
| return mcdnn.conv_forward( | |
| data, | |
| kernel, | |
| [pad_d, pad_h, pad_w], | |
| [stride_d, stride_h, stride_w], | |
| [dilation_d, dilation_h, dilation_w], | |
| conv_mode=1, | |
| tensor_format=tensor_format, | |
| algo=cfg["algo"].val, | |
| conv_dtype=out_dtype, | |
| ) |
| """Create an extern op that compute matrix mult of A and rhs with mcBLAS | ||
|
|
||
| Parameters | ||
| ---------- | ||
| lhs : Tensor | ||
| The left matrix operand | ||
| rhs : Tensor | ||
| The right matrix operand | ||
| transa : bool | ||
| Whether transpose lhs | ||
| transb : bool | ||
| Whether transpose rhs | ||
|
|
||
| Returns | ||
| ------- | ||
| C : Tensor | ||
| The result tensor. | ||
| """ | ||
| if n == 0: | ||
| n = lhs.shape[1] if transa else lhs.shape[0] | ||
| if m == 0: | ||
| m = rhs.shape[0] if transb else rhs.shape[1] | ||
| dtype = dtype if dtype is not None else lhs.dtype | ||
| return te.extern( | ||
| (n, m), | ||
| [lhs, rhs], | ||
| lambda ins, outs: tvm.tir.call_packed( | ||
| "tvm.contrib.mcblaslt.matmul", ins[0], ins[1], outs[0], transa, transb | ||
| ), | ||
| dtype=dtype, | ||
| name="C", |
There was a problem hiding this comment.
There are a couple of minor issues here for better clarity and maintainability:
- The docstring on line 23 says "mcBLAS" but this module is for "mcBLASlt". It should be updated to avoid confusion.
- The
namefor thete.externcall on line 53 is "C", which is very generic. It would be better to use a more descriptive name like "matmul_mcblaslt" to avoid potential name clashes and improve readability.
"""Create an extern op that compute matrix mult of A and rhs with mcBLASlt
Parameters
----------
lhs : Tensor
The left matrix operand
rhs : Tensor
The right matrix operand
transa : bool
Whether transpose lhs
transb : bool
Whether transpose rhs
Returns
-------
C : Tensor
The result tensor.
"""
if n == 0:
n = lhs.shape[1] if transa else lhs.shape[0]
if m == 0:
m = rhs.shape[0] if transb else rhs.shape[1]
dtype = dtype if dtype is not None else lhs.dtype
return te.extern(
(n, m),
[lhs, rhs],
lambda ins, outs: tvm.tir.call_packed(
"tvm.contrib.mcblaslt.matmul", ins[0], ins[1], outs[0], transa, transb
),
dtype=dtype,
name="matmul_mcblaslt",
)| /* | ||
| * 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. | ||
| */ | ||
|
|
||
| /*! | ||
| * \file src/relay/backend/contrib/mctlass/codegen.h | ||
| * \brief The 'custom' compilation pass for MCTLASS (invoked by the RelayToTIRTargetHook pass). | ||
| */ | ||
|
|
||
| #ifndef TVM_RELAY_BACKEND_CONTRIB_MCTLASS_CODEGEN_H_ | ||
| #define TVM_RELAY_BACKEND_CONTRIB_MCTLASS_CODEGEN_H_ | ||
|
|
||
| #include <tvm/ir/transform.h> | ||
|
|
||
| #include <string> | ||
| #include <vector> | ||
|
|
||
| #include "../codegen_c/codegen_c.h" | ||
|
|
||
| namespace tvm { | ||
| namespace relay { | ||
| namespace contrib { | ||
| namespace mctlass { | ||
|
|
||
| /*! | ||
| * \brief Returns the pass which replaces all calls to "Primitive" functions with "Compiler" | ||
| * attribute of "mctlass" with an call to an extern, and binds a \p runtime::StaticLibrary | ||
| * to the IRModule's "external_mods" attribute containing compiled implementations of | ||
| * those functions using the MCTLASS C++ template library. | ||
| */ | ||
| transform::Pass CompileForCutlass(); | ||
|
|
||
| // The rest is sparsely documented since they are exposed only for code sharing between Relay | ||
| // and Relax backend implementations. | ||
|
|
||
| /*! \brief Emit the function signature for a kernel */ | ||
| std::string EmitSignature(const std::vector<relay::contrib::Output>& out, | ||
| const std::string& func_id, const std::vector<std::string>& arg_names); | ||
|
|
||
| /*! \brief Generate the body of the kernel */ | ||
| GenerateBodyOutput GenerateBody(const std::string& func_name, const std::string& ext_func_id, | ||
| const std::vector<std::string>& output_types, | ||
| const Array<String>& func_args, const Map<String, ObjectRef>& attrs, | ||
| int* buf_idx); | ||
|
|
||
| /*! \brief Create a C-source module from the given kernel string */ | ||
| runtime::Module Finalize(const std::string& code, const Array<String>& func_names); | ||
|
|
||
| } // namespace mctlass | ||
| } // namespace contrib | ||
| } // namespace relay | ||
| } // namespace tvm | ||
|
|
||
| #endif // TVM_RELAY_BACKEND_CONTRIB_MCTLASS_CODEGEN_H_ |
There was a problem hiding this comment.
This file, along with other files under src/relay/backend/contrib/mctlass/, seems to have been accidentally included in this pull request. The PR is about mcblas and mcdnn, but these files are for a mctlass backend which is not wired into the build system in this PR. Please consider removing them to keep the PR focused and avoid adding unused code.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Change-Id: I1cae7e03deb667afb2564cd4c54c195c45f28038