Mousius commented on a change in pull request #9313: URL: https://github.com/apache/tvm/pull/9313#discussion_r737795247
########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) + << "SEScope annotations must include at least a device_type"; + Target target = FindPrimitiveTargetOrFail(se_scope->device_type()); + return se_scope_cache_.Unique( + SEScope(device_type, se_scope->virtual_device_id(), target, se_scope->memory_scope())); +} + +/*! + * \brief Returns the default \p SEScope for primitives and the \p SEScope for the host + * given vector of available \p targets. If necessary, add new \p Targets to \p targets + * to match the required devices. + */ +void CompilationConfigNode::EstablishDefaultSEScopes(const transform::PassContext& pass_ctx) { + // + // Gather the hints as to what our default device type for primitives should be. + // + DLDeviceType default_primitive_device_type; + Optional<Integer> opt_fallback_dev = pass_ctx->GetConfig<Integer>("relay.fallback_device_type"); + if (opt_fallback_dev) { + const int64_t v = opt_fallback_dev.value()->value; + if (v <= 0) { + LOG(FATAL) + << "The 'relay.fallback_device_type' pass attribute is set to an invalid device type " + << v; + default_primitive_device_type = kDLCPU; + } else { + default_primitive_device_type = static_cast<DLDeviceType>(v); + LOG(INFO) << "Using the 'relay.fallback_device_type' pass attribute " + << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + } else if (primitive_targets.size() == 1) { + // In the homogeneous case there's no free choice. + default_primitive_device_type = + static_cast<DLDeviceType>(primitive_targets.front()->kind->device_type); + LOG(INFO) << "Using the unique target '" << primitive_targets.front()->str() + << "' of device type " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } else { + // Fallback. + default_primitive_device_type = kDLCPU; + LOG(WARNING) << "Using " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + + // + // Establish the default primitive SEScope, choosing a known Target to match the device type. + // + default_primitive_se_scope = se_scope_cache_.Unique( + SEScope(default_primitive_device_type, + /*virtual_device_id=*/0, FindPrimitiveTargetOrFail(default_primitive_device_type))); + + // + // Gather the hints as to what our default device type for the 'host' should be. + // + DLDeviceType host_device_type; + if (host_target.defined()) { + host_device_type = static_cast<DLDeviceType>(host_target->kind->device_type); + if (host_device_type != kDLCPU) { + LOG(WARNING) << "Using the given host target '" << host_target->str() + << "' of non-CPU device type " << host_device_type + << " for all host operations and data"; + } else { + LOG(INFO) << "Using the given host target '" << host_target->str() << "' of device type " + << host_device_type << " for all host operations and data"; + } + } else if (primitive_targets.size() == 1 && + primitive_targets.front()->kind->device_type == kDLCPU) { + // In the homogenous case without an explicit host target just use the given target so long as + // it's a CPU. + host_device_type = kDLCPU; + host_target = + FindPrimitiveTargetOrFail(host_device_type); // ie just primitive_targets.front()! + LOG(INFO) << "Using the unique target '" << host_target->str() << "' of device type " + << host_device_type << " for all host operations and data"; + } else { + // Fallback. + host_device_type = kDLCPU; + // Even if the list of available targets already includes one for kDLCPU we won't use it + // since its options may not be appropriate for host code (eg shape functions). Instead, + // create a fresh default Target. + host_target = MakeDefaultTarget(host_device_type); + LOG(WARNING) << "Using the default host target '" << host_target->str() << "' of device type " + << host_device_type << " for all host operations and data"; + } + + // + // Establish the host SEScope. + // + host_se_scope = se_scope_cache_.Unique(SEScope(host_device_type, + /*virtual_device_id=*/0, host_target)); +} + +/* static */ Target CompilationConfigNode::MakeDefaultTarget(DLDeviceType device_type) { + std::string name = runtime::DeviceName(device_type); + if (name == "cpu") { + if (runtime::Registry::Get("codegen.LLVMModuleCreate")) { + // LLVM is available. + return Target("llvm"); + } else { + // LLVM is not available. + return Target("stackvm"); + } + } else { + return Target(name); + } +} + +Target CompilationConfigNode::FindPrimitiveTargetOrFail(DLDeviceType device_type) const { + auto itr = std::find_if( + primitive_targets.begin(), primitive_targets.end(), + [device_type](const Target& target) { return target->kind->device_type == device_type; }); + CHECK(itr != primitive_targets.end()) << "No target for device type " << device_type << " in the " Review comment: We should also replicate this in a test with `ASSERT_RAISES` ########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) + << "SEScope annotations must include at least a device_type"; + Target target = FindPrimitiveTargetOrFail(se_scope->device_type()); + return se_scope_cache_.Unique( + SEScope(device_type, se_scope->virtual_device_id(), target, se_scope->memory_scope())); +} + +/*! + * \brief Returns the default \p SEScope for primitives and the \p SEScope for the host + * given vector of available \p targets. If necessary, add new \p Targets to \p targets + * to match the required devices. + */ +void CompilationConfigNode::EstablishDefaultSEScopes(const transform::PassContext& pass_ctx) { + // + // Gather the hints as to what our default device type for primitives should be. + // + DLDeviceType default_primitive_device_type; + Optional<Integer> opt_fallback_dev = pass_ctx->GetConfig<Integer>("relay.fallback_device_type"); + if (opt_fallback_dev) { + const int64_t v = opt_fallback_dev.value()->value; + if (v <= 0) { + LOG(FATAL) + << "The 'relay.fallback_device_type' pass attribute is set to an invalid device type " + << v; + default_primitive_device_type = kDLCPU; + } else { + default_primitive_device_type = static_cast<DLDeviceType>(v); + LOG(INFO) << "Using the 'relay.fallback_device_type' pass attribute " + << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + } else if (primitive_targets.size() == 1) { + // In the homogeneous case there's no free choice. + default_primitive_device_type = + static_cast<DLDeviceType>(primitive_targets.front()->kind->device_type); + LOG(INFO) << "Using the unique target '" << primitive_targets.front()->str() + << "' of device type " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } else { + // Fallback. + default_primitive_device_type = kDLCPU; + LOG(WARNING) << "Using " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + + // + // Establish the default primitive SEScope, choosing a known Target to match the device type. + // + default_primitive_se_scope = se_scope_cache_.Unique( + SEScope(default_primitive_device_type, + /*virtual_device_id=*/0, FindPrimitiveTargetOrFail(default_primitive_device_type))); + + // + // Gather the hints as to what our default device type for the 'host' should be. + // + DLDeviceType host_device_type; + if (host_target.defined()) { + host_device_type = static_cast<DLDeviceType>(host_target->kind->device_type); + if (host_device_type != kDLCPU) { + LOG(WARNING) << "Using the given host target '" << host_target->str() + << "' of non-CPU device type " << host_device_type + << " for all host operations and data"; + } else { + LOG(INFO) << "Using the given host target '" << host_target->str() << "' of device type " + << host_device_type << " for all host operations and data"; + } + } else if (primitive_targets.size() == 1 && + primitive_targets.front()->kind->device_type == kDLCPU) { + // In the homogenous case without an explicit host target just use the given target so long as + // it's a CPU. + host_device_type = kDLCPU; + host_target = + FindPrimitiveTargetOrFail(host_device_type); // ie just primitive_targets.front()! + LOG(INFO) << "Using the unique target '" << host_target->str() << "' of device type " + << host_device_type << " for all host operations and data"; + } else { + // Fallback. + host_device_type = kDLCPU; + // Even if the list of available targets already includes one for kDLCPU we won't use it + // since its options may not be appropriate for host code (eg shape functions). Instead, + // create a fresh default Target. + host_target = MakeDefaultTarget(host_device_type); Review comment: I don't think any of our test cases cover this defaulting logic as all the tests preconfigure most things. ########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) + << "SEScope annotations must include at least a device_type"; + Target target = FindPrimitiveTargetOrFail(se_scope->device_type()); + return se_scope_cache_.Unique( + SEScope(device_type, se_scope->virtual_device_id(), target, se_scope->memory_scope())); +} + +/*! + * \brief Returns the default \p SEScope for primitives and the \p SEScope for the host + * given vector of available \p targets. If necessary, add new \p Targets to \p targets + * to match the required devices. + */ +void CompilationConfigNode::EstablishDefaultSEScopes(const transform::PassContext& pass_ctx) { + // + // Gather the hints as to what our default device type for primitives should be. + // + DLDeviceType default_primitive_device_type; + Optional<Integer> opt_fallback_dev = pass_ctx->GetConfig<Integer>("relay.fallback_device_type"); + if (opt_fallback_dev) { + const int64_t v = opt_fallback_dev.value()->value; + if (v <= 0) { + LOG(FATAL) + << "The 'relay.fallback_device_type' pass attribute is set to an invalid device type " + << v; + default_primitive_device_type = kDLCPU; + } else { + default_primitive_device_type = static_cast<DLDeviceType>(v); + LOG(INFO) << "Using the 'relay.fallback_device_type' pass attribute " + << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + } else if (primitive_targets.size() == 1) { + // In the homogeneous case there's no free choice. + default_primitive_device_type = Review comment: Test case please :smile_cat: ########## File path: include/tvm/target/compilation_config.h ########## @@ -0,0 +1,150 @@ +/* + * 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 tvm/target/compilation_config.h + * \brief A helper class to collect all the targets in canonical form necessary for compilation. + * CAUTION: Preliminary, currently only used to support device planning, very likely to change. + */ + +#ifndef TVM_TARGET_COMPILATION_CONFIG_H_ +#define TVM_TARGET_COMPILATION_CONFIG_H_ + +#include <tvm/target/se_scope.h> + +namespace tvm { + +/*! + * \brief Gathers the \p Targets and \p SEScopes in canonical form needed to compile a Relay module. + * + * CAUTION: This is a temporary class to help us bridge legacy and new target/device handling Review comment: Following the discussion around this offline, I believe we'd like to keep this longer term - so we should update this comment in future patches. ########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) + << "SEScope annotations must include at least a device_type"; + Target target = FindPrimitiveTargetOrFail(se_scope->device_type()); + return se_scope_cache_.Unique( + SEScope(device_type, se_scope->virtual_device_id(), target, se_scope->memory_scope())); +} + +/*! + * \brief Returns the default \p SEScope for primitives and the \p SEScope for the host + * given vector of available \p targets. If necessary, add new \p Targets to \p targets + * to match the required devices. + */ +void CompilationConfigNode::EstablishDefaultSEScopes(const transform::PassContext& pass_ctx) { + // + // Gather the hints as to what our default device type for primitives should be. + // + DLDeviceType default_primitive_device_type; + Optional<Integer> opt_fallback_dev = pass_ctx->GetConfig<Integer>("relay.fallback_device_type"); + if (opt_fallback_dev) { + const int64_t v = opt_fallback_dev.value()->value; + if (v <= 0) { + LOG(FATAL) + << "The 'relay.fallback_device_type' pass attribute is set to an invalid device type " + << v; + default_primitive_device_type = kDLCPU; + } else { + default_primitive_device_type = static_cast<DLDeviceType>(v); + LOG(INFO) << "Using the 'relay.fallback_device_type' pass attribute " + << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + } else if (primitive_targets.size() == 1) { + // In the homogeneous case there's no free choice. + default_primitive_device_type = + static_cast<DLDeviceType>(primitive_targets.front()->kind->device_type); + LOG(INFO) << "Using the unique target '" << primitive_targets.front()->str() + << "' of device type " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } else { + // Fallback. + default_primitive_device_type = kDLCPU; Review comment: This is a valid assumption for TVM today, I'm fairly sure introducing a new host type will require us to rethink this though - would be happy with the check else. ########## File path: src/target/se_scope.cc ########## @@ -0,0 +1,225 @@ +/* + * 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 tvm/target/se_scope.cc + * \brief Implementation of \p SEScope for representing a Storage or Execution scope. + */ +#include <tvm/node/reflection.h> +#include <tvm/runtime/device_api.h> +#include <tvm/target/se_scope.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(SEScopeNode); + +void SEScopeNode::VisitAttrs(AttrVisitor* v) { + int i = static_cast<int>(device_type_); + v->Visit("device_type", &i); + device_type_ = static_cast<DLDeviceType>(i); + v->Visit("virtual_device_id", &virtual_device_id_); + v->Visit("target", &target_); + v->Visit("memory_scope", &memory_scope_); +} + +bool SEScopeNode::SEqualReduce(const SEScopeNode* other, SEqualReducer equal) const { + return device_type_ == other->device_type_ && virtual_device_id_ == other->virtual_device_id_ && + // NOTE: Comparing targets by their str representations + target_->str() == other->target_->str() && memory_scope_ == other->memory_scope_; +} + +void SEScopeNode::SHashReduce(SHashReducer hash_reduce) const { + hash_reduce(device_type_); + hash_reduce(virtual_device_id_); + // NOTE: Reducing target to its str representation + hash_reduce(target_->str()); + hash_reduce(memory_scope_); +} + +TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) + .set_dispatch<SEScopeNode>([](const ObjectRef& ref, ReprPrinter* p) { + auto* node = ref.as<SEScopeNode>(); + p->stream << "SEScopeNode("; + if (node->is_fully_unconstrained()) { + p->stream << "?"; + } else { + bool need_sep = false; + if (node->device_type() != kInvalidDeviceType) { + p->stream << "device_type=" << node->device_type(); + need_sep = true; + } + if (node->virtual_device_id() >= 0) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "virtual_device_id=" << node->virtual_device_id(); + need_sep = true; + } + if (node->target().defined()) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "target='" << node->target()->str() << "'"; + need_sep = true; + } + if (!node->memory_scope().empty()) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "memory_scope='" << node->memory_scope() << "'"; + } + } + p->stream << ")"; + }); + +SEScope::SEScope(DLDeviceType device_type, int virtual_device_id, Target target, + String memory_scope) { + ICHECK(!target.defined() || device_type == target->kind->device_type) + << "target '" << target->str() << "' has device type " << target->kind->device_type + << " but scope has device type " << device_type; + auto node = make_object<SEScopeNode>(); + node->device_type_ = device_type; + node->virtual_device_id_ = virtual_device_id; + node->target_ = std::move(target); + node->memory_scope_ = std::move(memory_scope); + data_ = std::move(node); +} + +/* static */ SEScope SEScope::FullyUnconstrained() { + static const SEScope unconstrained{}; + return unconstrained; +} + +/* static */ +Optional<SEScope> SEScope::Join(const SEScope& lhs, const SEScope& rhs) { Review comment: There's a lot of logic in both `Join` and `Default` but not many test cases, could you see if we can bulk out the test suite a bit? ########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) + << "SEScope annotations must include at least a device_type"; + Target target = FindPrimitiveTargetOrFail(se_scope->device_type()); + return se_scope_cache_.Unique( + SEScope(device_type, se_scope->virtual_device_id(), target, se_scope->memory_scope())); +} + +/*! + * \brief Returns the default \p SEScope for primitives and the \p SEScope for the host + * given vector of available \p targets. If necessary, add new \p Targets to \p targets + * to match the required devices. + */ +void CompilationConfigNode::EstablishDefaultSEScopes(const transform::PassContext& pass_ctx) { + // + // Gather the hints as to what our default device type for primitives should be. + // + DLDeviceType default_primitive_device_type; + Optional<Integer> opt_fallback_dev = pass_ctx->GetConfig<Integer>("relay.fallback_device_type"); + if (opt_fallback_dev) { + const int64_t v = opt_fallback_dev.value()->value; + if (v <= 0) { + LOG(FATAL) Review comment: I think this should trigger an error we can test for? ########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) Review comment: Can we add a test for this with `ASSERT_RAISES` ? ########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) + << "SEScope annotations must include at least a device_type"; + Target target = FindPrimitiveTargetOrFail(se_scope->device_type()); + return se_scope_cache_.Unique( + SEScope(device_type, se_scope->virtual_device_id(), target, se_scope->memory_scope())); +} + +/*! + * \brief Returns the default \p SEScope for primitives and the \p SEScope for the host + * given vector of available \p targets. If necessary, add new \p Targets to \p targets + * to match the required devices. + */ +void CompilationConfigNode::EstablishDefaultSEScopes(const transform::PassContext& pass_ctx) { + // + // Gather the hints as to what our default device type for primitives should be. + // + DLDeviceType default_primitive_device_type; + Optional<Integer> opt_fallback_dev = pass_ctx->GetConfig<Integer>("relay.fallback_device_type"); + if (opt_fallback_dev) { + const int64_t v = opt_fallback_dev.value()->value; + if (v <= 0) { + LOG(FATAL) + << "The 'relay.fallback_device_type' pass attribute is set to an invalid device type " + << v; + default_primitive_device_type = kDLCPU; + } else { + default_primitive_device_type = static_cast<DLDeviceType>(v); + LOG(INFO) << "Using the 'relay.fallback_device_type' pass attribute " + << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + } else if (primitive_targets.size() == 1) { + // In the homogeneous case there's no free choice. + default_primitive_device_type = + static_cast<DLDeviceType>(primitive_targets.front()->kind->device_type); + LOG(INFO) << "Using the unique target '" << primitive_targets.front()->str() + << "' of device type " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } else { + // Fallback. + default_primitive_device_type = kDLCPU; + LOG(WARNING) << "Using " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + + // + // Establish the default primitive SEScope, choosing a known Target to match the device type. + // + default_primitive_se_scope = se_scope_cache_.Unique( + SEScope(default_primitive_device_type, + /*virtual_device_id=*/0, FindPrimitiveTargetOrFail(default_primitive_device_type))); + + // + // Gather the hints as to what our default device type for the 'host' should be. + // + DLDeviceType host_device_type; + if (host_target.defined()) { + host_device_type = static_cast<DLDeviceType>(host_target->kind->device_type); + if (host_device_type != kDLCPU) { + LOG(WARNING) << "Using the given host target '" << host_target->str() Review comment: Can we construct a test case for this? As it's supported here it makes sense to add the check @csullivan suggested when we default back to kDLCPU and raise an error if it's not found. ########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) + << "SEScope annotations must include at least a device_type"; + Target target = FindPrimitiveTargetOrFail(se_scope->device_type()); + return se_scope_cache_.Unique( + SEScope(device_type, se_scope->virtual_device_id(), target, se_scope->memory_scope())); +} + +/*! + * \brief Returns the default \p SEScope for primitives and the \p SEScope for the host + * given vector of available \p targets. If necessary, add new \p Targets to \p targets + * to match the required devices. + */ +void CompilationConfigNode::EstablishDefaultSEScopes(const transform::PassContext& pass_ctx) { + // + // Gather the hints as to what our default device type for primitives should be. + // + DLDeviceType default_primitive_device_type; + Optional<Integer> opt_fallback_dev = pass_ctx->GetConfig<Integer>("relay.fallback_device_type"); + if (opt_fallback_dev) { + const int64_t v = opt_fallback_dev.value()->value; + if (v <= 0) { + LOG(FATAL) + << "The 'relay.fallback_device_type' pass attribute is set to an invalid device type " + << v; + default_primitive_device_type = kDLCPU; + } else { + default_primitive_device_type = static_cast<DLDeviceType>(v); Review comment: I think because we consume `PassContext` we can actually test this with many values as well? Did something similar here: https://github.com/apache/tvm/pull/9338/files#diff-4608944c7ab774c4dddfef6e2eaada43112ca48be2b10c255c95839f37c991a5R43-R53 ########## File path: python/tvm/target/compilation_config.py ########## @@ -0,0 +1,22 @@ +# 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. +"""Python bindings for creating CompilationConfigs.""" +from . import _ffi_api + + +def make_compilation_config(ctxt, targets, host_target=None): Review comment: Ideally this would be a Python class which we can construct for `relay.build` and restoring the default behaviour with something similar to `CompilationConfig.from_target(target)`. Though that's not in scope of this patch, it'd be something like this example: https://github.com/apache/tvm/blob/37a8d7b2b7df647a5dac6fbda18ee54d902ce4e4/python/tvm/relay/backend/executor.py#L26-L33 ########## File path: src/target/se_scope.cc ########## @@ -0,0 +1,225 @@ +/* + * 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 tvm/target/se_scope.cc + * \brief Implementation of \p SEScope for representing a Storage or Execution scope. + */ +#include <tvm/node/reflection.h> +#include <tvm/runtime/device_api.h> +#include <tvm/target/se_scope.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(SEScopeNode); + +void SEScopeNode::VisitAttrs(AttrVisitor* v) { + int i = static_cast<int>(device_type_); + v->Visit("device_type", &i); + device_type_ = static_cast<DLDeviceType>(i); + v->Visit("virtual_device_id", &virtual_device_id_); + v->Visit("target", &target_); + v->Visit("memory_scope", &memory_scope_); +} + +bool SEScopeNode::SEqualReduce(const SEScopeNode* other, SEqualReducer equal) const { + return device_type_ == other->device_type_ && virtual_device_id_ == other->virtual_device_id_ && + // NOTE: Comparing targets by their str representations + target_->str() == other->target_->str() && memory_scope_ == other->memory_scope_; +} + +void SEScopeNode::SHashReduce(SHashReducer hash_reduce) const { + hash_reduce(device_type_); + hash_reduce(virtual_device_id_); + // NOTE: Reducing target to its str representation + hash_reduce(target_->str()); + hash_reduce(memory_scope_); +} + +TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) + .set_dispatch<SEScopeNode>([](const ObjectRef& ref, ReprPrinter* p) { + auto* node = ref.as<SEScopeNode>(); + p->stream << "SEScopeNode("; + if (node->is_fully_unconstrained()) { + p->stream << "?"; + } else { + bool need_sep = false; + if (node->device_type() != kInvalidDeviceType) { + p->stream << "device_type=" << node->device_type(); + need_sep = true; + } + if (node->virtual_device_id() >= 0) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "virtual_device_id=" << node->virtual_device_id(); + need_sep = true; + } + if (node->target().defined()) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "target='" << node->target()->str() << "'"; + need_sep = true; + } + if (!node->memory_scope().empty()) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "memory_scope='" << node->memory_scope() << "'"; + } + } + p->stream << ")"; + }); + +SEScope::SEScope(DLDeviceType device_type, int virtual_device_id, Target target, + String memory_scope) { + ICHECK(!target.defined() || device_type == target->kind->device_type) Review comment: an `ASSERT_RAISES` test would be good here! ########## File path: src/target/compilation_config.cc ########## @@ -0,0 +1,219 @@ +/* + * 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 tvm/target/compilation_config.cc + * \brief Implementation of \p CompilationConfig for collecting \p Targets. + */ + +#include <tvm/runtime/device_api.h> +#include <tvm/target/compilation_config.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(CompilationConfigNode); + +void CompilationConfigNode::VisitAttrs(AttrVisitor* v) { + v->Visit("legacy_target_map", &legacy_target_map); + v->Visit("host_target", &host_target); + v->Visit("primitive_targets", &primitive_targets); + v->Visit("default_primitive_se_scope", &default_primitive_se_scope); + v->Visit("host_se_scope", &host_se_scope); + v->Visit("optional_homogenous_target", &optional_homogeneous_target); + // NOTE: The se_scope_cache_ is not accessible via FFI. +} + +SEScope CompilationConfigNode::CanonicalSEScope(const SEScope& se_scope) const { + if (se_scope->target().defined()) { + return se_scope_cache_.Unique(se_scope); + } + DLDeviceType device_type = se_scope->device_type(); + // TODO(mbs): Proper diagnostics. + CHECK(device_type != kInvalidDeviceType) + << "SEScope annotations must include at least a device_type"; + Target target = FindPrimitiveTargetOrFail(se_scope->device_type()); + return se_scope_cache_.Unique( + SEScope(device_type, se_scope->virtual_device_id(), target, se_scope->memory_scope())); +} + +/*! + * \brief Returns the default \p SEScope for primitives and the \p SEScope for the host + * given vector of available \p targets. If necessary, add new \p Targets to \p targets + * to match the required devices. + */ +void CompilationConfigNode::EstablishDefaultSEScopes(const transform::PassContext& pass_ctx) { + // + // Gather the hints as to what our default device type for primitives should be. + // + DLDeviceType default_primitive_device_type; + Optional<Integer> opt_fallback_dev = pass_ctx->GetConfig<Integer>("relay.fallback_device_type"); + if (opt_fallback_dev) { + const int64_t v = opt_fallback_dev.value()->value; + if (v <= 0) { + LOG(FATAL) + << "The 'relay.fallback_device_type' pass attribute is set to an invalid device type " + << v; + default_primitive_device_type = kDLCPU; + } else { + default_primitive_device_type = static_cast<DLDeviceType>(v); + LOG(INFO) << "Using the 'relay.fallback_device_type' pass attribute " + << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + } else if (primitive_targets.size() == 1) { + // In the homogeneous case there's no free choice. + default_primitive_device_type = + static_cast<DLDeviceType>(primitive_targets.front()->kind->device_type); + LOG(INFO) << "Using the unique target '" << primitive_targets.front()->str() + << "' of device type " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } else { + // Fallback. + default_primitive_device_type = kDLCPU; + LOG(WARNING) << "Using " << default_primitive_device_type + << " as the default device type for all primitive operations"; + } + + // + // Establish the default primitive SEScope, choosing a known Target to match the device type. + // + default_primitive_se_scope = se_scope_cache_.Unique( + SEScope(default_primitive_device_type, + /*virtual_device_id=*/0, FindPrimitiveTargetOrFail(default_primitive_device_type))); + + // + // Gather the hints as to what our default device type for the 'host' should be. + // + DLDeviceType host_device_type; + if (host_target.defined()) { + host_device_type = static_cast<DLDeviceType>(host_target->kind->device_type); + if (host_device_type != kDLCPU) { + LOG(WARNING) << "Using the given host target '" << host_target->str() + << "' of non-CPU device type " << host_device_type + << " for all host operations and data"; + } else { + LOG(INFO) << "Using the given host target '" << host_target->str() << "' of device type " + << host_device_type << " for all host operations and data"; + } + } else if (primitive_targets.size() == 1 && + primitive_targets.front()->kind->device_type == kDLCPU) { + // In the homogenous case without an explicit host target just use the given target so long as + // it's a CPU. + host_device_type = kDLCPU; + host_target = + FindPrimitiveTargetOrFail(host_device_type); // ie just primitive_targets.front()! + LOG(INFO) << "Using the unique target '" << host_target->str() << "' of device type " + << host_device_type << " for all host operations and data"; + } else { + // Fallback. + host_device_type = kDLCPU; + // Even if the list of available targets already includes one for kDLCPU we won't use it + // since its options may not be appropriate for host code (eg shape functions). Instead, + // create a fresh default Target. + host_target = MakeDefaultTarget(host_device_type); + LOG(WARNING) << "Using the default host target '" << host_target->str() << "' of device type " + << host_device_type << " for all host operations and data"; + } + + // + // Establish the host SEScope. + // + host_se_scope = se_scope_cache_.Unique(SEScope(host_device_type, + /*virtual_device_id=*/0, host_target)); +} + +/* static */ Target CompilationConfigNode::MakeDefaultTarget(DLDeviceType device_type) { + std::string name = runtime::DeviceName(device_type); + if (name == "cpu") { + if (runtime::Registry::Get("codegen.LLVMModuleCreate")) { Review comment: It may be better to use `TargetKindRegistry::Global()->ListAllNames();` or ` TargetKindRegEntry::ListTargetKinds()` as a slightly more robust way of checking this than looking for a specific function? Although I know this is copied so it's worked so far :smile_cat: ########## File path: src/target/se_scope.cc ########## @@ -0,0 +1,225 @@ +/* + * 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 tvm/target/se_scope.cc + * \brief Implementation of \p SEScope for representing a Storage or Execution scope. + */ +#include <tvm/node/reflection.h> +#include <tvm/runtime/device_api.h> +#include <tvm/target/se_scope.h> + +namespace tvm { + +TVM_REGISTER_NODE_TYPE(SEScopeNode); + +void SEScopeNode::VisitAttrs(AttrVisitor* v) { + int i = static_cast<int>(device_type_); + v->Visit("device_type", &i); + device_type_ = static_cast<DLDeviceType>(i); + v->Visit("virtual_device_id", &virtual_device_id_); + v->Visit("target", &target_); + v->Visit("memory_scope", &memory_scope_); +} + +bool SEScopeNode::SEqualReduce(const SEScopeNode* other, SEqualReducer equal) const { + return device_type_ == other->device_type_ && virtual_device_id_ == other->virtual_device_id_ && + // NOTE: Comparing targets by their str representations + target_->str() == other->target_->str() && memory_scope_ == other->memory_scope_; +} + +void SEScopeNode::SHashReduce(SHashReducer hash_reduce) const { + hash_reduce(device_type_); + hash_reduce(virtual_device_id_); + // NOTE: Reducing target to its str representation + hash_reduce(target_->str()); + hash_reduce(memory_scope_); +} + +TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) + .set_dispatch<SEScopeNode>([](const ObjectRef& ref, ReprPrinter* p) { + auto* node = ref.as<SEScopeNode>(); + p->stream << "SEScopeNode("; + if (node->is_fully_unconstrained()) { + p->stream << "?"; + } else { + bool need_sep = false; + if (node->device_type() != kInvalidDeviceType) { + p->stream << "device_type=" << node->device_type(); + need_sep = true; + } + if (node->virtual_device_id() >= 0) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "virtual_device_id=" << node->virtual_device_id(); + need_sep = true; + } + if (node->target().defined()) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "target='" << node->target()->str() << "'"; + need_sep = true; + } + if (!node->memory_scope().empty()) { + if (need_sep) { + p->stream << ", "; + } + p->stream << "memory_scope='" << node->memory_scope() << "'"; + } + } + p->stream << ")"; + }); + +SEScope::SEScope(DLDeviceType device_type, int virtual_device_id, Target target, + String memory_scope) { + ICHECK(!target.defined() || device_type == target->kind->device_type) + << "target '" << target->str() << "' has device type " << target->kind->device_type + << " but scope has device type " << device_type; + auto node = make_object<SEScopeNode>(); + node->device_type_ = device_type; + node->virtual_device_id_ = virtual_device_id; + node->target_ = std::move(target); + node->memory_scope_ = std::move(memory_scope); + data_ = std::move(node); +} + +/* static */ SEScope SEScope::FullyUnconstrained() { + static const SEScope unconstrained{}; + return unconstrained; +} + +/* static */ +Optional<SEScope> SEScope::Join(const SEScope& lhs, const SEScope& rhs) { + if (lhs == rhs) { + return lhs; + } + DLDeviceType joined_device_type; + if (lhs->device_type_ != kInvalidDeviceType) { + joined_device_type = lhs->device_type_; + if (rhs->device_type_ != kInvalidDeviceType && lhs->device_type_ != rhs->device_type_) { + return {}; + } + } else { + joined_device_type = rhs->device_type_; + } + int joined_virtual_device_id; + if (lhs->virtual_device_id_ >= 0) { + joined_virtual_device_id = lhs->virtual_device_id_; + if (rhs->virtual_device_id_ >= 0 && lhs->virtual_device_id_ != rhs->virtual_device_id_) { + return {}; + } + } else { + joined_virtual_device_id = rhs->virtual_device_id_; + } + Target joined_target; + if (lhs->target_.defined()) { + joined_target = lhs->target_; + if (rhs->target_.defined() && lhs->target_ != rhs->target_) { + return {}; + } + } else { + joined_target = rhs->target_; + } + String joined_memory_scope; + if (!lhs->memory_scope_.empty()) { + joined_memory_scope = lhs->memory_scope_; + if (!rhs->memory_scope_.empty() && lhs->memory_scope_ != rhs->memory_scope_) { + return {}; + } + } else { + joined_memory_scope = rhs->memory_scope_; + } + return SEScope(joined_device_type, joined_virtual_device_id, joined_target, joined_memory_scope); +} + +/* static */ +SEScope SEScope::Default(const SEScope& lhs, const SEScope& rhs) { + if (lhs == rhs) { + return lhs; + } + DLDeviceType defaulted_device_type; + if (lhs->device_type_ != kInvalidDeviceType) { + defaulted_device_type = lhs->device_type_; + } else { + defaulted_device_type = rhs->device_type_; + } + int defaulted_virtual_device_id; + if (lhs->virtual_device_id_ >= 0) { + defaulted_virtual_device_id = lhs->virtual_device_id_; + } else { + defaulted_virtual_device_id = rhs->virtual_device_id_; + } + Target defaulted_target; + if (lhs->target_.defined()) { + defaulted_target = lhs->target_; + } else { + // We can only default to the rhs's target if it is consistent with the device type + if (rhs->target_.defined() && rhs->target_->kind->device_type == defaulted_device_type) { + defaulted_target = rhs->target_; + } + // else: leave as null + } + String defaulted_memory_scope; + if (!lhs->memory_scope_.empty()) { + defaulted_memory_scope = lhs->memory_scope_; + } else { + defaulted_memory_scope = rhs->memory_scope_; + } + return SEScope(defaulted_device_type, defaulted_virtual_device_id, defaulted_target, + defaulted_memory_scope); +} + +SEScope SEScopeCache::Make(DLDeviceType device_type, int virtual_device_id, Target target, + String memory_scope) { + // Not the most efficient, but reducing the key to a string seems to be the simplest. Review comment: Realistically I think being able to properly index `Target` is the fix here, it's starting to become a bit of a pain. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
