668: Push-down of arithmetics before source operations

src/ir/daphneir/Daphne.h

@@ -82,6 +82,7 @@ BoolOrUnknown stringToBoolOrUnknown(const std::string &str);
 #include <ir/daphneir/DaphneInferSparsityOpInterface.h>
 #include <ir/daphneir/DaphneInferSymmetricOpInterface.h>
 #include <ir/daphneir/DaphneInferTypesOpInterface.h>
+#include <ir/daphneir/DaphnePushDownTraits.h>
 
 #include <string>
 #include <utility>

src/ir/daphneir/DaphneOps.td

@@ -30,6 +30,7 @@ include "ir/daphneir/DaphneInferSparsityOpInterface.td"
 include "ir/daphneir/DaphneInferSparsityTraits.td"
 include "ir/daphneir/DaphneInferTypesOpInterface.td"
 include "ir/daphneir/DaphneInferTypesTraits.td"
+include "ir/daphneir/DaphnePushDownTraits.td"
 include "ir/daphneir/DaphneVectorizableOpInterface.td"
 include "ir/daphneir/DaphneTypeAdaptationTraits.td"
 include "ir/daphneir/CUDASupport.td"
@@ -119,7 +120,7 @@ def Daphne_ConvertCSRMatrixToRowOffsetsMemRef : Daphne_Op<"convertCSRMatrixToRow
 
 def Daphne_FillOp : Daphne_Op<"fill", [
     DataTypeMat, ValueTypeFromFirstArg,
-    NumRowsFromIthScalar<1>, NumColsFromIthScalar<2>, CUDASupport,
+    NumRowsFromIthScalar<1>, NumColsFromIthScalar<2>, CUDASupport, Pure,
     DeclareOpInterfaceMethods<InferSymmetricOpInterface>
 ]> {
     let arguments = (ins AnyScalar:$arg, Size:$numRows, Size:$numCols);
@@ -157,7 +158,7 @@ def Daphne_DiagMatrixOp : Daphne_Op<"diagMatrix", [
 def Daphne_RandMatrixOp : Daphne_Op<"randMatrix", [
     TypesMatchWith<"min and max arguments must be of the same type", "min", "max", "$_self">,
     NumRowsFromIthScalar<0>, NumColsFromIthScalar<1>, DeclareOpInterfaceMethods<InferTypesOpInterface>,
-    SparsityFromIthScalar<4>, CastArgsToResTypeRandMatrixOp
+    SparsityFromIthScalar<4>, CastArgsToResTypeRandMatrixOp, Pure
 ]> {
     let arguments = (ins Size:$numRows, Size:$numCols, AnyScalar:$min, AnyScalar:$max, F64:$sparsity, IntScalar:$seed);
     let results = (outs MatrixOrU:$res);
@@ -175,7 +176,8 @@ def Daphne_SeqOp : Daphne_Op<"seq", [
     DataTypeMat, ValueTypeFromArgs,
     OneCol, DeclareOpInterfaceMethods<InferNumRowsOpInterface>,
     CompletelyDense,
-    CastArgsToResType
+    CastArgsToResType,
+    Pure
 ]> {
     let arguments = (ins NumScalar:$from, NumScalar:$to, NumScalar:$inc);
     let results = (outs MatrixOf<[NumScalar]>);
@@ -254,15 +256,25 @@ class Daphne_EwUnaryOp<string name, Type scalarType, list<Trait> traits = []> :
 // ----------------------------------------------------------------------------
 
 // TODO EwMinusOp: Should an unsigned integer argument yield a signed integer result?
-def Daphne_EwMinusOp : Daphne_EwUnaryOp<"ewMinus", NumScalar, [ValueTypeFromFirstArg]> {
+def Daphne_EwMinusOp : Daphne_EwUnaryOp<"ewMinus", NumScalar, [ValueTypeFromFirstArg,PushDownLinear]> {
    let hasFolder = 1;
    let hasCanonicalizeMethod = 1;
 }
-def Daphne_EwAbsOp : Daphne_EwUnaryOp<"ewAbs", NumScalar, [ValueTypeFromFirstArg]>;
-def Daphne_EwSignOp : Daphne_EwUnaryOp<"ewSign", NumScalar, [ValueTypeFromFirstArg]>;
-def Daphne_EwExpOp : Daphne_EwUnaryOp<"ewExp", NumScalar, [ValueTypeFromArgsFP]>;
-def Daphne_EwLnOp : Daphne_EwUnaryOp<"ewLn", NumScalar, [ValueTypeFromArgsFP]>;
-def Daphne_EwSqrtOp : Daphne_EwUnaryOp<"ewSqrt", NumScalar, [ValueTypeFromArgsFP, DeclareOpInterfaceMethods<VectorizableOpInterface>]>;
+def Daphne_EwAbsOp : Daphne_EwUnaryOp<"ewAbs", NumScalar, [ValueTypeFromFirstArg]>{
+   let hasCanonicalizeMethod = 1;
+}
+def Daphne_EwSignOp : Daphne_EwUnaryOp<"ewSign", NumScalar, [ValueTypeFromFirstArg]>{
+   let hasCanonicalizeMethod = 1;
+}
+def Daphne_EwExpOp : Daphne_EwUnaryOp<"ewExp", NumScalar, [ValueTypeFromArgsFP]>{
+   let hasCanonicalizeMethod = 1;
+}
+def Daphne_EwLnOp : Daphne_EwUnaryOp<"ewLn", NumScalar, [ValueTypeFromArgsFP]>{
+   let hasCanonicalizeMethod = 1;
+}
+def Daphne_EwSqrtOp : Daphne_EwUnaryOp<"ewSqrt", NumScalar, [ValueTypeFromArgsFP, DeclareOpInterfaceMethods<VectorizableOpInterface>]>{
+   let hasCanonicalizeMethod = 1;
+}
 
 // ----------------------------------------------------------------------------
 // Logical
@@ -334,21 +346,27 @@ class Daphne_EwBinaryOp<string name, Type scalarType, list<Trait> traits = []>
 // ----------------------------------------------------------------------------
 
 // TODO Make EwAddOp Commutative again (see #449).
-def Daphne_EwAddOp    : Daphne_EwBinaryOp<"ewAdd", AnyScalar, [ValueTypeFromArgs, CastArgsToResType/*, Commutative*/, EwSparseIfBoth, CUDASupport]> {
+def Daphne_EwAddOp    : Daphne_EwBinaryOp<"ewAdd", AnyScalar, [ValueTypeFromArgs, CastArgsToResType/*, Commutative*/, EwSparseIfBoth, CUDASupport, PushDownLinear]> {
+    let hasCanonicalizeMethod = 1;
+}
+def Daphne_EwSubOp    : Daphne_EwBinaryOp<"ewSub", NumScalar, [ValueTypeFromArgs, CastArgsToResType, EwSparseIfBoth, CUDASupport, PushDownLinear]> {
+    let hasCanonicalizeMethod = 1;
+}
+def Daphne_EwMulOp    : Daphne_EwBinaryOp<"ewMul", NumScalar, [ValueTypeFromArgs, CastArgsToResType, Commutative, EwSparseIfEither, CUDASupport, PushDownLinear, PushDownIncrementUpdate]> {
+    let hasCanonicalizeMethod = 1;
+}
+def Daphne_EwDivOp    : Daphne_EwBinaryOp<"ewDiv", NumScalar, [ValueTypeFromArgs, CastArgsToResType, CUDASupport, PushDownLinear, PushDownIncrementUpdate]> {
     let hasCanonicalizeMethod = 1;
 }
-def Daphne_EwSubOp    : Daphne_EwBinaryOp<"ewSub", NumScalar, [ValueTypeFromArgs, CastArgsToResType, EwSparseIfBoth, CUDASupport]> {
+def Daphne_EwPowOp    : Daphne_EwBinaryOp<"ewPow", NumScalar, [ValueTypeFromArgs, CastArgsToResType, CUDASupport]> {
     let hasCanonicalizeMethod = 1;
 }
-def Daphne_EwMulOp    : Daphne_EwBinaryOp<"ewMul", NumScalar, [ValueTypeFromArgs, CastArgsToResType, Commutative, EwSparseIfEither, CUDASupport]> {
+def Daphne_EwModOp    : Daphne_EwBinaryOp<"ewMod", NumScalar, [ValueTypeFromArgs, CastArgsToResType]> {
     let hasCanonicalizeMethod = 1;
 }
-def Daphne_EwDivOp    : Daphne_EwBinaryOp<"ewDiv", NumScalar, [ValueTypeFromArgs, CastArgsToResType, CUDASupport]> {
+def Daphne_EwLogOp    : Daphne_EwBinaryOp<"ewLog", NumScalar, [ValueTypeFromArgsFP, CastArgsToResType]> {
     let hasCanonicalizeMethod = 1;
 }
-def Daphne_EwPowOp    : Daphne_EwBinaryOp<"ewPow", NumScalar, [ValueTypeFromArgs, CastArgsToResType, CUDASupport]>;
-def Daphne_EwModOp    : Daphne_EwBinaryOp<"ewMod", NumScalar, [ValueTypeFromArgs, CastArgsToResType]>;
-def Daphne_EwLogOp    : Daphne_EwBinaryOp<"ewLog", NumScalar, [ValueTypeFromArgsFP, CastArgsToResType]>;
 
 // ----------------------------------------------------------------------------
 // Min/max

src/ir/daphneir/DaphnePushDownTraits.h

@@ -0,0 +1,24 @@
+#ifndef SRC_IR_DAPHNEIR_DAPHNEPUSHDOWNTRAITS_H
+#define SRC_IR_DAPHNEIR_DAPHNEPUSHDOWNTRAITS_H
+
+namespace mlir::OpTrait {
+
+// ============================================================================
+// Traits determining advanced push-down rewrite possibilities
+// ============================================================================
+
+/**
+ * @brief This trait is for operations that are linear and might allow for push-downs into functions that
+ * act on e.g. ranges
+ */
+template <class ConcreteOp> class PushDownLinear : public TraitBase<ConcreteOp, PushDownLinear> {};
+
+/**
+ * @brief This trait is for operations that are using an increment value which
+ * needs to be accounted for during push-down optimizations
+ */
+template <class ConcreteOp> class PushDownIncrementUpdate : public TraitBase<ConcreteOp, PushDownIncrementUpdate> {};
+
+} // namespace mlir::OpTrait
+
+#endif // SRC_IR_DAPHNEIR_DAPHNEPUSHDOWNTRAITS_H

src/ir/daphneir/DaphnePushDownTraits.td

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2025 The DAPHNE Consortium
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef SRC_IR_DAPHNEIR_DAPHNEPUSHDOWNTRAITS_TD
+#define SRC_IR_DAPHNEIR_DAPHNEPUSHDOWNTRAITS_TD
+
+include "mlir/IR/OpBase.td"
+
+// ============================================================================
+// Traits determining pushdown rewrite possibilities
+// ============================================================================
+
+def PushDownLinear: NativeOpTrait<"PushDownLinear">;
+def PushDownIncrementUpdate: NativeOpTrait<"PushDownIncrementUpdate">;
+
+#endif //SRC_IR_DAPHNEIR_DAPHNEPUSHDOWNTRAITS_TD

test/CMakeLists.txt

@@ -67,7 +67,9 @@ set(TEST_SOURCES
         api/cli/codegen/TransposeTest.cpp
 
         ir/daphneir/InferTypesTest.cpp
+        ir/daphneir/PushdownArithmetics.cpp
         api/cli/operations/CanonicalizationConstantFoldingOpTest.cpp
+        api/cli/operations/CanonicalizationPushDownTest.cpp
 
         parser/config/ConfigParserTest.cpp
 

test/api/cli/distributed/distributed_4.daphne

@@ -1,3 +1,3 @@
 m1 = rand(10, 10, 0.0, 20.0, 1.0, 1);
-m2 = rand(10, 10, 0.0, 20.0, 1.0, 1);
-print(m1 @ m2);
+m2 = rand(10, 10, 0.0, 30.0, 1.0, 1);
+print(m1 @ m2);

test/api/cli/operations/CanonicalizationPushDownTest.cpp

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2021 The DAPHNE Consortium
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <api/cli/Utils.h>
+
+#include <tags.h>
+
+#include <catch.hpp>
+
+#include <sstream>
+#include <string>
+
+const std::string dirPath = "test/api/cli/operations/";
+
+#define MAKE_TEST_CASE(name, count)                                                                                    \
+    TEST_CASE(name, TAG_OPERATIONS) {                                                                                  \
+        for (unsigned i = 1; i <= count; i++) {                                                                        \
+            DYNAMIC_SECTION(name "_" << i << ".daphne") { compareDaphneToRefSimple(dirPath, name, i); }                \
+        }                                                                                                              \
+    }
+
+// These tests simply ensure that basic functionality stays
+// functional after pushdown optimizations are performed.
+// Check the dedicated "pushdown_arithmetics" testcase for more in depth testing
+// of the canonicalizations.
+MAKE_TEST_CASE("fillEwAdd", 2)
+MAKE_TEST_CASE("fillEwSub", 2)
+MAKE_TEST_CASE("fillEwMul", 2)
+MAKE_TEST_CASE("fillEwDiv", 2)
+MAKE_TEST_CASE("fillEwPow", 2)
+MAKE_TEST_CASE("fillEwMod", 2)
+MAKE_TEST_CASE("fillEwLog", 2)
+MAKE_TEST_CASE("fillEwAbs", 2)
+MAKE_TEST_CASE("fillEwSign", 2)
+MAKE_TEST_CASE("fillEwExp", 2)
+MAKE_TEST_CASE("fillEwLn", 1)
+MAKE_TEST_CASE("fillEwSqrt", 1)
+MAKE_TEST_CASE("fillEwMinus", 2)
+MAKE_TEST_CASE("randEwAdd", 2)
+MAKE_TEST_CASE("randEwSub", 1)
+MAKE_TEST_CASE("randEwMul", 3)
+MAKE_TEST_CASE("randEwDiv", 3)
+MAKE_TEST_CASE("randEwAbs", 3)
+MAKE_TEST_CASE("randEwMinus", 3)
+MAKE_TEST_CASE("seqEwAdd", 2)
+MAKE_TEST_CASE("seqEwSub", 2)
+MAKE_TEST_CASE("seqEwMul", 3)
+MAKE_TEST_CASE("seqEwDiv", 3)

test/api/cli/operations/fillEwAbs_1.daphne

@@ -0,0 +1,4 @@
+fillIntermediate = fill(-3,3,3);
+result = abs(fillIntermediate);
+
+print(result);

test/api/cli/operations/fillEwAbs_1.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, int64_t)
+3 3 3
+3 3 3
+3 3 3

test/api/cli/operations/fillEwAbs_2.daphne

@@ -0,0 +1,4 @@
+fillIntermediate = fill(3,3,3);
+result = abs(fillIntermediate);
+
+print(result);

test/api/cli/operations/fillEwAbs_2.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, int64_t)
+3 3 3
+3 3 3
+3 3 3

test/api/cli/operations/fillEwAdd_1.daphne

@@ -0,0 +1,3 @@
+result = fill(1,3,3) + 3;
+
+print(result);

test/api/cli/operations/fillEwAdd_1.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, int64_t)
+4 4 4
+4 4 4
+4 4 4

test/api/cli/operations/fillEwAdd_2.daphne

@@ -0,0 +1,2 @@
+result = 2 + fill(1.0,2,2);
+print(result);

test/api/cli/operations/fillEwAdd_2.txt

@@ -0,0 +1,3 @@
+DenseMatrix(2x2, double)
+3 3
+3 3

test/api/cli/operations/fillEwDiv_1.daphne

@@ -0,0 +1,3 @@
+result = fill(12.3,2,3) / 3;
+
+print(result);

test/api/cli/operations/fillEwDiv_1.txt

@@ -0,0 +1,3 @@
+DenseMatrix(2x3, double)
+4.1 4.1 4.1
+4.1 4.1 4.1

test/api/cli/operations/fillEwDiv_2.daphne

@@ -0,0 +1,3 @@
+result = fill(12,3,2) / 3;
+
+print(result);

test/api/cli/operations/fillEwDiv_2.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x2, int64_t)
+4 4
+4 4
+4 4

test/api/cli/operations/fillEwExp_1.daphne

@@ -0,0 +1,4 @@
+fillIntermediate = fill(3,3,3);
+result = exp(fillIntermediate);
+
+print(result);

test/api/cli/operations/fillEwExp_1.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, double)
+20.0855 20.0855 20.0855
+20.0855 20.0855 20.0855
+20.0855 20.0855 20.0855

test/api/cli/operations/fillEwExp_2.daphne

@@ -0,0 +1,4 @@
+fillIntermediate = fill(-3,3,3);
+result = exp(fillIntermediate);
+
+print(result);

test/api/cli/operations/fillEwExp_2.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, double)
+0.0497871 0.0497871 0.0497871
+0.0497871 0.0497871 0.0497871
+0.0497871 0.0497871 0.0497871

test/api/cli/operations/fillEwLn_1.daphne

@@ -0,0 +1,4 @@
+fillIntermediate = fill(3,3,3);
+result = ln(fillIntermediate);
+
+print(result);

test/api/cli/operations/fillEwLn_1.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, double)
+1.09861 1.09861 1.09861
+1.09861 1.09861 1.09861
+1.09861 1.09861 1.09861

test/api/cli/operations/fillEwLog_1.daphne

@@ -0,0 +1,4 @@
+a = fill(8,3,3);
+result = log(a,2);
+
+print(result);

test/api/cli/operations/fillEwLog_1.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, double)
+3 3 3
+3 3 3
+3 3 3

test/api/cli/operations/fillEwLog_2.daphne

@@ -0,0 +1,4 @@
+a = fill(16.0,3,3);
+result = log(a,2);
+
+print(result);

test/api/cli/operations/fillEwLog_2.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, double)
+4 4 4
+4 4 4
+4 4 4

test/api/cli/operations/fillEwMinus_1.daphne

@@ -0,0 +1,3 @@
+result = - fill(3,2,3);
+
+print(result);

test/api/cli/operations/fillEwMinus_1.txt

@@ -0,0 +1,3 @@
+DenseMatrix(2x3, int64_t)
+-3 -3 -3
+-3 -3 -3

test/api/cli/operations/fillEwMinus_2.daphne

@@ -0,0 +1,3 @@
+result = - fill(-3,2,3);
+
+print(result);

test/api/cli/operations/fillEwMinus_2.txt

@@ -0,0 +1,3 @@
+DenseMatrix(2x3, int64_t)
+3 3 3
+3 3 3

test/api/cli/operations/fillEwMod_1.daphne

@@ -0,0 +1,3 @@
+result = mod(fill(33,2,2),5);
+
+print(result);

test/api/cli/operations/fillEwMod_1.txt

@@ -0,0 +1,3 @@
+DenseMatrix(2x2, int64_t)
+3 3
+3 3

test/api/cli/operations/fillEwMod_2.daphne

@@ -0,0 +1,3 @@
+result = fill(33.0,2,2) % 5;
+
+print(result);

test/api/cli/operations/fillEwMod_2.txt

@@ -0,0 +1,3 @@
+DenseMatrix(2x2, double)
+3 3
+3 3

test/api/cli/operations/fillEwMul_1.daphne

@@ -0,0 +1,3 @@
+result = fill(1,3,3) * 3;
+
+print(result);

test/api/cli/operations/fillEwMul_1.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, int64_t)
+3 3 3
+3 3 3
+3 3 3

test/api/cli/operations/fillEwMul_2.daphne

@@ -0,0 +1,2 @@
+result = fill(2.1,3,3) * 4;
+print(result);

test/api/cli/operations/fillEwMul_2.txt

@@ -0,0 +1,4 @@
+DenseMatrix(3x3, double)
+8.4 8.4 8.4
+8.4 8.4 8.4
+8.4 8.4 8.4