preserve linalg.matvec as kernel.linear_transform in ConvertToCiphertextSemantics

This change adds a separate pattern PreserveLinalgMatvecAsLinearTransform with higher priority that identifies constant public matrix and secret vector inputs and preserves it as a generic kernel.linear_transform op.

PiperOrigin-RevId: 961045685

Author: j2kun

lib/Transforms/ConvertToCiphertextSemantics/BUILD

@@ -15,6 +15,7 @@ cc_library(
         ":TypeConversion",
         ":pass_inc_gen",
         "@heir//lib/Dialect:ModuleAttributes",
+        "@heir//lib/Dialect/Kernel/IR:Dialect",
         "@heir//lib/Dialect/Secret/IR:SecretAttributes",
         "@heir//lib/Dialect/Secret/IR:SecretPatterns",
         "@heir//lib/Dialect/TensorExt/IR:Dialect",
@@ -24,6 +25,7 @@ cc_library(
         "@heir//lib/Kernel:IRMaterializingVisitor",
         "@heir//lib/Kernel:KernelImplementation",
         "@heir//lib/Kernel:Utils",
+        "@heir//lib/Target/CompilationTarget",
         "@heir//lib/Transforms/DropUnitDims",
         "@heir//lib/Transforms/LayoutPropagation:Utils",
         "@heir//lib/Utils",

lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp

@@ -12,6 +12,7 @@
 #include <utility>
 #include <vector>
 
+#include "lib/Dialect/Kernel/IR/KernelOps.h"
 #include "lib/Dialect/ModuleAttributes.h"
 #include "lib/Dialect/Secret/IR/SecretAttributes.h"
 #include "lib/Dialect/Secret/IR/SecretDialect.h"
@@ -26,6 +27,7 @@
 #include "lib/Kernel/KernelImplementation.h"
 #include "lib/Kernel/KernelName.h"
 #include "lib/Kernel/Utils.h"
+#include "lib/Target/CompilationTarget/CompilationTarget.h"
 #include "lib/Transforms/ConvertToCiphertextSemantics/AssignLayout.h"
 #include "lib/Transforms/ConvertToCiphertextSemantics/TypeConversion.h"
 #include "lib/Transforms/DropUnitDims/DropUnitDims.h"
@@ -895,6 +897,125 @@ struct ConvertLinalgMatvecLayout : public ConversionBase<linalg::MatvecOp> {
   bool unrollKernels;
 };
 
+struct PreserveLinalgMatvecAsLinearTransform
+    : public ConversionBase<linalg::MatvecOp> {
+ public:
+  using ConversionBase<linalg::MatvecOp>::ConversionBase;
+
+  PreserveLinalgMatvecAsLinearTransform(
+      const ContextAwareTypeConverter& typeConverter, MLIRContext* context)
+      : ConversionBase<linalg::MatvecOp>(typeConverter, context,
+                                         /*benefit=*/20) {}
+
+  LogicalResult matchAndRewrite(
+      linalg::MatvecOp op, OpAdaptor adaptor,
+      ContextAwareConversionPatternRewriter& rewriter) const final {
+    auto target = getTargetConfig(op->getParentOfType<ModuleOp>());
+    if (failed(target) || !target->has_kernel_linear_transform) {
+      return rewriter.notifyMatchFailure(op, "linear transform not enabled");
+    }
+
+    Value matrixOperand = op.getInputs()[0];
+    LayoutAttr matrixLayout = getLayoutAttr(matrixOperand);
+    if (!matrixLayout) {
+      return rewriter.notifyMatchFailure(op, "missing layout for matrix");
+    }
+
+    Value matrix = matrixOperand;
+    if (auto assignLayoutOp =
+            matrix.getDefiningOp<tensor_ext::AssignLayoutOp>()) {
+      matrix = assignLayoutOp.getValue();
+    }
+    auto constantMatrixOp = matrix.getDefiningOp<arith::ConstantOp>();
+    if (!constantMatrixOp) {
+      return rewriter.notifyMatchFailure(op, "matrix is not a constant");
+    }
+    auto denseAttr = dyn_cast<DenseElementsAttr>(constantMatrixOp.getValue());
+    if (!denseAttr) {
+      return rewriter.notifyMatchFailure(op,
+                                         "matrix is not a DenseElementsAttr");
+    }
+
+    auto matrixType = cast<RankedTensorType>(matrix.getType());
+    auto convertedMatrixType = cast<ShapedType>(
+        getTypeConverter()->convertType(matrixType, matrixLayout));
+    if (!convertedMatrixType) {
+      return rewriter.notifyMatchFailure(op, "failed to convert matrix type");
+    }
+
+    int64_t numDiagonals = convertedMatrixType.getShape()[0];
+    int64_t slots = convertedMatrixType.getShape()[1];
+    auto elementType = matrixType.getElementType();
+
+    Attribute zeroAttr = rewriter.getZeroAttr(elementType);
+    std::vector<Attribute> diagonalValues(numDiagonals * slots, zeroAttr);
+
+    auto matrixRelation = matrixLayout.getIntegerRelation();
+    PointPairCollector collector(2, 2);
+    enumeratePoints(matrixRelation, collector);
+
+    int64_t numCols = matrixType.getDimSize(1);
+    for (const auto& pointPair : collector.points) {
+      int64_t row = pointPair.first[0];
+      int64_t col = pointPair.first[1];
+      int64_t d = pointPair.second[0];
+      int64_t s = pointPair.second[1];
+
+      int64_t flatIndex = row * numCols + col;
+      Attribute val = denseAttr.getValues<Attribute>()[flatIndex];
+      diagonalValues[d * slots + s] = val;
+    }
+
+    std::vector<int64_t> nonZeroDiagonalIndices;
+    std::vector<Attribute> nonZeroDiagonalValues;
+    for (int64_t d = 0; d < numDiagonals; ++d) {
+      bool isZero = true;
+      for (int64_t s = 0; s < slots; ++s) {
+        if (diagonalValues[d * slots + s] != zeroAttr) {
+          isZero = false;
+          break;
+        }
+      }
+      if (!isZero) {
+        nonZeroDiagonalIndices.push_back(d);
+        for (int64_t s = 0; s < slots; ++s) {
+          nonZeroDiagonalValues.push_back(diagonalValues[d * slots + s]);
+        }
+      }
+    }
+
+    auto diagonalsType = RankedTensorType::get(
+        {static_cast<int64_t>(nonZeroDiagonalIndices.size()), slots},
+        elementType);
+    auto diagonalsAttr =
+        DenseElementsAttr::get(diagonalsType, nonZeroDiagonalValues);
+    auto diagonalIndicesAttr =
+        rewriter.getDenseI64ArrayAttr(nonZeroDiagonalIndices);
+
+    auto resultLayout = findAttributeAssociatedWith(
+        op.getResult(0), tensor_ext::TensorExtDialect::kLayoutAttrName);
+    if (failed(resultLayout)) {
+      return rewriter.notifyMatchFailure(op, "missing output layout");
+    }
+
+    auto outputType = op.getResult(0).getType();
+    auto convertedOutputType =
+        getTypeConverter()->convertType(outputType, resultLayout.value());
+
+    rewriter.setInsertionPointAfter(op);
+    auto linearTransformOp = rewriter.create<kernel::LinearTransformOp>(
+        op.getLoc(), convertedOutputType, adaptor.getInputs()[1], diagonalsAttr,
+        diagonalIndicesAttr, /*bsgs_ratio=*/nullptr);
+
+    setMaterializedAttr(linearTransformOp);
+    linearTransformOp->setAttr(kLayoutAttrName, resultLayout.value());
+
+    addBiasAndReplace(rewriter, op, linearTransformOp.getResult(),
+                      adaptor.getOutputs()[0], resultLayout.value());
+    return success();
+  }
+};
+
 struct ConvertLinalgConv1D : public ConversionBase<linalg::Conv1DOp> {
  public:
   using ConversionBase<linalg::Conv1DOp>::ConversionBase;
@@ -2761,13 +2882,14 @@ struct ConvertToCiphertextSemantics
       return isa<ModuleOp>(op) || hasMaterializedAttr(op);
     });
 
-    patterns.add<ConvertAnyAddingMaterializedAttr, ConvertConvertLayout,
-                 ConvertFunc, ConvertLinalgMatmul, ConvertLinalgBatchMatmul,
-                 ConvertLinalgReduce, ConvertLinalgDot, ConvertSecretGeneric,
-                 ConvertTensorCollapseShape, ConvertTensorExpandShape,
-                 ConvertTensorExtractLayout, ConvertTensorExtractSlice,
-                 ConvertTensorPad, ConvertTensorInsertLayout,
-                 ConvertTensorInsertSlice>(typeConverter, context);
+    patterns.add<
+        ConvertAnyAddingMaterializedAttr, ConvertConvertLayout, ConvertFunc,
+        ConvertLinalgMatmul, ConvertLinalgBatchMatmul, ConvertLinalgReduce,
+        ConvertLinalgDot, ConvertSecretGeneric, ConvertTensorCollapseShape,
+        ConvertTensorExpandShape, ConvertTensorExtractLayout,
+        ConvertTensorExtractSlice, ConvertTensorPad, ConvertTensorInsertLayout,
+        ConvertTensorInsertSlice, PreserveLinalgMatvecAsLinearTransform>(
+        typeConverter, context);
     patterns.add<ConvertLinalgMatvecLayout, ConvertLinalgConv1D,
                  ConvertLinalgConv2D, ConvertLinalgConv2DNchwFchw,
                  ConvertLinalgConv1DNcwFcw>(typeConverter, context,

lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.td

@@ -48,6 +48,7 @@ def ConvertToCiphertextSemantics : Pass<"convert-to-ciphertext-semantics"> {
   }];
   let dependentDialects = [
     "mlir::heir::tensor_ext::TensorExtDialect",
+    "mlir::heir::kernel::KernelDialect",
     "mlir::linalg::LinalgDialect",
     "mlir::tensor::TensorDialect",
     "mlir::scf::SCFDialect"

tests/Transforms/convert_to_ciphertext_semantics/linear_transform.mlir

@@ -0,0 +1,33 @@
+// RUN: heir-opt %s --convert-to-ciphertext-semantics=min-slot-count=4 | FileCheck %s
+
+// CHECK: module
+module attributes {
+  backend.openfhe,
+  backend.config_override = {has_kernel_linear_transform = true}
+} {
+  // CHECK: @main
+  func.func @main(%arg0: !secret.secret<tensor<4xf32>> {tensor_ext.layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 3 and 0 <= slot <= 3 }">}) -> (!secret.secret<tensor<2xf32>> {tensor_ext.layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 1 and 0 <= slot <= 3 }">}) {
+    %cst = arith.constant dense<0.0> : tensor<2xf32>
+    %cst_mat = arith.constant dense<[[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]> : tensor<2x4xf32>
+    %0 = tensor_ext.assign_layout %cst_mat {
+      layout = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : exists (e0: ct + slot - i1 - 4e0 = 0 and 0 <= ct <= 3) and slot - i0 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 3 and 0 <= slot <= 3 }">,
+      tensor_ext.layout = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : exists (e0: ct + slot - i1 - 4e0 = 0 and 0 <= ct <= 3) and slot - i0 = 0 and 0 <= i0 <= 1 and 0 <= i1 <= 3 and 0 <= slot <= 3 }">
+    } : tensor<2x4xf32>
+    %1 = tensor_ext.assign_layout %cst {
+      layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 1 and 0 <= slot <= 3 }">,
+      tensor_ext.layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 1 and 0 <= slot <= 3 }">
+    } : tensor<2xf32>
+    %2 = secret.generic(%arg0 : !secret.secret<tensor<4xf32>> {tensor_ext.layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 3 and 0 <= slot <= 3 }">}) {
+    ^body(%input: tensor<4xf32>):
+      // CHECK: kernel.linear_transform
+      // CHECK-SAME: diagonal_indices = array<i64: 0, 1, 2, 3>
+      // CHECK-SAME: diagonals = dense<{{\[\[}}1.000000e+00, 6.000000e+00, 0.000000e+00, 0.000000e+00], [2.000000e+00, 7.000000e+00, 0.000000e+00, 0.000000e+00], [3.000000e+00, 8.000000e+00, 0.000000e+00, 0.000000e+00], [4.000000e+00, 5.000000e+00, 0.000000e+00, 0.000000e+00]]> : tensor<4x4xf32>
+      %3 = linalg.matvec {
+        secret.kernel = #secret.kernel<name = "MatvecDiagonal", force = false>,
+        tensor_ext.layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 1 and 0 <= slot <= 3 }">
+      } ins(%0, %input : tensor<2x4xf32>, tensor<4xf32>) outs(%1 : tensor<2xf32>) -> tensor<2xf32>
+      secret.yield %3 : tensor<2xf32>
+    } -> (!secret.secret<tensor<2xf32>> {tensor_ext.layout = #tensor_ext.layout<"{ [i0] -> [ct, slot] : ct = 0 and slot = i0 and 0 <= i0 <= 1 and 0 <= slot <= 3 }">})
+    return %2 : !secret.secret<tensor<2xf32>>
+  }
+}