Raise nested reductions and lower tensor products to cuTensorNet

Author: arpitj1

generic_solver/kernel_library_phase2.mlir

@@ -1486,6 +1486,18 @@ module {
     kernel.yield %C : tensor<?x2xf32>
   }
 
+  // Separable 3D tensor product: ai,bj,ck,ijk->abc. The operands are the
+  // rank-6 broadcast/submap views produced by raising; ABI lowering unwraps
+  // them to the shared psi buffer, the u buffer, and the output buffer.
+  kernel.defn @cutensornetTensorProduct3D_f32_tensor(
+      %psiA: tensor<?x?x?x?x?x?xf32>,
+      %psiB: tensor<?x?x?x?x?x?xf32>,
+      %psiC: tensor<?x?x?x?x?x?xf32>,
+      %u: tensor<?x?x?x?x?x?xf32>,
+      %out: tensor<?x?x?x?x?x?xf32>) -> tensor<?x?x?x?x?x?xf32> {
+    kernel.yield %out : tensor<?x?x?x?x?x?xf32>
+  }
+
   kernel.defn @cudnnConvolution2D_9tap_f16(
       %A0: memref<?x?xf16, strided<[?, 1], offset: ?>>,
       %A1: memref<?x?xf16, strided<[?, 1], offset: ?>>,

issues/tensor_product_3d_cutensornet_harness.c

@@ -0,0 +1,41 @@
+#include <math.h>
+#include <stdint.h>
+#include <stdio.h>
+
+#define KP 4
+#define KQ 5
+
+// Bare-pointer memref ABI for three memref<?xf32> arguments.
+extern void tensor_product_3d(
+    float *, float *, int64_t, int64_t, int64_t,
+    float *, float *, int64_t, int64_t, int64_t,
+    float *, float *, int64_t, int64_t, int64_t);
+
+int main(void) {
+  float psi[KQ * KP], u[KP * KP * KP], out[KQ * KQ * KQ];
+  for (int i = 0; i < KQ * KP; ++i)
+    psi[i] = (float)(i - 7) / 13.0f;
+  for (int i = 0; i < KP * KP * KP; ++i)
+    u[i] = (float)(i % 11 - 5) / 9.0f;
+
+  tensor_product_3d(
+      psi, psi, 0, KQ * KP, 1,
+      u, u, 0, KP * KP * KP, 1,
+      out, out, 0, KQ * KQ * KQ, 1);
+
+  float maxErr = 0.0f;
+  for (int a = 0; a < KQ; ++a)
+    for (int b = 0; b < KQ; ++b)
+      for (int c = 0; c < KQ; ++c) {
+        float ref = 0.0f;
+        for (int i = 0; i < KP; ++i)
+          for (int j = 0; j < KP; ++j)
+            for (int k = 0; k < KP; ++k)
+              ref += psi[a * KP + i] * psi[b * KP + j] *
+                     psi[c * KP + k] * u[(i * KP + j) * KP + k];
+        float err = fabsf(out[(a * KQ + b) * KQ + c] - ref);
+        if (err > maxErr) maxErr = err;
+      }
+  printf("cutensornet tensor product max_err=%g\n", maxErr);
+  return maxErr <= 2.0e-5f ? 0 : 1;
+}

lib/polygeist/Passes/LowerKernelLaunchToCuBLAS.cpp

@@ -124,6 +124,8 @@ static StringRef shimSymbolFor(StringRef libSym) {
     return "polygeist_cufft_z2z_1d";
   if (libSym == "cufftC2C_1D_tensor")
     return "polygeist_cufft_c2c_1d";
+  if (libSym == "cutensornetTensorProduct3D_f32_tensor")
+    return "polygeist_cutensornet_tensor_product_3d_f32";
   // NOTE: cudnnConvolution2D_9tap_i{8,16} are intentionally absent — those
   // launches route to PVA Solutions' libpva_operator and are lowered by
   // a separate pass (see LowerKernelLaunchToPVA.cpp). cuDNN itself has
@@ -1007,6 +1009,71 @@ static LogicalResult lowerCufftC2C1DTensor(LaunchOp launch, ModuleOp module,
   return success();
 }
 
+// Tensor-product contraction
+//   out[a,b,c] = sum(i,j,k) psi[a,i] psi[b,j] psi[c,k] u[i,j,k]
+// reaches the matcher as five rank-6 submap views over three flat buffers.
+// The first three views share the psi base.  Unwrap all views and let the
+// cuTensorNet shim express the full Einstein contraction directly.
+static LogicalResult lowerCutensornetTensorProduct3DF32(LaunchOp launch,
+                                                        ModuleOp module) {
+  if (launch.getNumOperands() != 5 || launch.getNumResults() != 1)
+    return launch.emitError(
+        "cuTensorNet tensor product: expected 5 operands and 1 result");
+
+  for (Value operand : launch.getOperands()) {
+    auto ty = dyn_cast<RankedTensorType>(operand.getType());
+    if (!ty || ty.getRank() != 6 || !ty.getElementType().isF32())
+      return launch.emitError(
+          "cuTensorNet tensor product: operands must be rank-6 f32 tensors");
+    auto submap = operand.getDefiningOp<polygeist::SubmapOp>();
+    if (!submap || submap.getSizes().size() != 6)
+      return launch.emitError(
+          "cuTensorNet tensor product: operands must be rank-6 submaps");
+  }
+
+  Value psi0 = resolveSubmapBase(launch.getOperand(0));
+  Value psi1 = resolveSubmapBase(launch.getOperand(1));
+  Value psi2 = resolveSubmapBase(launch.getOperand(2));
+  Value u = resolveSubmapBase(launch.getOperand(3));
+  Value out = resolveSubmapBase(launch.getOperand(4));
+  if (psi0 != psi1 || psi0 != psi2)
+    return launch.emitError(
+        "cuTensorNet tensor product: first three views must share psi base");
+
+  auto psiTy = dyn_cast<RankedTensorType>(psi0.getType());
+  auto uTy = dyn_cast<RankedTensorType>(u.getType());
+  auto outTy = dyn_cast<RankedTensorType>(out.getType());
+  if (!psiTy || !uTy || !outTy || !psiTy.getElementType().isF32() ||
+      !uTy.getElementType().isF32() || !outTy.getElementType().isF32())
+    return launch.emitError(
+        "cuTensorNet tensor product: submap bases must be f32 tensors");
+
+  auto firstView = launch.getOperand(0).getDefiningOp<polygeist::SubmapOp>();
+  OpBuilder b(launch);
+  Location loc = launch.getLoc();
+  Value KQ = valueAsI32(b, loc, firstView.getSizes()[0]);
+  Value KP = valueAsI32(b, loc, firstView.getSizes()[3]);
+  Value psiMr = tensorToMemref(b, loc, psi0);
+  Value uMr = tensorToMemref(b, loc, u);
+  Value outMr = tensorToMemref(b, loc, out);
+  Value psiPtr = memrefBasePtr(b, loc, psiMr);
+  Value uPtr = memrefBasePtr(b, loc, uMr);
+  Value outPtr = memrefBasePtr(b, loc, outMr);
+
+  auto ptrTy = LLVM::LLVMPointerType::get(b.getContext());
+  SmallVector<Type> argTypes = {b.getI32Type(), b.getI32Type(), ptrTy,
+                                ptrTy, ptrTy};
+  func::FuncOp shim = ensureShimDecl(
+      module, "polygeist_cutensornet_tensor_product_3d_f32", argTypes, b);
+  b.create<func::CallOp>(loc, shim,
+                         ValueRange{KQ, KP, psiPtr, uPtr, outPtr});
+
+  Value updatedOut = memrefToTensor(b, loc, outMr, out.getType());
+  rewireLaunchResult(launch, updatedOut);
+  launch.erase();
+  return success();
+}
+
 // Darknet im2col+GEMM reaches the matcher as rank-3 broadcasted submaps:
 //   A(m, k, n) -> weights[m, k]
 //   B(m, k, n) -> workspace[k, n]
@@ -2974,6 +3041,8 @@ struct LowerKernelLaunchToCuBLASPass
       } else if (libSym == "cufftZ2Z_1D_tensor" ||
                  libSym == "cufftC2C_1D_tensor") {
         r = lowerCufftC2C1DTensor(launch, module, shim);
+      } else if (libSym == "cutensornetTensorProduct3D_f32_tensor") {
+        r = lowerCutensornetTensorProduct3DF32(launch, module);
       } else if (libSym == "cudnnConvolutionFwd_batched") {
         r = lowerCudnnConv2dBatched(launch, module);
       } else if (libSym == "cudnnConvolutionFwd_im2col_gemm") {