[mcBlas] Support mcblas contrib for MACA target

[Five-HZ]

• add USE_MCBLAS to enable compile with mcblas;
• add tvm.contrib.mcblas.matmul and tvm.contrib.mcblas.batch_matmul to call mcblas library;
• add tvm.conrtib.mcblaslt.matmul to call mcblaslt library;
• add mcblas-supproted subgraphs parition rules.

[gemini-code-assist]

Code Review

This pull request introduces support for the mcBLAS and mcBLASLt libraries on MACA devices in TVM, adding CMake configurations, Python interfaces, Relax pattern matching, and JSON runtime integration. The review feedback highlights several critical issues, including potential compiler and runtime crashes due to missing type/attribute checks, multi-GPU safety bugs where thread-local handles do not respect the active device, performance overhead from redundant handle creation, and incorrect test gating that skips FP8 tests on MACA.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The pattern checker accesses zero_point.data and scale.ty directly without verifying if they are actually Constants. If the model uses dynamic quantization where scale or zero_point are Vars, this will raise an AttributeError and crash the compiler. We should defensively check if they are instances of tvm.relax.Constant first.

Suggested change

	if "scale" in context.annotated_expr and "zp" in context.annotated_expr:
	scale = context.annotated_expr["scale"]
	zero_point = context.annotated_expr["zp"]
	# Only scalar values for scale and zero_point are supported.
	if scale.ty.ndim != 0 or zero_point.ty.ndim != 0:
	return False
	# Only zero_point == 0.0 is supported.
	if zero_point.data.numpy()[()].item() != 0.0:
	return False
	if "scale" in context.annotated_expr and "zp" in context.annotated_expr:
	scale = context.annotated_expr["scale"]
	zero_point = context.annotated_expr["zp"]
	if not isinstance(scale, tvm.relax.Constant) or not isinstance(zero_point, tvm.relax.Constant):
	return False
	if scale.ty.ndim != 0 or zero_point.ty.ndim != 0:
	return False
	if zero_point.data.numpy()[()].item() != 0.0:
	return False

[gemini-code-assist]

The flattening logic in CallMcblasLt restricts flattening to transa == false. However, if A is transposed (transa == true), the batch axes and spatial axis M are still contiguous in memory and can be safely flattened into Batch * M. Restricting this causes a runtime crash at TVM_FFI_ICHECK_EQ(batch_count_A, batch_count_B) because the pattern checker _check_matmul allows this configuration.

Suggested change

	if (A->ndim > 2 && B->ndim == 2 && transa == false) {
	if (A->ndim > 2 && B->ndim == 2) {

[gemini-code-assist]

Creating and destroying a mcblasLt handle via mcblasLtCreate and mcblasLtDestroy on every single invocation of tvm.contrib.mcblaslt.matmul introduces massive performance overhead. We should reuse the cached thread-local McBlasLtThreadEntry instead.

Suggested change

	mcblasLtHandle_t ltHandle;
	CHECK_MCBLAS_ERROR(mcblasLtCreate(&ltHandle));
	mcStream_t stream = static_cast<mcStream_t>(TVMFFIEnvGetStream(kDLMACA, A->device.device_id));
	CallLtIgemm(args, ret, ltHandle, stream);
	CHECK_MCBLAS_ERROR(mcblasLtDestroy(ltHandle));
	auto* entry_ptr = McBlasLtThreadEntry::ThreadLocal(A->device);
	TVM_FFI_ICHECK((A->dtype == DLDataType{kDLInt, 8, 1})) << "Expects dtype to be int8\n";
	mcStream_t stream = static_cast<mcStream_t>(TVMFFIEnvGetStream(kDLMACA, A->device.device_id));
	CallLtIgemm(args, ret, entry_ptr->handle, stream);

[gemini-code-assist]

If op_name contains "dequantize", the runtime unconditionally attempts to retrieve the "dq_scale" attribute. However, if the scale is not a constant, "dq_scale" is not serialized in codegen.cc, which will cause a crash. We should check if the attribute exists first.

Suggested change

	std::optional<float> dq_scale = std::nullopt;
	if (op_name.find("dequantize") != std::string::npos) {
	dq_scale = static_cast<float>(node.GetAttr<double>("dq_scale"));
	}
	std::optional<float> dq_scale = std::nullopt;
	if (op_name.find("dequantize") != std::string::npos && node.HasAttr("dq_scale")) {
	dq_scale = static_cast<float>(node.GetAttr<double>("dq_scale"));
	}

[gemini-code-assist]

A mcblasHandle_t is associated with the active device at the time of creation. In multi-GPU environments, a single thread might switch devices, or the handle might be bound to the wrong device if the active device was not set correctly during construction. We should use a thread-local map of device IDs to instances. Please also add #include <unordered_map> and #include <memory> at the top of the file.

Suggested change

	McBlasThreadEntry* McBlasThreadEntry::ThreadLocal(DLDevice curr_device) {
	static thread_local McBlasThreadEntry inst;
	McBlasThreadEntry* retval = &inst;
	mcStream_t stream =
	static_cast<mcStream_t>(TVMFFIEnvGetStream(curr_device.device_type, curr_device.device_id));
	CHECK_MCBLAS_ERROR(mcblasSetStream(retval->handle, stream));
	return retval;
	}
	McBlasThreadEntry* McBlasThreadEntry::ThreadLocal(DLDevice curr_device) {
	thread_local std::unordered_map<int, std::unique_ptr<McBlasThreadEntry>> instances;
	auto& inst = instances[curr_device.device_id];
	if (!inst) {
	int old_device;
	MACA_CALL(mcGetDevice(&old_device));
	MACA_CALL(mcSetDevice(curr_device.device_id));
	inst = std::make_unique<McBlasThreadEntry>();
	MACA_CALL(mcSetDevice(old_device));
	}
	mcStream_t stream =
	static_cast<mcStream_t>(TVMFFIEnvGetStream(curr_device.device_type, curr_device.device_id));
	CHECK_MCBLAS_ERROR(mcblasSetStream(inst->handle, stream));
	return inst.get();
	}

[gemini-code-assist]

Using a single static thread_local instance for McBlasLtThreadEntry completely ignores curr_device. In multi-GPU environments, this will cause the workspace pointer to be allocated on the wrong device, leading to crashes or silent data corruption. We should use a thread-local map of device IDs to instances. Please also add #include <unordered_map> and #include <memory> at the top of the file.

McBlasLtThreadEntry* McBlasLtThreadEntry::ThreadLocal(DLDevice curr_device) {
  thread_local std::unordered_map<int, std::unique_ptr<McBlasLtThreadEntry>> instances;
  auto& inst = instances[curr_device.device_id];
  if (!inst) {
    int old_device;
    MACA_CALL(mcGetDevice(&old_device));
    MACA_CALL(mcSetDevice(curr_device.device_id));
    inst = std::make_unique<McBlasLtThreadEntry>();
    MACA_CALL(mcSetDevice(old_device));
  }
  return inst.get();
}

[gemini-code-assist]

The test uses env.has_cuda_compute(9) to gate FP8 tests on MACA. Since MACA is not a CUDA device, has_cuda_compute will return False, causing these tests to be silently skipped on MACA even though MACA supports FP8. It should use env.has_maca() or env.has_maca_compute(...). The same applies to lines 350 and 377.

Suggested change

	@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
	@pytest.mark.skipif(not env.has_maca(), reason="need maca")

[gemini-code-assist]

If USE_MCBLAS is enabled but the mcBLAS library is not found in the SDK, MACA_MCBLAS_LIBRARY will be set to MACA_MCBLAS_LIBRARY-NOTFOUND. Linking against this directly will cause a confusing CMake or linker error. It is safer to check if the library was successfully found first.

  if(MACA_MCBLAS_LIBRARY STREQUAL "MACA_MCBLAS_LIBRARY-NOTFOUND")
    message(FATAL_ERROR "mcBLAS library not found, but USE_MCBLAS is enabled.")
  endif()
  target_link_libraries(tvm_runtime_extra PRIVATE tvm_mcblas_objs ${MACA_MCBLAS_LIBRARY})

[gemini-code-assist]

Using assert in the pattern check function _check_matmul will crash the entire compilation pipeline with an AssertionError if the matched binding is unexpected. Instead of asserting, the function should gracefully return False to indicate that the pattern is not supported for offloading.

Suggested change

	if matmul_rhs_var in context.matched_bindings:
	matmul_rhs_call = context.matched_bindings[matmul_rhs_var]
	assert (
	isinstance(matmul_rhs_call, tvm.relax.Call)
	and matmul_rhs_call.op.name == "relax.permute_dims"
	)
	rhs_transposed = True
	if matmul_rhs_var in context.matched_bindings:
	matmul_rhs_call = context.matched_bindings[matmul_rhs_var]
	if not (
	isinstance(matmul_rhs_call, tvm.relax.Call)
	and matmul_rhs_call.op.name == "relax.permute_dims"
	):
	return False
	rhs_transposed = True

[JoieAli]

lgtm