From b362f86038b1a5dbb31a659e82b622799697b585 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Thu, 14 Mar 2024 09:57:05 -0500 Subject: [PATCH 1/3] [TIR] LowerTVMBuiltin may use device_type from PrimFunc annotation If an allocation occurs within a host function, it may not have a device/host split. --- src/tir/transforms/lower_tvm_builtin.cc | 36 +++++++++++++----- .../test_tir_transform_lower_tvm_builtin.py | 37 +++++++++++++++++-- 2 files changed, 60 insertions(+), 13 deletions(-) diff --git a/src/tir/transforms/lower_tvm_builtin.cc b/src/tir/transforms/lower_tvm_builtin.cc index 6da2f873b728..486830e90a50 100644 --- a/src/tir/transforms/lower_tvm_builtin.cc +++ b/src/tir/transforms/lower_tvm_builtin.cc @@ -38,6 +38,19 @@ namespace tir { // These information are needed during codegen. class BuiltinLower : public StmtExprMutator { public: + static PrimFunc Build(PrimFunc func) { + Optional device_type = NullOpt; + if (auto target = func->GetAttr(tvm::attr::kTarget)) { + device_type = Integer(target.value()->kind->default_device_type); + } + + BuiltinLower mutator(device_type); + func.CopyOnWrite()->body = mutator.VisitBodyAndRealizeAlloca(func->body); + return func; + } + + BuiltinLower(Optional device_type = NullOpt) : device_type_(device_type) {} + // NOTE: Right now, we make the following scoping requirement // for memory allocated by the following primitives // - tvm_stack_make_array @@ -284,13 +297,17 @@ class BuiltinLower : public StmtExprMutator { Stmt VisitStmt_(const AttrStmtNode* op) final { if (op->attr_key == attr::device_id) { - ICHECK(!device_id_); + auto cache = device_id_; device_id_ = op->value; - return this->VisitStmt(op->body); + Stmt out = this->VisitStmt(op->body); + device_id_ = cache; + return out; } else if (op->attr_key == attr::device_type) { - ICHECK(!device_type_); + auto cache = device_type_; device_type_ = op->value; - return this->VisitStmt(op->body); + Stmt out = this->VisitStmt(op->body); + device_type_ = cache; + return out; } else { return StmtExprMutator::VisitStmt_(op); } @@ -656,13 +673,12 @@ class BuiltinLower : public StmtExprMutator { namespace transform { Pass LowerTVMBuiltin() { - auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) { - if (IsHostFunc(f).value_or(false)) { - auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol); - f.CopyOnWrite()->body = BuiltinLower().Build(f->body); - VLOG(2) << "LowerTVMBuiltin: " << f; + auto pass_func = [](PrimFunc func, IRModule m, PassContext ctx) { + if (IsHostFunc(func).value_or(false)) { + func = BuiltinLower::Build(func); + VLOG(2) << "LowerTVMBuiltin: " << func; } - return f; + return func; }; return CreatePrimFuncPass(pass_func, 0, "tir.LowerTVMBuiltin", {}); } diff --git a/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py b/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py index de1020ef2078..754ce032404d 100644 --- a/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py +++ b/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py @@ -260,11 +260,13 @@ def expected(): class TestLowerAllocateRequiresDeviceID(tvm.testing.CompareBeforeAfter): + """If device id is missing, error.""" + transform = tvm.tir.transform.LowerTVMBuiltin() def before(): T.func_attr({"target": T.target("llvm")}) - T.attr("dummy", "device_id", 0) + T.attr("dummy", "device_type", 2) # kDLCuda ptr = T.allocate([16], "float32") buf = T.decl_buffer(16, "float32", data=ptr) buf[0] = 0.0 @@ -273,16 +275,45 @@ def before(): class TestLowerAllocateRequiresDeviceType(tvm.testing.CompareBeforeAfter): + """If device type is missing, error. + + The device type can be inferred either from the `"device_type"` + statement attribute, or from the `"target"` function attribute. + Here, we provide neither. The `"tir.is_host_func"` attribute is + provided as otherwise the function would be skipped altogether by + LowerTVMBuiltin. + """ + transform = tvm.tir.transform.LowerTVMBuiltin() def before(): - T.func_attr({"target": T.target("llvm")}) + T.func_attr({"tir.is_host_func": True}) T.attr("dummy", "device_id", 0) + ptr = T.allocate([1024 * 1024], "float32") + buf = T.decl_buffer(1024 * 1024, "float32", data=ptr) + buf[0] = 0.0 + + expected = tvm.TVMError + + +class TestLowerCPUAllocWithFunctionAttr(tvm.testing.CompareBeforeAfter): + """CPU allocations can be handled at codegen time + + Like `TestLowerCPUAllocation`, but the device type is taken from + the function attribute. The `AttrStmt` can override the device + type for allocations within its scope, but it defaults to the + function's target. + """ + + transform = tvm.tir.transform.LowerTVMBuiltin() + + def before(): + T.func_attr({"target": T.target("llvm")}) ptr = T.allocate([16], "float32") buf = T.decl_buffer(16, "float32", data=ptr) buf[0] = 0.0 - expected = tvm.TVMError + expected = before if __name__ == "__main__": From 3145b353b8b98eb231b1a74d314c16dd300cc1a9 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Fri, 15 Mar 2024 14:53:16 -0500 Subject: [PATCH 2/3] lint fix --- src/tir/transforms/lower_tvm_builtin.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tir/transforms/lower_tvm_builtin.cc b/src/tir/transforms/lower_tvm_builtin.cc index 486830e90a50..1a3888a7cd48 100644 --- a/src/tir/transforms/lower_tvm_builtin.cc +++ b/src/tir/transforms/lower_tvm_builtin.cc @@ -49,7 +49,7 @@ class BuiltinLower : public StmtExprMutator { return func; } - BuiltinLower(Optional device_type = NullOpt) : device_type_(device_type) {} + explicit BuiltinLower(Optional device_type = NullOpt) : device_type_(device_type) {} // NOTE: Right now, we make the following scoping requirement // for memory allocated by the following primitives From c2be6006a9577aa5190b6e0fedb376078f195931 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 26 Mar 2024 09:22:48 -0500 Subject: [PATCH 3/3] Remove magic number from unit test --- .../tir-transform/test_tir_transform_lower_tvm_builtin.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py b/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py index 754ce032404d..08e641a61d1b 100644 --- a/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py +++ b/tests/python/tir-transform/test_tir_transform_lower_tvm_builtin.py @@ -217,7 +217,7 @@ class TestLowerDeviceAllocate(tvm.testing.CompareBeforeAfter): def before(): T.func_attr({"target": T.target("llvm")}) - T.attr("dummy", "device_type", 2) # kDLCuda + T.attr("dummy", "device_type", tvm.runtime.Device.kDLCUDA) T.attr("dummy", "device_id", 0) ptr = T.allocate([16], "float32") buf = T.decl_buffer(16, "float32", data=ptr) @@ -246,7 +246,7 @@ class TestLowerCPUAllocation(tvm.testing.CompareBeforeAfter): def before(): T.func_attr({"target": T.target("llvm")}) - T.attr("dummy", "device_type", 1) # kDLCPU + T.attr("dummy", "device_type", tvm.runtime.Device.kDLCPU) T.attr("dummy", "device_id", 0) ptr = T.allocate([16], "float32") buf = T.decl_buffer(16, "float32", data=ptr) @@ -266,7 +266,7 @@ class TestLowerAllocateRequiresDeviceID(tvm.testing.CompareBeforeAfter): def before(): T.func_attr({"target": T.target("llvm")}) - T.attr("dummy", "device_type", 2) # kDLCuda + T.attr("dummy", "device_type", tvm.runtime.Device.kDLCUDA) ptr = T.allocate([16], "float32") buf = T.decl_buffer(16, "float32", data=ptr) buf[0] = 0.0