Address code review comments

Signed-off-by: Tiotto, Ettore <[email protected]>

Author: etiotto

third_party/intel/backend/compiler.py

@@ -235,7 +235,7 @@ def make_ttgir(mod, metadata, opt, properties):
         intel.passes.ttgpuir.add_accelerate_matmul(pm)
         intel.passes.ttgpuir.add_remove_layout_conversions(pm)
         intel.passes.ttgpuir.add_materialize_block_pointer(pm)
-        intel.passes.ttgpuir.add_rewrite_tensor_pointer(pm)
+        #        intel.passes.ttgpuir.add_rewrite_tensor_pointer(pm)
         intel.passes.ttgpuir.add_pipeline(pm, opt.num_stages, False)
 
         intel.passes.ttgpuir.add_coalesce(pm)

third_party/intel/include/Analysis/AxisInfo.h

@@ -11,7 +11,6 @@
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 
 #include <optional>
-#include <type_traits>
 
 namespace mlir::triton::intel {
 

third_party/intel/include/Dialect/TritonIntelGPU/IR/Utils.h

@@ -46,7 +46,6 @@ inline unsigned getNumElementsPerThread(
 inline bool applyTransposedReduction() {
   return tools::getBoolEnv("TRITON_INTEL_REDUCE_TRANSPOSE");
 }
-
 } // namespace mlir::triton::gpu::intel
 
 #endif // TRITON_DIALECT_TRITON_INTEL_GPU_IR_UTILS_H

third_party/intel/lib/Analysis/AxisInfo.cpp

@@ -1,12 +1,10 @@
+#include "intel/include/Analysis/AxisInfo.h"
 #include "mlir/Analysis/DataFlowFramework.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
-#include "intel/include/Analysis/AxisInfo.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "triton/Dialect/Triton/IR/Dialect.h"
-
 #define DEBUG_TYPE "intel-axis-info"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")

third_party/intel/lib/TritonIntelGPUTransforms/Coalesce.cpp

@@ -1,19 +1,18 @@
 #include "intel/include/Analysis/AxisInfo.h"
-#include "intel/include/Dialect/TritonIntelGPU/IR/Dialect.h"
+// #include "intel/include/Dialect/TritonIntelGPU/IR/Dialect.h"
 #include "intel/include/Dialect/TritonIntelGPU/IR/Utils.h"
 #include "intel/include/Dialect/TritonIntelGPU/Transforms/Passes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Support/LLVM.h"
-#include "triton/Dialect/Triton/IR/Dialect.h"
-#include "triton/Dialect/Triton/IR/Types.h"
+// #include "triton/Dialect/Triton/IR/Dialect.h"
+// #include "triton/Dialect/Triton/IR/Types.h"
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
 #include "triton/Tools/StrUtil.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
-#include <variant>
 
 #define DEBUG_TYPE "tritonintelgpu-coalesce"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
@@ -104,8 +103,8 @@ struct CoalescePass
     perThread = std::min<int>(perThread, std::max(numElems / numThreads, 1));
     LDBG("perThread: " << perThread);
 
-    if (perThread <= 1)
-      return;
+    //    if (perThread <= 1)
+    //      return;
 
     if (!dyn_cast<triton::LoadOp>(op)) {
       // For ops that can result in a global memory write, we should enforce
@@ -299,7 +298,6 @@ struct CoalescePass
     LDBG("Coalescing op: " << *op);
 
     OpBuilder builder(op);
-    IRRewriter rewriter(builder);
 
     // Convert operands
     // Note: for load/store with a blocked pointers argument we cannot change
@@ -312,14 +310,15 @@ struct CoalescePass
       if (tensorType &&
           !isa<ttg::SharedEncodingAttr>(tensorType.getEncoding())) {
         RankedTensorType newType = getNewType(tensorType, encoding);
-        newArgs.push_back(rewriter.create<ttg::ConvertLayoutOp>(
+        newArgs.push_back(builder.create<ttg::ConvertLayoutOp>(
             op->getLoc(), newType, operand));
       } else {
         assert(isa<tt::PointerType>(operand.getType()) &&
                "Expecting operand to have blocked pointer type");
         auto defOp = findDefiningMakeTensorPtrOp(operand);
         assert(defOp && "Expected a make_tensor_ptr operation");
         LDBG("Found make_tensor_ptr definition: " << *defOp);
+        IRRewriter rewriter(builder);
         changeAndPropagateLayout(*defOp, encoding, rewriter);
         newArgs.push_back(operand);
       }
@@ -335,14 +334,14 @@ struct CoalescePass
 
     // Construct new op with the new encoding.
     Operation *newOp =
-        rewriter.create(op->getLoc(), op->getName().getIdentifier(), newArgs,
-                        newTypes, op->getAttrs());
+        builder.create(op->getLoc(), op->getName().getIdentifier(), newArgs,
+                       newTypes, op->getAttrs());
 
     // Cast the results back to the original layout.
     for (size_t i = 0; i < op->getNumResults(); i++) {
       Value newResult = newOp->getResult(i);
       if (newTypes[i] != op->getResultTypes()[i]) {
-        newResult = rewriter.create<ttg::ConvertLayoutOp>(
+        newResult = builder.create<ttg::ConvertLayoutOp>(
             op->getLoc(), op->getResult(i).getType(), newResult);
       }
       op->getResult(i).replaceAllUsesWith(newResult);
@@ -400,11 +399,7 @@ struct CoalescePass
       coalesceOp(layout, op);
     }
 
-    // Verify the module's functions after the transformation.
-    for (auto op : moduleOp.getOps<tt::FuncOp>()) {
-      for (Operation &op1 : op.getOps())
-        assert(succeeded(verify(&op1)));
-    }
+    assert(succeeded(verify(moduleOp)) && "Module verification failed");
   }
 };