[ADRENO] Revive and consolicate Adreno features

cmake/modules/contrib/CLML.cmake

@@ -31,13 +31,15 @@ if(USE_CLML)
             string(REGEX MATCH "CL_QCOM_ML_OPS_H_MAJOR_VERSION ([0-9]*)" _ ${ver})
             set(CLML_VERSION_MAJOR ${CMAKE_MATCH_1})
         else()
-            set(CLML_VERSION_MAJOR "2")
+            set(CLML_VERSION_MAJOR "4")
         endif()
     else()
-        set(CLML_VERSION_MAJOR "2")
+        set(CLML_VERSION_MAJOR "4")
     endif()
     add_definitions(-DTVM_CLML_VERSION=${CLML_VERSION_MAJOR})
     message(STATUS "CLML SDK Version :" ${CLML_VERSION_MAJOR})
+else()
+    set(CLML_VERSION_MAJOR "4")
 endif()
 
 if(USE_CLML_GRAPH_EXECUTOR)

python/tvm/relax/backend/adreno/clml.py

@@ -18,10 +18,12 @@
 """Pattern table for CLML backend"""
 
 import tvm
-from tvm import IRModule, relax
+from tvm import IRModule, relax, tir
 from tvm.ir.transform import PassContext, module_pass
 from tvm.relax import transform
 from tvm.relax.dpl.pattern import (
+    GlobalVarPattern,
+    TuplePattern,
     is_const,
     is_op,
     is_tuple_get_item,
@@ -611,3 +613,102 @@ def transform_module(self, mod: IRModule, ctx: PassContext) -> IRModule:
         )
         mod = seq(mod)
         return mod
+
+
+def _check_dequantize_matmul(ctx: relax.transform.PatternCheckContext) -> bool:
+    _input = ctx.annotated_expr["lhs"]
+    root = ctx.annotated_expr["root"]
+    wdq = ctx.annotated_expr["w_decoded"]
+    w_pack = ctx.annotated_expr["w_encoded"]
+
+    if ctx.annotated_expr["lhs"].struct_info.dtype != "float16":
+        return False
+    if not isinstance(wdq, relax.Call):
+        return False
+    g_var = wdq.args[0]
+    if not (isinstance(g_var, relax.GlobalVar) and "dequantize" in g_var.name_hint):
+        return False
+
+    if not (
+        (len(root.struct_info.shape) == 3)
+        and isinstance(root.struct_info.shape[0], tir.IntImm)
+        and (root.struct_info.dtype == "float16")
+        and (root.struct_info.shape[0] == 1)
+    ):
+        return False
+
+    if not (
+        (len(wdq.struct_info.shape) == 2)
+        and (w_pack.struct_info.shape[-1] == root.struct_info.shape[-1])
+        and (wdq.struct_info.shape[-2] == _input.struct_info.shape[-1])
+    ):
+        return False
+
+    return True
+
+
+def dequantize_matmul_patterns():
+    """Returns a list of supported decode -> matmul patterns."""
+
+    def _dequantize_matmul_pattern(name):
+        scales = wildcard()
+        x = wildcard()
+        w_packed = wildcard()
+
+        w_decoded = is_op("relax.call_tir")(
+            GlobalVarPattern(),
+            TuplePattern([w_packed, scales]),
+        )
+        matmul = is_op("relax.matmul")(x, w_decoded)
+
+        annotations = {
+            "root": matmul,
+            "lhs": x,
+            "w_encoded": w_packed,
+            "w_decoded": w_decoded,
+            "scales": scales,
+        }
+
+        return name, matmul, annotations, _check_dequantize_matmul
+
+    return [
+        _dequantize_matmul_pattern("openclml.dequant_matmul"),
+    ]
+
+
+clml_llm_patterns = [
+    *dequantize_matmul_patterns(),
+]
+register_patterns(clml_llm_patterns)
+
+
+@tvm.transform.module_pass(opt_level=0, name="OpenCLMLOffLoadForLLM")
+class OpenCLMLOffLoadForLLM:
+    """A compiler pass that partition the graph with dequant Matmul to CLML backend offload."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        """Initializer.
+        Parameters
+        ----------
+        target : tvm.target.Target
+            Target device.
+        """
+        self.target = target
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """Apply required passed to transform"""
+
+        if "adreno" in self.target.keys and (clml_sdk_version() >= 5):
+            mod = tvm.transform.Sequential(
+                [
+                    transform.Normalize(),
+                    transform.FuseOpsByPattern(clml_llm_patterns, annotate_codegen=True),
+                    transform.RunCodegen(),
+                ]
+            )(mod)
+
+        return mod

python/tvm/relax/backend/adreno/pipeline.py

@@ -24,7 +24,10 @@
 def library_dispatch_passes(target: tvm.target.Target):  # pylint: disable=unused-argument
     """The default library dispatch passes for Adreno GPU backend."""
     if "clml" in target.keys:
-        return [relax.backend.adreno.clml.OpenCLMLOffLoad()]
+        return [
+            relax.backend.adreno.clml.OpenCLMLOffLoadForLLM(target),
+            relax.backend.adreno.clml.OpenCLMLOffLoad(),
+        ]
     else:
         return []
 

src/relax/backend/contrib/clml/codegen.cc

@@ -168,9 +168,6 @@ class OpenCLMLJSONSerializer : public JSONSerializer {
       // Transfer attributes from the collector's node to the final node.
       node->CaptureAttrs(*collector.node_);
 
-      // Capture global settings on the JSON node.
-      SaveGlobalAttributes(node);
-
       VLOG(1) << name << " has " << node->GetInputs().size() << " inputs";
     }
 

src/relax/backend/contrib/codegen_json/codegen_json.h

@@ -232,7 +232,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator<NodeEntries> {
    * \param constant_names The names of all constants in the original module.
    */
   explicit JSONSerializer(const ffi::Map<Constant, ffi::String>& constant_names)
-      : constant_names_(constant_names) {}
+      : constant_names_(std::move(constant_names)) {}
 
   void serialize(Function func) {
     // First we convert all the parameters into input nodes.
@@ -493,7 +493,7 @@ class JSONSerializer : public relax::MemoizedExprTranslator<NodeEntries> {
   /*! \brief The list of required constants, ordered. */
   ffi::Array<ffi::String> constants_used_;
   /*! \brief The names of all constants in the original module. */
-  const ffi::Map<Constant, ffi::String>& constant_names_;
+  const ffi::Map<Constant, ffi::String> constant_names_;
 };
 
 }  // namespace contrib

src/runtime/contrib/clml/clml_memory_planner.cc

@@ -115,9 +115,9 @@ size_t PingPongAllocate(CachedLayer* layer, const std::map<size_t, size_t>& segm
    * for most of the graphs.
    *
    */
-  ssize_t free_start;
-  ssize_t free_size;
-  ssize_t last_found_size = CLMLWorkspace::Global()->onchip_mem_size + 1;
+  size_t free_start;
+  size_t free_size;
+  size_t last_found_size = CLMLWorkspace::Global()->onchip_mem_size + 1;
 
   for (auto it = segments.begin(); it != segments.end(); it++) {
     if (it->second < last_found_size) {