|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, |
| 13 | + * software distributed under the License is distributed on an |
| 14 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | + * KIND, either express or implied. See the License for the |
| 16 | + * specific language governing permissions and limitations |
| 17 | + * under the License. |
| 18 | + */ |
| 19 | + |
| 20 | +#include <tvm/arith/analyzer.h> |
| 21 | +#include <tvm/arith/iter_affine_map.h> |
| 22 | +#include <tvm/runtime/registry.h> |
| 23 | +#include <tvm/tir/analysis.h> |
| 24 | +#include <tvm/tir/op.h> |
| 25 | +#include <tvm/tir/stmt.h> |
| 26 | +#include <tvm/tir/stmt_functor.h> |
| 27 | +#include <tvm/tir/transform.h> |
| 28 | + |
| 29 | +#include "../../arith/const_fold.h" |
| 30 | +#include "../../arith/pattern_match.h" |
| 31 | + |
| 32 | +namespace tvm { |
| 33 | +namespace tir { |
| 34 | + |
| 35 | +class PTXRewriter : public StmtMutator { |
| 36 | + public: |
| 37 | + Stmt VisitStmt_(const AllocateNode* allocate) final { |
| 38 | + if (!has_buffer_1) { |
| 39 | + has_buffer_1 = true; |
| 40 | + // addr[0] -> global_addr / addr[1] -> local_addr |
| 41 | + addr_buffer = decl_buffer({IntImm(DataType::Int(32), 2)}, DataType::Int(32), "addr", "local"); |
| 42 | + predicate_buffer = |
| 43 | + decl_buffer({IntImm(DataType::Int(32), 1)}, DataType::Bool(1), "predicate", "local"); |
| 44 | + } |
| 45 | + Stmt result = StmtMutator::VisitStmt_(allocate); |
| 46 | + if (!has_buffer_2) { |
| 47 | + has_buffer_2 = true; |
| 48 | + result = |
| 49 | + Allocate(addr_buffer->data, addr_buffer->dtype, addr_buffer->shape, Bool(true), result); |
| 50 | + result = Allocate(predicate_buffer->data, predicate_buffer->dtype, predicate_buffer->shape, |
| 51 | + Bool(true), result); |
| 52 | + } |
| 53 | + return result; |
| 54 | + } |
| 55 | + |
| 56 | + Stmt VisitStmt_(const BufferStoreNode* store) final { |
| 57 | + Stmt result = StmtMutator::VisitStmt_(store); |
| 58 | + Buffer load_buffer = store->buffer; |
| 59 | + PrimExpr load_value = store->value; |
| 60 | + // const BufferLoadNode* gload = load_value.as<BufferLoadNode>(); // take |
| 61 | + // the place of instance of |
| 62 | + const CallNode* call = load_value.as<CallNode>(); |
| 63 | + if (call != nullptr) { |
| 64 | + const OpNode* op = call->op.as<OpNode>(); |
| 65 | + if (op != nullptr && op->name == "tir.if_then_else") { |
| 66 | + const PrimExpr& predicate = call->args[0]; |
| 67 | + const PrimExpr& lhs = call->args[1]; |
| 68 | + const PrimExpr& rhs = call->args[2]; |
| 69 | + PrimExpr global_addr, local_addr; |
| 70 | + const BufferLoadNode* load = lhs.as<BufferLoadNode>(); |
| 71 | + PrimExpr imm_value = rhs; |
| 72 | + if (load == nullptr) { |
| 73 | + load = rhs.as<BufferLoadNode>(); |
| 74 | + imm_value = lhs; |
| 75 | + if (load == nullptr) { |
| 76 | + return result; |
| 77 | + } |
| 78 | + } |
| 79 | + global_addr = load->indices[0]; |
| 80 | + const RampNode* ramp = global_addr.as<RampNode>(); |
| 81 | + if (ramp != nullptr) { |
| 82 | + return result; |
| 83 | + } |
| 84 | + local_addr = store->indices[0]; |
| 85 | + BufferStore addr_store(addr_buffer, global_addr, {IntImm(DataType::Int(32), 0)}); |
| 86 | + BufferStore local_addr_store(addr_buffer, local_addr, {IntImm(DataType::Int(32), 1)}); |
| 87 | + BufferStore predicate_store(predicate_buffer, predicate, {IntImm(DataType::Int(32), 0)}); |
| 88 | + PrimExpr new_lhs, new_rhs, new_predicate, new_indice; |
| 89 | + new_lhs = |
| 90 | + BufferLoad(load->buffer, {BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 0)})}); |
| 91 | + new_rhs = IntImm(DataType::Int(32), 0); |
| 92 | + new_predicate = BufferLoad(predicate_buffer, {IntImm(DataType::Int(32), 0)}); |
| 93 | + new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); |
| 94 | + BufferStore value_store(store->buffer, imm_value, {new_indice}); |
| 95 | + Evaluate ptx_load(Call(store->buffer->dtype, tvm::tir::builtin::ptx_ldg32(), |
| 96 | + {store->buffer->data, new_predicate, new_lhs, new_indice})); |
| 97 | + Array<Stmt> tmp_seq = {addr_store, local_addr_store, predicate_store, value_store, |
| 98 | + ptx_load}; |
| 99 | + SeqStmt seq_stmt = SeqStmt(tmp_seq); |
| 100 | + return seq_stmt; |
| 101 | + } |
| 102 | + } |
| 103 | + return result; |
| 104 | + } |
| 105 | + |
| 106 | + bool has_buffer_1 = false, has_buffer_2 = false; |
| 107 | + Buffer addr_buffer, predicate_buffer; |
| 108 | +}; |
| 109 | + |
| 110 | +namespace transform { |
| 111 | + |
| 112 | +Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) { |
| 113 | + auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m, PassContext ctx) { |
| 114 | + if (enable_inject_ptx_intrin) { |
| 115 | + auto* n = f.CopyOnWrite(); |
| 116 | + n->body = PTXRewriter()(n->body); |
| 117 | + // inject ptx |
| 118 | + } |
| 119 | + return f; |
| 120 | + }; |
| 121 | + return CreatePrimFuncPass(pass_func, 0, "tir.InjectPTXLDG32", {}); |
| 122 | +} |
| 123 | + |
| 124 | +// The pass can now be invoked via the pass infrastructure, but we also add a |
| 125 | +// Python binding for it |
| 126 | +TVM_REGISTER_GLOBAL("tir.transform.InjectPTXLDG32").set_body_typed(InjectPTXLDG32); |
| 127 | + |
| 128 | +} // namespace transform |
| 129 | +} // namespace tir |
| 130 | +} // namespace tvm |
0 commit comments