[BYOC] Enable mcBlas/mcDnn

cmake/config.cmake

@@ -494,5 +494,6 @@ SET(CMAKE_VS_PLATFORM_TOOLSET_HOST_ARCHITECTURE "x64")
 
 # Enable using flash-attn as a BYOC backend
 # Need to have USE_MACA=ON
-set(USE_FLASHATTN ON)
+set(USE_FLASHATTN OFF)
 set(USE_MCDNN ON)
+set(USE_MCBLAS ON)

cmake/modules/LibInfo.cmake

@@ -119,6 +119,8 @@ function(add_lib_info src_file)
     TVM_INFO_USE_HIPBLAS="${USE_HIPBLAS}"
     TVM_INFO_USE_ROCM="${USE_ROCM}"
     TVM_INFO_USE_MACA="${USE_MACA}"
+    TVM_INFO_USE_MCBLAS="${USE_MCBLAS}"
+    TVM_INFO_USE_MCDNN="${USE_MCDNN}"
     TVM_INFO_USE_RCCL="${USE_RCCL}"
     TVM_INFO_USE_RPC="${USE_RPC}"
     TVM_INFO_USE_RTTI="${USE_RTTI}"

cmake/modules/MACA.cmake

@@ -47,6 +47,29 @@ if(USE_MACA)
     list(APPEND RUNTIME_SRCS ${CONTRIB_FLASHATTN_SRCS})
     list(APPEND TVM_RUNTIME_LINKER_LIBS ${MACA_FLASHATTN_LIBRARY})
   endif(USE_FLASHATTN)
+
+  if(USE_MCBLAS)
+    message(STATUS "Build with mcBLAS support")
+    tvm_file_glob(GLOB MCBLAS_CONTRIB_SRC src/relay/backend/contrib/mcblas/*.cc src/relax/backend/contrib/mcblas/*.cc)
+    list(APPEND COMPILER_SRCS ${MCBLAS_CONTRIB_SRC})
+    tvm_file_glob(GLOB CONTRIB_MCBLAS_SRCS src/runtime/contrib/mcblas/*.cc)
+    list(APPEND RUNTIME_SRCS ${CONTRIB_MCBLAS_SRCS})
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${MACA_MCBLAS_LIBRARY})
+    if(NOT MACA_MCBLASLT_LIBRARY STREQUAL "MACA_MCBLASLT_LIBRARY-NOTFOUND")
+      list(APPEND TVM_RUNTIME_LINKER_LIBS ${MACA_MCBLASLT_LIBRARY})
+    endif()
+  endif(USE_MCBLAS)
+
+  if(USE_MCDNN)
+    message(STATUS "Build with mcdnn support")
+    include_directories(SYSTEM ${MACA_INCLUDE_DIRS}/mcdnn)
+    tvm_file_glob(GLOB MCDNN_RELAY_CONTRIB_SRC src/relay/backend/contrib/mcdnn/*.cc src/relax/backend/contrib/mcdnn/*.cc)
+    list(APPEND COMPILER_SRCS ${MCDNN_RELAY_CONTRIB_SRC})
+    tvm_file_glob(GLOB CONTRIB_MCDNN_SRCS src/runtime/contrib/mcdnn/*.cc)
+    list(APPEND RUNTIME_SRCS ${CONTRIB_MCDNN_SRCS})
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${MACA_MCDNN_LIBRARY})
+  endif(USE_MCDNN)
+
 else(USE_MACA)
   list(APPEND COMPILER_SRCS src/target/opt/build_maca_off.cc)
 endif(USE_MACA)

cmake/utils/FindMACA.cmake

@@ -48,7 +48,9 @@ macro(find_maca use_maca)
     find_library(MACA_MACAMCC_LIBRARY mcruntime ${__maca_sdk}/lib)
     find_library(MACA_HCA_LIBRARY mxc-runtime64 ${__maca_sdk}/lib)
     find_library(MACA_FLASHATTN_LIBRARY mcFlashAttn ${__maca_sdk}/lib)
-
+    find_library(MACA_MCBLAS_LIBRARY mcblas ${__maca_sdk}/lib)
+    find_library(MACA_MCBLASLT_LIBRARY mcblasLt ${__maca_sdk}/lib)
+    find_library(MACA_MCDNN_LIBRARY mcdnn ${__maca_sdk}/lib)
     if(MACA_MACAMCC_LIBRARY)
       set(MACA_FOUND TRUE)
     endif()
@@ -57,6 +59,8 @@ macro(find_maca use_maca)
     message(STATUS "Found MACA_INCLUDE_DIRS=" ${MACA_INCLUDE_DIRS})
     message(STATUS "Found MACA_MACAMCC_LIBRARY=" ${MACA_MACAMCC_LIBRARY})
     message(STATUS "Found MACA_FLASHATTN_LIBRARY=" ${MACA_FLASHATTN_LIBRARY})
+    message(STATUS "Found MACA_MCBLAS_LIBRARY=" ${MACA_MCBLAS_LIBRARY})
+    message(STATUS "Found MACA_MCDNN_LIBRARY=" ${MACA_MCDNN_LIBRARY})
   endif(MACA_FOUND)
 endmacro(find_maca)
 

include/tvm/topi/contrib/mcblas.h

@@ -0,0 +1,88 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \brief External function interface to mcBLAS libraries
+ * \file mcblas.h
+ */
+#ifndef TVM_TOPI_CONTRIB_MCBLAS_H_
+#define TVM_TOPI_CONTRIB_MCBLAS_H_
+
+#include <tvm/te/operation.h>
+#include <tvm/topi/detail/extern.h>
+
+namespace tvm {
+namespace topi {
+namespace contrib {
+
+using namespace tvm::te;
+using namespace topi::detail;
+/*!
+ * \brief Create an op that multiplies lhs and rhs with mcBLAS
+ *
+ * \param lhs The left matrix operand
+ * \param rhs The right matrix operand
+ * \param transa Whether to transpose lhs
+ * \param transb Whether to transpose rhs
+ *
+ * \return The output tensor
+ */
+inline Tensor mcblas_matmul(const Tensor& lhs, const Tensor& rhs, bool transa, bool transb) {
+  auto n = transa ? lhs->shape[1] : lhs->shape[0];
+  auto m = transb ? rhs->shape[0] : rhs->shape[1];
+
+  return make_extern(
+      {{n, m}}, {lhs->dtype}, {lhs, rhs},
+      [&](Array<Buffer> ins, Array<Buffer> outs) {
+        return call_packed({StringImm("tvm.contrib.mcblas.matmul"), pack_buffer(ins[0]),
+                            pack_buffer(ins[1]), pack_buffer(outs[0]), transa, transb});
+      },
+      "C", "", {})[0];
+}
+
+/*!
+ * \brief Create an op that multiplies batch matrices
+ *        lhs and rhs with mcBLAS
+ *
+ * \param lhs The left matrix operand
+ * \param rhs The right matrix operand
+ * \param transa Whether to transpose lhs
+ * \param transb Whether to transpose rhs
+ *
+ * \return The output tensor
+ */
+inline Tensor mcblas_batch_matmul(const Tensor& lhs, const Tensor& rhs, bool transa, bool transb) {
+  auto b = lhs->shape[0];
+  auto n = transa ? lhs->shape[2] : lhs->shape[1];
+  auto m = transb ? rhs->shape[1] : rhs->shape[2];
+
+  return make_extern(
+      {{b, n, m}}, {lhs->dtype}, {lhs, rhs},
+      [&](Array<Buffer> ins, Array<Buffer> outs) {
+        return call_packed({StringImm("tvm.contrib.mcblas.batch_matmul"), pack_buffer(ins[0]),
+                            pack_buffer(ins[1]), pack_buffer(outs[0]), transa, transb});
+      },
+      "C", "", {})[0];
+}
+
+}  // namespace contrib
+}  // namespace topi
+}  // namespace tvm
+
+#endif  // TVM_TOPI_CONTRIB_MCBLAS_H_

include/tvm/topi/maca/dense.h

@@ -0,0 +1,99 @@
+/*
+ * 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 maca/dense.h
+ * \brief MACA schedule for dense operation
+ */
+#ifndef TVM_TOPI_MACA_DENSE_H_
+#define TVM_TOPI_MACA_DENSE_H_
+
+#include <tvm/target/generic_func.h>
+#include <tvm/te/operation.h>
+#include <tvm/te/schedule_pass.h>
+#include <tvm/topi/contrib/mcblas.h>
+#include <tvm/topi/cuda/dense.h>
+#include <tvm/topi/detail/array_utils.h>
+#include <tvm/topi/generic/extern.h>
+#include <tvm/topi/nn/dense.h>
+#include <tvm/topi/tags.h>
+
+namespace tvm {
+namespace topi {
+
+using namespace tvm::te;
+
+namespace maca {
+/*!
+ * \brief Implementation of dense for MACA backend
+ *
+ * \param target The target device
+ * \param data Tensor with shape [batch, in_dim]
+ * \param weight Tensor with shape [out_dim, in_dim]
+ * \param bias Tensor with shape [out_dim]. Optional; to omit bias, pass Tensor()
+ * \param out_dtype Output data type. Used for mixed precision.
+ *
+ * \return Tensor with shape [batch, out_dim]
+ */
+inline tvm::te::Tensor dense_maca(const Target& target, const tvm::te::Tensor& data,
+                                  const tvm::te::Tensor& weight, const tvm::te::Tensor& bias,
+                                  const DataType& out_dtype) {
+  ICHECK_EQ(data->shape.size(), 2) << "dense requires 2-D data";
+  ICHECK_EQ(weight->shape.size(), 2) << "dense requires 2-D weight";
+  if (bias.defined()) {
+    ICHECK_EQ(bias->shape.size(), 1) << "dense requires 1-D bias";
+  }
+
+  auto batch = data->shape[0];
+  auto in_dim = data->shape[1];
+  auto out_dim = weight->shape[0];
+
+  if (target->GetLibs().count("mcblas")) {
+    ICHECK_EQ(data->dtype, out_dtype) << "Mixed precision not supported.";
+    auto mm = topi::contrib::mcblas_matmul(data, weight, false, true);
+    if (bias.defined()) {
+      mm = tvm::te::compute(
+          {batch, out_dim}, [&](Var i, Var j) { return mm(i, j) + bias(j); }, "tensor", kBroadcast);
+    }
+
+    return mm;
+  } else {
+    return topi::nn::dense(data, weight, bias, out_dtype);
+  }
+}
+
+/*!
+ * \brief Create a MACA schedule for dense
+ *
+ * \param target The target to generate a schedule for.
+ * \param outs The output tensors.
+ *
+ * \return A schedule for the given ops.
+ */
+inline Schedule schedule_dense(const Target& target, const Array<Tensor>& outs) {
+  if (target->kind->name == "maca" && target->GetLibs().count("mcblas")) {
+    return topi::generic::schedule_extern(target, outs);
+  }
+  return topi::cuda::schedule_dense(target, outs);
+}
+
+}  // namespace maca
+}  // namespace topi
+}  // namespace tvm
+#endif  // TVM_TOPI_MACA_DENSE_H_

python/tvm/contrib/mcblas.py

@@ -0,0 +1,86 @@
+# 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.
+"""External function interface to mcBLAS libraries."""
+import tvm
+from tvm import te
+
+
+def matmul(lhs, rhs, transa=False, transb=False, dtype=None):
+    """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.
+    """
+    n = lhs.shape[1] if transa else lhs.shape[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.mcblas.matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        dtype=dtype,
+        name="matmul_mcblas",
+    )
+
+
+def batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None):
+    """Create an extern op that compute batch 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.
+    """
+    b = lhs.shape[0]
+    n = lhs.shape[2] if transa else lhs.shape[1]
+    m = rhs.shape[1] if transb else rhs.shape[2]
+    dtype = dtype if dtype is not None else lhs.dtype
+    return te.extern(
+        (b, n, m),
+        [lhs, rhs],
+        lambda ins, outs: tvm.tir.call_packed(
+            "tvm.contrib.mcblas.batch_matmul", ins[0], ins[1], outs[0], transa, transb
+        ),
+        dtype=dtype,
+        name="batch_matmul_mcblas",
+    )

python/tvm/contrib/mcblaslt.py

@@ -0,0 +1,54 @@
+# 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.
+"""External function interface to mcBLASlt libraries."""
+import tvm
+from tvm import te
+
+
+def matmul(lhs, rhs, transa=False, transb=False, n=0, m=0, dtype=None):
+    """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",
+    )