OpenXLA-specific changes.

- Add `BUILD` files. - Remove `python/` folder. - Disable short pointer. - Fixup expensiveLoadOrStore().
openxla · May 30, 2023 · 39ff695 · 39ff695
1 parent f29838a
commit 39ff695
Show file tree

Hide file tree

Showing 90 changed files with 681 additions and 40,488 deletions.
diff --git a/BUILD b/BUILD
diff --git a/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/lib/Dialect/TritonGPU/Transforms/Utility.cpp
@@ -93,6 +93,11 @@ bool expensiveLoadOrStore(Operation *op, Attribute &targetEncoding) {
   // same
   if (isSingleValue(op->getOperand(0)))
     return false;
+  // TODO(manany): Investigate with Openai why the change here
+  // https://github.com/openai/triton/commit/640f3c392184cd14291c1bca6a4795eb0f32a61a
+  // which introduces Case 2 causes breakage to this test
+  // //third_party/py/jax_triton/tests:pallas_test_sm80 --test_filter=test_fused_attention_bwd
+  return true;
   // Case 2: Tensor of pointers has more threads than elements
   // we can presume a high hit-rate that makes it cheap to load
   auto ptrType = op->getOperand(0).getType().cast<RankedTensorType>();

diff --git a/lib/Target/PTX/PTXTranslation.cpp b/lib/Target/PTX/PTXTranslation.cpp
@@ -49,7 +49,7 @@ std::string translateLLVMIRToPTX(llvm::Module &module, int cc, int version) {
   auto *shortPtr =
       static_cast<llvm::cl::opt<bool> *>(options["nvptx-short-ptr"]);
   assert(shortPtr);
-  shortPtr->setValue(true);
+  shortPtr->setValue(false);
   std::string sm = cc == 90 ? "sm_90a" : "sm_" + std::to_string(cc);
   // max PTX version
   int ptxMajor = maxPTX / 10;

diff --git a/python/MANIFEST.in b/python/MANIFEST.in
diff --git a/python/README.md b/python/README.md
diff --git a/python/examples/copy_strided.py b/python/examples/copy_strided.py
diff --git a/python/examples/empty.py b/python/examples/empty.py
diff --git a/python/pyproject.toml b/python/pyproject.toml
diff --git a/python/setup.py b/python/setup.py
diff --git a/python/src/extra/cuda.ll b/python/src/extra/cuda.ll