leeexyz commented on a change in pull request #8802:
URL: https://github.com/apache/tvm/pull/8802#discussion_r693698609
##########
File path: src/relay/backend/te_compiler.h
##########
@@ -166,30 +166,73 @@ void UpdateFunctionMetadata(Function relay_func,
/*!
* \brief Obtain the Target from the device type.
* If homogenous compilation, this will return the only target.
- * If heteregenous compilation, this will select associated using the targets_
Map.
+ * If heterogeneous compilation, this will select the associated target using
the
+ * targets_ Map.
*
* \param dev_type
* \return Target
*/
Target GetTargetFromInteger(DLDeviceType dev_type, TargetMap targets);
+/*! \brief Utility to convert a LoweredModule to an IRModule.
+ *
+ * This function takes all the target specific modules in LoweredModule and
+ * annotates their functions with the correct target, and puts all those
functions
+ * in one IRModule.
+ * The purpose of this utility is to allow us to slowly remove LoweredModule
from the codebase.
+ *
+ * \param mod The LoweredModule to convert.
+ * \return The IRModule form of the input LoweredModule.
+ */
+IRModule LoweredModuleToIRModule(LoweredModule mod);
+
+/*! \brief Utility to convert an IRModule to a LoweredModule.
+ *
+ * This function takes all the functions in the IRModule and moves them into
target-specific
+ * IRModules stored inside a LoweredModule.
+ * The purpose of this utility is to allow us to slowly remove LoweredModule
from the codebase.
+ * \param mod The IRModule to convert.
+ * \return The LoweredModule form of the input IRModule.
+ */
+LoweredModule IRModuleToLoweredModule(IRModule mod);
+
/*! \brief Lower an IRModule's primitive functions to TIR.
*
* This is the "back half" of the Relay compiler which lowers "primitive
functions"
* to TE expressions, schedules them, and then to TIR.
*
- * \param compiler The TE-to-TIR compliler (which caches lowered functions)
* \param module The IRModule.
* \param targets The mapping for devices to targets.
* \param device_map An analysis result mapping each sub-expression to a
device.
+ * \param memory_plan The memory plan used during lowering
+ * \param module_name The name of this module
+ * \param process_fn Callback allowing one-level up code generators to process
+ * each function that we lower
* \return The lowered module, see above.
*/
-// TODO(@electriclilies): Not sure if this default initialization is correct...
LoweredModule LowerTE(
const IRModule& module, TargetMap targets, DeviceMap device_map,
backend::StaticMemoryPlan memory_plan, const String& module_name,
ProcessFn process_fn = [](Function f) {});
+/*! \brief Pass to lower an IRModule's primitive functions to TIR.
+ *
+ * This is the "back half" of the Relay compiler which lowers "primitive
functions"
+ * to TE expressions, schedules them, and then to TIR. This Pass calls
LowerTE, and
+ * uses LoweredModuleToIRModule utility to convert the output LowerTE's output
+ * LoweredModule into an IRModule before returning it.
+ *
+ * \param targets The mapping for devices to targets.
+ * \param device_map An analysis result mapping each sub-expression to a
device.
Review comment:
`device_map ` -> `device_context_map`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
