guan404ming commented on code in PR #19873: URL: https://github.com/apache/tvm/pull/19873#discussion_r3458060656
########## src/tirx/transform/lower_bool_buffer.cc: ########## @@ -0,0 +1,134 @@ +/* + * 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 lower_bool_buffer.cc + * \brief Replace boolean buffers with an int8 backing array. + */ +#include <tvm/ffi/cast.h> +#include <tvm/ffi/reflection/registry.h> +#include <tvm/tirx/op.h> +#include <tvm/tirx/stmt_functor.h> +#include <tvm/tirx/transform.h> + +#include <unordered_map> + +namespace tvm { +namespace tirx { + +/*! + * \brief Boolean tensors are stored in an int8 backing array. This pass rewrites bool-typed + * buffers to int8 and inserts the casts on the surrounding load/store, so the rest of + * lowering and codegen only ever sees int8 storage. Only the function body is rewritten, + * leaving buffer_map arguments boolean for argument validation in MakePackedAPI. + */ +class BoolBufferLegalizer : public StmtExprMutator { + public: + static PrimFunc Legalize(PrimFunc func) { + auto pass = BoolBufferLegalizer(); + auto* n = func.CopyOnWrite(); + n->body = pass.VisitStmt(std::move(n->body)); + return func; + } + + private: + Buffer GetRemappedBuffer(Buffer buf) { + auto it = buffer_remap_.find(buf); + if (it != buffer_remap_.end()) { + return it->second; + } + Buffer new_buf = buf; + if (buf->dtype.is_bool()) { + new_buf.CopyOnWrite()->dtype = DataType::Int(8).with_lanes(buf->dtype.lanes()); + } + buffer_remap_[buf] = new_buf; + return new_buf; + } + + Stmt VisitStmt_(const AllocBufferNode* op) final { Review Comment: Confirmed reachable via the `default` pipeline (no opaque-block lowering before `FlattenBuffer`). Added `VisitStmt_(SBlockNode)` + `MutateBufferRegion` mirroring `FlattenBuffer`, with a regression test. Done. ########## src/tirx/transform/lower_bool_buffer.cc: ########## @@ -0,0 +1,134 @@ +/* + * 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 lower_bool_buffer.cc + * \brief Replace boolean buffers with an int8 backing array. + */ +#include <tvm/ffi/cast.h> +#include <tvm/ffi/reflection/registry.h> +#include <tvm/tirx/op.h> +#include <tvm/tirx/stmt_functor.h> +#include <tvm/tirx/transform.h> + +#include <unordered_map> + +namespace tvm { +namespace tirx { + +/*! + * \brief Boolean tensors are stored in an int8 backing array. This pass rewrites bool-typed + * buffers to int8 and inserts the casts on the surrounding load/store, so the rest of + * lowering and codegen only ever sees int8 storage. Only the function body is rewritten, + * leaving buffer_map arguments boolean for argument validation in MakePackedAPI. + */ +class BoolBufferLegalizer : public StmtExprMutator { + public: + static PrimFunc Legalize(PrimFunc func) { + auto pass = BoolBufferLegalizer(); + auto* n = func.CopyOnWrite(); + n->body = pass.VisitStmt(std::move(n->body)); + return func; + } + + private: + Buffer GetRemappedBuffer(Buffer buf) { + auto it = buffer_remap_.find(buf); + if (it != buffer_remap_.end()) { + return it->second; + } + Buffer new_buf = buf; + if (buf->dtype.is_bool()) { + new_buf.CopyOnWrite()->dtype = DataType::Int(8).with_lanes(buf->dtype.lanes()); + } + buffer_remap_[buf] = new_buf; + return new_buf; + } + + Stmt VisitStmt_(const AllocBufferNode* op) final { + auto node = StmtExprMutator::VisitStmt_(op).as_or_throw<AllocBuffer>(); + Buffer new_buf = GetRemappedBuffer(node->buffer); + if (!new_buf.same_as(node->buffer)) { + node.CopyOnWrite()->buffer = new_buf; + } + return std::move(node); + } + + Stmt VisitStmt_(const DeclBufferNode* op) final { + auto node = StmtExprMutator::VisitStmt_(op).as_or_throw<DeclBuffer>(); + Buffer new_buf = GetRemappedBuffer(node->buffer); + if (!new_buf.same_as(node->buffer)) { + node.CopyOnWrite()->buffer = new_buf; + } + return std::move(node); + } + + Stmt VisitStmt_(const BufferStoreNode* op) final { + BufferStore store = StmtExprMutator::VisitStmt_(op).as_or_throw<BufferStore>(); + bool store_returns_bool = op->value.dtype().is_bool(); + Buffer new_buf = GetRemappedBuffer(store->buffer); + if (new_buf.same_as(store->buffer) && !store_returns_bool) { + return std::move(store); + } + auto writer = store.CopyOnWrite(); + writer->buffer = new_buf; + if (store_returns_bool) { + writer->value = + tvm::cast(DataType::Int(8).with_lanes(store->value.dtype().lanes()), store->value); + } Review Comment: Done. Ties the store cast to the remapped buffer's storage type instead of a duplicated literal. ########## src/tirx/transform/lower_bool_buffer.cc: ########## @@ -0,0 +1,134 @@ +/* + * 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 lower_bool_buffer.cc + * \brief Replace boolean buffers with an int8 backing array. + */ +#include <tvm/ffi/cast.h> +#include <tvm/ffi/reflection/registry.h> +#include <tvm/tirx/op.h> +#include <tvm/tirx/stmt_functor.h> +#include <tvm/tirx/transform.h> + +#include <unordered_map> + +namespace tvm { +namespace tirx { + +/*! + * \brief Boolean tensors are stored in an int8 backing array. This pass rewrites bool-typed + * buffers to int8 and inserts the casts on the surrounding load/store, so the rest of + * lowering and codegen only ever sees int8 storage. Only the function body is rewritten, + * leaving buffer_map arguments boolean for argument validation in MakePackedAPI. + */ +class BoolBufferLegalizer : public StmtExprMutator { + public: + static PrimFunc Legalize(PrimFunc func) { + auto pass = BoolBufferLegalizer(); + auto* n = func.CopyOnWrite(); + n->body = pass.VisitStmt(std::move(n->body)); + return func; + } + + private: + Buffer GetRemappedBuffer(Buffer buf) { + auto it = buffer_remap_.find(buf); + if (it != buffer_remap_.end()) { + return it->second; + } + Buffer new_buf = buf; + if (buf->dtype.is_bool()) { + new_buf.CopyOnWrite()->dtype = DataType::Int(8).with_lanes(buf->dtype.lanes()); + } + buffer_remap_[buf] = new_buf; + return new_buf; + } + + Stmt VisitStmt_(const AllocBufferNode* op) final { + auto node = StmtExprMutator::VisitStmt_(op).as_or_throw<AllocBuffer>(); + Buffer new_buf = GetRemappedBuffer(node->buffer); + if (!new_buf.same_as(node->buffer)) { + node.CopyOnWrite()->buffer = new_buf; + } + return std::move(node); + } + + Stmt VisitStmt_(const DeclBufferNode* op) final { + auto node = StmtExprMutator::VisitStmt_(op).as_or_throw<DeclBuffer>(); + Buffer new_buf = GetRemappedBuffer(node->buffer); + if (!new_buf.same_as(node->buffer)) { + node.CopyOnWrite()->buffer = new_buf; + } + return std::move(node); + } + + Stmt VisitStmt_(const BufferStoreNode* op) final { + BufferStore store = StmtExprMutator::VisitStmt_(op).as_or_throw<BufferStore>(); + bool store_returns_bool = op->value.dtype().is_bool(); + Buffer new_buf = GetRemappedBuffer(store->buffer); + if (new_buf.same_as(store->buffer) && !store_returns_bool) { + return std::move(store); + } + auto writer = store.CopyOnWrite(); + writer->buffer = new_buf; + if (store_returns_bool) { + writer->value = + tvm::cast(DataType::Int(8).with_lanes(store->value.dtype().lanes()), store->value); + } + return std::move(store); + } + + PrimExpr VisitExpr_(const BufferLoadNode* op) final { + bool load_returns_bool = op->dtype.is_bool(); + BufferLoad load = StmtExprMutator::VisitExpr_(op).as_or_throw<BufferLoad>(); + Buffer new_buf = GetRemappedBuffer(load->buffer); + if (load_returns_bool) { + int lanes = op->dtype.lanes(); + auto writer = load.CopyOnWrite(); + writer->buffer = new_buf; + writer->dtype = DataType::Int(8).with_lanes(lanes); + return tvm::cast(DataType::Bool().with_lanes(lanes), load); + } Review Comment: Done, same as above for the load side. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
