[MatMul][WIP] MatMul epilog support: share memory reuse

torch/csrc/jit/codegen/cuda/codegen.cpp

@@ -617,18 +617,25 @@ class CudaKernelGenerator : private OptOutConstDispatch {
  // Utility function to emit a cp.async intrinsic
  void genCpAsync(const LoadStoreOp* ldst, int vec_size) {
  auto dtype = ldst->in()->getDataType().value();
+ bool is_cg = ldst->opType() == LoadStoreOpType::CpAsyncCg;
+
+ if (is_cg) {
+ indent() << "Ampere::cpAsyncCg";
+ } else {
+ indent() << "Ampere::cpAsync";
+ }
 
  if (ldst->predicate() == nullptr) {
  // Out of line predicate variant
- indent() << "Ampere::cpAsync<" << dtype << "," << vec_size << ">("
-  << genMaybeHoistedPointer(ldst->out()) << ","
-  << genMaybeHoistedPointer(ldst->in()) << ");\n";
+ code_ << "<" << dtype << "," << vec_size << ">("
+ << genMaybeHoistedPointer(ldst->out()) << ","
+ << genMaybeHoistedPointer(ldst->in()) << ");\n";
  } else {
  // Inline predicate variant
- indent() << "Ampere::cpAsync<" << dtype << "," << vec_size << ">("
-  << genMaybeHoistedPointer(ldst->out()) << ","
-  << genMaybeHoistedPointer(ldst->in()) << ","
-  << genInline(ldst->predicate()) << ");\n";
+ code_ << "<" << dtype << "," << vec_size << ">("
+ << genMaybeHoistedPointer(ldst->out()) << ","
+ << genMaybeHoistedPointer(ldst->in()) << ","
+ << genInline(ldst->predicate()) << ");\n";
  }
  }
 
@@ -1402,6 +1409,7 @@ class CudaKernelGenerator : private OptOutConstDispatch {
  genLdMatrix(ldst, vector_word_size);
  break;
  case LoadStoreOpType::CpAsync:
+ case LoadStoreOpType::CpAsyncCg:
  genCpAsync(ldst, vector_word_size);
  break;
  default:
@@ -2618,6 +2626,18 @@ class CudaKernelGenerator : private OptOutConstDispatch {
  indent() << sync_call << ";\n";
  }
 
+ void handle(const kir::Deallocate* dealloc) final {
+ const auto alloc = dealloc->buffer();
+ const auto tv = alloc->buffer()->as<TensorView>();
+ const auto size = alloc->size();
+ const auto buffer_dtype = alloc->buffer()->dtype();
+
+ TORCH_INTERNAL_ASSERT(size != nullptr);
+ indent() << "// de-alloc " << varName(tv) << "\n";
+ indent() << "offset -= (" << genInline(size) << " * sizeof(" << buffer_dtype
+ << "));\n";
+ }
+
  void handle(const kir::InitMagicZero*) final {
  indent() << "NVFUSER_DEFINE_MAGIC_ZERO\n";
  }

torch/csrc/jit/codegen/cuda/dispatch.cpp

@@ -154,6 +154,9 @@ void Expr::dispatch(T handler, Expr* expr) {
  case ExprType::Allocate:
  ptr(handler)->handle(expr->as<kir::Allocate>());
  return;
+ case ExprType::DeAllocate:
+ ptr(handler)->handle(expr->as<kir::Deallocate>());
+ return;
  case ExprType::BlockSync:
  ptr(handler)->handle(expr->as<kir::BlockSync>());
  return;
@@ -331,6 +334,9 @@ void Expr::constDispatch(T handler, const Expr* expr) {
  case ExprType::Allocate:
  ptr(handler)->handle(expr->as<kir::Allocate>());
  return;
+ case ExprType::DeAllocate:
+ ptr(handler)->handle(expr->as<kir::Deallocate>());
+ return;
  case ExprType::BlockSync:
  ptr(handler)->handle(expr->as<kir::BlockSync>());
  return;
@@ -516,6 +522,9 @@ void Expr::mutatorDispatch(T mutator, Expr* expr) {
  case ExprType::Allocate:
  ptr(mutator)->mutate(expr->as<kir::Allocate>());
  return;
+ case ExprType::DeAllocate:
+ ptr(mutator)->mutate(expr->as<kir::Deallocate>());
+ return;
  case ExprType::BlockSync:
  ptr(mutator)->mutate(expr->as<kir::BlockSync>());
  return;
@@ -766,6 +775,9 @@ void OptOutConstDispatch::handle(const ViewOp* stmt) {
 void OptOutConstDispatch::handle(const kir::Allocate* stmt) {
  unhandled(stmt);
 }
+void OptOutConstDispatch::handle(const kir::Deallocate* stmt) {
+ unhandled(stmt);
+}
 void OptOutConstDispatch::handle(const kir::BlockSync* stmt) {
  unhandled(stmt);
 }
@@ -913,6 +925,9 @@ void OptOutDispatch::handle(ViewOp* stmt) {
 void OptOutDispatch::handle(kir::Allocate* stmt) {
  unhandled(stmt);
 }
+void OptOutDispatch::handle(kir::Deallocate* stmt) {
+ unhandled(stmt);
+}
 void OptOutDispatch::handle(kir::BlockSync* stmt) {
  unhandled(stmt);
 }

torch/csrc/jit/codegen/cuda/dispatch.h

@@ -95,6 +95,7 @@ class TensorIndex;
 class IntPair;
 
 class Allocate;
+class Deallocate;
 class BlockSync;
 class GridSync;
 class CpAsyncWait;
@@ -162,6 +163,7 @@ class TORCH_CUDA_CU_API OptOutConstDispatch : public PolymorphicBase {
  virtual void handle(const ViewOp* stmt);
 
  virtual void handle(const kir::Allocate*);
+ virtual void handle(const kir::Deallocate*);
  virtual void handle(const kir::BlockSync*);
  virtual void handle(const kir::GridSync*);
  virtual void handle(const kir::CpAsyncWait*);
@@ -226,6 +228,7 @@ class TORCH_CUDA_CU_API OptOutDispatch : public PolymorphicBase {
  virtual void handle(ViewOp* stmt);
 
  virtual void handle(kir::Allocate* stmt);
+ virtual void handle(kir::Deallocate* stmt);
  virtual void handle(kir::BlockSync* stmt);
  virtual void handle(kir::GridSync* stmt);
  virtual void handle(kir::CpAsyncWait* stmt);
@@ -331,6 +334,7 @@ class TORCH_CUDA_CU_API OptOutMutator : public PolymorphicBase {
  virtual void mutate(ViewOp*);
 
  virtual void mutate(kir::Allocate*);
+ virtual void mutate(kir::Deallocate*);
  virtual void mutate(kir::BlockSync*);
  virtual void mutate(kir::GridSync*);
  virtual void mutate(kir::CpAsyncWait*);

torch/csrc/jit/codegen/cuda/executor.cpp

@@ -442,6 +442,46 @@ uint64_t FusionExecutor::computeSharedMemory(
  return total;
 }
 
+uint64_t FusionExecutor::computeDynamicSharedMemory(
+ kir::ExpressionEvaluator& expr_eval,
+ const kir::AllocateRecord& records,
+ uint64_t total) {
+ FUSER_PERF_SCOPE("computeDynamicSharedMemory");
+ uint64_t max_total = total;
+ int N = records.actions.size();
+
+ for (int idx : c10::irange(N)) {
+ auto size = records.sizes[idx];
+ auto alloc = records.allocs[idx];
+
+ // If this buffer aliases another buffer,
+ // then do not allocate memory for this buffer.
+ const auto inferred_val = expr_eval.evaluate(size);
+ TORCH_INTERNAL_ASSERT(
+ inferred_val.has_value(),
+ "Failed to evaluate the size ",
+ size,
+ " of shared memory buffer ",
+ alloc->buffer()->toString());
+
+ const uint64_t data_size = dataTypeSize(alloc->buffer()->dtype());
+ const uint64_t buffer_size = data_size * inferred_val.value();
+
+ if (records.actions[idx] == kir::SmemAllocAction::Allocate) {
+ // Allocate:
+ const int align_size = 16; // always align to 16B/128b.
+ total = ceilDiv(total, align_size) * align_size;
+ total += buffer_size;
+ max_total = std::max(total, max_total);
+ } else {
+ // Deallocate:
+ total -= buffer_size;
+ }
+ }
+
+ return max_total;
+}
+
 LaunchParams FusionExecutor::computeLaunchParams(
  const LaunchParams& launch_constraints,
  kir::ExpressionEvaluator& expr_eval,
@@ -615,10 +655,9 @@ LaunchParams FusionExecutor::computeLaunchParams(
  }
 
  // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
- const uint64_t dynamic_smem_size = computeSharedMemory(
+ const uint64_t dynamic_smem_size = computeDynamicSharedMemory(
  expr_eval,
- kernel_summary.dynamic_smem_allocations,
- true,
+ kernel_summary.allocation_record,
  reduction_broadcast_workspace);
 
  // Check that requested smem size can be dynamically allocated.

torch/csrc/jit/codegen/cuda/executor.h

@@ -5,6 +5,7 @@
 #include <torch/csrc/jit/codegen/cuda/ir_all_nodes.h>
 #include <torch/csrc/jit/codegen/cuda/ir_cloner.h>
 #include <torch/csrc/jit/codegen/cuda/ir_printer.h>
+#include <torch/csrc/jit/codegen/cuda/kernel.h>
 #include <torch/csrc/jit/codegen/cuda/kernel_expr_evaluator.h>
 #include <torch/csrc/jit/codegen/cuda/lower2device.h>
 #include <torch/csrc/jit/codegen/cuda/utils.h>
@@ -179,6 +180,11 @@ class TORCH_CUDA_CU_API FusionExecutor : public NonCopyable {
  bool align_padding = false,
  uint64_t total = 0);
 
+ uint64_t computeDynamicSharedMemory(
+ kir::ExpressionEvaluator& expr_eval,
+ const kir::AllocateRecord& records,
+ uint64_t total = 0);
+
  // return a pair of vector of tensors, where tensors in the first vector are
  // not initialized, while the second vector contains zero-initiliazed tensors
  GlobalBuffers allocGlobalVals(kir::ExpressionEvaluator& expr_eval);

torch/csrc/jit/codegen/cuda/ir_iostream.cpp

@@ -629,6 +629,12 @@ void IrPrinter::handle(const kir::GridSync* node) {
  os_ << ")\n";
 }
 
+void IrPrinter::handle(const kir::Deallocate* node) {
+ indent() << "DeAllocate(";
+ handle(node->buffer());
+ os_ << ")\n";
+}
+
 void IrPrinter::handle(const kir::ForLoop* node) {
  indent() << "FOR ";
  handle(node->index());

torch/csrc/jit/codegen/cuda/ir_iostream.h

@@ -109,6 +109,7 @@ class TORCH_CUDA_CU_API IrPrinter : public OptInConstDispatch {
  void handle(const kir::ForLoop*) final;
  void handle(const kir::IfThenElse*) final;
  void handle(const kir::Allocate*) final;
+ void handle(const kir::Deallocate*) final;
  void handle(const kir::BlockSync*) final;
  void handle(const kir::GridSync*) final;
  void handle(const kir::CpAsyncWait*) final;

torch/csrc/jit/codegen/cuda/ir_nodes.cpp

@@ -1402,11 +1402,11 @@ void IterDomain::parallelize(ParallelType t) {
  // to make copies of the iterdomains. We might eventually just want
  // to lock these parallel types and not allowing any changes once
  // they are swizzled.
- TORCH_CHECK(
- t == ParallelType::Vectorize || t == ParallelType::TIDx ||
- t == ParallelType::Serial || t == ParallelType::Mma,
- "Parallel type other than serial, tidx, vectorize not allowed for mma swizzled ids",
- t);
+ // TORCH_CHECK(
+ //  t == ParallelType::Vectorize || t == ParallelType::TIDx ||
+ //  t == ParallelType::Serial || t == ParallelType::Mma,
+ //  "Parallel type other than serial, tidx, vectorize not allowed for mma
+ //  swizzled ids", t);
  }
 
  parallel_type_ = t;

torch/csrc/jit/codegen/cuda/kernel.cpp

@@ -70,6 +70,12 @@ class KernelIrScanner : private IrVisitor {
  break;
  case MemoryType::Shared:
  summary_.dynamic_smem_allocations.push_back(allocate);
+ if (!allocate->alias()) {
+ summary_.allocation_record.actions.push_back(
+ SmemAllocAction::Allocate);
+ summary_.allocation_record.sizes.push_back(allocate->size());
+ summary_.allocation_record.allocs.push_back(allocate);
+ }
  break;
  case MemoryType::Local:
  if (!ExpressionEvaluator::isConst(allocate->size())) {
@@ -80,6 +86,12 @@ class KernelIrScanner : private IrVisitor {
  }
  }
 
+ void handle(kir::Deallocate* deallocate) final {
+ summary_.allocation_record.actions.push_back(SmemAllocAction::DeAllocate);
+ summary_.allocation_record.sizes.push_back(deallocate->buffer()->size());
+ summary_.allocation_record.allocs.push_back(deallocate->buffer());
+ }
+
  void handle(UnaryOp* unary_op) final {
  if (unary_op->getUnaryOpType() == UnaryOpType::RandLike) {
  summary_.max_rng_offsets =

torch/csrc/jit/codegen/cuda/kernel.h

@@ -22,6 +22,20 @@ namespace fuser {
 namespace cuda {
 namespace kir {
 
+// Smem allocation-deallocation trace:
+enum class SmemAllocAction { Allocate = 0, DeAllocate };
+
+struct AllocateRecord {
+ // Vector of allocate or deallocate actions in a kernel
+ std::vector<SmemAllocAction> actions;
+
+ // Vector of sizes in corresponding allocations.
+ std::vector<Val*> sizes;
+
+ // Vector of alloc expressions.
+ std::vector<const kir::Allocate*> allocs;
+};
+
 //! Summary of interesting facts about the kernel
 // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
 struct KernelSummary {
@@ -96,6 +110,9 @@ struct KernelSummary {
 
  //! Track information on vectorized set operations for runtime validation
  std::vector<VectorizedSetInfo> vectorized_set_info;
+
+ //! Records allocation and deallocation of dynamic shared memory space.
+ AllocateRecord allocation_record;
 };
 
 class TORCH_CUDA_CU_API KernelPerformanceProfile {

torch/csrc/jit/codegen/cuda/kernel_ir.cpp

@@ -592,6 +592,19 @@ Allocate::Allocate(
  "IR type only valid for Kernel container.");
 }
 
+Deallocate::Deallocate(IrBuilderPasskey passkey, Allocate* buffer)
+ : Expr(passkey, ExprType::DeAllocate), buffer_(buffer) {
+ TORCH_INTERNAL_ASSERT(
+ buffer->buffer()->isA<TensorView>() &&
+ // Only shared mem is supported for now.
+ buffer->buffer()->as<TensorView>()->getMemoryType() ==
+ MemoryType::Shared);
+
+ TORCH_INTERNAL_ASSERT(
+ passkey.ir_container_->isA<kir::Kernel>(),
+ "IR type only valid for Kernel container.");
+}
+
 GridReduction::GridReduction(
  IrBuilderPasskey passkey,
  BinaryOpType reduction_op_type,

torch/csrc/jit/codegen/cuda/kernel_ir.h

@@ -270,6 +270,19 @@ class TORCH_CUDA_CU_API Allocate final : public Expr {
  const Allocate* alias_ = nullptr;
 };
 
+//! Deallocate a space that has been occupied by a buffer.
+class TORCH_CUDA_CU_API Deallocate final : public Expr {
+ public:
+ explicit Deallocate(IrBuilderPasskey passkey, Allocate* buffer);
+
+ auto buffer() const {
+ return buffer_;
+ }
+
+ private:
+ const Allocate* buffer_ = nullptr;
+};
+
 // Sync represents __syncthreads barrier for block level coordination.
 //
 // TODO(kir): change name to SyncThreads as we could have other barriers.