Merge denzel and ethan work, and squash

Author: denzel-farmer

.gitignore

@@ -1,3 +1,6 @@
+
+.working/
+output/
 # Triton builds
 build/
 build-*/

README.md

@@ -6,6 +6,10 @@
 |-------------------- | -------------------- |
 | [![Documentation](https://github.com/triton-lang/triton/actions/workflows/documentation.yml/badge.svg)](https://triton-lang.org/) | [![Wheels](https://github.com/triton-lang/triton/actions/workflows/wheels.yml/badge.svg?branch=release/2.0.x)](https://github.com/triton-lang/triton/actions/workflows/wheels.yml) |
 
+# COMS E6998 Runtime Analysis Project
+
+See [our readme](runtime-analysis/README.md) in the `runtime-analysis` folder for details.
+
 # Triton
 
 This is the development repository of Triton, a language and compiler for writing highly efficient custom Deep-Learning primitives. The aim of Triton is to provide an open-source environment to write fast code at higher productivity than CUDA, but also with higher flexibility than other existing DSLs.

bin/benchmark-lower.cpp

@@ -0,0 +1,11 @@
+#include "./RegisterTritonDialects.h"
+
+#include "mlir/Tools/mlir-reduce/MlirReduceMain.h"
+
+int main(int argc, char **argv) {
+  mlir::DialectRegistry registry;
+  registerTritonDialects(registry);
+
+  mlir::MLIRContext context(registry);
+  return mlir::failed(mlir::mlirReduceMain(argc, argv, context));
+}

commands.md

@@ -9,6 +9,7 @@ namespace mlir {
 std::unique_ptr<Pass> createLLVMDIScopePass();
 
 std::unique_ptr<Pass> createLLVMRuntimeAnalysisPass();
+std::unique_ptr<Pass> createLLVMRuntimeAnalysisPass(const std::string &analysisResultFile);
 
 /// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION

empty_kernel.cpp

@@ -20,6 +20,11 @@ def LLVMRuntimeAnalysis: Pass<"runtime-analysis", "mlir::ModuleOp"> {
   }];
 
   let constructor = "mlir::createLLVMRuntimeAnalysisPass()";
+
+  let options = [
+    Option<"result_file", "result_file", "std::string", /*default*/"\"\"",
+           "path to place analysis results">,
+  ];
 }
 
 #endif

include/triton/Target/LLVMIR/Passes.h

@@ -0,0 +1,9 @@
+// Denzel / Ethan Project, Analyze per-kernel Runtime 
+
+
+
+// Components: 
+// - kernel launch overhead, depends on parameters for kernel launch (number of threads, number of blocks per thread)
+//      - Do some profiling, build a linear regression 
+// - compute-bound instruction latency 
+//     - Do some profiling, build a linear regression

include/triton/Target/LLVMIR/Passes.td

@@ -21,6 +21,7 @@ add_triton_library(TritonGPUToLLVM
     SPMDOpToLLVM.cpp
     DecomposeUnsupportedConversions.cpp
     PrintOpToLLVM.cpp
+    LoopHintPass.cpp
 
     DEPENDS
     TritonGPUConversionPassIncGen

lib/Analysis/Runtime.cpp

@@ -0,0 +1,77 @@
+#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+/// Matches scf -> cf loop pattern:
+///   %cmp = arith.cmpi slt, %lhs, %const : i32
+///   cf.cond_br %cmp, ^loop, ^exit
+///
+/// and annotates the `cf.cond_br` with loop iteration count.
+
+namespace {
+
+struct AnnotateLoopCondBrPattern : public OpRewritePattern<cf::CondBranchOp> {
+  using OpRewritePattern<cf::CondBranchOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(cf::CondBranchOp condBrOp,
+                                PatternRewriter &rewriter) const override {
+
+    auto cmpOp = condBrOp.getCondition().getDefiningOp<arith::CmpIOp>();
+    if (!cmpOp)
+      return failure();
+
+    if (cmpOp.getPredicate() != arith::CmpIPredicate::slt)
+      return failure();
+
+    auto rhsOp = cmpOp.getRhs().getDefiningOp<arith::ConstantOp>();
+    if (!rhsOp)
+      return failure();
+
+    auto intAttr = mlir::dyn_cast<mlir::IntegerAttr>(rhsOp.getValue());
+    if (!intAttr)
+      return failure();
+
+    int64_t loopCount = intAttr.getInt();
+
+    rewriter.setInsertionPointAfter(condBrOp);
+    condBrOp->setAttr("loop_count", rewriter.getI64IntegerAttr(loopCount));
+
+    return success();
+  }
+};
+
+/// A pass that runs on a function and looks for any cf.cond_br that
+/// implements a loop condition (in the pattern above), then annotates
+/// it with `loop_count`.
+struct AnnotateLoopCountPass
+    : public PassWrapper<AnnotateLoopCountPass, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AnnotateLoopCountPass)
+
+  StringRef getArgument() const final { return "annotate-loop-count"; }
+  StringRef getDescription() const final {
+    return "Annotate cf.cond_br ops with loop iteration counts.";
+  }
+
+  void runOnOperation() override {
+    auto function = getOperation();
+
+    RewritePatternSet patterns(&getContext());
+    patterns.add<AnnotateLoopCondBrPattern>(&getContext());
+
+    if (failed(applyPatternsAndFoldGreedily(function.getBody(),
+                                            std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createAnnotateLoopCountPass() {
+  return std::make_unique<AnnotateLoopCountPass>();
+}