Hzfengsy commented on code in PR #14673: URL: https://github.com/apache/tvm/pull/14673#discussion_r1241582434
########## src/tir/transforms/inject_permuted_layout.cc: ########## @@ -0,0 +1,266 @@ +/* + * 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 inject_permuted_layout.cc + * \brief The pass for inject permuted layout. + */ + +#include <tvm/arith/analyzer.h> +#include <tvm/tir/function.h> +#include <tvm/tir/op.h> +#include <tvm/tir/stmt_functor.h> +#include <tvm/tir/transform.h> + +#include "../../support/utils.h" +#include "../ir/functor_common.h" +#include "ir_utils.h" + +namespace tvm { +namespace tir { + +using tir::Block; +using tir::BlockRealize; +using tir::Call; +using tir::For; + +class PermutedLayoutInjector : public StmtExprMutator { + public: + PermutedLayoutInjector() {} + + private: + Array<PrimExpr> GetNewIndices(PrimExpr s0, PrimExpr s1, int smem_width) { + // index after vectorize(8) + PrimExpr i = s0, j = floordiv(s1, 8), v = floormod(s1, 8); + PrimExpr permuted_j; + // In the following comments, each number represent a 8 * fp16 load + // which is correspond to a index (i, j) in line 50's PrimExpr + // Each 8 number correspond to 32 memory bank (every bank has 32 bit): + // 8 * 8 * 16bit = 32 * 32bit + // And we have 32 banks in total, so all loads in one column share + // same memory bank + if (smem_width % 64 == 0) { + // use 8 * 8 permuted + // 0 1 2 3 4 5 6 7 ==> 0 1 2 3 4 5 6 7 + // 0 1 2 3 4 5 6 7 ==> 1 0 3 2 5 4 7 6 + // 0 1 2 3 4 5 6 7 ==> 2 3 0 1 6 7 4 5 + // 0 1 2 3 4 5 6 7 ==> 3 2 1 0 7 6 5 4 + // 0 1 2 3 4 5 6 7 ==> 4 5 6 7 0 1 2 3 + // 0 1 2 3 4 5 6 7 ==> 5 4 7 6 1 0 3 2 + // 0 1 2 3 4 5 6 7 ==> 6 7 4 5 2 3 0 1 + // 0 1 2 3 4 5 6 7 ==> 7 6 5 4 3 2 1 0 + PrimExpr permuted_j_mod_8 = (floormod(j, 8) ^ floormod(i, 8)); + permuted_j = floordiv(j, 8) * 8 + permuted_j_mod_8; + } else { + // use 8 * 4 permuted + // 0 1 2 3 ==> 0 1 2 3 + // 0 1 2 3 ==> 0 1 2 3 + // 0 1 2 3 ==> 1 0 3 2 + // 0 1 2 3 ==> 1 0 3 2 + // 0 1 2 3 ==> 2 3 0 1 + // 0 1 2 3 ==> 2 3 0 1 + // 0 1 2 3 ==> 3 2 1 0 + // 0 1 2 3 ==> 3 2 1 0 + // in 8 number each line view: + // 0 1 2 3 4 0 1 2 3 ==> 0 1 2 3 0 1 2 3 + // 0 1 2 3 4 0 1 2 3 ==> 1 0 3 2 1 0 3 2 + // 0 1 2 3 4 0 1 2 3 ==> 2 3 0 1 2 3 0 1 + // 0 1 2 3 4 0 1 2 3 ==> 3 2 1 0 3 2 1 0 + permuted_j = floormod(j, 4) ^ floordiv(floormod(i, 8), 2); + } + return {s0, permuted_j * 8 + v}; + } + + Stmt VisitStmt_(const BlockRealizeNode* _op) final { + BlockRealize br = Downcast<BlockRealize>(StmtExprMutator::VisitStmt_(_op)); + BlockRealizeNode* op = br.CopyOnWrite(); + if (op->block->annotations.count("permuted_layout") == 0) { + return br; + } + String val = Downcast<String>(op->block->annotations.at("permuted_layout")); + if (val.empty()) return br; + Block blk = op->block; + Stmt body = blk->body; + if (support::StartsWith(val, "g2s")) { + // Case 1. Rewrite global to share.dyn + + // Step 1.1. Handle case when have local stage + // Block with local stage is like + // body { + // SeqStmt { + // seq[0]: local <- global + // seq[1]: shared.dyn <- local + // } + // } + // We only need to rewrite seq[1] + bool have_local_stage = (body.as<SeqStmtNode>() != nullptr); + Stmt upper_loop; + if (have_local_stage) { + SeqStmt seq = Downcast<SeqStmt>(body); + ICHECK(seq->size() == 2); + upper_loop = seq->seq[0]; + body = seq->seq[1]; + } + + // Step 1.2. get inner loop body + std::vector<const ForNode*> loops; + while (const ForNode* loop = body.as<ForNode>()) { + loops.push_back(loop); + body = loop->body; + } + Optional<PrimExpr> if_then_else_condition = NullOpt; + const BufferStoreNode* store = body.as<BufferStoreNode>(); + if (!store) { + // Case 1.2.1. IfThenElse generated by reverse_compute_inline + // It is always like + // if condition: + // loop_body + // We just extract the inner loop body inside IfThenElseNode + const IfThenElseNode* if_then_else = body.as<IfThenElseNode>(); + store = if_then_else->then_case.as<BufferStoreNode>(); + ICHECK(!if_then_else->else_case); + if_then_else_condition = if_then_else->condition; + } + ICHECK(store) << body; + + // Step 1.3. Get smem width and refuse to make any difference if invalid + auto smem_width = store->buffer->shape[1].as<IntImmNode>()->value; + if (smem_width % 32 != 0) { + LOG(INFO) << "Permuted Layout for " << op->block->name_hint + << " is not supported since its second dimension is not divisible by 32"; + return br; + } + if (smem_width % 64 == 32) { + if (store->buffer->shape[0].as<IntImmNode>()->value % 2 != 0) { + LOG(INFO) << "Permuted Layout for " << op->block->name_hint Review Comment: ```suggestion LOG(WARNING) << "Permuted Layout for " << op->block->name_hint ``` ########## src/tir/transforms/inject_permuted_layout.cc: ########## @@ -0,0 +1,266 @@ +/* + * 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 inject_permuted_layout.cc + * \brief The pass for inject permuted layout. + */ + +#include <tvm/arith/analyzer.h> +#include <tvm/tir/function.h> +#include <tvm/tir/op.h> +#include <tvm/tir/stmt_functor.h> +#include <tvm/tir/transform.h> + +#include "../../support/utils.h" +#include "../ir/functor_common.h" +#include "ir_utils.h" + +namespace tvm { +namespace tir { + +using tir::Block; +using tir::BlockRealize; +using tir::Call; +using tir::For; + +class PermutedLayoutInjector : public StmtExprMutator { + public: + PermutedLayoutInjector() {} + + private: + Array<PrimExpr> GetNewIndices(PrimExpr s0, PrimExpr s1, int smem_width) { + // index after vectorize(8) + PrimExpr i = s0, j = floordiv(s1, 8), v = floormod(s1, 8); + PrimExpr permuted_j; + // In the following comments, each number represent a 8 * fp16 load + // which is correspond to a index (i, j) in line 50's PrimExpr + // Each 8 number correspond to 32 memory bank (every bank has 32 bit): + // 8 * 8 * 16bit = 32 * 32bit + // And we have 32 banks in total, so all loads in one column share + // same memory bank + if (smem_width % 64 == 0) { + // use 8 * 8 permuted + // 0 1 2 3 4 5 6 7 ==> 0 1 2 3 4 5 6 7 + // 0 1 2 3 4 5 6 7 ==> 1 0 3 2 5 4 7 6 + // 0 1 2 3 4 5 6 7 ==> 2 3 0 1 6 7 4 5 + // 0 1 2 3 4 5 6 7 ==> 3 2 1 0 7 6 5 4 + // 0 1 2 3 4 5 6 7 ==> 4 5 6 7 0 1 2 3 + // 0 1 2 3 4 5 6 7 ==> 5 4 7 6 1 0 3 2 + // 0 1 2 3 4 5 6 7 ==> 6 7 4 5 2 3 0 1 + // 0 1 2 3 4 5 6 7 ==> 7 6 5 4 3 2 1 0 + PrimExpr permuted_j_mod_8 = (floormod(j, 8) ^ floormod(i, 8)); + permuted_j = floordiv(j, 8) * 8 + permuted_j_mod_8; + } else { + // use 8 * 4 permuted + // 0 1 2 3 ==> 0 1 2 3 + // 0 1 2 3 ==> 0 1 2 3 + // 0 1 2 3 ==> 1 0 3 2 + // 0 1 2 3 ==> 1 0 3 2 + // 0 1 2 3 ==> 2 3 0 1 + // 0 1 2 3 ==> 2 3 0 1 + // 0 1 2 3 ==> 3 2 1 0 + // 0 1 2 3 ==> 3 2 1 0 + // in 8 number each line view: + // 0 1 2 3 4 0 1 2 3 ==> 0 1 2 3 0 1 2 3 + // 0 1 2 3 4 0 1 2 3 ==> 1 0 3 2 1 0 3 2 + // 0 1 2 3 4 0 1 2 3 ==> 2 3 0 1 2 3 0 1 + // 0 1 2 3 4 0 1 2 3 ==> 3 2 1 0 3 2 1 0 + permuted_j = floormod(j, 4) ^ floordiv(floormod(i, 8), 2); + } + return {s0, permuted_j * 8 + v}; + } + + Stmt VisitStmt_(const BlockRealizeNode* _op) final { + BlockRealize br = Downcast<BlockRealize>(StmtExprMutator::VisitStmt_(_op)); + BlockRealizeNode* op = br.CopyOnWrite(); + if (op->block->annotations.count("permuted_layout") == 0) { + return br; + } + String val = Downcast<String>(op->block->annotations.at("permuted_layout")); + if (val.empty()) return br; + Block blk = op->block; + Stmt body = blk->body; + if (support::StartsWith(val, "g2s")) { + // Case 1. Rewrite global to share.dyn + + // Step 1.1. Handle case when have local stage + // Block with local stage is like + // body { + // SeqStmt { + // seq[0]: local <- global + // seq[1]: shared.dyn <- local + // } + // } + // We only need to rewrite seq[1] + bool have_local_stage = (body.as<SeqStmtNode>() != nullptr); + Stmt upper_loop; + if (have_local_stage) { + SeqStmt seq = Downcast<SeqStmt>(body); + ICHECK(seq->size() == 2); + upper_loop = seq->seq[0]; + body = seq->seq[1]; + } + + // Step 1.2. get inner loop body + std::vector<const ForNode*> loops; + while (const ForNode* loop = body.as<ForNode>()) { + loops.push_back(loop); + body = loop->body; + } + Optional<PrimExpr> if_then_else_condition = NullOpt; + const BufferStoreNode* store = body.as<BufferStoreNode>(); + if (!store) { + // Case 1.2.1. IfThenElse generated by reverse_compute_inline + // It is always like + // if condition: + // loop_body + // We just extract the inner loop body inside IfThenElseNode + const IfThenElseNode* if_then_else = body.as<IfThenElseNode>(); + store = if_then_else->then_case.as<BufferStoreNode>(); + ICHECK(!if_then_else->else_case); + if_then_else_condition = if_then_else->condition; + } + ICHECK(store) << body; + + // Step 1.3. Get smem width and refuse to make any difference if invalid + auto smem_width = store->buffer->shape[1].as<IntImmNode>()->value; + if (smem_width % 32 != 0) { + LOG(INFO) << "Permuted Layout for " << op->block->name_hint Review Comment: ```suggestion LOG(WARNING) << "Permuted Layout for " << op->block->name_hint ``` -- 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]
