Fix zero-rank tensor handling in shape functions (#731)

- Fix getHostConstantTensor to handle zero-rank tensors correctly
- Create rank-1 tensor with zeros when values.size() is 0
- Add test for zero-rank tensor handling in shape function creation

GitOrigin-RevId: dc1ed6d94ce37d6d47661030273b042fbee85acc

Author: christopherbate

mlir-tensorrt/compiler/lib/Dialect/Plan/Transforms/CreateShapeFuncs.cpp

@@ -398,13 +398,17 @@ static MemorySpaceAttr getHostSpace(RewriterBase &rewriter) {
 
 static Value getHostConstantTensor(RewriterBase &rewriter, Location loc,
                                    ArrayRef<int64_t> values) {
-  auto rtt =
-      RankedTensorType::get({static_cast<int64_t>(values.size())},
-                            rewriter.getIndexType(), getHostSpace(rewriter));
+  auto rtt = RankedTensorType::get(
+      {std::max<int64_t>(1, static_cast<int64_t>(values.size()))},
+      rewriter.getIndexType(), getHostSpace(rewriter));
 
   // Create the DenseIntElementsAttr with host space
-  auto attr = DenseIntElementsAttr::get(rtt, values);
-  return rewriter.create<arith::ConstantOp>(loc, rtt, attr);
+  auto attrData = [&]() -> TypedAttr {
+    if (static_cast<int64_t>(values.size()) == rtt.getRank())
+      return DenseIntElementsAttr::get(rtt, values);
+    return rewriter.getZeroAttr(rtt);
+  }();
+  return rewriter.create<arith::ConstantOp>(loc, rtt, attrData);
 }
 
 static SmallVector<Value> createConstantIndices(RewriterBase &rewriter,
@@ -523,20 +527,29 @@ static FailureOr<func::FuncOp> createAggregateShapeFunc(
 
   SmallVector<Value> argValues = getArguments(rewriter, func);
   SmallVector<Type> argTypes;
+
   for (auto [idx, argValue] : llvm::enumerate(argValues)) {
     Type t = argValue.getType();
     auto rtt = dyn_cast<RankedTensorType>(t);
-    if (!rtt)
-      argTypes.push_back(t);
-    const TensorKindLattice *lattice =
-        solver.lookupState<TensorKindLattice>(argValue);
-    if (!lattice || lattice->getValue().isUninitialized())
-      return failure();
-    if (lattice->getValue().isHostVisible()) {
-      argTypes.push_back(t);
+
+    bool isHostTensorArg = [&, idx = idx, argValue = argValue]() {
+      if (!rtt)
+        return false;
+      if (func.getArgAttr(idx, plan::PlanDialect::kValueBoundsAttrName))
+        return true;
+      if (auto memSpace = func.getArgAttrOfType<plan::MemorySpaceAttr>(
+              idx, plan::PlanDialect::kMemorySpaceConstraintAttrName))
+        return memSpace.isHostVisible();
+      const TensorKindLattice *lattice =
+          solver.lookupState<TensorKindLattice>(argValue);
+      return lattice && !lattice->getValue().isUninitialized() &&
+             lattice->getValue().isHostVisible();
+    }();
+    if (!isHostTensorArg) {
+      argTypes.push_back(!rtt ? t : getShapeTensorType(rtt));
       continue;
     }
-    argTypes.push_back(getShapeTensorType(rtt));
+    argTypes.push_back(rtt);
   }
 
   FailureOr<SmallVector<Value>> returnedValues = getReturnedValues(func);
@@ -712,8 +725,6 @@ class CreateShapeFuncsPass
       (*aggShapeFunc)
           ->setAttr(PlanDialect::kShapeFuncMarkerAttrName,
                     UnitAttr::get(rewriter.getContext()));
-      if (failed(aggShapeFunc))
-        continue;
 
       // Add the symbol to the original func.
       symbolTable.insert(*aggShapeFunc);

mlir-tensorrt/compiler/test/Dialect/Plan/create-shape-funcs-zero-rank.mlir

@@ -0,0 +1,30 @@
+// RUN: mlir-tensorrt-opt %s -split-input-file -plan-create-shape-funcs="abi-version=1" -inline | FileCheck %s
+
+// Test zero-rank tensor handling in shape function creation
+// The fix ensures that when creating constant tensors with zero values,
+// we create a rank-1 tensor with zeros instead of a zero-rank tensor.
+
+func.func public @test_zero_rank_shape(%arg0: !executor.ptr<host> {executor.abi = #executor.arg<byval, tensor<f32>>},
+                                %arg1: i32,
+                                %arg2: !executor.ptr<host> {executor.abi = #executor.arg<byref, tensor<?xf32>>},
+                                %arg3: !executor.ptr<host> {executor.abi = #executor.arg<byref, tensor<f32>>})
+    attributes {executor.func_abi = (tensor<f32>, i32) -> (tensor<?xf32>, tensor<f32>)} {
+  %0 = executor.abi.recv %arg0 : tensor<f32>
+  %1 = stablehlo.exponential %0 : tensor<f32>
+  %2 = plan.with_shape %1() : (tensor<f32>) -> tensor<f32>
+  %size = tensor.from_elements %arg1 : tensor<1xi32>
+  %3 = stablehlo.dynamic_broadcast_in_dim %2, %size, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
+  %4 = plan.with_shape %3(%arg1) : (tensor<?xf32>, i32) -> tensor<?xf32>
+  executor.abi.send %4 to %arg2 : tensor<?xf32>
+  executor.abi.send %2 to %arg3 : tensor<f32>
+  return
+}
+
+// CHECK-LABEL: func.func public @test_zero_rank_shape(
+// CHECK-LABEL: func.func public @test_zero_rank_shape_get_shapes(
+//   CHECK: %[[cst:.+]] = arith.constant dense<0> : tensor<1xindex, #plan.memory_space<host>>
+//   CHECK-DAG: %[[v0:.+]] = arith.index_cast %{{.+}} : i32 to index
+//   CHECK-DAG: %[[v1:.+]] = tensor.from_elements %[[v0]] : tensor<1xindex, #plan.memory_space<host>>
+//   CHECK: executor.abi.send %[[v1]] to %{{.+}}
+//   CHECK: executor.abi.send %[[cst]] to %{{.+}}
+//   CHECK: return