From 0850d2b4f4afff0231a5db258e831602e8cd00c0 Mon Sep 17 00:00:00 2001 From: andy_yang Date: Mon, 13 Feb 2023 19:35:34 +0800 Subject: [PATCH 01/10] add ptx load global 32bit --- include/tvm/tir/builtin.h | 7 + include/tvm/tir/transform.h | 6 + src/driver/driver_api.cc | 12 ++ src/target/source/codegen_cuda.cc | 41 +++++ src/tir/op/builtin.cc | 3 + src/tir/transforms/inject_ptx_ldg32.cc | 160 ++++++++++++++++++ .../python/unittest/test_inject_ptx_ldg32.py | 70 ++++++++ 7 files changed, 299 insertions(+) create mode 100644 src/tir/transforms/inject_ptx_ldg32.cc create mode 100644 tests/python/unittest/test_inject_ptx_ldg32.py diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h index d830ea579aa7..e984c7e2ab6d 100644 --- a/include/tvm/tir/builtin.h +++ b/include/tvm/tir/builtin.h @@ -610,6 +610,13 @@ TVM_DLL const Op& tvm_store_matrix_sync(); */ TVM_DLL const Op& ptx_mma(); + +/*! + * \brief tvm intrinsic for ptx predicate load with 32-bit data type. + * + */ +TVM_DLL const Op& ptx_pred_ldg32(); + /*! * \brief tvm intrinsic for sparse tensor core ptx instructions. * diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h index be7589b04bf5..6d2a9e18da68 100644 --- a/include/tvm/tir/transform.h +++ b/include/tvm/tir/transform.h @@ -677,6 +677,12 @@ TVM_DLL Pass Filter(runtime::TypedPackedFunc fcond); */ TVM_DLL Pass InjectPTXAsyncCopy(); +/*! + * \brief Pass to rewrite global to local memory copy on CUDA with ldg32 instruction. + * \return The pass. + */ +TVM_DLL Pass InjectPTXLDG32(bool enable_inject_ptx_ldg32 = true); + /*! * \brief Remove the weight layout rewrite block * \param skip_ndarray_rewrite If True, exact rewrite of NDArray, according to the given index map, diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc index 288ac7b92a2c..a8349c67bc33 100644 --- a/src/driver/driver_api.cc +++ b/src/driver/driver_api.cc @@ -55,6 +55,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.use_async_copy", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.merge_async_commit_queue_scope", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.instrument_lwp", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.vtcm_capacity", Integer); +TVM_REGISTER_PASS_CONFIG_OPTION("tir.ptx_pred_ldg32", Bool); // WARNING: May cause coherency issues resulting data miscompares // Experimental feature that, when enabled by the runtime, bypasses the cache when using DMA. When @@ -159,6 +160,8 @@ Array CreatePassList(bool disable_loop_partition) { bool enable_equiv_terms_in_cse_tir = pass_ctx->GetConfig("tir.enable_equiv_terms_in_cse_tir", Bool(false)).value(); + bool ptx_pred_ldg32 = pass_ctx->GetConfig("tir.ptx_pred_ldg32", Bool(false)).value(); + // Get any user-added passes Array> add_lower_pass = pass_ctx->GetConfig>>("tir.add_lower_pass", Array>()) @@ -257,6 +260,10 @@ Array CreatePassList(bool disable_loop_partition) { pass_list.push_back(tir::transform::InstrumentBoundCheckers()); } + if(ptx_pred_ldg32){ + pass_list.push_back(tir::transform::InjectPTXLDG32(true)); + } + pass_list.push_back( tir::transform::CommonSubexprElimTIR(!disable_cse_tir, enable_equiv_terms_in_cse_tir)); @@ -584,6 +591,11 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target) mixed_pass_list.push_back(tir::transform::InjectPTXAsyncCopy()); } + bool ptx_pred_ldg32 = pass_ctx->GetConfig("tir.ptx_pred_ldg32", Bool(false)).value(); + if (ptx_pred_ldg32) { + mixed_pass_list.push_back(tir::transform::InjectPTXLDG32()); + } + bool unpacked_api = mixed_mod->GetAttr(tvm::attr::kExecutor) .value_or(relay::Executor::Create("graph", {})) ->GetAttr("unpacked-api") diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc index c891ec5a28cf..3b964fb8f08b 100644 --- a/src/target/source/codegen_cuda.cc +++ b/src/target/source/codegen_cuda.cc @@ -920,6 +920,47 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { } else if (op->op.same_as(builtin::ptx_wait_group())) { std::string N = this->PrintExpr(op->args[0]); this->stream << "__asm__ __volatile__(\"cp.async.wait_group " + N + ";\");\n\n"; + } else if (op->op.same_as(builtin::ptx_pred_ldg32())){ + /* + asm volatile ( + "{.reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + // " @p ld.global.nc.f32 %0, [%1];}\n"t + " @p ld.global.nc.L2::128B.f32 %0, [%1];}\n" + : "=f"(reg) + : "l"(addr), "r"((int)guard) + ); + */ + + // get local + std::string reg = this->PrintExpr(op->args[0]); + // get guard + std::string guard = this->PrintExpr(op->args[1]); + // std::string lhs = this->PrintExpr(op->args[2]); + const BufferLoadNode* addr_buffer = op->args[2].as(); + + std::string global_addr = this->PrintExpr(addr_buffer->indices[0]); + std::string global_buffer = this->PrintExpr(addr_buffer->buffer->data); + + std::string local_addr = this->PrintExpr(op->args[3]); + + + this->stream << "asm volatile (\n" ; + this->PrintIndent(); + stream << "\"{.reg .pred p;\\n\"\n" ; + this->PrintIndent(); + stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n" ; + this->PrintIndent(); + stream << "\" @!p mov.b32 %0, 0;\\n\"\n"; + this->PrintIndent(); + stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n" ; + // stream << "\" @p ld.global.nc.L2::128B.f32 %0, [%1];}\\n\"\n" ; + this->PrintIndent(); + stream << ": \"=f\"(" << reg << "[" << local_addr << "]" << ")\n" ; + this->PrintIndent(); + stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr << ")), \"r\"((int)" << guard << ")\n" ; + this->PrintIndent(); + stream << ");\n" ; } else { CodeGenC::VisitExpr_(op, os); } diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc index dc3208f484e3..fdcb687b36cd 100644 --- a/src/tir/op/builtin.cc +++ b/src/tir/op/builtin.cc @@ -251,6 +251,9 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync) TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); +TIR_DEFINE_BUILTIN_FUNC(ptx_pred_ldg32) + .set_num_inputs(4) + .set_attr("TCallEffectKind", Integer(CallEffectKind::kPure)); TIR_DEFINE_BUILTIN_FUNC(ptx_mma_sp) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); diff --git a/src/tir/transforms/inject_ptx_ldg32.cc b/src/tir/transforms/inject_ptx_ldg32.cc new file mode 100644 index 000000000000..4597d477e848 --- /dev/null +++ b/src/tir/transforms/inject_ptx_ldg32.cc @@ -0,0 +1,160 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../../arith/const_fold.h" +#include "../../arith/pattern_match.h" + +namespace tvm { +namespace tir { + + + +class PTXRewriter : public StmtMutator{ + +public: + explicit PTXRewriter() {} + + Stmt VisitStmt_(const AllocateNode* allocate) final{ + + if (!has_buffer_1){ + has_buffer_1 = true; + // addr[0] -> global_addr / addr[1] -> local_addr + addr_buffer = decl_buffer({IntImm(DataType::Int(32), 2)}, DataType::Int(32), "addr", "local"); + predicate_buffer = decl_buffer({IntImm(DataType::Int(32), 1)}, DataType::Bool(1), "predicate", "local"); + } + + Stmt result = StmtMutator::VisitStmt_(allocate); + if (!has_buffer_2){ + has_buffer_2 = true; + result = Allocate(addr_buffer->data, addr_buffer->dtype, addr_buffer->shape, Bool(true), result); + result = Allocate(predicate_buffer->data, predicate_buffer->dtype, predicate_buffer->shape, Bool(true), result); + } + return result; + } + + + Stmt VisitStmt_(const BufferStoreNode* store) final{ + + Stmt result = StmtMutator::VisitStmt_(store); + Buffer load_buffer = store->buffer; + PrimExpr load_value = store->value; + + + const CallNode* call = load_value.as(); + + if (call != nullptr){ + + const OpNode* op = call->op.as(); + if (op != nullptr && op->name == "tir.if_then_else"){ + const PrimExpr& predicate = call->args[0]; + const PrimExpr& lhs = call->args[1]; + const PrimExpr& rhs = call->args[2]; + + PrimExpr global_addr, local_addr; + const BufferLoadNode* load = lhs.as(); + + + if (load == nullptr){ + load = rhs.as(); + } + + for ( auto each : load->indices){ + + } + global_addr = load->indices[0]; + const RampNode* ramp = global_addr.as(); + if (ramp != nullptr){ + return result; + + } + + local_addr = store->indices[0]; + + + BufferStore addr_store(addr_buffer, global_addr, {IntImm(DataType::Int(32), 0)}); + BufferStore local_addr_store(addr_buffer,local_addr,{IntImm(DataType::Int(32), 1)}); + BufferStore predicate_store(predicate_buffer, predicate, {IntImm(DataType::Int(32), 0)}); + + PrimExpr new_lhs, new_rhs, new_predicate, new_indice; + + new_lhs = BufferLoad(load->buffer, {BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 0)})}); + new_rhs = IntImm(DataType::Int(32), 0); + new_predicate = BufferLoad(predicate_buffer, {IntImm(DataType::Int(32), 0)}); + new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); + + + Evaluate ptx_load(Call(store->buffer->dtype, tvm::tir::builtin::ptx_pred_ldg32(),{store->buffer->data, new_predicate , new_lhs,new_indice})); + + Array tmp_seq = {addr_store, local_addr_store, predicate_store, ptx_load}; + + SeqStmt seq_stmt = SeqStmt(tmp_seq); + + return seq_stmt; + + } + } + + + return result; + + } + + + bool has_buffer_1 = false, has_buffer_2 = false ; + Buffer addr_buffer , predicate_buffer ; + + +}; + + +namespace transform{ + + +Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) { + auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m, PassContext ctx) { + if (enable_inject_ptx_intrin) { + auto* n = f.CopyOnWrite(); + n->body = PTXRewriter()(n->body); + // inject ptx + } + return f; + }; + return CreatePrimFuncPass(pass_func, 0, "tir.InjectPTXLDG32", {}); +} + +// The pass can now be invoked via the pass infrastructure, but we also add a Python binding for it +TVM_REGISTER_GLOBAL("tir.transform.InjectPTXLDG32").set_body_typed(InjectPTXLDG32); + + + +} // namespace transform +} // namespace tir +} // namespace tvm + + + diff --git a/tests/python/unittest/test_inject_ptx_ldg32.py b/tests/python/unittest/test_inject_ptx_ldg32.py new file mode 100644 index 000000000000..f20032c164eb --- /dev/null +++ b/tests/python/unittest/test_inject_ptx_ldg32.py @@ -0,0 +1,70 @@ +# 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. +import tvm +from tvm.script import tir as T +import numpy as np +import tvm.testing + + +@T.prim_func +def vector_add(A: T.Buffer[(16), "float32"], B: T.Buffer[(32), "float32"]) -> None: + T.func_attr({"global_symbol": "default_function", "tir.noalias": True}) + bx = T.env_thread("blockIdx.x") + tx = T.env_thread("threadIdx.x") + T.launch_thread(bx, 1) + T.launch_thread(tx, 32) + A_local = T.alloc_buffer([32], "float32", scope="local") + + with T.block(): + T.reads(A[0:16]) + T.writes(A_local[0:32]) + A_local[tx] = T.if_then_else(tx%2==0, A[tx/2], T.float32(0), dtype="float32") + B[tx] = A_local[tx] + 1.0 + + + + +@tvm.testing.requires_cuda +def test_inject_ptx_intrin(): + f = vector_add + arch = tvm.contrib.nvcc.get_target_compute_version() + major, _ = tvm.contrib.nvcc.parse_compute_version(arch) + if major < 8: + # Require at least SM80 + return + with tvm.transform.PassContext(config={"tir.ptx_pred_ldg32": True}): + mod = tvm.build(f, target="cuda") + A_np = np.random.rand(16).astype("float32") + B_np = np.zeros((32)).astype("float32") + dev = tvm.cuda(0) + A_nd = tvm.nd.array(A_np, device=dev) + B_nd = tvm.nd.array(B_np, device=dev) + mod(A_nd, B_nd) + + C_np = np.zeros((32)).astype("float32") + + + for i in range(32): + if i % 2 == 0: + C_np[i] = A_np[i//2] + C_np[i] += 1.0 + + tvm.testing.assert_allclose(B_nd.numpy(), C_np) + + +if __name__ == "__main__": + test_inject_ptx_intrin() From 38c64df49ec41773b1e1a41425b0621ac3c3e03c Mon Sep 17 00:00:00 2001 From: andy_yang Date: Mon, 13 Feb 2023 21:44:11 +0800 Subject: [PATCH 02/10] Update inject_ptx_ldg32.cc --- src/tir/transforms/inject_ptx_ldg32.cc | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/src/tir/transforms/inject_ptx_ldg32.cc b/src/tir/transforms/inject_ptx_ldg32.cc index 4597d477e848..201308e66aa6 100644 --- a/src/tir/transforms/inject_ptx_ldg32.cc +++ b/src/tir/transforms/inject_ptx_ldg32.cc @@ -54,6 +54,11 @@ class PTXRewriter : public StmtMutator{ result = Allocate(addr_buffer->data, addr_buffer->dtype, addr_buffer->shape, Bool(true), result); result = Allocate(predicate_buffer->data, predicate_buffer->dtype, predicate_buffer->shape, Bool(true), result); } + + + + + return result; } @@ -61,9 +66,12 @@ class PTXRewriter : public StmtMutator{ Stmt VisitStmt_(const BufferStoreNode* store) final{ Stmt result = StmtMutator::VisitStmt_(store); + + Buffer load_buffer = store->buffer; PrimExpr load_value = store->value; + // const BufferLoadNode* gload = load_value.as(); // take the place of instance of const CallNode* call = load_value.as(); @@ -77,20 +85,22 @@ class PTXRewriter : public StmtMutator{ PrimExpr global_addr, local_addr; const BufferLoadNode* load = lhs.as(); + PrimExpr imm_value = rhs; if (load == nullptr){ load = rhs.as(); + imm_value = lhs; + if (load == nullptr){ + return result; + } } - for ( auto each : load->indices){ - - } global_addr = load->indices[0]; + const RampNode* ramp = global_addr.as(); if (ramp != nullptr){ return result; - } local_addr = store->indices[0]; @@ -99,6 +109,7 @@ class PTXRewriter : public StmtMutator{ BufferStore addr_store(addr_buffer, global_addr, {IntImm(DataType::Int(32), 0)}); BufferStore local_addr_store(addr_buffer,local_addr,{IntImm(DataType::Int(32), 1)}); BufferStore predicate_store(predicate_buffer, predicate, {IntImm(DataType::Int(32), 0)}); + PrimExpr new_lhs, new_rhs, new_predicate, new_indice; @@ -108,14 +119,17 @@ class PTXRewriter : public StmtMutator{ new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); + BufferStore value_store(store->buffer, imm_value, {new_indice}); + + Evaluate ptx_load(Call(store->buffer->dtype, tvm::tir::builtin::ptx_pred_ldg32(),{store->buffer->data, new_predicate , new_lhs,new_indice})); - Array tmp_seq = {addr_store, local_addr_store, predicate_store, ptx_load}; + Array tmp_seq = {addr_store, local_addr_store, predicate_store, value_store , ptx_load}; SeqStmt seq_stmt = SeqStmt(tmp_seq); return seq_stmt; - + } } From 690db54db03104bcd1b5beeefc6694384c433ad1 Mon Sep 17 00:00:00 2001 From: andy_yang Date: Thu, 16 Feb 2023 14:16:31 +0800 Subject: [PATCH 03/10] test use clang-format --- include/tvm/tir/builtin.h | 2 +- src/driver/driver_api.cc | 10 +- src/target/source/codegen_cuda.cc | 2 +- src/tir/op/builtin.cc | 2 +- src/tir/transforms/inject_ptx_ldg32.cc | 195 +++++++----------- .../python/unittest/test_inject_ptx_ldg32.py | 2 +- 6 files changed, 89 insertions(+), 124 deletions(-) diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h index e984c7e2ab6d..d3972beb7a24 100644 --- a/include/tvm/tir/builtin.h +++ b/include/tvm/tir/builtin.h @@ -615,7 +615,7 @@ TVM_DLL const Op& ptx_mma(); * \brief tvm intrinsic for ptx predicate load with 32-bit data type. * */ -TVM_DLL const Op& ptx_pred_ldg32(); +TVM_DLL const Op& inject_ptx_ldg32(); /*! * \brief tvm intrinsic for sparse tensor core ptx instructions. diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc index a8349c67bc33..41a155aa1ec5 100644 --- a/src/driver/driver_api.cc +++ b/src/driver/driver_api.cc @@ -55,7 +55,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.use_async_copy", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.merge_async_commit_queue_scope", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.instrument_lwp", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.vtcm_capacity", Integer); -TVM_REGISTER_PASS_CONFIG_OPTION("tir.ptx_pred_ldg32", Bool); +TVM_REGISTER_PASS_CONFIG_OPTION("tir.inject_ptx_ldg32", Bool); // WARNING: May cause coherency issues resulting data miscompares // Experimental feature that, when enabled by the runtime, bypasses the cache when using DMA. When @@ -160,7 +160,7 @@ Array CreatePassList(bool disable_loop_partition) { bool enable_equiv_terms_in_cse_tir = pass_ctx->GetConfig("tir.enable_equiv_terms_in_cse_tir", Bool(false)).value(); - bool ptx_pred_ldg32 = pass_ctx->GetConfig("tir.ptx_pred_ldg32", Bool(false)).value(); + bool inject_ptx_ldg32 = pass_ctx->GetConfig("tir.inject_ptx_ldg32", Bool(false)).value(); // Get any user-added passes Array> add_lower_pass = @@ -260,7 +260,7 @@ Array CreatePassList(bool disable_loop_partition) { pass_list.push_back(tir::transform::InstrumentBoundCheckers()); } - if(ptx_pred_ldg32){ + if(inject_ptx_ldg32){ pass_list.push_back(tir::transform::InjectPTXLDG32(true)); } @@ -591,8 +591,8 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target) mixed_pass_list.push_back(tir::transform::InjectPTXAsyncCopy()); } - bool ptx_pred_ldg32 = pass_ctx->GetConfig("tir.ptx_pred_ldg32", Bool(false)).value(); - if (ptx_pred_ldg32) { + bool inject_ptx_ldg32 = pass_ctx->GetConfig("tir.inject_ptx_ldg32", Bool(false)).value(); + if (inject_ptx_ldg32) { mixed_pass_list.push_back(tir::transform::InjectPTXLDG32()); } diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc index 3b964fb8f08b..e0b8dc2b7e45 100644 --- a/src/target/source/codegen_cuda.cc +++ b/src/target/source/codegen_cuda.cc @@ -920,7 +920,7 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { } else if (op->op.same_as(builtin::ptx_wait_group())) { std::string N = this->PrintExpr(op->args[0]); this->stream << "__asm__ __volatile__(\"cp.async.wait_group " + N + ";\");\n\n"; - } else if (op->op.same_as(builtin::ptx_pred_ldg32())){ + } else if (op->op.same_as(builtin::inject_ptx_ldg32())){ /* asm volatile ( "{.reg .pred p;\n" diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc index fdcb687b36cd..ec6250c6269c 100644 --- a/src/tir/op/builtin.cc +++ b/src/tir/op/builtin.cc @@ -251,7 +251,7 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync) TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); -TIR_DEFINE_BUILTIN_FUNC(ptx_pred_ldg32) +TIR_DEFINE_BUILTIN_FUNC(inject_ptx_ldg32) .set_num_inputs(4) .set_attr("TCallEffectKind", Integer(CallEffectKind::kPure)); diff --git a/src/tir/transforms/inject_ptx_ldg32.cc b/src/tir/transforms/inject_ptx_ldg32.cc index 201308e66aa6..265abd1c4b3b 100644 --- a/src/tir/transforms/inject_ptx_ldg32.cc +++ b/src/tir/transforms/inject_ptx_ldg32.cc @@ -32,125 +32,93 @@ namespace tvm { namespace tir { - - -class PTXRewriter : public StmtMutator{ - -public: - explicit PTXRewriter() {} - - Stmt VisitStmt_(const AllocateNode* allocate) final{ - - if (!has_buffer_1){ - has_buffer_1 = true; - // addr[0] -> global_addr / addr[1] -> local_addr - addr_buffer = decl_buffer({IntImm(DataType::Int(32), 2)}, DataType::Int(32), "addr", "local"); - predicate_buffer = decl_buffer({IntImm(DataType::Int(32), 1)}, DataType::Bool(1), "predicate", "local"); - } - - Stmt result = StmtMutator::VisitStmt_(allocate); - if (!has_buffer_2){ - has_buffer_2 = true; - result = Allocate(addr_buffer->data, addr_buffer->dtype, addr_buffer->shape, Bool(true), result); - result = Allocate(predicate_buffer->data, predicate_buffer->dtype, predicate_buffer->shape, Bool(true), result); - } - - - - - - return result; +class PTXRewriter : public StmtMutator { + public: + Stmt VisitStmt_(const AllocateNode* allocate) final { + if (!has_buffer_1) { + has_buffer_1 = true; + // addr[0] -> global_addr / addr[1] -> local_addr + addr_buffer = decl_buffer({IntImm(DataType::Int(32), 2)}, + DataType::Int(32), "addr", "local"); + predicate_buffer = decl_buffer({IntImm(DataType::Int(32), 1)}, + DataType::Bool(1), "predicate", "local"); } - - - Stmt VisitStmt_(const BufferStoreNode* store) final{ - - Stmt result = StmtMutator::VisitStmt_(store); - - - Buffer load_buffer = store->buffer; - PrimExpr load_value = store->value; - - // const BufferLoadNode* gload = load_value.as(); // take the place of instance of - - const CallNode* call = load_value.as(); - - if (call != nullptr){ - - const OpNode* op = call->op.as(); - if (op != nullptr && op->name == "tir.if_then_else"){ - const PrimExpr& predicate = call->args[0]; - const PrimExpr& lhs = call->args[1]; - const PrimExpr& rhs = call->args[2]; - - PrimExpr global_addr, local_addr; - const BufferLoadNode* load = lhs.as(); - PrimExpr imm_value = rhs; - - - if (load == nullptr){ - load = rhs.as(); - imm_value = lhs; - if (load == nullptr){ - return result; - } - } - - global_addr = load->indices[0]; - - const RampNode* ramp = global_addr.as(); - if (ramp != nullptr){ + Stmt result = StmtMutator::VisitStmt_(allocate); + if (!has_buffer_2) { + has_buffer_2 = true; + result = Allocate(addr_buffer->data, addr_buffer->dtype, + addr_buffer->shape, Bool(true), result); + result = Allocate(predicate_buffer->data, predicate_buffer->dtype, + predicate_buffer->shape, Bool(true), result); + } + return result; + } + + Stmt VisitStmt_(const BufferStoreNode* store) final { + Stmt result = StmtMutator::VisitStmt_(store); + Buffer load_buffer = store->buffer; + PrimExpr load_value = store->value; + // const BufferLoadNode* gload = load_value.as(); // take + // the place of instance of + const CallNode* call = load_value.as(); + if (call != nullptr) { + const OpNode* op = call->op.as(); + if (op != nullptr && op->name == "tir.if_then_else") { + const PrimExpr& predicate = call->args[0]; + const PrimExpr& lhs = call->args[1]; + const PrimExpr& rhs = call->args[2]; + PrimExpr global_addr, local_addr; + const BufferLoadNode* load = lhs.as(); + PrimExpr imm_value = rhs; + if (load == nullptr) { + load = rhs.as(); + imm_value = lhs; + if (load == nullptr) { return result; } - - local_addr = store->indices[0]; - - - BufferStore addr_store(addr_buffer, global_addr, {IntImm(DataType::Int(32), 0)}); - BufferStore local_addr_store(addr_buffer,local_addr,{IntImm(DataType::Int(32), 1)}); - BufferStore predicate_store(predicate_buffer, predicate, {IntImm(DataType::Int(32), 0)}); - - - PrimExpr new_lhs, new_rhs, new_predicate, new_indice; - - new_lhs = BufferLoad(load->buffer, {BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 0)})}); - new_rhs = IntImm(DataType::Int(32), 0); - new_predicate = BufferLoad(predicate_buffer, {IntImm(DataType::Int(32), 0)}); - new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); - - - BufferStore value_store(store->buffer, imm_value, {new_indice}); - - - Evaluate ptx_load(Call(store->buffer->dtype, tvm::tir::builtin::ptx_pred_ldg32(),{store->buffer->data, new_predicate , new_lhs,new_indice})); - - Array tmp_seq = {addr_store, local_addr_store, predicate_store, value_store , ptx_load}; - - SeqStmt seq_stmt = SeqStmt(tmp_seq); - - return seq_stmt; - } + global_addr = load->indices[0]; + const RampNode* ramp = global_addr.as(); + if (ramp != nullptr) { + return result; + } + local_addr = store->indices[0]; + BufferStore addr_store(addr_buffer, global_addr, + {IntImm(DataType::Int(32), 0)}); + BufferStore local_addr_store(addr_buffer, local_addr, + {IntImm(DataType::Int(32), 1)}); + BufferStore predicate_store(predicate_buffer, predicate, + {IntImm(DataType::Int(32), 0)}); + PrimExpr new_lhs, new_rhs, new_predicate, new_indice; + new_lhs = BufferLoad( + load->buffer, + {BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 0)})}); + new_rhs = IntImm(DataType::Int(32), 0); + new_predicate = + BufferLoad(predicate_buffer, {IntImm(DataType::Int(32), 0)}); + new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); + BufferStore value_store(store->buffer, imm_value, {new_indice}); + Evaluate ptx_load( + Call(store->buffer->dtype, tvm::tir::builtin::inject_ptx_ldg32(), + {store->buffer->data, new_predicate, new_lhs, new_indice})); + Array tmp_seq = {addr_store, local_addr_store, predicate_store, + value_store, ptx_load}; + SeqStmt seq_stmt = SeqStmt(tmp_seq); + return seq_stmt; } - - - return result; - } + return result; + } - - bool has_buffer_1 = false, has_buffer_2 = false ; - Buffer addr_buffer , predicate_buffer ; - - + bool has_buffer_1 = false, has_buffer_2 = false; + Buffer addr_buffer, predicate_buffer; }; - -namespace transform{ - +namespace transform { Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) { - auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m, PassContext ctx) { + auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m, + PassContext ctx) { if (enable_inject_ptx_intrin) { auto* n = f.CopyOnWrite(); n->body = PTXRewriter()(n->body); @@ -161,14 +129,11 @@ Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) { return CreatePrimFuncPass(pass_func, 0, "tir.InjectPTXLDG32", {}); } -// The pass can now be invoked via the pass infrastructure, but we also add a Python binding for it -TVM_REGISTER_GLOBAL("tir.transform.InjectPTXLDG32").set_body_typed(InjectPTXLDG32); - - +// The pass can now be invoked via the pass infrastructure, but we also add a +// Python binding for it +TVM_REGISTER_GLOBAL("tir.transform.InjectPTXLDG32") + .set_body_typed(InjectPTXLDG32); } // namespace transform } // namespace tir } // namespace tvm - - - diff --git a/tests/python/unittest/test_inject_ptx_ldg32.py b/tests/python/unittest/test_inject_ptx_ldg32.py index f20032c164eb..4946f0acf77d 100644 --- a/tests/python/unittest/test_inject_ptx_ldg32.py +++ b/tests/python/unittest/test_inject_ptx_ldg32.py @@ -46,7 +46,7 @@ def test_inject_ptx_intrin(): if major < 8: # Require at least SM80 return - with tvm.transform.PassContext(config={"tir.ptx_pred_ldg32": True}): + with tvm.transform.PassContext(config={"tir.inject_ptx_ldg32": True}): mod = tvm.build(f, target="cuda") A_np = np.random.rand(16).astype("float32") B_np = np.zeros((32)).astype("float32") From 78141e3558fc7f7cd86208ab68b262636207cce0 Mon Sep 17 00:00:00 2001 From: andy_yang Date: Thu, 16 Feb 2023 14:34:53 +0800 Subject: [PATCH 04/10] remove printIndent --- src/target/source/codegen_cuda.cc | 20 ++++---------------- 1 file changed, 4 insertions(+), 16 deletions(-) diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc index e0b8dc2b7e45..9d45e1c25919 100644 --- a/src/target/source/codegen_cuda.cc +++ b/src/target/source/codegen_cuda.cc @@ -936,30 +936,18 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { std::string reg = this->PrintExpr(op->args[0]); // get guard std::string guard = this->PrintExpr(op->args[1]); - // std::string lhs = this->PrintExpr(op->args[2]); const BufferLoadNode* addr_buffer = op->args[2].as(); - std::string global_addr = this->PrintExpr(addr_buffer->indices[0]); std::string global_buffer = this->PrintExpr(addr_buffer->buffer->data); - std::string local_addr = this->PrintExpr(op->args[3]); - - this->stream << "asm volatile (\n" ; - this->PrintIndent(); - stream << "\"{.reg .pred p;\\n\"\n" ; - this->PrintIndent(); - stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n" ; - this->PrintIndent(); - stream << "\" @!p mov.b32 %0, 0;\\n\"\n"; - this->PrintIndent(); - stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n" ; + this->stream << "\"{.reg .pred p;\\n\"\n" ; + this->stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n" ; + this->stream << "\" @!p mov.b32 %0, 0;\\n\"\n"; + this->stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n" ; // stream << "\" @p ld.global.nc.L2::128B.f32 %0, [%1];}\\n\"\n" ; - this->PrintIndent(); stream << ": \"=f\"(" << reg << "[" << local_addr << "]" << ")\n" ; - this->PrintIndent(); stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr << ")), \"r\"((int)" << guard << ")\n" ; - this->PrintIndent(); stream << ");\n" ; } else { CodeGenC::VisitExpr_(op, os); From bb7a70a6e4dbeb9bb9357cc0b8ab120c4742f017 Mon Sep 17 00:00:00 2001 From: andy_yang Date: Fri, 17 Feb 2023 14:03:00 +0800 Subject: [PATCH 05/10] change inject_ptx_ldg32 to ptx_ldg32 --- include/tvm/tir/builtin.h | 2 +- include/tvm/tir/transform.h | 2 +- src/driver/driver_api.cc | 10 +++++----- src/target/source/codegen_cuda.cc | 2 +- src/tir/op/builtin.cc | 2 +- src/tir/transforms/inject_ptx_ldg32.cc | 2 +- tests/python/unittest/test_inject_ptx_ldg32.py | 2 +- 7 files changed, 11 insertions(+), 11 deletions(-) diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h index d3972beb7a24..aefe19de2a3a 100644 --- a/include/tvm/tir/builtin.h +++ b/include/tvm/tir/builtin.h @@ -615,7 +615,7 @@ TVM_DLL const Op& ptx_mma(); * \brief tvm intrinsic for ptx predicate load with 32-bit data type. * */ -TVM_DLL const Op& inject_ptx_ldg32(); +TVM_DLL const Op& ptx_ldg32(); /*! * \brief tvm intrinsic for sparse tensor core ptx instructions. diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h index 6d2a9e18da68..85b381a52950 100644 --- a/include/tvm/tir/transform.h +++ b/include/tvm/tir/transform.h @@ -681,7 +681,7 @@ TVM_DLL Pass InjectPTXAsyncCopy(); * \brief Pass to rewrite global to local memory copy on CUDA with ldg32 instruction. * \return The pass. */ -TVM_DLL Pass InjectPTXLDG32(bool enable_inject_ptx_ldg32 = true); +TVM_DLL Pass InjectPTXLDG32(bool enable_ptx_ldg32 = true); /*! * \brief Remove the weight layout rewrite block diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc index 41a155aa1ec5..aa6f169c59a7 100644 --- a/src/driver/driver_api.cc +++ b/src/driver/driver_api.cc @@ -55,7 +55,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.use_async_copy", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.merge_async_commit_queue_scope", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.instrument_lwp", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.vtcm_capacity", Integer); -TVM_REGISTER_PASS_CONFIG_OPTION("tir.inject_ptx_ldg32", Bool); +TVM_REGISTER_PASS_CONFIG_OPTION("tir.ptx_ldg32", Bool); // WARNING: May cause coherency issues resulting data miscompares // Experimental feature that, when enabled by the runtime, bypasses the cache when using DMA. When @@ -160,7 +160,7 @@ Array CreatePassList(bool disable_loop_partition) { bool enable_equiv_terms_in_cse_tir = pass_ctx->GetConfig("tir.enable_equiv_terms_in_cse_tir", Bool(false)).value(); - bool inject_ptx_ldg32 = pass_ctx->GetConfig("tir.inject_ptx_ldg32", Bool(false)).value(); + bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value(); // Get any user-added passes Array> add_lower_pass = @@ -260,7 +260,7 @@ Array CreatePassList(bool disable_loop_partition) { pass_list.push_back(tir::transform::InstrumentBoundCheckers()); } - if(inject_ptx_ldg32){ + if(ptx_ldg32){ pass_list.push_back(tir::transform::InjectPTXLDG32(true)); } @@ -591,8 +591,8 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target) mixed_pass_list.push_back(tir::transform::InjectPTXAsyncCopy()); } - bool inject_ptx_ldg32 = pass_ctx->GetConfig("tir.inject_ptx_ldg32", Bool(false)).value(); - if (inject_ptx_ldg32) { + bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value(); + if (ptx_ldg32) { mixed_pass_list.push_back(tir::transform::InjectPTXLDG32()); } diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc index 9d45e1c25919..6701d3b3ce40 100644 --- a/src/target/source/codegen_cuda.cc +++ b/src/target/source/codegen_cuda.cc @@ -920,7 +920,7 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { } else if (op->op.same_as(builtin::ptx_wait_group())) { std::string N = this->PrintExpr(op->args[0]); this->stream << "__asm__ __volatile__(\"cp.async.wait_group " + N + ";\");\n\n"; - } else if (op->op.same_as(builtin::inject_ptx_ldg32())){ + } else if (op->op.same_as(builtin::ptx_ldg32())){ /* asm volatile ( "{.reg .pred p;\n" diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc index ec6250c6269c..c47e3b264bf8 100644 --- a/src/tir/op/builtin.cc +++ b/src/tir/op/builtin.cc @@ -251,7 +251,7 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync) TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); -TIR_DEFINE_BUILTIN_FUNC(inject_ptx_ldg32) +TIR_DEFINE_BUILTIN_FUNC(ptx_ldg32) .set_num_inputs(4) .set_attr("TCallEffectKind", Integer(CallEffectKind::kPure)); diff --git a/src/tir/transforms/inject_ptx_ldg32.cc b/src/tir/transforms/inject_ptx_ldg32.cc index 265abd1c4b3b..52726c25cf7b 100644 --- a/src/tir/transforms/inject_ptx_ldg32.cc +++ b/src/tir/transforms/inject_ptx_ldg32.cc @@ -99,7 +99,7 @@ class PTXRewriter : public StmtMutator { new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); BufferStore value_store(store->buffer, imm_value, {new_indice}); Evaluate ptx_load( - Call(store->buffer->dtype, tvm::tir::builtin::inject_ptx_ldg32(), + Call(store->buffer->dtype, tvm::tir::builtin::ptx_ldg32(), {store->buffer->data, new_predicate, new_lhs, new_indice})); Array tmp_seq = {addr_store, local_addr_store, predicate_store, value_store, ptx_load}; diff --git a/tests/python/unittest/test_inject_ptx_ldg32.py b/tests/python/unittest/test_inject_ptx_ldg32.py index 4946f0acf77d..1c5bf1f4fb21 100644 --- a/tests/python/unittest/test_inject_ptx_ldg32.py +++ b/tests/python/unittest/test_inject_ptx_ldg32.py @@ -46,7 +46,7 @@ def test_inject_ptx_intrin(): if major < 8: # Require at least SM80 return - with tvm.transform.PassContext(config={"tir.inject_ptx_ldg32": True}): + with tvm.transform.PassContext(config={"tir.ptx_ldg32": True}): mod = tvm.build(f, target="cuda") A_np = np.random.rand(16).astype("float32") B_np = np.zeros((32)).astype("float32") From 36f9ad9fe76f93e1292100af05bec4161ae43c54 Mon Sep 17 00:00:00 2001 From: andy_yang Date: Mon, 13 Feb 2023 19:35:34 +0800 Subject: [PATCH 06/10] [TIR] Introduce Pass InjectPTXLDG32 This PR introduces a new pass InjectPTXLDG32 to change the `if_then_else` call node to `ptx_pred_ldg32` call node. When the store buffer is local and the load value is global, the pass can change the if_then_else pattern to a ptx pattern. Test the pass with: ```python with tvm.transform.PassContext(config={"tir.ptx_pred_ldg32": True}): mod = tvm.build(f, target="cuda") ```` --- include/tvm/tir/builtin.h | 7 + include/tvm/tir/transform.h | 6 + src/driver/driver_api.cc | 12 ++ src/target/source/codegen_cuda.cc | 29 ++++ src/tir/op/builtin.cc | 3 + src/tir/transforms/inject_ptx_ldg32.cc | 139 ++++++++++++++++++ .../python/unittest/test_inject_ptx_ldg32.py | 67 +++++++++ 7 files changed, 263 insertions(+) create mode 100644 src/tir/transforms/inject_ptx_ldg32.cc create mode 100644 tests/python/unittest/test_inject_ptx_ldg32.py diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h index d830ea579aa7..aefe19de2a3a 100644 --- a/include/tvm/tir/builtin.h +++ b/include/tvm/tir/builtin.h @@ -610,6 +610,13 @@ TVM_DLL const Op& tvm_store_matrix_sync(); */ TVM_DLL const Op& ptx_mma(); + +/*! + * \brief tvm intrinsic for ptx predicate load with 32-bit data type. + * + */ +TVM_DLL const Op& ptx_ldg32(); + /*! * \brief tvm intrinsic for sparse tensor core ptx instructions. * diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h index be7589b04bf5..85b381a52950 100644 --- a/include/tvm/tir/transform.h +++ b/include/tvm/tir/transform.h @@ -677,6 +677,12 @@ TVM_DLL Pass Filter(runtime::TypedPackedFunc fcond); */ TVM_DLL Pass InjectPTXAsyncCopy(); +/*! + * \brief Pass to rewrite global to local memory copy on CUDA with ldg32 instruction. + * \return The pass. + */ +TVM_DLL Pass InjectPTXLDG32(bool enable_ptx_ldg32 = true); + /*! * \brief Remove the weight layout rewrite block * \param skip_ndarray_rewrite If True, exact rewrite of NDArray, according to the given index map, diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc index 288ac7b92a2c..aa6f169c59a7 100644 --- a/src/driver/driver_api.cc +++ b/src/driver/driver_api.cc @@ -55,6 +55,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.use_async_copy", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.merge_async_commit_queue_scope", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.instrument_lwp", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.vtcm_capacity", Integer); +TVM_REGISTER_PASS_CONFIG_OPTION("tir.ptx_ldg32", Bool); // WARNING: May cause coherency issues resulting data miscompares // Experimental feature that, when enabled by the runtime, bypasses the cache when using DMA. When @@ -159,6 +160,8 @@ Array CreatePassList(bool disable_loop_partition) { bool enable_equiv_terms_in_cse_tir = pass_ctx->GetConfig("tir.enable_equiv_terms_in_cse_tir", Bool(false)).value(); + bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value(); + // Get any user-added passes Array> add_lower_pass = pass_ctx->GetConfig>>("tir.add_lower_pass", Array>()) @@ -257,6 +260,10 @@ Array CreatePassList(bool disable_loop_partition) { pass_list.push_back(tir::transform::InstrumentBoundCheckers()); } + if(ptx_ldg32){ + pass_list.push_back(tir::transform::InjectPTXLDG32(true)); + } + pass_list.push_back( tir::transform::CommonSubexprElimTIR(!disable_cse_tir, enable_equiv_terms_in_cse_tir)); @@ -584,6 +591,11 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target) mixed_pass_list.push_back(tir::transform::InjectPTXAsyncCopy()); } + bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value(); + if (ptx_ldg32) { + mixed_pass_list.push_back(tir::transform::InjectPTXLDG32()); + } + bool unpacked_api = mixed_mod->GetAttr(tvm::attr::kExecutor) .value_or(relay::Executor::Create("graph", {})) ->GetAttr("unpacked-api") diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc index 9bf0109cace1..48c2c6be44e9 100644 --- a/src/target/source/codegen_cuda.cc +++ b/src/target/source/codegen_cuda.cc @@ -926,6 +926,35 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { } else if (op->op.same_as(builtin::ptx_wait_group())) { std::string N = this->PrintExpr(op->args[0]); this->stream << "__asm__ __volatile__(\"cp.async.wait_group " + N + ";\");\n\n"; + } else if (op->op.same_as(builtin::ptx_ldg32())){ + /* + asm volatile ( + "{.reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + // " @p ld.global.nc.f32 %0, [%1];}\n"t + " @p ld.global.nc.L2::128B.f32 %0, [%1];}\n" + : "=f"(reg) + : "l"(addr), "r"((int)guard) + ); + */ + + // get local + std::string reg = this->PrintExpr(op->args[0]); + // get guard + std::string guard = this->PrintExpr(op->args[1]); + const BufferLoadNode* addr_buffer = op->args[2].as(); + std::string global_addr = this->PrintExpr(addr_buffer->indices[0]); + std::string global_buffer = this->PrintExpr(addr_buffer->buffer->data); + std::string local_addr = this->PrintExpr(op->args[3]); + this->stream << "asm volatile (\n" ; + this->stream << "\"{.reg .pred p;\\n\"\n" ; + this->stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n" ; + this->stream << "\" @!p mov.b32 %0, 0;\\n\"\n"; + this->stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n" ; + // stream << "\" @p ld.global.nc.L2::128B.f32 %0, [%1];}\\n\"\n" ; + stream << ": \"=f\"(" << reg << "[" << local_addr << "]" << ")\n" ; + stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr << ")), \"r\"((int)" << guard << ")\n" ; + stream << ");\n" ; } else { CodeGenC::VisitExpr_(op, os); } diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc index dc3208f484e3..c47e3b264bf8 100644 --- a/src/tir/op/builtin.cc +++ b/src/tir/op/builtin.cc @@ -251,6 +251,9 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync) TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); +TIR_DEFINE_BUILTIN_FUNC(ptx_ldg32) + .set_num_inputs(4) + .set_attr("TCallEffectKind", Integer(CallEffectKind::kPure)); TIR_DEFINE_BUILTIN_FUNC(ptx_mma_sp) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); diff --git a/src/tir/transforms/inject_ptx_ldg32.cc b/src/tir/transforms/inject_ptx_ldg32.cc new file mode 100644 index 000000000000..52726c25cf7b --- /dev/null +++ b/src/tir/transforms/inject_ptx_ldg32.cc @@ -0,0 +1,139 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../../arith/const_fold.h" +#include "../../arith/pattern_match.h" + +namespace tvm { +namespace tir { + +class PTXRewriter : public StmtMutator { + public: + Stmt VisitStmt_(const AllocateNode* allocate) final { + if (!has_buffer_1) { + has_buffer_1 = true; + // addr[0] -> global_addr / addr[1] -> local_addr + addr_buffer = decl_buffer({IntImm(DataType::Int(32), 2)}, + DataType::Int(32), "addr", "local"); + predicate_buffer = decl_buffer({IntImm(DataType::Int(32), 1)}, + DataType::Bool(1), "predicate", "local"); + } + Stmt result = StmtMutator::VisitStmt_(allocate); + if (!has_buffer_2) { + has_buffer_2 = true; + result = Allocate(addr_buffer->data, addr_buffer->dtype, + addr_buffer->shape, Bool(true), result); + result = Allocate(predicate_buffer->data, predicate_buffer->dtype, + predicate_buffer->shape, Bool(true), result); + } + return result; + } + + Stmt VisitStmt_(const BufferStoreNode* store) final { + Stmt result = StmtMutator::VisitStmt_(store); + Buffer load_buffer = store->buffer; + PrimExpr load_value = store->value; + // const BufferLoadNode* gload = load_value.as(); // take + // the place of instance of + const CallNode* call = load_value.as(); + if (call != nullptr) { + const OpNode* op = call->op.as(); + if (op != nullptr && op->name == "tir.if_then_else") { + const PrimExpr& predicate = call->args[0]; + const PrimExpr& lhs = call->args[1]; + const PrimExpr& rhs = call->args[2]; + PrimExpr global_addr, local_addr; + const BufferLoadNode* load = lhs.as(); + PrimExpr imm_value = rhs; + if (load == nullptr) { + load = rhs.as(); + imm_value = lhs; + if (load == nullptr) { + return result; + } + } + global_addr = load->indices[0]; + const RampNode* ramp = global_addr.as(); + if (ramp != nullptr) { + return result; + } + local_addr = store->indices[0]; + BufferStore addr_store(addr_buffer, global_addr, + {IntImm(DataType::Int(32), 0)}); + BufferStore local_addr_store(addr_buffer, local_addr, + {IntImm(DataType::Int(32), 1)}); + BufferStore predicate_store(predicate_buffer, predicate, + {IntImm(DataType::Int(32), 0)}); + PrimExpr new_lhs, new_rhs, new_predicate, new_indice; + new_lhs = BufferLoad( + load->buffer, + {BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 0)})}); + new_rhs = IntImm(DataType::Int(32), 0); + new_predicate = + BufferLoad(predicate_buffer, {IntImm(DataType::Int(32), 0)}); + new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); + BufferStore value_store(store->buffer, imm_value, {new_indice}); + Evaluate ptx_load( + Call(store->buffer->dtype, tvm::tir::builtin::ptx_ldg32(), + {store->buffer->data, new_predicate, new_lhs, new_indice})); + Array tmp_seq = {addr_store, local_addr_store, predicate_store, + value_store, ptx_load}; + SeqStmt seq_stmt = SeqStmt(tmp_seq); + return seq_stmt; + } + } + return result; + } + + bool has_buffer_1 = false, has_buffer_2 = false; + Buffer addr_buffer, predicate_buffer; +}; + +namespace transform { + +Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) { + auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m, + PassContext ctx) { + if (enable_inject_ptx_intrin) { + auto* n = f.CopyOnWrite(); + n->body = PTXRewriter()(n->body); + // inject ptx + } + return f; + }; + return CreatePrimFuncPass(pass_func, 0, "tir.InjectPTXLDG32", {}); +} + +// The pass can now be invoked via the pass infrastructure, but we also add a +// Python binding for it +TVM_REGISTER_GLOBAL("tir.transform.InjectPTXLDG32") + .set_body_typed(InjectPTXLDG32); + +} // namespace transform +} // namespace tir +} // namespace tvm diff --git a/tests/python/unittest/test_inject_ptx_ldg32.py b/tests/python/unittest/test_inject_ptx_ldg32.py new file mode 100644 index 000000000000..c92bdd8c6488 --- /dev/null +++ b/tests/python/unittest/test_inject_ptx_ldg32.py @@ -0,0 +1,67 @@ +# 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. +import tvm +from tvm.script import tir as T +import numpy as np +import tvm.testing + + +@T.prim_func +def vector_add(A: T.Buffer[(16), "float32"], B: T.Buffer[(32), "float32"]) -> None: + T.func_attr({"global_symbol": "default_function", "tir.noalias": True}) + bx = T.env_thread("blockIdx.x") + tx = T.env_thread("threadIdx.x") + T.launch_thread(bx, 1) + T.launch_thread(tx, 32) + A_local = T.alloc_buffer([32], "float32", scope="local") + + with T.block(): + T.reads(A[0:16]) + T.writes(A_local[0:32]) + A_local[tx] = T.if_then_else(tx % 2 == 0, A[tx / 2], T.float32(0), dtype="float32") + B[tx] = A_local[tx] + 1.0 + + +@tvm.testing.requires_cuda +def test_inject_ptx_intrin(): + f = vector_add + arch = tvm.contrib.nvcc.get_target_compute_version() + major, _ = tvm.contrib.nvcc.parse_compute_version(arch) + if major < 8: + # Require at least SM80 + return + with tvm.transform.PassContext(config={"tir.ptx_ldg32": True}): + mod = tvm.build(f, target="cuda") + A_np = np.random.rand(16).astype("float32") + B_np = np.zeros((32)).astype("float32") + dev = tvm.cuda(0) + A_nd = tvm.nd.array(A_np, device=dev) + B_nd = tvm.nd.array(B_np, device=dev) + mod(A_nd, B_nd) + + C_np = np.zeros((32)).astype("float32") + + for i in range(32): + if i % 2 == 0: + C_np[i] = A_np[i // 2] + C_np[i] += 1.0 + + tvm.testing.assert_allclose(B_nd.numpy(), C_np) + + +if __name__ == "__main__": + test_inject_ptx_intrin() From 3b04bf35ff0834665b8ebb819cf6a976063a8b80 Mon Sep 17 00:00:00 2001 From: andy_yang Date: Mon, 13 Feb 2023 19:35:34 +0800 Subject: [PATCH 07/10] [TIR] Introduce Pass InjectPTXLDG32 This PR introduces a new pass InjectPTXLDG32 to change the `if_then_else` call node to `ptx_pred_ldg32` call node. When the store buffer is local and the load value is global, the pass can change the if_then_else pattern to a ptx pattern. Test the pass with: ```python with tvm.transform.PassContext(config={"tir.ptx_pred_ldg32": True}): mod = tvm.build(f, target="cuda") ```` --- include/tvm/tir/builtin.h | 7 + include/tvm/tir/transform.h | 6 + src/driver/driver_api.cc | 12 ++ src/target/source/codegen_cuda.cc | 31 ++++ src/tir/op/builtin.cc | 3 + src/tir/transforms/inject_ptx_ldg32.cc | 139 ++++++++++++++++++ .../python/unittest/test_inject_ptx_ldg32.py | 67 +++++++++ 7 files changed, 265 insertions(+) create mode 100644 src/tir/transforms/inject_ptx_ldg32.cc create mode 100644 tests/python/unittest/test_inject_ptx_ldg32.py diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h index d830ea579aa7..aefe19de2a3a 100644 --- a/include/tvm/tir/builtin.h +++ b/include/tvm/tir/builtin.h @@ -610,6 +610,13 @@ TVM_DLL const Op& tvm_store_matrix_sync(); */ TVM_DLL const Op& ptx_mma(); + +/*! + * \brief tvm intrinsic for ptx predicate load with 32-bit data type. + * + */ +TVM_DLL const Op& ptx_ldg32(); + /*! * \brief tvm intrinsic for sparse tensor core ptx instructions. * diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h index be7589b04bf5..85b381a52950 100644 --- a/include/tvm/tir/transform.h +++ b/include/tvm/tir/transform.h @@ -677,6 +677,12 @@ TVM_DLL Pass Filter(runtime::TypedPackedFunc fcond); */ TVM_DLL Pass InjectPTXAsyncCopy(); +/*! + * \brief Pass to rewrite global to local memory copy on CUDA with ldg32 instruction. + * \return The pass. + */ +TVM_DLL Pass InjectPTXLDG32(bool enable_ptx_ldg32 = true); + /*! * \brief Remove the weight layout rewrite block * \param skip_ndarray_rewrite If True, exact rewrite of NDArray, according to the given index map, diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc index 288ac7b92a2c..aa6f169c59a7 100644 --- a/src/driver/driver_api.cc +++ b/src/driver/driver_api.cc @@ -55,6 +55,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.use_async_copy", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.merge_async_commit_queue_scope", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.instrument_lwp", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.vtcm_capacity", Integer); +TVM_REGISTER_PASS_CONFIG_OPTION("tir.ptx_ldg32", Bool); // WARNING: May cause coherency issues resulting data miscompares // Experimental feature that, when enabled by the runtime, bypasses the cache when using DMA. When @@ -159,6 +160,8 @@ Array CreatePassList(bool disable_loop_partition) { bool enable_equiv_terms_in_cse_tir = pass_ctx->GetConfig("tir.enable_equiv_terms_in_cse_tir", Bool(false)).value(); + bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value(); + // Get any user-added passes Array> add_lower_pass = pass_ctx->GetConfig>>("tir.add_lower_pass", Array>()) @@ -257,6 +260,10 @@ Array CreatePassList(bool disable_loop_partition) { pass_list.push_back(tir::transform::InstrumentBoundCheckers()); } + if(ptx_ldg32){ + pass_list.push_back(tir::transform::InjectPTXLDG32(true)); + } + pass_list.push_back( tir::transform::CommonSubexprElimTIR(!disable_cse_tir, enable_equiv_terms_in_cse_tir)); @@ -584,6 +591,11 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target) mixed_pass_list.push_back(tir::transform::InjectPTXAsyncCopy()); } + bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value(); + if (ptx_ldg32) { + mixed_pass_list.push_back(tir::transform::InjectPTXLDG32()); + } + bool unpacked_api = mixed_mod->GetAttr(tvm::attr::kExecutor) .value_or(relay::Executor::Create("graph", {})) ->GetAttr("unpacked-api") diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc index 9bf0109cace1..20607254370e 100644 --- a/src/target/source/codegen_cuda.cc +++ b/src/target/source/codegen_cuda.cc @@ -926,6 +926,37 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { } else if (op->op.same_as(builtin::ptx_wait_group())) { std::string N = this->PrintExpr(op->args[0]); this->stream << "__asm__ __volatile__(\"cp.async.wait_group " + N + ";\");\n\n"; + } else if (op->op.same_as(builtin::ptx_ldg32())) { + /* + asm volatile ( + "{.reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + // " @p ld.global.nc.f32 %0, [%1];}\n"t + " @p ld.global.nc.L2::128B.f32 %0, [%1];}\n" + : "=f"(reg) + : "l"(addr), "r"((int)guard) + ); + */ + + // get local + std::string reg = this->PrintExpr(op->args[0]); + // get guard + std::string guard = this->PrintExpr(op->args[1]); + const BufferLoadNode* addr_buffer = op->args[2].as(); + std::string global_addr = this->PrintExpr(addr_buffer->indices[0]); + std::string global_buffer = this->PrintExpr(addr_buffer->buffer->data); + std::string local_addr = this->PrintExpr(op->args[3]); + this->stream << "asm volatile (\n"; + this->stream << "\"{.reg .pred p;\\n\"\n"; + this->stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n"; + this->stream << "\" @!p mov.b32 %0, 0;\\n\"\n"; + this->stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n"; + // stream << "\" @p ld.global.nc.L2::128B.f32 %0, [%1];}\\n\"\n" ; + stream << ": \"=f\"(" << reg << "[" << local_addr << "]" + << ")\n"; + stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr << ")), \"r\"((int)" + << guard << ")\n"; + stream << ");\n"; } else { CodeGenC::VisitExpr_(op, os); } diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc index dc3208f484e3..c47e3b264bf8 100644 --- a/src/tir/op/builtin.cc +++ b/src/tir/op/builtin.cc @@ -251,6 +251,9 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync) TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); +TIR_DEFINE_BUILTIN_FUNC(ptx_ldg32) + .set_num_inputs(4) + .set_attr("TCallEffectKind", Integer(CallEffectKind::kPure)); TIR_DEFINE_BUILTIN_FUNC(ptx_mma_sp) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); diff --git a/src/tir/transforms/inject_ptx_ldg32.cc b/src/tir/transforms/inject_ptx_ldg32.cc new file mode 100644 index 000000000000..52726c25cf7b --- /dev/null +++ b/src/tir/transforms/inject_ptx_ldg32.cc @@ -0,0 +1,139 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../../arith/const_fold.h" +#include "../../arith/pattern_match.h" + +namespace tvm { +namespace tir { + +class PTXRewriter : public StmtMutator { + public: + Stmt VisitStmt_(const AllocateNode* allocate) final { + if (!has_buffer_1) { + has_buffer_1 = true; + // addr[0] -> global_addr / addr[1] -> local_addr + addr_buffer = decl_buffer({IntImm(DataType::Int(32), 2)}, + DataType::Int(32), "addr", "local"); + predicate_buffer = decl_buffer({IntImm(DataType::Int(32), 1)}, + DataType::Bool(1), "predicate", "local"); + } + Stmt result = StmtMutator::VisitStmt_(allocate); + if (!has_buffer_2) { + has_buffer_2 = true; + result = Allocate(addr_buffer->data, addr_buffer->dtype, + addr_buffer->shape, Bool(true), result); + result = Allocate(predicate_buffer->data, predicate_buffer->dtype, + predicate_buffer->shape, Bool(true), result); + } + return result; + } + + Stmt VisitStmt_(const BufferStoreNode* store) final { + Stmt result = StmtMutator::VisitStmt_(store); + Buffer load_buffer = store->buffer; + PrimExpr load_value = store->value; + // const BufferLoadNode* gload = load_value.as(); // take + // the place of instance of + const CallNode* call = load_value.as(); + if (call != nullptr) { + const OpNode* op = call->op.as(); + if (op != nullptr && op->name == "tir.if_then_else") { + const PrimExpr& predicate = call->args[0]; + const PrimExpr& lhs = call->args[1]; + const PrimExpr& rhs = call->args[2]; + PrimExpr global_addr, local_addr; + const BufferLoadNode* load = lhs.as(); + PrimExpr imm_value = rhs; + if (load == nullptr) { + load = rhs.as(); + imm_value = lhs; + if (load == nullptr) { + return result; + } + } + global_addr = load->indices[0]; + const RampNode* ramp = global_addr.as(); + if (ramp != nullptr) { + return result; + } + local_addr = store->indices[0]; + BufferStore addr_store(addr_buffer, global_addr, + {IntImm(DataType::Int(32), 0)}); + BufferStore local_addr_store(addr_buffer, local_addr, + {IntImm(DataType::Int(32), 1)}); + BufferStore predicate_store(predicate_buffer, predicate, + {IntImm(DataType::Int(32), 0)}); + PrimExpr new_lhs, new_rhs, new_predicate, new_indice; + new_lhs = BufferLoad( + load->buffer, + {BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 0)})}); + new_rhs = IntImm(DataType::Int(32), 0); + new_predicate = + BufferLoad(predicate_buffer, {IntImm(DataType::Int(32), 0)}); + new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); + BufferStore value_store(store->buffer, imm_value, {new_indice}); + Evaluate ptx_load( + Call(store->buffer->dtype, tvm::tir::builtin::ptx_ldg32(), + {store->buffer->data, new_predicate, new_lhs, new_indice})); + Array tmp_seq = {addr_store, local_addr_store, predicate_store, + value_store, ptx_load}; + SeqStmt seq_stmt = SeqStmt(tmp_seq); + return seq_stmt; + } + } + return result; + } + + bool has_buffer_1 = false, has_buffer_2 = false; + Buffer addr_buffer, predicate_buffer; +}; + +namespace transform { + +Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) { + auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m, + PassContext ctx) { + if (enable_inject_ptx_intrin) { + auto* n = f.CopyOnWrite(); + n->body = PTXRewriter()(n->body); + // inject ptx + } + return f; + }; + return CreatePrimFuncPass(pass_func, 0, "tir.InjectPTXLDG32", {}); +} + +// The pass can now be invoked via the pass infrastructure, but we also add a +// Python binding for it +TVM_REGISTER_GLOBAL("tir.transform.InjectPTXLDG32") + .set_body_typed(InjectPTXLDG32); + +} // namespace transform +} // namespace tir +} // namespace tvm diff --git a/tests/python/unittest/test_inject_ptx_ldg32.py b/tests/python/unittest/test_inject_ptx_ldg32.py new file mode 100644 index 000000000000..c92bdd8c6488 --- /dev/null +++ b/tests/python/unittest/test_inject_ptx_ldg32.py @@ -0,0 +1,67 @@ +# 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. +import tvm +from tvm.script import tir as T +import numpy as np +import tvm.testing + + +@T.prim_func +def vector_add(A: T.Buffer[(16), "float32"], B: T.Buffer[(32), "float32"]) -> None: + T.func_attr({"global_symbol": "default_function", "tir.noalias": True}) + bx = T.env_thread("blockIdx.x") + tx = T.env_thread("threadIdx.x") + T.launch_thread(bx, 1) + T.launch_thread(tx, 32) + A_local = T.alloc_buffer([32], "float32", scope="local") + + with T.block(): + T.reads(A[0:16]) + T.writes(A_local[0:32]) + A_local[tx] = T.if_then_else(tx % 2 == 0, A[tx / 2], T.float32(0), dtype="float32") + B[tx] = A_local[tx] + 1.0 + + +@tvm.testing.requires_cuda +def test_inject_ptx_intrin(): + f = vector_add + arch = tvm.contrib.nvcc.get_target_compute_version() + major, _ = tvm.contrib.nvcc.parse_compute_version(arch) + if major < 8: + # Require at least SM80 + return + with tvm.transform.PassContext(config={"tir.ptx_ldg32": True}): + mod = tvm.build(f, target="cuda") + A_np = np.random.rand(16).astype("float32") + B_np = np.zeros((32)).astype("float32") + dev = tvm.cuda(0) + A_nd = tvm.nd.array(A_np, device=dev) + B_nd = tvm.nd.array(B_np, device=dev) + mod(A_nd, B_nd) + + C_np = np.zeros((32)).astype("float32") + + for i in range(32): + if i % 2 == 0: + C_np[i] = A_np[i // 2] + C_np[i] += 1.0 + + tvm.testing.assert_allclose(B_nd.numpy(), C_np) + + +if __name__ == "__main__": + test_inject_ptx_intrin() From a4587079a91d8a2a32524c693a65a8c208d1487d Mon Sep 17 00:00:00 2001 From: andy_yang Date: Mon, 13 Feb 2023 19:35:34 +0800 Subject: [PATCH 08/10] [TIR] Introduce Pass InjectPTXLDG32 This PR introduces a new pass InjectPTXLDG32 to change the `if_then_else` call node to `ptx_pred_ldg32` call node. When the store buffer is local and the load value is global, the pass can change the if_then_else pattern to a ptx pattern. Test the pass with: ```python with tvm.transform.PassContext(config={"tir.ptx_pred_ldg32": True}): mod = tvm.build(f, target="cuda") ```` --- include/tvm/tir/builtin.h | 6 + include/tvm/tir/transform.h | 6 + src/driver/driver_api.cc | 12 ++ src/target/source/codegen_cuda.cc | 31 +++++ src/tir/op/builtin.cc | 2 + src/tir/transforms/inject_ptx_ldg32.cc | 130 ++++++++++++++++++ .../python/unittest/test_inject_ptx_ldg32.py | 67 +++++++++ 7 files changed, 254 insertions(+) create mode 100644 src/tir/transforms/inject_ptx_ldg32.cc create mode 100644 tests/python/unittest/test_inject_ptx_ldg32.py diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h index d830ea579aa7..5f6315001102 100644 --- a/include/tvm/tir/builtin.h +++ b/include/tvm/tir/builtin.h @@ -610,6 +610,12 @@ TVM_DLL const Op& tvm_store_matrix_sync(); */ TVM_DLL const Op& ptx_mma(); +/*! + * \brief tvm intrinsic for ptx predicate load with 32-bit data type. + * + */ +TVM_DLL const Op& ptx_ldg32(); + /*! * \brief tvm intrinsic for sparse tensor core ptx instructions. * diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h index be7589b04bf5..85b381a52950 100644 --- a/include/tvm/tir/transform.h +++ b/include/tvm/tir/transform.h @@ -677,6 +677,12 @@ TVM_DLL Pass Filter(runtime::TypedPackedFunc fcond); */ TVM_DLL Pass InjectPTXAsyncCopy(); +/*! + * \brief Pass to rewrite global to local memory copy on CUDA with ldg32 instruction. + * \return The pass. + */ +TVM_DLL Pass InjectPTXLDG32(bool enable_ptx_ldg32 = true); + /*! * \brief Remove the weight layout rewrite block * \param skip_ndarray_rewrite If True, exact rewrite of NDArray, according to the given index map, diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc index 288ac7b92a2c..da1bbc296a49 100644 --- a/src/driver/driver_api.cc +++ b/src/driver/driver_api.cc @@ -55,6 +55,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.use_async_copy", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.merge_async_commit_queue_scope", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.instrument_lwp", Bool); TVM_REGISTER_PASS_CONFIG_OPTION("tir.vtcm_capacity", Integer); +TVM_REGISTER_PASS_CONFIG_OPTION("tir.ptx_ldg32", Bool); // WARNING: May cause coherency issues resulting data miscompares // Experimental feature that, when enabled by the runtime, bypasses the cache when using DMA. When @@ -159,6 +160,8 @@ Array CreatePassList(bool disable_loop_partition) { bool enable_equiv_terms_in_cse_tir = pass_ctx->GetConfig("tir.enable_equiv_terms_in_cse_tir", Bool(false)).value(); + bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value(); + // Get any user-added passes Array> add_lower_pass = pass_ctx->GetConfig>>("tir.add_lower_pass", Array>()) @@ -257,6 +260,10 @@ Array CreatePassList(bool disable_loop_partition) { pass_list.push_back(tir::transform::InstrumentBoundCheckers()); } + if (ptx_ldg32) { + pass_list.push_back(tir::transform::InjectPTXLDG32(true)); + } + pass_list.push_back( tir::transform::CommonSubexprElimTIR(!disable_cse_tir, enable_equiv_terms_in_cse_tir)); @@ -584,6 +591,11 @@ transform::Sequential MixedModulePassManager(IRModule mixed_mod, Target target) mixed_pass_list.push_back(tir::transform::InjectPTXAsyncCopy()); } + bool ptx_ldg32 = pass_ctx->GetConfig("tir.ptx_ldg32", Bool(false)).value(); + if (ptx_ldg32) { + mixed_pass_list.push_back(tir::transform::InjectPTXLDG32()); + } + bool unpacked_api = mixed_mod->GetAttr(tvm::attr::kExecutor) .value_or(relay::Executor::Create("graph", {})) ->GetAttr("unpacked-api") diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc index 9bf0109cace1..20607254370e 100644 --- a/src/target/source/codegen_cuda.cc +++ b/src/target/source/codegen_cuda.cc @@ -926,6 +926,37 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { } else if (op->op.same_as(builtin::ptx_wait_group())) { std::string N = this->PrintExpr(op->args[0]); this->stream << "__asm__ __volatile__(\"cp.async.wait_group " + N + ";\");\n\n"; + } else if (op->op.same_as(builtin::ptx_ldg32())) { + /* + asm volatile ( + "{.reg .pred p;\n" + " setp.ne.b32 p, %2, 0;\n" + // " @p ld.global.nc.f32 %0, [%1];}\n"t + " @p ld.global.nc.L2::128B.f32 %0, [%1];}\n" + : "=f"(reg) + : "l"(addr), "r"((int)guard) + ); + */ + + // get local + std::string reg = this->PrintExpr(op->args[0]); + // get guard + std::string guard = this->PrintExpr(op->args[1]); + const BufferLoadNode* addr_buffer = op->args[2].as(); + std::string global_addr = this->PrintExpr(addr_buffer->indices[0]); + std::string global_buffer = this->PrintExpr(addr_buffer->buffer->data); + std::string local_addr = this->PrintExpr(op->args[3]); + this->stream << "asm volatile (\n"; + this->stream << "\"{.reg .pred p;\\n\"\n"; + this->stream << "\" setp.ne.b32 p, %2, 0;\\n\"\n"; + this->stream << "\" @!p mov.b32 %0, 0;\\n\"\n"; + this->stream << "\" @p ld.global.nc.f32 %0, [%1];}\\n\"\n"; + // stream << "\" @p ld.global.nc.L2::128B.f32 %0, [%1];}\\n\"\n" ; + stream << ": \"=f\"(" << reg << "[" << local_addr << "]" + << ")\n"; + stream << ": \"l\"((void*)(" << global_buffer << "+" << global_addr << ")), \"r\"((int)" + << guard << ")\n"; + stream << ");\n"; } else { CodeGenC::VisitExpr_(op, os); } diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc index dc3208f484e3..680202751f12 100644 --- a/src/tir/op/builtin.cc +++ b/src/tir/op/builtin.cc @@ -251,6 +251,8 @@ TIR_DEFINE_BUILTIN_FUNC(tvm_store_matrix_sync) TIR_DEFINE_BUILTIN_FUNC(ptx_mma).set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); +TIR_DEFINE_BUILTIN_FUNC(ptx_ldg32).set_num_inputs(4).set_attr( + "TCallEffectKind", Integer(CallEffectKind::kPure)); TIR_DEFINE_BUILTIN_FUNC(ptx_mma_sp) .set_attr("TCallEffectKind", Integer(CallEffectKind::kOpaque)); diff --git a/src/tir/transforms/inject_ptx_ldg32.cc b/src/tir/transforms/inject_ptx_ldg32.cc new file mode 100644 index 000000000000..b4c398bd17eb --- /dev/null +++ b/src/tir/transforms/inject_ptx_ldg32.cc @@ -0,0 +1,130 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../../arith/const_fold.h" +#include "../../arith/pattern_match.h" + +namespace tvm { +namespace tir { + +class PTXRewriter : public StmtMutator { + public: + Stmt VisitStmt_(const AllocateNode* allocate) final { + if (!has_buffer_1) { + has_buffer_1 = true; + // addr[0] -> global_addr / addr[1] -> local_addr + addr_buffer = decl_buffer({IntImm(DataType::Int(32), 2)}, DataType::Int(32), "addr", "local"); + predicate_buffer = + decl_buffer({IntImm(DataType::Int(32), 1)}, DataType::Bool(1), "predicate", "local"); + } + Stmt result = StmtMutator::VisitStmt_(allocate); + if (!has_buffer_2) { + has_buffer_2 = true; + result = + Allocate(addr_buffer->data, addr_buffer->dtype, addr_buffer->shape, Bool(true), result); + result = Allocate(predicate_buffer->data, predicate_buffer->dtype, predicate_buffer->shape, + Bool(true), result); + } + return result; + } + + Stmt VisitStmt_(const BufferStoreNode* store) final { + Stmt result = StmtMutator::VisitStmt_(store); + Buffer load_buffer = store->buffer; + PrimExpr load_value = store->value; + // const BufferLoadNode* gload = load_value.as(); // take + // the place of instance of + const CallNode* call = load_value.as(); + if (call != nullptr) { + const OpNode* op = call->op.as(); + if (op != nullptr && op->name == "tir.if_then_else") { + const PrimExpr& predicate = call->args[0]; + const PrimExpr& lhs = call->args[1]; + const PrimExpr& rhs = call->args[2]; + PrimExpr global_addr, local_addr; + const BufferLoadNode* load = lhs.as(); + PrimExpr imm_value = rhs; + if (load == nullptr) { + load = rhs.as(); + imm_value = lhs; + if (load == nullptr) { + return result; + } + } + global_addr = load->indices[0]; + const RampNode* ramp = global_addr.as(); + if (ramp != nullptr) { + return result; + } + local_addr = store->indices[0]; + BufferStore addr_store(addr_buffer, global_addr, {IntImm(DataType::Int(32), 0)}); + BufferStore local_addr_store(addr_buffer, local_addr, {IntImm(DataType::Int(32), 1)}); + BufferStore predicate_store(predicate_buffer, predicate, {IntImm(DataType::Int(32), 0)}); + PrimExpr new_lhs, new_rhs, new_predicate, new_indice; + new_lhs = + BufferLoad(load->buffer, {BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 0)})}); + new_rhs = IntImm(DataType::Int(32), 0); + new_predicate = BufferLoad(predicate_buffer, {IntImm(DataType::Int(32), 0)}); + new_indice = BufferLoad(addr_buffer, {IntImm(DataType::Int(32), 1)}); + BufferStore value_store(store->buffer, imm_value, {new_indice}); + Evaluate ptx_load(Call(store->buffer->dtype, tvm::tir::builtin::ptx_ldg32(), + {store->buffer->data, new_predicate, new_lhs, new_indice})); + Array tmp_seq = {addr_store, local_addr_store, predicate_store, value_store, + ptx_load}; + SeqStmt seq_stmt = SeqStmt(tmp_seq); + return seq_stmt; + } + } + return result; + } + + bool has_buffer_1 = false, has_buffer_2 = false; + Buffer addr_buffer, predicate_buffer; +}; + +namespace transform { + +Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) { + auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m, PassContext ctx) { + if (enable_inject_ptx_intrin) { + auto* n = f.CopyOnWrite(); + n->body = PTXRewriter()(n->body); + // inject ptx + } + return f; + }; + return CreatePrimFuncPass(pass_func, 0, "tir.InjectPTXLDG32", {}); +} + +// The pass can now be invoked via the pass infrastructure, but we also add a +// Python binding for it +TVM_REGISTER_GLOBAL("tir.transform.InjectPTXLDG32").set_body_typed(InjectPTXLDG32); + +} // namespace transform +} // namespace tir +} // namespace tvm diff --git a/tests/python/unittest/test_inject_ptx_ldg32.py b/tests/python/unittest/test_inject_ptx_ldg32.py new file mode 100644 index 000000000000..c92bdd8c6488 --- /dev/null +++ b/tests/python/unittest/test_inject_ptx_ldg32.py @@ -0,0 +1,67 @@ +# 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. +import tvm +from tvm.script import tir as T +import numpy as np +import tvm.testing + + +@T.prim_func +def vector_add(A: T.Buffer[(16), "float32"], B: T.Buffer[(32), "float32"]) -> None: + T.func_attr({"global_symbol": "default_function", "tir.noalias": True}) + bx = T.env_thread("blockIdx.x") + tx = T.env_thread("threadIdx.x") + T.launch_thread(bx, 1) + T.launch_thread(tx, 32) + A_local = T.alloc_buffer([32], "float32", scope="local") + + with T.block(): + T.reads(A[0:16]) + T.writes(A_local[0:32]) + A_local[tx] = T.if_then_else(tx % 2 == 0, A[tx / 2], T.float32(0), dtype="float32") + B[tx] = A_local[tx] + 1.0 + + +@tvm.testing.requires_cuda +def test_inject_ptx_intrin(): + f = vector_add + arch = tvm.contrib.nvcc.get_target_compute_version() + major, _ = tvm.contrib.nvcc.parse_compute_version(arch) + if major < 8: + # Require at least SM80 + return + with tvm.transform.PassContext(config={"tir.ptx_ldg32": True}): + mod = tvm.build(f, target="cuda") + A_np = np.random.rand(16).astype("float32") + B_np = np.zeros((32)).astype("float32") + dev = tvm.cuda(0) + A_nd = tvm.nd.array(A_np, device=dev) + B_nd = tvm.nd.array(B_np, device=dev) + mod(A_nd, B_nd) + + C_np = np.zeros((32)).astype("float32") + + for i in range(32): + if i % 2 == 0: + C_np[i] = A_np[i // 2] + C_np[i] += 1.0 + + tvm.testing.assert_allclose(B_nd.numpy(), C_np) + + +if __name__ == "__main__": + test_inject_ptx_intrin() From d86324ad3725ccfb3a408e0ec94c5bbe7336bc75 Mon Sep 17 00:00:00 2001 From: andy_yang Date: Sat, 18 Feb 2023 10:18:23 +0800 Subject: [PATCH 09/10] Update test_inject_ptx_ldg32.py --- tests/python/unittest/test_inject_ptx_ldg32.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/unittest/test_inject_ptx_ldg32.py b/tests/python/unittest/test_inject_ptx_ldg32.py index c92bdd8c6488..81c6e89ad921 100644 --- a/tests/python/unittest/test_inject_ptx_ldg32.py +++ b/tests/python/unittest/test_inject_ptx_ldg32.py @@ -21,13 +21,13 @@ @T.prim_func -def vector_add(A: T.Buffer[(16), "float32"], B: T.Buffer[(32), "float32"]) -> None: +def vector_add(A: T.Buffer((16), "float32"), B: T.Buffer((32), "float32")) -> None: T.func_attr({"global_symbol": "default_function", "tir.noalias": True}) bx = T.env_thread("blockIdx.x") tx = T.env_thread("threadIdx.x") T.launch_thread(bx, 1) T.launch_thread(tx, 32) - A_local = T.alloc_buffer([32], "float32", scope="local") + A_local = T.Buffer((32), "float32", scope="local") with T.block(): T.reads(A[0:16]) From b72d775c08624078e8687957772abb2b972552c1 Mon Sep 17 00:00:00 2001 From: andy_yang Date: Sat, 18 Feb 2023 13:18:32 +0800 Subject: [PATCH 10/10] Update builtin.h --- include/tvm/tir/builtin.h | 1 - 1 file changed, 1 deletion(-) diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h index 12d789676cf1..708abde2cd31 100644 --- a/include/tvm/tir/builtin.h +++ b/include/tvm/tir/builtin.h @@ -610,7 +610,6 @@ TVM_DLL const Op& tvm_store_matrix_sync(); */ TVM_DLL const Op& ptx_mma(); - /*! * \brief tvm intrinsic for ptx predicate load with 32-bit data type. *