gemini-code-assist[bot] commented on code in PR #19873: URL: https://github.com/apache/tvm/pull/19873#discussion_r3457897946
########## 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:  Instead of hardcoding `DataType::Int(8)`, use `new_buf->dtype` to determine the target backing storage type. This is more robust, self-documenting, and prevents potential type mismatches if the backing storage type is ever changed or parameterized in the future. ```suggestion if (load_returns_bool) { int lanes = op->dtype.lanes(); auto writer = load.CopyOnWrite(); writer->buffer = new_buf; writer->dtype = new_buf->dtype.with_lanes(lanes); return tvm::cast(DataType::Bool().with_lanes(lanes), load); } ``` ########## 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:  The `BoolBufferLegalizer` currently does not handle `SBlockNode`. Since `LowerBoolBuffer` runs immediately after `FlattenBuffer` (which explicitly handles `SBlockNode`), block-structured TIR may still be present. If a boolean buffer is allocated inside an `SBlock` (`alloc_buffers`) or referenced in its `reads`/`writes` regions, those references will not be remapped to `int8`, leading to a mismatch with the remapped loads/stores in the block body. Consider overriding `VisitStmt_(const SBlockNode* op)` and implementing a helper to mutate `BufferRegion`s to ensure all buffer references within blocks are consistently remapped. ```c BufferRegion MutateBufferRegion(BufferRegion region) { Buffer new_buf = GetRemappedBuffer(region->buffer); if (new_buf.same_as(region->buffer)) { return region; } auto writer = region.CopyOnWrite(); writer->buffer = new_buf; return region; } Stmt VisitStmt_(const SBlockNode* op) final { SBlock block = ffi::GetRef<SBlock>(op); ffi::Array<Buffer> alloc_buffers = op->alloc_buffers; alloc_buffers.MutateByApply([this](Buffer buf) { return GetRemappedBuffer(buf); }); if (!alloc_buffers.same_as(op->alloc_buffers)) { block.CopyOnWrite()->alloc_buffers = alloc_buffers; } ffi::Array<BufferRegion> reads = op->reads; reads.MutateByApply([this](BufferRegion region) { return MutateBufferRegion(region); }); if (!reads.same_as(op->reads)) { block.CopyOnWrite()->reads = reads; } ffi::Array<BufferRegion> writes = op->writes; writes.MutateByApply([this](BufferRegion region) { return MutateBufferRegion(region); }); if (!writes.same_as(op->writes)) { block.CopyOnWrite()->writes = writes; } return StmtExprMutator::VisitStmt_(block.get()); } Stmt VisitStmt_(const AllocBufferNode* op) final { ``` ########## 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:  Instead of hardcoding `DataType::Int(8)`, use `new_buf->dtype` to determine the target backing storage type. This is more robust, self-documenting, and prevents potential type mismatches if the backing storage type is ever changed or parameterized in the future. ```suggestion if (store_returns_bool) { writer->value = tvm::cast(new_buf->dtype.with_lanes(store->value.dtype().lanes()), store->value); } ``` -- 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]
