chore: import upstream snapshot with attribution

python/tvm/s_tir/backend/__init__.py

@@ -0,0 +1,19 @@
+# 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.
+"""S-TIR backend compilation pipeline and other passes."""
+
+from . import adreno

python/tvm/s_tir/backend/adreno/__init__.py

@@ -0,0 +1,21 @@
+# isort: skip_file
+# 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.
+"""The S-TIR Adreno pipeline passes"""
+
+from . import pipeline
+from . import transform

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

@@ -0,0 +1,124 @@
+# 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
+"""The TIR backend compilation pipeline for Adreno"""
+
+import tvm
+from tvm import s_tir, tirx
+from tvm.tirx import compilation_pipeline as tir_pipeline
+
+
+def default_tir_pipeline():
+    """The default tirx pipeline used in tvm.tirx.build"""
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+        """The default lowering passes for TIR backend."""
+        pass_ctx = tvm.transform.PassContext.current()
+        config = pass_ctx.config
+        passes = [
+            s_tir.backend.adreno.transform.TextureFlatten(),
+            s_tir.transform.CanonicalizeLoop(),
+            s_tir.transform.LowerCrossThreadReduction(),
+            s_tir.transform.LowerInitBlock(),
+            s_tir.transform.PlanAndUpdateBufferAllocationLocation(),
+            s_tir.transform.ConvertBlocksToOpaque(),
+            s_tir.transform.LiftThreadBinding(),
+            s_tir.transform.ManifestSharedMemoryLocalStage(),
+            s_tir.transform.CompactBufferAllocation(),
+            s_tir.transform.LowerAutoCopy(),
+            s_tir.transform.UnifyThreadBinding(),
+            s_tir.transform.LowerMatchBuffer(),
+            tirx.transform.StmtSimplify(),
+            s_tir.transform.InjectPermutedLayout(),
+            s_tir.transform.AnnotateIrregularLoop(),
+            s_tir.transform.InjectSoftwarePipeline(),
+            s_tir.transform.TransformMmaBufferLayout(),
+            s_tir.transform.LowerOpaqueBlock(),
+            s_tir.backend.adreno.transform.InjectTextureAlloc(),
+            tirx.transform.FlattenBuffer(),
+            tirx.transform.BF16ComputeLegalize(),
+            tirx.transform.NarrowDataType(32),
+            s_tir.transform.LoopPartition(),
+            tirx.transform.VectorizeLoop(not bool(config.get("tirx.disable_vectorize", False))),
+            s_tir.transform.InjectVirtualThread(),
+            s_tir.transform.InjectDoubleBuffer(),
+        ]
+        if not bool(config.get("tirx.disable_storage_rewrite", False)):
+            passes.append(tirx.transform.StorageRewrite())
+        if config.get("tirx.use_async_copy", False):
+            passes.append(s_tir.transform.LowerAsyncDMA())
+        passes.extend(
+            [
+                s_tir.transform.HoistIfThenElse(),
+                tirx.transform.UnrollLoop(),
+                s_tir.transform.RenormalizeSplitPattern(),
+                tirx.transform.StmtSimplify(),
+                tirx.transform.RemoveNoOp(),
+                s_tir.transform.RewriteUnsafeSelect(),
+            ]
+        )
+        # Additional passes based on configuration.
+        if bool(config.get("tirx.instrument_bound_checkers", False)):
+            passes.append(s_tir.transform.InstrumentBoundCheckers())
+        if bool(config.get("tirx.ptx.ldg32", False)):
+            passes.append(s_tir.transform.InjectPTXLDG32(True))
+        if not bool(config.get("tirx.disable_cse_tir", False)):
+            passes.append(tirx.transform.CommonSubexprElim())
+        if bool(config.get("tirx.instrument_lwp", False)):
+            passes.append(s_tir.transform.InstrumentProfileIntrinsics())
+        passes.extend(
+            [
+                # Bind the target first so that target-specific attributes are available.
+                tirx.transform.FP8ComputeLegalize(),
+                # VerifyVTCMLimit must occur before LowerVtcmAlloc.
+                s_tir.transform.VerifyVTCMLimit(),
+                s_tir.transform.LowerVtcmAlloc(),
+                tirx.transform.VerifyMemory(),
+                tirx.transform.AnnotateEntryFunc(),
+            ]
+        )
+        passes.extend(
+            [
+                s_tir.transform.ThreadSync("shared"),
+                s_tir.transform.ThreadSync("shared.dyn"),
+                s_tir.transform.ThreadSync("warp"),
+                s_tir.transform.InferFragment(),
+                s_tir.transform.LowerThreadAllreduce(),
+            ]
+        )
+        if bool(config.get("tirx.use_async_copy", False)):
+            passes.append(s_tir.transform.InjectPTXAsyncCopy())
+        if bool(config.get("tirx.ptx.ldg32", False)):
+            passes.append(s_tir.transform.InjectPTXLDG32())
+        passes.extend(
+            [
+                s_tir.transform.MergeSharedMemoryAllocations(),
+                tirx.transform.SplitHostDevice(),
+                tirx.transform.MakePackedAPI(),
+                tirx.transform.FP8StorageLegalize(),
+                tirx.transform.BF16StorageLegalize(),
+            ]
+        )
+        mod = tvm.ir.transform.Sequential(passes)(mod)
+        return mod
+
+    return _pipeline
+
+
+tir_pipeline.PIPELINE_MAP["adreno"] = default_tir_pipeline

python/tvm/s_tir/backend/adreno/transform/__init__.py

@@ -0,0 +1,23 @@
+# isort: skip_file
+# 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.
+"""Adreno TIR transformations."""
+
+from .transform import (
+    InjectTextureAlloc,
+    TextureFlatten,
+)

python/tvm/s_tir/backend/adreno/transform/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI APIs for tvm.s_tir.backend.adreno.transform"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("s_tir.backend.adreno.transform", __name__)

python/tvm/s_tir/backend/adreno/transform/transform.py

@@ -0,0 +1,46 @@
+# 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.
+"""Wrapping existing transformations."""
+# pylint: disable=invalid-name, unsupported-binary-operation
+
+from . import _ffi_api
+
+
+def InjectTextureAlloc():
+    """Inject Texture Allocation Intrinsic to make sure appropriate lowering
+    via alloc_nd/alloc_free calls
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.InjectTextureAlloc()  # type: ignore
+
+
+def TextureFlatten():
+    """Flatten the multi-dimensional read/write to 2D.
+
+
+    Parameters
+    ----------
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.TextureFlatten()  # type: ignore