b2b793dfb5
Most notably, this PR: - changes the traits (and assembly format) of addptr so it can handle offsets that have arbitrary integer width. - adds support for `cat`
49 lines
1.9 KiB
TableGen
49 lines
1.9 KiB
TableGen
#ifndef TRITON_PATTERNS
|
|
#define TRITON_PATTERNS
|
|
|
|
include "mlir/Dialect/StandardOps/IR/Ops.td"
|
|
include "mlir/Dialect/Arithmetic/IR/ArithmeticOps.td"
|
|
include "triton/Dialect/Triton/IR/TritonOps.td"
|
|
|
|
|
|
// AddIOp(DotOp(a, b, c), d) and c==0 => DotOp(a, b, d)
|
|
// AddFOp(DotOp(a, b, c), d) and c==0 => DotOp(a, b, d)
|
|
|
|
// AddIOp(d, DotOp(a, b, c)) and c==0 => DotOp(a, b, d)
|
|
// AddFOp(d, DotOp(a, b, c)) and c==0 => DotOp(a, b, d)
|
|
def CombineDotAddIPattern : Pat<
|
|
(Arith_AddIOp $d, (TT_DotOp:$res $a, $b, $c, $allowTF32)),
|
|
(TT_DotOp $a, $b, $d, $allowTF32),
|
|
[(Constraint<CPred<"isZero($0)">> $c)]>;
|
|
def CombineDotAddFPattern : Pat<
|
|
(Arith_AddFOp $d, (TT_DotOp:$res $a, $b, $c, $allowTF32)),
|
|
(TT_DotOp $a, $b, $d, $allowTF32),
|
|
[(Constraint<CPred<"isZero($0)">> $c)]>;
|
|
|
|
def CombineDotAddIRevPattern : Pat<
|
|
(Arith_AddIOp (TT_DotOp:$res $a, $b, $c, $allowTF32), $d),
|
|
(TT_DotOp $a, $b, $d, $allowTF32),
|
|
[(Constraint<CPred<"isZero($0)">> $c)]>;
|
|
def CombineDotAddFRevPattern : Pat<
|
|
(Arith_AddFOp (TT_DotOp:$res $a, $b, $c, $allowTF32), $d),
|
|
(TT_DotOp $a, $b, $d, $allowTF32),
|
|
[(Constraint<CPred<"isZero($0)">> $c)]>;
|
|
|
|
// TODO: this fails for addptr(addptr(ptr, i32), i64)
|
|
// Commented out until fixed
|
|
// addptr(addptr(%ptr, %idx0), %idx1) => addptr(%ptr, AddI(%idx0, %idx1))
|
|
// Note: leave (sub %c0, %c0) canceling to ArithmeticDialect
|
|
// (ref: ArithmeticCanonicalization.td)
|
|
// def CombineAddPtrPattern : Pat<
|
|
// (TT_AddPtrOp (TT_AddPtrOp $ptr, $idx0), $idx1),
|
|
// (TT_AddPtrOp $ptr, (Arith_AddIOp $idx0, $idx1))>;
|
|
|
|
// broadcast(cst) => cst
|
|
def getConstantValue : NativeCodeCall<"getConstantValue($_builder, $0, $1)">;
|
|
def CombineBroadcastConstantPattern : Pat<
|
|
(TT_BroadcastOp:$bcast_res (Arith_ConstantOp $value)),
|
|
(Arith_ConstantOp (getConstantValue $value, $bcast_res)),
|
|
[(Constraint<CPred<"isBroadcastConstantCombinable($0)">> $value)]>;
|
|
|
|
#endif
|