[BYOC-OpenCLML] OpenCLML integration with TVM.

CMakeLists.txt

@@ -110,6 +110,8 @@ tvm_option(USE_TENSORRT_RUNTIME "Build with TensorRT runtime" OFF)
 tvm_option(USE_RUST_EXT "Build with Rust based compiler extensions, STATIC, DYNAMIC, or OFF" OFF)
 tvm_option(USE_VITIS_AI "Build with VITIS-AI Codegen support" OFF)
 tvm_option(SUMMARIZE "Print CMake option summary after configuring" OFF)
+tvm_option(USE_CLML "Build with CLML Codegen support" OFF)
+tvm_option(USE_CLML_GRAPH_EXECUTOR "Build with CLML graph runtime" OFF)
 
 # include directories
 include_directories(${CMAKE_INCLUDE_PATH})
@@ -492,6 +494,7 @@ include(cmake/modules/contrib/ArmComputeLib.cmake)
 include(cmake/modules/contrib/TensorRT.cmake)
 include(cmake/modules/contrib/VitisAI.cmake)
 include(cmake/modules/contrib/Verilator.cmake)
+include(cmake/modules/contrib/CLML.cmake)
 include(cmake/modules/Git.cmake)
 include(cmake/modules/LibInfo.cmake)
 include(cmake/modules/RustExt.cmake)

cmake/config.cmake

@@ -262,6 +262,11 @@ set(USE_VITIS_AI OFF)
 # Build Verilator codegen and runtime
 set(USE_VERILATOR OFF)
 
+#Whether to use CLML codegen
+set(USE_CLML OFF)
+# USE_CLML_GRAPH_EXECUTOR - CLML SDK PATH or ON or OFF
+set(USE_CLML_GRAPH_EXECUTOR OFF)
+
 # Build ANTLR parser for Relay text format
 # Possible values:
 # - ON: enable ANTLR by searching default locations (cmake find_program for antlr4 and /usr/local for jar)

cmake/modules/LibInfo.cmake

@@ -113,6 +113,8 @@ function(add_lib_info src_file)
     TVM_INFO_USE_THRUST="${USE_THRUST}"
     TVM_INFO_USE_VITIS_AI="${USE_VITIS_AI}"
     TVM_INFO_USE_VULKAN="${USE_VULKAN}"
+    TVM_INFO_USE_CLML="${USE_CLML}"
+    TVM_INFO_USE_CLML_GRAPH_EXECUTOR="${USE_CLML_GRAPH_EXECUTOR}"
   )
 
 endfunction()

cmake/modules/contrib/CLML.cmake

@@ -0,0 +1,58 @@
+# 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.
+
+if(USE_CLML)
+    file(GLOB CLML_RELAY_CONTRIB_SRC src/relay/backend/contrib/clml/*.cc)
+    file(GLOB CLML_RUNTIME_MODULE src/runtime/contrib/clml/clml_runtime.cc)
+    list(APPEND COMPILER_SRCS ${CLML_RELAY_CONTRIB_SRC})
+    if(NOT USE_CLML_GRAPH_EXECUTOR)
+        list(APPEND COMPILER_SRCS ${CLML_RUNTIME_MODULE})
+    endif()
+    message(STATUS "Build with CLML support...")
+endif()
+
+if(USE_CLML_GRAPH_EXECUTOR)
+    set(CLML_PATH ${CMAKE_CURRENT_SOURCE_DIR}/clml)
+    # Detect custom CLML path.
+    if (NOT USE_CLML_GRAPH_EXECUTOR STREQUAL "ON")
+        set(CLML_PATH ${USE_CLML_GRAPH_EXECUTOR})
+    endif()
+
+    file(GLOB CLML_CONTRIB_SRC src/runtime/contrib/clml/*)
+
+    # Cmake needs to find clml library, include and support directories
+    # in the path specified by CLML_PATH.
+    set(CLML_INCLUDE_DIRS ${CLML_PATH}/include ${CLML_PATH})
+    include_directories(${CLML_INCLUDE_DIRS})
+    find_library(EXTERN_CLML_COMPUTE_LIB
+          NAMES OpenCL libOpenCL
+          HINTS "${CLML_PATH}" "${CLML_PATH}/lib64"
+          )
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${EXTERN_CLML_COMPUTE_LIB})
+    list(APPEND RUNTIME_SRCS ${CLML_CONTRIB_SRC})
+    message(STATUS "Build with CLML graph runtime support: "
+            ${EXTERN_CLML_COMPUTE_LIB})
+
+    # Set flag to detect CLML graph runtime support.
+    add_definitions(-DTVM_GRAPH_EXECUTOR_CLML)
+
+    message(STATUS "Enable OpenCL as fallback to CLML")
+    file(GLOB RUNTIME_OPENCL_SRCS src/runtime/opencl/*.cc)
+    list(APPEND RUNTIME_SRCS ${RUNTIME_OPENCL_SRCS})
+    set(USE_OPENCL ON)
+
+endif()

python/tvm/relay/op/contrib/__init__.py

@@ -26,3 +26,4 @@
 from .libtorch import *
 from .tensorrt import *
 from .cutlass import *
+from .clml import *

python/tvm/relay/op/contrib/clml.py

@@ -0,0 +1,247 @@
+# 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.
+# pylint: disable=invalid-name, unused-argument
+"""CLML Library supported operators."""
+import tvm
+
+from tvm import relay
+from tvm._ffi import register_func
+from tvm.relay import transform
+from tvm.relay.build_module import bind_params_by_name
+
+from ...dataflow_pattern import wildcard, is_op, is_constant, is_tuple_get_item
+from .register import register_pattern_table
+from ..strategy.generic import is_depthwise_conv2d
+
+
+def is_clml_runtime_enabled():
+    """Check if the CLML graph runtime is present.
+
+    Returns
+    -------
+    ret: bool
+        True if present, False if not.
+    """
+    check_enabled = tvm.get_global_func("relay.op.is_clml_runtime_enabled", True)
+    if check_enabled:
+        return check_enabled()
+    return False
+
+
+def partition_for_clml(mod, params=None):
+    """Partition the graph greedily offloading supported
+    operators to CLML Library.
+
+    Parameters
+    ----------
+    mod : Module
+        The module to run passes on.
+    params : Optional[Dict[str, NDArray]]
+        Constant input parameters.
+
+    Returns
+    -------
+    ret : annotated and partitioned module.
+    """
+
+    if params:
+        mod["main"] = bind_params_by_name(mod["main"], params)
+
+    seq = tvm.transform.Sequential(
+        [
+            transform.InferType(),
+            transform.FoldConstant(),
+            transform.MergeComposite(clml_pattern_table()),
+            transform.AnnotateTarget("clml", False),
+            transform.MergeCompilerRegions(),
+            transform.PartitionGraph(),
+        ]
+    )
+
+    result_mod = seq(mod)
+    return result_mod
+
+
+@register_func("relay.ext.clml.optimize")
+def preprocess_module(mod):
+    """
+    Pre-process a module containing functions ready for CLML codegen. For now we enforce OIHW
+    kernel layout and fold the transforms away.
+
+    Parameters
+    ----------
+    mod : Module
+        The module to run passes on.
+
+    Returns
+    -------
+    preprocessed_mod : The processed module.
+    """
+
+    def convert_layout_conv2d(conv2d_function):
+        def convert_conv(attrs, inputs, tinfos, desired_layouts):
+            new_attrs = dict(attrs)
+            data_info = tinfos[0]
+            weight_info = tinfos[1]
+            desired_data_layout, desired_kernel_layout = map(str, desired_layouts)
+            new_attrs["data_layout"] = desired_data_layout
+            new_attrs["kernel_layout"] = desired_kernel_layout
+
+            if is_depthwise_conv2d(
+                data_info.shape,
+                attrs["data_layout"],
+                weight_info.shape,
+                attrs["kernel_layout"],
+                attrs["groups"],
+            ):
+                dkl = desired_kernel_layout
+                new_attrs["kernel_layout"] = dkl[1] + dkl[0] + dkl[2] + dkl[3]
+            return conv2d_function(*inputs, **new_attrs)
+
+        return convert_conv
+
+    with OpAttrContext(
+        "nn.conv2d", "FTVMConvertOpLayout", convert_layout_conv2d(tvm.relay.nn.conv2d)
+    ):
+        seq = tvm.transform.Sequential(
+            [
+                transform.ConvertLayout({"nn.conv2d": ["NCHW", "OIHW"]}),
+                transform.FoldConstant(),
+            ]
+        )
+        preprocessed_mod = seq(mod)
+    return preprocessed_mod
+
+
+@register_pattern_table("clml")
+def clml_pattern_table():
+    """Get the CLML pattern table."""
+
+    def conv_pattern():
+        """Create a convolution pattern."""
+        pattern = is_op("nn.conv2d")(wildcard(), is_constant())
+        pattern = pattern.optional(lambda x: is_op("nn.bias_add")(x, is_constant()))
+        pattern = pattern.optional(
+            lambda x: is_op("nn.batch_norm")(
+                x, is_constant(), is_constant(), is_constant(), is_constant()
+            )
+        )
+        pattern = pattern.optional(is_tuple_get_item)
+        pattern = pattern.optional(is_op("nn.relu"))
+        return pattern
+
+    def batch_norm_pattern():
+        """Create a batch norm pattern."""
+        pattern = is_op("nn.batch_norm")(
+            wildcard(), is_constant(), is_constant(), is_constant(), is_constant()
+        )
+        pattern = is_tuple_get_item(pattern)
+        return pattern
+
+    def dense_pattern():
+        """Create a dense pattern."""
+        pattern = is_op("nn.dense")(wildcard(), is_constant())
+        pattern = pattern.optional(lambda x: is_op("add")(x, is_constant()))
+        return pattern
+
+    def pad_pattern():
+        """Create a pad pattern."""
+        pattern = is_op("nn.pad")(wildcard(), wildcard())
+        return pattern
+
+    def check_conv(extract):
+        """Check conv pattern is supported by CLML."""
+        call = extract
+        if isinstance(call, tvm.relay.expr.TupleGetItem):
+            call = call.tuple_value
+        elif call.op.name == "nn.relu":
+            call = call.args[0]
+            if isinstance(call, tvm.relay.expr.TupleGetItem):
+                call = call.tuple_value
+        while call.op.name != "nn.conv2d":
+            call = call.args[0]
+        attrs, args = call.attrs, call.args
+        if attrs.data_layout != "NCHW":
+            return False
+        data_typ = args[0].checked_type
+        kernel_typ = args[1].checked_type
+        is_depthwise = is_depthwise_conv2d(
+            data_typ.shape,
+            attrs["data_layout"],
+            kernel_typ.shape,
+            attrs["kernel_layout"],
+            attrs["groups"],
+        )
+        if attrs.groups != 1 and not is_depthwise:
+            return False
+        return True
+
+    return [
+        ("clml.conv2d", conv_pattern(), check_conv),
+        ("clml.dense", dense_pattern()),
+        ("clml.pad", pad_pattern()),
+        ("clml.batch_norm", batch_norm_pattern()),
+    ]
+
+
+def _register_external_op_helper(op_name, supported=True):
+    @tvm.ir.register_op_attr(op_name, "target.clml")
+    def _func_wrapper(expr):
+        return supported
+
+    return _func_wrapper
+
+
+_register_external_op_helper("clip")
+_register_external_op_helper("relu")
+_register_external_op_helper("nn.global_avg_pool2d")
+_register_external_op_helper("nn.global_max_pool2d")
+_register_external_op_helper("nn.softmax")
+_register_external_op_helper("reshape")
+
+
+class OpAttrContext(object):
+    """Temporarily changes the attr of an op."""
+
+    def __init__(self, op_name, attr_key, attr_value):
+        """Saves the required info for RAII pattern usage.
+
+        Parameters
+        ----------
+        op_name : str
+            The op name.
+
+        attr_key : str
+            The attribute name.
+
+        attr_value : object
+            The attribute value.
+        """
+        self.op = relay.op.get(op_name)
+        self.attr_key = attr_key
+        self.attr_value = attr_value
+
+    def __enter__(self):
+        self.older_attr = self.op.get_attr(self.attr_key)
+        self.op.reset_attr(self.attr_key)
+        self.op.set_attr(self.attr_key, self.attr_value)
+        return self
+
+    def __exit__(self, ptype, value, trace):
+        self.op.reset_attr(self.attr_key)
+        if self.older_attr:
+            self.op.set_attr(self.attr_key, self.older_attr)