Skip to content

[MetaXGPU][FFI] Using tvm-ffi interface for MACA target#15

Merged
Five-HZ merged 1 commit into
MetaX-MACA:mainfrom
Five-HZ:verify
Feb 6, 2026
Merged

[MetaXGPU][FFI] Using tvm-ffi interface for MACA target#15
Five-HZ merged 1 commit into
MetaX-MACA:mainfrom
Five-HZ:verify

Conversation

@Five-HZ

@Five-HZ Five-HZ commented Feb 6, 2026

Copy link
Copy Markdown
Member

Fix some APIs based the lastest code.

@gemini-code-assist

Copy link
Copy Markdown

Summary of Changes

Hello @Five-HZ, 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 modernizes the MetaXGPU (MACA) target's integration with TVM by migrating its Foreign Function Interface (FFI) and module system to the latest tvm-ffi standards. The changes streamline how the MACA backend interacts with TVM's core, improving consistency and leveraging updated FFI capabilities for better device and module management.

Highlights

  • FFI Interface Migration: The MetaXGPU (MACA) backend has been updated to use the new tvm-ffi interface, replacing older tvm.ffi calls with tvm_ffi and tvm.ffi.reflection for global function registrations and module handling.
  • Python API Updates: Python files now import tvm_ffi and tvm directly, and functions previously registered with tvm.ffi.register_func are now registered using tvm_ffi.register_global_func.
  • C++ FFI and Module System Refactor: Significant changes were made in C++ files to align with the new FFI. This includes updating includes, changing base classes for module nodes (MACAModuleNode now inherits ffi::ModuleObj), and modifying method signatures for module serialization (SaveToBytes, WriteToFile) and function retrieval (GetFunction). Global FFI registrations now use TVM_FFI_STATIC_INIT_BLOCK.
  • MACA Device API Enhancements: The MACADeviceAPI now includes explicit stream management functions (CreateStream, FreeStream, SyncStreamFromTo) and refactored attribute retrieval logic. The MACATimerNode has also been updated to use the new stream handling mechanisms.
  • Type System Consistency: Various C++ code sections have been updated to use ffi::String and ffi::Array types for better consistency with the new FFI framework.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • python/tvm/contrib/mxcc.py
    • Updated imports from tvm.ffi, tvm.runtime, tvm.target to tvm_ffi, tvm, tvm.target.Target.
    • Added warnings and typing.Tuple imports.
    • Changed tvm.get_global_func to tvm_ffi.get_global_func.
    • Replaced tvm.ffi.register_func with tvm_ffi.register_global_func for several functions.
  • src/meta_schedule/schedule_rule/schedule_rule.cc
    • Updated String to ffi::String and Array<String> to ffi::Array<ffi::String> for parameters in ScheduleRule::DefaultMACA.
  • src/runtime/maca/maca_device_api.cc
    • Added new FFI-related includes (mcc/mcc_global.h, tvm/ffi/extra/c_env_api.h, tvm/ffi/reflection/registry.h, tvm/ffi/base_details.h).
    • Refactored GetAttr logic for architecture attributes and added kImagePitchAlignment.
    • Introduced CreateStream, FreeStream, SyncStreamFromTo methods for stream management.
    • Simplified GPUCopy to always use asynchronous memory copy.
    • Migrated FFI global function registrations from TVM_FFI_REGISTER_GLOBAL to TVM_FFI_STATIC_INIT_BLOCK.
    • Updated MACATimerNode to use TVMFFIEnvGetStream and the new object info macro.
  • src/runtime/maca/maca_module.cc
    • Added new FFI-related includes (dmlc/memory_io.h, tvm/ffi/extra/c_env_api.h, tvm/ffi/reflection/registry.h).
    • Changed MACAModuleNode to inherit from ffi::ModuleObj.
    • Updated module property methods (type_key to kind, GetPropertyMask to use ffi::Module constants).
    • Modified module serialization and source inspection methods (SaveToFile to WriteToFile, SaveToBinary to SaveToBytes, GetSource to InspectSource).
    • Updated GetFunction signature to return ffi::Optional<ffi::Function>.
    • Migrated module creation and loading functions to return/accept ffi::Module and ffi::Bytes.
    • Replaced multiple TVM_FFI_REGISTER_GLOBAL calls with a single TVM_FFI_STATIC_INIT_BLOCK for module loading.
  • src/runtime/maca/maca_module.h
    • Updated MACAModuleCreate function signature to return ffi::Module.
  • src/target/opt/build_maca_on.cc
    • Added include for tvm/ffi/reflection/registry.h.
    • Updated BuildMACA function signature to return ffi::Module and use ffi::Map.
    • Migrated FFI global function registration for target.build.maca to TVM_FFI_STATIC_INIT_BLOCK.
    • Updated TVM_REGISTER_PASS_CONFIG_OPTION to use ffi::String.
  • src/target/source/codegen_maca.cc
    • Updated String to ffi::String for scope and name variables in DeclBufferNode handling.
    • Updated String to ffi::String and Array<PrimExpr> to ffi::Array<PrimExpr> in PrintCallExtern.
  • src/target/source/codegen_maca.h
    • Updated String to ffi::String and Array<PrimExpr> to ffi::Array<PrimExpr> in PrintCallExtern signature.
    • Updated std::unordered_map<String, uint32_t> to std::unordered_map<ffi::String, uint32_t> for shared buffer alignments.
  • src/target/source/intrin_rule_maca.cc
    • Updated Array<PrimExpr> to ffi::Array<PrimExpr> for maca_args in DispatchMACAShuffle.
  • src/target/tag.cc
    • Updated String to ffi::String and Array<String> to ffi::Array<ffi::String> in TVM_REGISTER_MACA_TAG macro.
  • src/target/target_kind.cc
    • Updated String to ffi::String for mcpu in UpdateMACAAttrs.
    • Updated String to ffi::String and Array<String> to ffi::Array<ffi::String> in TVM_REGISTER_TARGET_KIND("maca") definition.
Activity
  • The pull request author, Five-HZ, has submitted changes to fix APIs based on the latest code, specifically migrating to the new tvm-ffi interface for the MetaXGPU (MACA) target.
Using Gemini Code Assist

The 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 /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.

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

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@Five-HZ Five-HZ merged commit e4274b0 into MetaX-MACA:main Feb 6, 2026
0 of 5 checks passed

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

The pull request migrates the MACA target to the new tvm-ffi interface. The migration is generally well-executed, but there are a few regressions in the runtime logic, specifically regarding synchronous memory copies and source inspection. Additionally, some FFI return types could be more idiomatic by using std::nullopt instead of null objects.

Comment on lines 250 to 253
static void GPUCopy(const void* from, void* to, size_t size, mcMemcpyKind kind,
mcStream_t stream) {
if (stream != 0) {
MACA_CALL(mcMemcpyAsync(to, from, size, kind, stream));
} else {
MACA_CALL(mcMemcpy(to, from, size, kind));
}
MACA_CALL(mcMemcpyAsync(to, from, size, kind, stream));
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The removal of the synchronous mcMemcpy path is a regression. In TVM, a null stream handle indicates that the operation should be synchronous with respect to the host. Always using mcMemcpyAsync makes these operations asynchronous, which can lead to race conditions if the caller expects the copy to be completed before the function returns (e.g., when copying data back to the host for immediate use).

  static void GPUCopy(const void* from, void* to, size_t size, mcMemcpyKind kind,
                      mcStream_t stream) {
    if (stream != nullptr) {
      MACA_CALL(mcMemcpyAsync(to, from, size, kind, stream));
    } else {
      MACA_CALL(mcMemcpy(to, from, size, kind));
    }
  }

Comment on lines +99 to 107
ffi::String InspectSource(const ffi::String& format) const final {
if (format == fmt_) return data_;
if (maca_source_.length() != 0) {
return maca_source_;
} else {
if (fmt_ == "fatbin") return data_;
return "";
}
return "";
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The implementation of InspectSource incorrectly ignores the format argument when maca_source_ is present. It should only return the source code if the requested format is compatible (e.g., "c" or an empty string). Returning C source when a different format like "asm" is requested is misleading and violates the InspectSource contract. Additionally, returning binary data (data_) as source code when fmt_ == "fatbin" is generally incorrect for this API.

  ffi::String InspectSource(const ffi::String& format) const final {
    if (format == fmt_) return data_;
    if (format == "c" || format == "") {
      return maca_source_;
    }
    return "";
  }

ICHECK_EQ(sptr_to_self.get(), this);
ICHECK_NE(name, symbol::tvm_module_main) << "Device function do not have main";
auto it = fmap_.find(name);
if (it == fmap_.end()) return ffi::Function();

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

When a function is not found in the module's function map, GetFunction should return std::nullopt instead of a null ffi::Function object. This is more consistent with the ffi::Optional return type and the standard pattern used in other TVM modules migrated to the new FFI.

  if (it == fmap_.end()) return std::nullopt;

Comment on lines +278 to +283
virtual void Start() {
int device_id;
MACA_CALL(mcGetDevice(&device_id));
stream_ = TVMFFIEnvGetStream(kDLMACA, device_id);
MACA_CALL(mcEventRecord(start_, static_cast<mcStream_t>(stream_)));
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Using mcGetDevice to determine the device ID for fetching the stream is potentially incorrect. A Timer is created for a specific Device, and its operations should be tied to that device regardless of the current thread's active device. The MACATimerNode should store the Device it was created for (passed in the factory function) and use its device_id to fetch the correct stream.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant