[PASS] Refactor a couple of TIR passes - BindTarget, AnnotateEntryFunc, Filter, LowerInitBlock

include/tvm/tir/transform.h

@@ -24,6 +24,7 @@
 #ifndef TVM_TIR_TRANSFORM_H_
 #define TVM_TIR_TRANSFORM_H_
 
+#include <tvm/driver/driver_api.h>
 #include <tvm/ir/transform.h>
 #include <tvm/tir/expr.h>
 #include <tvm/tir/function.h>
@@ -624,6 +625,24 @@ TVM_DLL Pass ExtractPrimFuncConstants();
  */
 TVM_DLL Pass RenormalizeSplitPattern();
 
+/*!
+ * \brief Annotate a PrimFunc with a given target.
+ * \return The pass.
+ */
+TVM_DLL Pass BindTarget(Target target);
+
+/*!
+ * \brief Set a PrimFunc as the entry point if it is only function in IRModule.
+ * \return The pass.
+ */
+TVM_DLL Pass AnnotateEntryFunc();
+
+/*!
+ * \brief Filter PrimFuncs with a given condition.
+ * \return The pass.
+ */
+TVM_DLL Pass Filter(runtime::TypedPackedFunc<bool(PrimFunc)> fcond);
+
 }  // namespace transform
 }  // namespace tir
 }  // namespace tvm

python/tvm/tir/transform/transform.py

@@ -16,7 +16,8 @@
 # under the License.
 """Wrapping existing transformations."""
 # pylint: disable=invalid-name
-from typing import Optional
+from typing import Optional, Callable
+from tvm.target import Target
 from . import _ffi_api
 from . import function_pass as _fpass
 
@@ -43,26 +44,6 @@ def _transform(func, mod, ctx):
     return _fpass.prim_func_pass(_transform, opt_level=0, name="Apply")  # type: ignore
 
 
-def Filter(fcond):
-    """Filter functions by the calling convention attribute.
-
-    Parameters
-    ----------
-    fcond : tvm.tir.PrimFunc -> bool
-        The condition of the filtering.
-
-    Returns
-    -------
-    fpass : tvm.transform.Pass
-        The result pass
-    """
-    # pylint: disable=unused-argument
-    def _transform(func, mod, ctx):
-        return func if fcond(func) else None
-
-    return _fpass.prim_func_pass(_transform, opt_level=0, name="Filter")  # type: ignore
-
-
 def InjectPrefetch():
     """Inject prefetch instructions into stmt.
 
@@ -806,3 +787,37 @@ def RenormalizeSplitPattern():
         The result pass
     """
     return _ffi_api.RenormalizeSplitPattern()  # type: ignore
+
+
+def BindTarget(target: Target):
+    """Annotate a PrimFunc with a given target.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.BindTarget(target)  # type: ignore
+
+
+def AnnotateEntryFunc():
+    """Set a PrimFunc as the entry point if it is only function in IRModule.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.AnnotateEntryFunc()  # type: ignore
+
+
+def Filter(fcond: Callable):
+    """Filter out PrimFuncs that does not satisfy the given condition.
+    `fcond` should be a function that takes a primfunc and returns boolean.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.Filter(fcond)  # type: ignore

src/driver/driver_api.cc

@@ -163,32 +163,6 @@ TVM_REGISTER_GLOBAL("driver.get_binds")
       return out_arr;
     });
 
-transform::Pass BindTarget(Target target) {
-  auto fpass = [target](tir::PrimFunc f, IRModule m, transform::PassContext ctx) {
-    return WithAttr(std::move(f), tvm::attr::kTarget, target);
-  };
-  return tir::transform::CreatePrimFuncPass(fpass, 0, "BindTarget", {});
-}
-
-static transform::Pass AnnotateEntryFunc(bool b) {
-  auto fpass = [](tir::PrimFunc f, IRModule m, transform::PassContext ctx) {
-    return WithAttr(std::move(f), tir::attr::kIsEntryFunc, Bool(true));
-  };
-  return tir::transform::CreatePrimFuncPass(fpass, 0, "AnnotateEntryFunc", {});
-}
-
-template <typename FCond>
-transform::Pass Filter(FCond fcond) {
-  auto fpass = [fcond](tir::PrimFunc f, IRModule m, transform::PassContext ctx) {
-    if (fcond(f)) {
-      return f;
-    } else {
-      return tir::PrimFunc(nullptr);
-    }
-  };
-  return tir::transform::CreatePrimFuncPass(fpass, 0, "Filter", {});
-}
-
 Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
   transform::PassContext pass_ctx = transform::PassContext::Current();
 
@@ -560,12 +534,12 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target)
 
   Array<Pass> mixed_pass_list;
 
-  mixed_pass_list.push_back(BindTarget(target));
+  mixed_pass_list.push_back(tir::transform::BindTarget(target));
 
   mixed_pass_list.push_back(tir::transform::VerifyMemory());
 
   if (ShouldAnnotateEntryFunc(mixed_mod)) {
-    mixed_pass_list.push_back(AnnotateEntryFunc(true));
+    mixed_pass_list.push_back(tir::transform::AnnotateEntryFunc());
   }
 
   bool detect_global_barrier =
@@ -602,14 +576,16 @@ TVM_REGISTER_GLOBAL("driver.mixed_mod_passes")
 
 transform::Sequential HostModulePassManager(IRModule mixed_mod, Target target_host) {
   Array<tvm::transform::Pass> host_pass_list;
-  host_pass_list.push_back(Filter([](const tir::PrimFunc& f) {
+
+  runtime::TypedPackedFunc<bool(tir::PrimFunc)> fcond = [](const tir::PrimFunc& f) {
     return f->GetAttr<Integer>(tvm::attr::kCallingConv, Integer(CallingConv::kDefault)) !=
            CallingConv::kDeviceKernelLaunch;
-  }));
+  };
+  host_pass_list.push_back(tir::transform::Filter(fcond));
 
   ICHECK(mixed_mod.defined()) << "This module must be defined";
 
-  host_pass_list.push_back(BindTarget(target_host));
+  host_pass_list.push_back(tir::transform::BindTarget(target_host));
 
   host_pass_list.push_back(tir::transform::LowerTVMBuiltin());
   host_pass_list.push_back(tir::transform::LowerCustomDatatypes());
@@ -627,12 +603,13 @@ TVM_REGISTER_GLOBAL("driver.host_mod_passes")
 
 transform::Sequential DeviceModulePassManager(IRModule mixed_mod, Target target) {
   Array<Pass> device_pass_list;
-  device_pass_list.push_back(Filter([](const tir::PrimFunc& f) {
+  runtime::TypedPackedFunc<bool(tir::PrimFunc)> fcond = [](const tir::PrimFunc& f) {
     return f->GetAttr<Integer>(tvm::attr::kCallingConv, Integer(CallingConv::kDefault)) ==
            CallingConv::kDeviceKernelLaunch;
-  }));
+  };
+  device_pass_list.push_back(tir::transform::Filter(fcond));
 
-  device_pass_list.push_back(BindTarget(target));
+  device_pass_list.push_back(tir::transform::BindTarget(target));
 
   device_pass_list.push_back(tir::transform::LowerWarpMemory());
   device_pass_list.push_back(tir::transform::Simplify());

src/tir/transforms/helpers.cc

@@ -0,0 +1,63 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file helpers.cc
+ * \brief Passes that serve as helper functions.
+ */
+
+#include <tvm/driver/driver_api.h>
+#include <tvm/tir/transform.h>
+
+namespace tvm {
+namespace tir {
+namespace transform {
+transform::Pass BindTarget(Target target) {
+  auto fpass = [target](tir::PrimFunc f, IRModule m, transform::PassContext ctx) {
+    return WithAttr(std::move(f), tvm::attr::kTarget, target);
+  };
+  return tir::transform::CreatePrimFuncPass(fpass, 0, "tir.BindTarget", {});
+}
+
+transform::Pass AnnotateEntryFunc() {
+  auto fpass = [](tir::PrimFunc f, IRModule m, transform::PassContext ctx) {
+    ICHECK(m->functions.size() == 1);
+    return WithAttr(std::move(f), tir::attr::kIsEntryFunc, Bool(true));
+  };
+  return tir::transform::CreatePrimFuncPass(fpass, 0, "tir.AnnotateEntryFunc", {});
+}
+
+transform::Pass Filter(runtime::TypedPackedFunc<bool(PrimFunc)> fcond) {
+  auto fpass = [fcond](tir::PrimFunc f, IRModule m, transform::PassContext ctx) {
+    if (fcond(f)) {
+      return f;
+    } else {
+      return tir::PrimFunc(nullptr);
+    }
+  };
+  return tir::transform::CreatePrimFuncPass(fpass, 0, "tir.Filter", {});
+}
+
+TVM_REGISTER_GLOBAL("tir.transform.BindTarget").set_body_typed(BindTarget);
+TVM_REGISTER_GLOBAL("tir.transform.AnnotateEntryFunc").set_body_typed(AnnotateEntryFunc);
+TVM_REGISTER_GLOBAL("tir.transform.Filter").set_body_typed(Filter);
+
+}  // namespace transform
+}  // namespace tir
+}  // namespace tvm

src/tir/transforms/lower_init_block.cc

@@ -81,7 +81,7 @@ Pass LowerInitBlock() {
   auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
     return LowerInitBlock(std::move(f));
   };
-  return CreatePrimFuncPass(pass_func, 0, "tir.LowerReduction", {});
+  return CreatePrimFuncPass(pass_func, 0, "tir.LowerInitBlock", {});
 }
 
 TVM_REGISTER_GLOBAL("tir.transform.LowerInitBlock").set_body_typed(LowerInitBlock);

src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc

@@ -60,7 +60,7 @@ class PoolAllocationToOffsetConverter : public StmtExprMutator {
       PoolInfo pool_info = pool_allocation->pool_info;
       int byte_pool_offset = pool_allocation->byte_offset->value;
       int required_pool_size_for_allocation =
-          byte_pool_offset + CalculateExtentsSize(allocate_node.operator->());
+          byte_pool_offset + static_cast<int>(CalculateExtentsSize(allocate_node.operator->()));
       if (all_pools_sizes_.find(pool_info) == all_pools_sizes_.end()) {
         all_pools_sizes_[pool_info] = required_pool_size_for_allocation;
       } else {

tests/python/unittest/test_tir_transform_helpers.py

@@ -0,0 +1,123 @@
+# 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.
+import pytest
+
+import tvm
+from tvm.script import tir as T
+import tvm.testing
+
+
+def test_annotate_entry_func_single_primfunc():
+    @tvm.script.ir_module
+    class MockModule:
+        @T.prim_func
+        def func1(A: T.Buffer[(16,), "float32"]):
+            for i in T.serial(16):
+                if i == 5:
+                    if i == 5:
+                        A[i] = 0.0
+
+    mod = MockModule
+    assert mod
+    assert mod["func1"].attrs is None
+    after = tvm.tir.transform.AnnotateEntryFunc()(mod)
+    assert (
+        after["func1"].attrs
+        and "tir.is_entry_func" in after["func1"].attrs
+        and after["func1"].attrs["tir.is_entry_func"]
+    )
+
+
+# Test module
+@tvm.script.ir_module
+class MockModule:
+    @T.prim_func
+    def func1(A: T.Buffer[(16,), "float32"]):
+        for i in T.serial(16):
+            if i == 5:
+                if i == 5:
+                    A[i] = 0.0
+
+    @T.prim_func
+    def func2(A: T.Buffer[(32,), "float32"]):
+        for i in T.serial(32):
+            if i == 15:
+                if i == 15:
+                    A[i] = 0.0
+
+
+@pytest.mark.xfail
+def test_annotate_entry_func_multiple_primfunc():
+    mod = MockModule
+    assert mod
+    assert mod["func1"].attrs is None
+    assert mod["func2"].attrs is None
+    # This should fail
+    after = tvm.tir.transform.AnnotateEntryFunc()(mod)
+
+
+def test_bind_target():
+    mod = MockModule
+    assert mod
+
+    target = tvm.target.Target("cuda")
+    assert mod["func1"].attrs is None
+    assert mod["func2"].attrs is None
+    after = tvm.tir.transform.BindTarget(target)(mod)
+
+    assert after["func1"].attrs and "target" in after["func1"].attrs
+    assert after["func1"].attrs["target"] == target
+    assert after["func2"].attrs and "target" in after["func2"].attrs
+    assert after["func2"].attrs["target"] == target
+
+
+def test_filter_primfunc():
+    mod = MockModule
+    assert mod
+    # Annotate each function for testing
+    mod["func1"] = mod["func1"].with_attr("temp", "test1")
+    mod["func2"] = mod["func2"].with_attr("temp", "test2")
+
+    # Test condition that does not filter out anything
+    def checker_filter_out_none(func: tvm.tir.PrimFunc):
+        return (func.attrs is not None) and ("temp" in func.attrs)
+
+    after = tvm.tir.transform.Filter(checker_filter_out_none)(mod)
+    assert len(after.functions) == 2
+    # Filtered functions should satisfy the given condition.
+    assert checker_filter_out_none(after["func1"])
+    assert checker_filter_out_none(after["func2"])
+
+    # Test condition that selectively filters out primfuncs
+    def checker_filter_out_one(func: tvm.tir.PrimFunc):
+        return (func.attrs is not None) and ("temp" in func.attrs) and func.attrs["temp"] == "test1"
+
+    after = tvm.tir.transform.Filter(checker_filter_out_one)(mod)
+    assert len(after.functions) == 1
+    # Filtered functions should satisfy the given condition.
+    assert checker_filter_out_one(after["func1"])
+
+    # Test condition that filters out everything
+    def checker_filter_out_both(func: tvm.tir.PrimFunc):
+        return (func.attrs is not None) and ("invalid_attr" in func.attrs)
+
+    after = tvm.tir.transform.Filter(checker_filter_out_both)(mod)
+    assert len(after.functions) == 0
+
+
+if __name__ == "__main__":
+    tvm.testing.main()