Skip lowering triton bitcast on tensor of pointers (#250)

We currently only skip lowering tt.bitcast if the input is of triton
pointer type. tt.bitcast can also take tensor of pointers, so we need to
skip lowering in that case as well.

Author: nhat-nguyen

include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp

@@ -12,6 +12,7 @@
 #include "triton-shared/Analysis/OpFoldResultUtils.h"
 #include "triton-shared/Analysis/PtrAnalysis.h"
 #include "triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtDialect.h"
+#include "triton-shared/Utils/Utils.h"
 
 #include "triton/Dialect/Triton/IR/Dialect.h"
 
@@ -853,7 +854,7 @@ struct BitcastConverter : public OpConversionPattern<triton::BitcastOp> {
   matchAndRewrite(triton::BitcastOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // arith::bitcast does not support casting pointers
-    if (isa<triton::PointerType>(op.getSrc().getType())) {
+    if (triton::isPtrTypeLike(op.getType())) {
       return failure();
     }
 
@@ -1577,7 +1578,7 @@ class ArgMinMaxBaseConverter : public OpConversionPattern<triton::ReduceOp> {
     } else {
       return failure();
     }
-    
+
     auto loc = op.getLoc();
 
     auto elemTypes = op.getElementTypes();

include/triton-shared/Utils/Utils.h

@@ -0,0 +1,14 @@
+#ifndef TRITON_SHARED_UTILITY_H
+#define TRITON_SHARED_UTILITY_H
+
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+namespace mlir {
+namespace triton {
+// Return true if the input type is a triton pointer or a tensor of triton pointers
+bool isPtrTypeLike(Type t);
+} // namespace triton
+
+} // namespace mlir
+
+#endif // TRITON_SHARED_UTILITY_H

lib/CMakeLists.txt

@@ -2,3 +2,4 @@ add_subdirectory(Analysis)
 add_subdirectory(AnalysisStructured)
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
+add_subdirectory(Utils)

lib/Conversion/TritonArithToLinalg/CMakeLists.txt

@@ -19,4 +19,5 @@ add_triton_library(TritonArithToLinalg
   TritonTransforms
   TritonTilingExtIR
   TritonStructuredIR
+  TritonSharedUtils
 )

lib/Conversion/TritonArithToLinalg/TritonArithToLinalgPass.cpp

@@ -10,6 +10,7 @@
 #include "triton-shared/Conversion/TritonArithToLinalg/TritonArithToLinalg.h"
 #include "triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.h"
 #include "triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtDialect.h"
+#include "triton-shared/Utils/Utils.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
@@ -123,7 +124,7 @@ class TritonArithToLinalgPass
     target.addLegalOp<triton::FuncOp, triton::ReturnOp>();
 
     target.addDynamicallyLegalOp<triton::BitcastOp>([](triton::BitcastOp op) {
-      return isa<triton::PointerType>(op.getSrc().getType());
+      return triton::isPtrTypeLike(op.getType());
     });
 
     target.addDynamicallyLegalDialect<arith::ArithDialect, math::MathDialect>(

lib/Utils/CMakeLists.txt

@@ -0,0 +1,6 @@
+add_triton_library(TritonSharedUtils
+  Utils.cpp
+
+  LINK_LIBS PUBLIC
+  TritonIR
+)

lib/Utils/Utils.cpp

@@ -0,0 +1,13 @@
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+namespace mlir {
+namespace triton {
+bool isPtrTypeLike(Type t) {
+  if (auto tensorType = dyn_cast<RankedTensorType>(t)) {
+    return isa<triton::PointerType>(tensorType.getElementType());
+  }
+  return isa<triton::PointerType>(t);
+}
+} // namespace triton
+
+} // namespace mlir